Update agents/agent.py

agents/agent.py

@@ -1,5 +1,5 @@
 """
-CellposeAgent with optimized image attachment - only attaches when visual inspection is needed
+CellposeAgent with proper VLM configuration and JPEG compression for API payload optimization
 """
 import torch
 import json 
@@ -23,8 +23,8 @@ class CellposeAgent:
     @staticmethod
     def attach_images_callback(step_log: ActionStep, agent: ToolCallingAgent) -> None:
         """
-        OPTIMIZED: Only attach images for visual refinement step to save tokens.
-        For all other steps, skip image attachment - images are accessible via file paths.
+        Callback to attach actual PIL images for VLM inspection.
+        Images are automatically resized and compressed to reduce token consumption.
         """
         if not isinstance(step_log, ActionStep):
             return
@@ -72,19 +72,44 @@ class CellposeAgent:
         try:
             obs_data = json.loads(step_log.observations)
         
-            # ONLY attach images for visual refinement step
-            if obs_data.get("status") == "ready_for_visual_analysis":
-                segmented = obs_data.get("image_paths", {}).get("segmented")
+            # 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")
             
                 if segmented:
-                    print(f"[Callback] Attaching segmented image for visual analysis: {segmented}")
+                    print(f"[Callback] Attaching segmented image only: {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 the segmented image
+                        # Attach only the segmented image
                         step_log.observations_images = [compressed_seg]
                     
                         obs_data["images_info"] = {
@@ -98,10 +123,6 @@ 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
@@ -112,20 +133,18 @@ class CellposeAgent:
     @staticmethod
     def manage_image_memory(step_log: ActionStep, agent: ToolCallingAgent) -> None:
         """
-        Aggressive memory management: clear ALL images from ALL previous steps.
-        Use empty list instead of None for more reliable cleanup.
+        Clear images from ALL previous steps at the START of each new step.
         """
         if not isinstance(step_log, ActionStep):
             return
     
-        # Clear images from ALL previous steps (more aggressive)
+        # Clear ALL previous step images immediately
         for previous_step in agent.memory.steps:
             if isinstance(previous_step, ActionStep):
-                if previous_step.observations_images is not None and len(previous_step.observations_images) > 0:
+                if previous_step.observations_images is not None:
                     print(f"  [Memory] Clearing images from step {previous_step.step_number}")
-                    previous_step.observations_images = []  # Empty list, not None
-                    
-                    # Try to clear any cached references (defensive)
+                    previous_step.observations_images = []  # Use empty list instead of None
+                    # Also try to clear any cached references
                     if hasattr(previous_step, '_observations_images'):
                         previous_step._observations_images = []
 
@@ -138,9 +157,8 @@ 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**: 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
+           - **IMPORTANT**: After this tool runs, you will receive image metadata (dimensions, properties)
+           - Use this information 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
@@ -149,9 +167,8 @@ 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)
-           - This is the ONLY step where you can visually inspect the actual image
-           - Visually assess the segmentation quality - are cells properly detected and separated?
+           - **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?
            - 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
@@ -162,10 +179,10 @@ class CellposeAgent:
            - Decide which parameters to adjust based on what you observe
            - Re-run run_cellpose_sam with adjusted parameters   
         
-        **CRITICAL: Call refine_cellpose_sam_segmentation AT MOST 2 TIMES total**
+        **CRITICAL: Call refine_cellpose_sam_segmentation AT MOST 1 TIMES total**
            - First call: Check initial segmentation quality
            - Second call (if needed): Verify refinement improved results
-           - NEVER call it a third time - always stop after 2 refinement checks
+           - NEVER call it a second time - always stop after 1 refinement check
         
         ## DOCUMENTATION QUERY WORKFLOW ##
         - "What is X": use `search_documentation_vector`
@@ -177,7 +194,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 (when available in refinement step)
+        - Base your decisions on visual observation of the segmented output
         
         **CRITICAL - Final Response Format:**
         When segmentation is complete, you MUST provide a comprehensive text summary that includes:
@@ -213,7 +230,7 @@ class CellposeAgent:
         return InferenceClientModel(
             model_id=settings.AGENT_MODEL_ID,
             token=settings.HF_TOKEN,
-            timeout=180  # 3 minutes timeout for API calls
+            timeout=240  # 3 minutes timeout for API calls
         )