Update tools/segmentation.py

tools/segmentation.py

@@ -1,9 +1,13 @@
 """
 Segmentation tools for cellpose-sam pipeline with proper smolagents VLM integration.
+
+Key change: Tools now resolve image paths from global context when the provided path
+is invalid or empty, preventing LLM path corruption issues.
 """
 import base64
 import json
 import re
+from pathlib import Path
 from typing import Any, Dict, TYPE_CHECKING
 import numpy as np
 import cv2
@@ -26,6 +30,52 @@ from config import settings
 langfuse = get_client()
 
 
+# =============================================================================
+# PATH RESOLUTION HELPER
+# =============================================================================
+def resolve_image_path(provided_path: str, context_type: str = "image") -> str:
+    """
+    Resolve the actual image path, falling back to global context if needed.
+    
+    This function handles the case where the LLM corrupts the file path by:
+    1. Checking if the provided path exists
+    2. If not, retrieving the correct path from global context
+    
+    Args:
+        provided_path: The path provided by the LLM (may be corrupted)
+        context_type: Either "image" for input or "output" for segmentation result
+    
+    Returns:
+        The resolved, valid path
+        
+    Raises:
+        FileNotFoundError: If no valid path can be resolved
+    """
+    # Import here to avoid circular imports
+    from agents.agent import get_current_image_path, get_current_output_path
+    
+    # First, check if the provided path is valid
+    if provided_path and Path(provided_path).exists():
+        print(f"[Path Resolution] Using provided path: {provided_path}")
+        return provided_path
+    
+    # Path is invalid - try to get from context
+    if context_type == "image":
+        context_path = get_current_image_path()
+    else:
+        context_path = get_current_output_path()
+    
+    if context_path and Path(context_path).exists():
+        print(f"[Path Resolution] Provided path invalid, using context: {context_path}")
+        print(f"[Path Resolution] (LLM provided: '{provided_path}')")
+        return context_path
+    
+    # Neither worked
+    error_msg = f"Could not resolve {context_type} path. Provided: '{provided_path}', Context: '{context_path}'"
+    print(f"[Path Resolution] ERROR: {error_msg}")
+    raise FileNotFoundError(error_msg)
+
+
 # --- Global State and Caching ---
 _image_cache: Dict[str, tuple[str, str]] = {}
 _cellpose_model = None
@@ -50,12 +100,14 @@ def get_sam_predictor():
         _sam_predictor = SamPredictor(sam)
     return _sam_predictor
 
+
 def _get_cached_image(image_path: str) -> tuple[str, str] | None:
     """Helper to retrieve an image from the cache."""
     if image_path in _image_cache:
         return _image_cache[image_path]
     return None
 
+
 def _load_and_cache_image(image_path: str) -> tuple[str, str]:
     """Helper to load, encode, and cache an image."""
     image_base64, media_type = resize_and_encode_image(image_path)
@@ -94,26 +146,35 @@ def parse_parameters_from_text(param_text: str) -> dict:
 
 
 @tool
-def get_segmentation_parameters(image_path: str, agent: Any = None) -> str:
+def get_segmentation_parameters(image_path: str = "", agent: Any = None) -> str:
     """
     Finds the best cellpose-sam segmentation parameters for an image using vector similarity.
     The image will be visible to the VLM for visual analysis.
     
+    NOTE: If image_path is empty or invalid, the tool will automatically use the 
+    current image from the system context.
+    
     Args:
-        image_path (str): Path to the image file to segment.
+        image_path (str): Path to the image file (optional - uses context if invalid).
         agent (Any, optional): The agent instance, passed automatically by smol-agents.
     
     Returns:
         str: JSON string containing recommended parameters and analysis context
-             (NO base64 to avoid GPU OOM)
     """
     print(f"\n--- TOOL CALLED: get_segmentation_parameters for '{image_path}' ---")
 
+    # Resolve the actual image path
+    try:
+        actual_path = resolve_image_path(image_path, context_type="image")
+    except FileNotFoundError as e:
+        return json.dumps({"error": str(e)})
+    
+    image_path = actual_path  # Use resolved path from here on
+
     try:
         # Load and cache image (for internal use)
         image_base64, media_type = _get_cached_image(image_path) or _load_and_cache_image(image_path)
         
-        
     except Exception as e:
         print(f"Warning: Could not read/resize image: {e}")
         return json.dumps({"error": f"Could not read image: {e}"})
@@ -204,7 +265,7 @@ def get_segmentation_parameters(image_path: str, agent: Any = None) -> str:
                              f"- min_size: {params['min_size']}\n\n"
                              f"Image stats: {image_shape[0]}x{image_shape[1]} pixels, "
                              f"mean intensity {stats['mean_intensity']:.1f}\n\n"
-                             f"To run segmentation, use: run_cellpose_sam(image_path='{image_path}', "
+                             f"To run segmentation, use: run_cellpose_sam(image_path='', "
                              f"diameter={params['diameter']}, flow_threshold={params['flow_threshold']}, "
                              f"cellprob_threshold={params['cellprob_threshold']}, min_size={params['min_size']})"
         }
