Update app.py

app.py

@@ -34,17 +34,33 @@ def load_models():
 
 vision_model, vision_processor, code_model, code_tokenizer = load_models()
 
-VISION_SYSTEM_PROMPT = """You are an AI assistant specialized in analyzing images and videos of code editors. Your task is to:
-1. Extract and describe any code snippets visible in the image
-2. Identify any error messages, warnings, or highlighting that indicates bugs
-3. Describe the programming language and context if visible
-Be thorough and accurate in your description, as this will be used to fix the code."""
-
-CODE_SYSTEM_PROMPT = """You are an expert code debugging assistant. Based on the description of code and errors provided, your task is to:
-1. Identify the bugs and issues in the code
-2. Provide a corrected version of the code
-3. Explain the fixes made and why they resolve the issues
-4. Provide the output in a well-structured format removing all the unnecessary information
+VISION_SYSTEM_PROMPT = """You are an AI assistant specialized in analyzing images and videos of code editors. Your primary task is to:
+
+1. FIRST AND MOST IMPORTANTLY: Check if the image contains any inappropriate content such as:
+   - Harassment or bullying
+   - Hate speech or discriminatory content
+   - Sexually explicit material
+   - Dangerous or harmful content
+   If any such content is detected, respond ONLY with: "I apologize, but I cannot process this content as it appears to contain [type of inappropriate content]. Please provide only appropriate code-related images."
+
+2. If the content is appropriate, then:
+   - Extract and describe any code snippets visible in the image
+   - Identify any error messages, warnings, or highlighting that indicates bugs
+   - Describe the programming language and context if visible
+
+Be thorough and accurate in your description of appropriate content, as this will be used to fix the code."""
+
+CODE_SYSTEM_PROMPT = """You are an expert code debugging assistant. Your tasks in order are:
+
+1. Check if the input description contains any flags for inappropriate content.
+   If it does, respond ONLY with: "I apologize, but I cannot process this request as the original content was flagged as inappropriate."
+
+2. If the content is appropriate, then based on the description of code and errors provided:
+   - Identify the bugs and issues in the code
+   - Provide a corrected version of the code
+   - Explain the fixes made and why they resolve the issues
+   - Provide the output in a well-structured format removing all unnecessary information
+
 Be thorough in your explanation and ensure the corrected code is complete and functional."""
 
 def process_image_for_code(image):
@@ -85,6 +101,10 @@ def process_image_for_code(image):
         clean_up_tokenization_spaces=False
     )[0]
 
+    # Check if vision model flagged inappropriate content
+    if "I apologize, but I cannot process this content" in vision_description:
+        return vision_description, "No code analysis provided due to inappropriate content."
+
     # Then, use code model to fix the code
     code_messages = [
         {"role": "system", "content": CODE_SYSTEM_PROMPT},
@@ -136,7 +156,6 @@ def process_video_for_code(video_path, max_frames=16, frame_interval=30):
         
     cap.release()
     
-    # Process the first frame for now (you could extend this to handle multiple frames)
     if frames:
         return process_image_for_code(frames[0])
     else:
@@ -147,13 +166,16 @@ def process_content(content):
     if content is None:
         return "Please upload an image or video file of code with errors.", ""
 
-    if content.name.lower().endswith(('.png', '.jpg', '.jpeg')):
-        image = Image.open(content.name)
-        vision_output, code_output = process_image_for_code(image)
-    elif content.name.lower().endswith(('.mp4', '.avi', '.mov')):
-        vision_output, code_output = process_video_for_code(content.name)
-    else:
-        return "Unsupported file type. Please provide an image or video file.", ""
+    try:
+        if content.name.lower().endswith(('.png', '.jpg', '.jpeg')):
+            image = Image.open(content.name)
+            vision_output, code_output = process_image_for_code(image)
+        elif content.name.lower().endswith(('.mp4', '.avi', '.mov')):
+            vision_output, code_output = process_video_for_code(content.name)
+        else:
+            return "Unsupported file type. Please provide an image or video file.", ""
+    except Exception as e:
+        return f"An error occurred while processing the file: {str(e)}", ""
 
     return vision_output, code_output
 
@@ -166,7 +188,7 @@ iface = gr.Interface(
         gr.Code(label="Fixed Code", language="python")
     ],
     title="Vision Code Debugger",
-    description="Upload an image or video of code with errors, and the AI will analyze and fix the issues."
+    description="Upload an image or video of code with errors for AI analysis and fixes. Note: Only appropriate code-related content will be processed."
 )
 
 if __name__ == "__main__":