Update agents/image_agent.py

agents/image_agent.py

@@ -18,45 +18,65 @@ class ImageAgent:
 
     def auto_analyze(self, image_path):
         """
-        STEP 1: Always understand image first
+        STEP 1: Understand image
         """
         if image_path:
             analysis = analyze_image(image_path)
+
+            # 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.last_used_image = image_path
+
             return analysis
+
         return None
 
     def run(self, user_input: str):
 
+        # ---------------- ROUTE ----------------
         task = self.router.route(user_input)
         steps = self.planner.plan(task)
 
         result = None
 
-        # 🔥 STEP 0: ALWAYS pick image first
+        # ---------------- IMAGE PICK ----------------
         image = None
-
         if self.memory.uploaded_files:
             image = self.memory.uploaded_files[0]
 
-        # 🔥 STEP 1: ANALYZE FIRST (IMPORTANT)
-        if image:
-            analysis = self.auto_analyze(image)
-            self.memory.analysis = analysis
+        # ---------------- AUTO ANALYZE ----------------
+        if image and not self.memory.image_context:
+            analysis = analyze_image(image)
+            self.memory.image_context = analysis["description"]
+
+        # ---------------- COMBINE CONTEXT ----------------
+        combined_prompt = user_input
+
+        if self.memory.image_context:
+            combined_prompt = f"""
+Image Context:
+{self.memory.image_context}
+
+User Request:
+{user_input}
+"""
 
-        # 🔥 STEP 2: EXECUTE TASK
+        # ---------------- EXECUTE STEPS ----------------
         for step in steps:
 
             if step == "generate":
-                result = generate_image(user_input)
+                result = generate_image(combined_prompt)
 
             elif step == "edit":
-                image = self.memory.last_used_image or image
-                result = edit_image(image, user_input)
-                self.memory.last_used_image = image
+                img = self.memory.last_used_image or image
+                result = edit_image(img, combined_prompt)
+                self.memory.last_used_image = img
 
             elif step == "analyze":
-                image = self.memory.last_used_image or image
                 result = analyze_image(image)
 
         return result