Update agents/image_agent.py

agents/image_agent.py

@@ -11,72 +11,106 @@ class ImageAgent:
         self.memory = memory
         self.planner = planner
 
+    # --------------------------
+    # Upload Files
+    # --------------------------
     def upload(self, files):
-        paths = save_uploads(files)
-        self.memory.add_files(paths)
-        return paths
 
-    def auto_analyze(self, image_path):
-        """
-        STEP 1: Understand image
-        """
-        if image_path:
-            analysis = analyze_image(image_path)
+        uploaded = save_uploads(files)
 
-            # store context safely
-            if hasattr(self.memory, "set_image_context"):
-                self.memory.set_image_context(analysis["description"])
-            else:
-                self.memory.image_context = analysis["description"]
+        self.memory.uploaded_files = uploaded
 
-            self.memory.last_used_image = image_path
+        # Automatically choose first uploaded image
+        for file in uploaded:
+            if file["type"] and file["type"].startswith("image/"):
+                self.memory.last_used_image = file["path"]
+                break
 
-            return analysis
+        return uploaded
 
-        return None
+    # --------------------------
+    # Analyze Image
+    # --------------------------
+    def auto_analyze(self):
 
-    def run(self, user_input: str):
+        image = self.memory.last_used_image
+
+        if not image:
+            return None
+
+        analysis = analyze_image(image)
+
+        self.memory.analysis = analysis
+        self.memory.image_context = analysis
+
+        return analysis
+
+    # --------------------------
+    # Main Agent
+    # --------------------------
+    def run(self, user_input):
 
-        # ---------------- ROUTE ----------------
         task = self.router.route(user_input)
+
         steps = self.planner.plan(task)
 
         result = None
 
-        # ---------------- IMAGE PICK ----------------
-        image = None
-        if self.memory.uploaded_files:
-            image = self.memory.uploaded_files[0]
+        image = self.memory.last_used_image
+
+        # Automatically analyze uploaded image once
+        if image and self.memory.image_context is None:
+            self.auto_analyze()
+
+        analysis = self.memory.image_context
+
+        context = ""
 
-        # ---------------- AUTO ANALYZE ----------------
-        if image and not self.memory.image_context:
-            analysis = analyze_image(image)
-            self.memory.image_context = analysis["description"]
+        if analysis:
 
-        # ---------------- COMBINE CONTEXT ----------------
-        combined_prompt = user_input
+            context = f"""
+Image Description:
+{analysis["description"]}
 
-        if self.memory.image_context:
-            combined_prompt = f"""
-Image Context:
-{self.memory.image_context}
+Objects:
+{", ".join(analysis.get("objects", []))}
+
+Style:
+{analysis.get("style", "")}
+
+Lighting:
+{analysis.get("lighting", "")}
+
+OCR:
+{analysis.get("text", "")}
+"""
+
+        prompt = f"""
+{context}
 
 User Request:
 {user_input}
 """
 
-        # ---------------- EXECUTE STEPS ----------------
         for step in steps:
 
             if step == "generate":
-                result = generate_image(combined_prompt)
+
+                result = generate_image(prompt)
 
             elif step == "edit":
-                img = self.memory.last_used_image or image
-                result = edit_image(img, combined_prompt)
-                self.memory.last_used_image = img
+
+                result = edit_image(
+                    image_path=image,
+                    prompt=prompt
+                )
+
+                # Remember edited image
+                if result:
+                    self.memory.last_used_image = result
 
             elif step == "analyze":
-                result = analyze_image(image)
+
+                result = self.auto_analyze()
 
         return result