Update agents/image_agent.py

agents/image_agent.py

@@ -11,32 +11,42 @@ class ImageAgent:
         self.memory = memory
         self.planner = planner
 
-    # --------------------------
+    # --------------------------------------------------
     # Upload Files
-    # --------------------------
+    # --------------------------------------------------
+
     def upload(self, files):
 
         uploaded = save_uploads(files)
 
         self.memory.uploaded_files = uploaded
 
-        # Automatically choose first uploaded image
+        self.memory.image_context = None
+        self.memory.analysis = None
+
         for file in uploaded:
-            if file["type"] and file["type"].startswith("image/"):
+
+            if file.get("type", "").startswith("image/"):
+
                 self.memory.last_used_image = file["path"]
+
                 break
 
         return uploaded
 
-    # --------------------------
+    # --------------------------------------------------
     # Analyze Image
-    # --------------------------
+    # --------------------------------------------------
+
     def auto_analyze(self, image_path=None):
 
-        image = self.memory.last_used_image
+        image = image_path or self.memory.last_used_image
 
         if not image:
-            return None
+
+            return {
+                "description": "No uploaded image found."
+            }
 
         analysis = analyze_image(image)
 
@@ -45,46 +55,123 @@ class ImageAgent:
 
         return analysis
 
-    # --------------------------
-    # Main Agent
-    # --------------------------
-    def run(self, user_input):
+    # --------------------------------------------------
+    # Convert object list safely
+    # --------------------------------------------------
 
-        task = self.router.route(user_input)
+    def format_objects(self, objects):
 
-        steps = self.planner.plan(task)
+        if not objects:
+            return ""
 
-        result = None
+        names = []
 
-        image = self.memory.last_used_image
+        for obj in objects:
 
-        # Automatically analyze uploaded image once
-        if image and self.memory.image_context is None:
-            self.auto_analyze()
+            if isinstance(obj, str):
 
-        analysis = self.memory.image_context
+                names.append(obj)
+
+            elif isinstance(obj, dict):
+
+                if "name" in obj:
+
+                    names.append(str(obj["name"]))
+
+                elif "label" in obj:
+
+                    names.append(str(obj["label"]))
+
+                else:
+
+                    names.append(str(obj))
+
+            else:
+
+                names.append(str(obj))
+
+        return ", ".join(names)
+
+    # --------------------------------------------------
+    # Build Prompt
+    # --------------------------------------------------
+
+    def build_context(self, analysis):
+
+        if not analysis:
+            return ""
+
+        objects = self.format_objects(
+            analysis.get("objects", [])
+        )
+
+        colors = analysis.get("colors", [])
 
-        context = ""
+        if isinstance(colors, list):
+            colors = ", ".join(colors)
 
-        if analysis:
+        suggestions = analysis.get("suggestions", [])
 
-            context = f"""
+        if isinstance(suggestions, list):
+            suggestions = ", ".join(suggestions)
+
+        context = f"""
 Image Description:
-{analysis["description"]}
+{analysis.get("description","")}
 
 Objects:
-{", ".join(analysis.get("objects", []))}
+{objects}
+
+People:
+{analysis.get("people","")}
 
 Style:
-{analysis.get("style", "")}
+{analysis.get("style","")}
 
 Lighting:
-{analysis.get("lighting", "")}
+{analysis.get("lighting","")}
+
+Colors:
+{colors}
+
+Background:
+{analysis.get("background","")}
 
-OCR:
-{analysis.get("text", "")}
+Camera Angle:
+{analysis.get("camera_angle","")}
+
+Quality:
+{analysis.get("quality","")}
+
+Editing Prompt:
+{analysis.get("editing_prompt","")}
+
+Suggestions:
+{suggestions}
 """
 
+        return context
+
+    # --------------------------------------------------
+    # Main Agent
+    # --------------------------------------------------
+
+    def run(self, user_input):
+
+        task = self.router.route(user_input)
+
+        steps = self.planner.plan(task)
+
+        image = self.memory.last_used_image
+
+        if image and self.memory.image_context is None:
+
+            self.auto_analyze(image)
+
+        analysis = self.memory.image_context
+
+        context = self.build_context(analysis)
+
         prompt = f"""
 {context}
 
@@ -92,6 +179,8 @@ User Request:
 {user_input}
 """
 
+        result = None
+
         for step in steps:
 
             if step == "generate":
@@ -105,12 +194,20 @@ User Request:
                     prompt=prompt
                 )
 
-                # Remember edited image
-                if result:
+                # If edit_image returns a file path
+                if isinstance(result, str):
+
                     self.memory.last_used_image = result
 
+                # If edit_image returns a PIL.Image,
+                # app.py will display it directly.
+
             elif step == "analyze":
 
-                result = self.auto_analyze()
+                result = self.auto_analyze(image)
+
+            elif step == "chat":
+
+                result = analysis
 
         return result