@@ -217,7 +278,7 @@ def get_segmentation_parameters(image_path: str, agent: Any = None) -> str:
 
 @tool
 def run_cellpose_sam(
-    image_path: str,
+    image_path: str = "",
     diameter: int = None,
     flow_threshold: float = None,
     cellprob_threshold: float = None,
@@ -230,8 +291,11 @@ def run_cellpose_sam(
     Runs cellpose-sam segmentation pipeline on an image with specified parameters.
     Returns results WITHOUT base64 images to prevent GPU memory issues.
     
+    NOTE: If image_path is empty or invalid, the tool will automatically use the 
+    current image from the system context.
+    
     Args:
-        image_path (str): Path to the image file to segment
+        image_path (str): Path to the image file (optional - uses context if invalid)
         diameter (int): Expected diameter of cells in pixels
         flow_threshold (float): Flow error threshold (range: 0-1)
         cellprob_threshold (float): Cell probability threshold (range: -6 to 6)
@@ -245,6 +309,14 @@ def run_cellpose_sam(
     """
     print(f"\n--- TOOL CALLED: run_cellpose_sam for '{image_path}' ---")
     
+    # Resolve the actual image path
+    try:
+        actual_path = resolve_image_path(image_path, context_type="image")
+    except FileNotFoundError as e:
+        return json.dumps({"error": str(e)})
+    
+    image_path = actual_path  # Use resolved path from here on
+    
     try:
         # Load and cache input image
         input_image_base64, input_media_type = _get_cached_image(image_path) or _load_and_cache_image(image_path)
@@ -337,6 +409,10 @@ def run_cellpose_sam(
         # Save output
         cv2.imwrite(output_path, cv2.cvtColor(colored_overlay.astype(np.uint8), cv2.COLOR_RGB2BGR))
         
+        # Store output path in context for later tools
+        from agents.agent import set_current_output_path
+        set_current_output_path(output_path)
+        
         # Load and cache output image
         output_image_base64, output_media_type = _load_and_cache_image(output_path)
         
@@ -394,9 +470,9 @@ def run_cellpose_sam(
 
 @tool
 def refine_cellpose_sam_segmentation(
-    original_image_path: str,
-    segmentation_output_path: str,
-    current_parameters: dict,
+    original_image_path: str = "",
+    segmentation_output_path: str = "",
+    current_parameters: dict = None,
     agent: Any = None,
 ) -> str:
     """
@@ -406,6 +482,8 @@ def refine_cellpose_sam_segmentation(
     Use this tool after run_cellpose_sam to check segmentation quality. The tool attaches
     both images to the current step so you can visually compare them.
     
+    NOTE: If paths are empty or invalid, the tool will automatically use paths from context.
+    
     Before calling, consider using search_knowledge_graph or hybrid_search to refresh
     your understanding of how cellpose parameters affect segmentation.
     
@@ -416,8 +494,8 @@ def refine_cellpose_sam_segmentation(
     - Too many false positives: increase cellprob_threshold or min_size
     
     Args:
-        original_image_path: Path to the original input image
-        segmentation_output_path: Path to the segmented overlay image
+        original_image_path: Path to the original input image (optional - uses context)
+        segmentation_output_path: Path to the segmented overlay image (optional - uses context)
         current_parameters: Dict with current diameter, flow_threshold, cellprob_threshold, min_size
         agent: The agent instance (passed automatically)
     
@@ -425,10 +503,27 @@ def refine_cellpose_sam_segmentation(
         str: JSON with guidance for VLM analysis (NO base64 images)
     """
     print(f"\n--- TOOL CALLED: refine_cellpose_sam_segmentation ---")
-    print(f"Original image: {original_image_path}")
-    print(f"Segmented image: {segmentation_output_path}")
+    print(f"Original image (provided): {original_image_path}")
+    print(f"Segmented image (provided): {segmentation_output_path}")
     print(f"Current parameters: {current_parameters}")
     
+    # Resolve paths from context if needed
+    try:
+        actual_original = resolve_image_path(original_image_path, context_type="image")
+    except FileNotFoundError as e:
+        return json.dumps({"error": f"Could not resolve original image: {e}"})
+    
+    try:
+        actual_segmented = resolve_image_path(segmentation_output_path, context_type="output")
+    except FileNotFoundError as e:
+        return json.dumps({"error": f"Could not resolve segmented image: {e}"})
+    
+    original_image_path = actual_original
+    segmentation_output_path = actual_segmented
+    
+    print(f"Resolved original: {original_image_path}")
+    print(f"Resolved segmented: {segmentation_output_path}")
+    
     try:
         # Load both images (for cache)
         original_b64, original_type = _get_cached_image(original_image_path) or _load_and_cache_image(original_image_path)
@@ -529,4 +624,4 @@ def refine_cellpose_sam_segmentation(
             "error": str(e),
             "message": "Could not load images for refinement. Check that both file paths are valid."
         }
-        return json.dumps(error_result, indent=2)
+        return json.dumps(error_result, indent=2)