Update app.py

app.py

@@ -2,10 +2,12 @@
 Gradio web interface for CellposeAgent
 """
 import os
+import json
 import gradio as gr
 from pathlib import Path
 from langfuse import get_client
 from openinference.instrumentation.smolagents import SmolagentsInstrumentor
+from smolagents.agents import ActionStep
 
 from config import settings
 from agents.agent import CellposeAgent
@@ -95,32 +97,33 @@ def process_image_task(image_path: str, task_text: str, agent: CellposeAgent) ->
         result = agent.run(task)
         get_client().flush()
         
-        # Extract output image path from the agent's response if segmentation was run
+        # Extract output image path from agent's memory (more reliable than parsing text)
         output_image_path = None
         try:
-            # Check if the response contains a segmentation output path
-            if "output_path" in result or "_cellpose_sam_overlay.png" in result:
-                # Parse the response to find the output path
-                import json
-                import re
-                
-                # Try to find JSON in the response
-                json_match = re.search(r'\{[^{}]*"output_path"[^{}]*\}', result)
-                if json_match:
-                    result_data = json.loads(json_match.group())
-                    output_image_path = result_data.get("output_path")
-                else:
-                    # Try to find the path in the text
-                    path_match = re.search(r'([^\s]+_cellpose_sam_overlay\.png)', result)
-                    if path_match:
-                        output_image_path = path_match.group(1)
-                
-                # Verify the file exists
-                if output_image_path and not Path(output_image_path).exists():
-                    output_image_path = None
+            # Search through agent's memory steps in reverse order (most recent first)
+            for step in reversed(agent.agent.memory.steps):
+                if isinstance(step, ActionStep) and step.observations:
+                    try:
+                        obs_data = json.loads(step.observations)
+                        
+                        # Check if this step contains segmentation output
+                        if obs_data.get("status") == "success" and "output_path" in obs_data:
+                            candidate_path = obs_data["output_path"]
+                            
+                            # Verify the file exists
+                            if Path(candidate_path).exists():
+                                output_image_path = candidate_path
+                                print(f"✓ Found segmentation output: {output_image_path}")
+                                break
+                                
+                    except (json.JSONDecodeError, Exception) as parse_error:
+                        continue
+            
+            if output_image_path is None:
+                print("ℹ️ No segmentation output found in agent memory")
                     
         except Exception as e:
-            print(f"Could not extract output image path: {e}")
+            print(f"Warning: Could not extract output image from agent memory: {e}")
         
         return result, output_image_path
         
@@ -235,4 +238,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()