Update agents/agent.py

agents/agent.py

@@ -1,5 +1,5 @@
 """
-CellposeAgent with proper VLM configuration and JPEG compression for API payload optimization
+CellposeAgent with optimized image attachment - only attaches when visual inspection is needed
 """
 import torch
 import json 
@@ -23,8 +23,8 @@ class CellposeAgent:
     @staticmethod
     def attach_images_callback(step_log: ActionStep, agent: ToolCallingAgent) -> None:
         """
-        Callback to attach actual PIL images for VLM inspection.
-        Images are automatically resized and compressed to reduce token consumption.
+        OPTIMIZED: Only attach images for visual refinement step to save tokens.
+        For all other steps, skip image attachment - images are accessible via file paths.
         """
         if not isinstance(step_log, ActionStep):
             return
@@ -72,44 +72,19 @@ class CellposeAgent:
         try:
             obs_data = json.loads(step_log.observations)
         
-            # Pattern 1: Single image from get_segmentation_parameters
-            if obs_data.get("status") == "success" and "image_path" in obs_data:
-                image_path = obs_data["image_path"]
-                print(f"[Callback] Attaching image: {image_path}")
-            
-                try:
-                    img = Image.open(image_path)
-                    compressed_img = resize_and_compress_image(img, max_size=512, quality=75)
-                
-                    # Attach compressed PIL Image
-                    step_log.observations_images = [compressed_img]
-                    
-                    # Keep metadata for context
-                    obs_data["image_info"] = {
-                        "original_dimensions": f"{img.size[0]}x{img.size[1]} pixels",
-                        "processed_dimensions": f"{compressed_img.size[0]}x{compressed_img.size[1]} pixels",
-                        "mode": compressed_img.mode,
-                        "note": "Image compressed for API efficiency (JPEG quality=75)"
-                    }
-                    step_log.observations = json.dumps(obs_data, indent=2)
-                    print(f"[Callback] ✓ Attached compressed image for VLM inspection")
-                except Exception as e:
-                    print(f"[Callback] Error attaching image: {e}")
-        
-            # Pattern 2: Segmented image ONLY from refine_segmentation
-            elif obs_data.get("status") == "ready_for_visual_analysis":
-                paths = obs_data.get("image_paths", {})
-                segmented = paths.get("segmented")
+            # ONLY attach images for visual refinement step
+            if obs_data.get("status") == "ready_for_visual_analysis":
+                segmented = obs_data.get("image_paths", {}).get("segmented")
             
                 if segmented:
-                    print(f"[Callback] Attaching segmented image only: {segmented}")
+                    print(f"[Callback] Attaching segmented image for visual analysis: {segmented}")
                     try:
                         seg_img = Image.open(segmented)
                     
                         # Compress the segmented image
                         compressed_seg = resize_and_compress_image(seg_img, max_size=512, quality=75)
                     
-                        # Attach only the segmented image
+                        # Attach the segmented image
                         step_log.observations_images = [compressed_seg]
                     
                         obs_data["images_info"] = {
@@ -123,6 +98,10 @@ class CellposeAgent:
                         print(f"[Callback] ✓ Attached compressed segmented image for VLM inspection")
                     except Exception as e:
                         print(f"[Callback] Error attaching segmented image: {e}")
+            else:
+                # For all other steps, explicitly skip image attachment to save tokens
+                step_log.observations_images = []
+                print(f"[Callback] Skipped image attachment (not a refinement step) to save tokens")
     
         except json.JSONDecodeError:
             pass
@@ -133,21 +112,22 @@ class CellposeAgent:
     @staticmethod
     def manage_image_memory(step_log: ActionStep, agent: ToolCallingAgent) -> None:
         """
-        Aggressive memory management: keep ONLY the last step's images.
-        All previous steps have their images cleared immediately.
+        Aggressive memory management: clear ALL images from ALL previous steps.
+        Use empty list instead of None for more reliable cleanup.
         """
         if not isinstance(step_log, ActionStep):
             return
     
-        current_step = step_log.step_number
-    
-        # Clear images from ALL previous steps (keeping only current)
+        # Clear images from ALL previous steps (more aggressive)
         for previous_step in agent.memory.steps:
-            if isinstance(previous_step, ActionStep) and \
-               previous_step.step_number < current_step:
-                if previous_step.observations_images is not None:
+            if isinstance(previous_step, ActionStep):
+                if previous_step.observations_images is not None and len(previous_step.observations_images) > 0:
                     print(f"  [Memory] Clearing images from step {previous_step.step_number}")
-                    previous_step.observations_images = None
+                    previous_step.observations_images = []  # Empty list, not None
+                    
+                    # Try to clear any cached references (defensive)
+                    if hasattr(previous_step, '_observations_images'):
+                        previous_step._observations_images = []
 
                     
     def __init__(self):
@@ -158,8 +138,9 @@ class CellposeAgent:
         When a user provides an image:
         1. use appropriate tools to review which cellpose-sam parameters are available. 
         2. use the tool: `get_segmentation_parameters`
-           - **IMPORTANT**: After this tool runs, you will receive image metadata (dimensions, properties)
-           - Use this information to reason about appropriate parameter values
+           - **IMPORTANT**: You will receive image metadata (dimensions, properties, statistics)
+           - The actual image file is accessible via the file path in the response
+           - Use the metadata to reason about appropriate parameter values
         3. carefully analyze the image metadata and matched parameters:
            - consider cell density based on image dimensions
            - compare matched parameter values to image characteristics
@@ -168,8 +149,9 @@ class CellposeAgent:
         5. Provide your final parameter recommendations in a clear, structured format 
         6. Use the parameters to run cellpose_sam through the tool: run_cellpose_sam
         7. after run_cellpose_sam, call the tool: refine_cellpose_sam_segmentation
-           - **IMPORTANT**: After this tool runs, you will see the SEGMENTED image (colored masks overlay)
-           - Visually inspect the segmentation quality - are cells properly detected and separated?
+           - **IMPORTANT**: After this tool runs, you WILL SEE the SEGMENTED image (colored masks overlay)
+           - This is the ONLY step where you can visually inspect the actual image
+           - Visually assess the segmentation quality - are cells properly detected and separated?
            - Use the visual analysis checklist provided in the tool output
         8. Based on visual analysis of the segmented image:
            - Assess if cell boundaries are accurate
@@ -195,7 +177,7 @@ class CellposeAgent:
         - Be concise and actionable
         - Always explain your reasoning when adjusting parameters
         - If keeping original matched parameters, briefly confirm why it's appropriate
-        - Base your decisions on visual observation of the segmented output
+        - Base your decisions on visual observation of the segmented output (when available in refinement step)
         
         **CRITICAL - Final Response Format:**
         When segmentation is complete, you MUST provide a comprehensive text summary that includes: