Upload 41 files

agents/agent.py

@@ -0,0 +1,225 @@
+"""
+CellposeAgent with proper VLM configuration
+"""
+import torch
+import json 
+from datetime import datetime
+from PIL import Image
+from smolagents import ToolCallingAgent, InferenceClientModel
+from smolagents.agents import ActionStep
+from langfuse import get_client, observe
+
+from config import settings
+from utils.gpu import clear_gpu_cache
+from tools import all_tools
+
+
+langfuse = get_client()
+
+
+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 to reduce token consumption.
+        """
+        if not isinstance(step_log, ActionStep):
+            return
+    
+        if not step_log.observations:
+            return
+    
+        def resize_image(img: Image.Image, max_size: int = 1024) -> Image.Image:
+            """Resize image maintaining aspect ratio, max dimension = max_size."""
+            if max(img.size) <= max_size:
+                return img
+        
+            ratio = max_size / max(img.size)
+            new_size = tuple(int(dim * ratio) for dim in img.size)
+            resized = img.resize(new_size, Image.Resampling.LANCZOS)
+            print(f"    Resized {img.size} → {resized.size}")
+            return resized
+    
+        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)
+                    resized_img = resize_image(img)
+                
+                    # Attach resized PIL Image
+                    step_log.observations_images = [resized_img]
+                    
+                    # Keep metadata for context
+                    obs_data["image_info"] = {
+                        "original_dimensions": f"{img.size[0]}x{img.size[1]} pixels",
+                        "resized_dimensions": f"{resized_img.size[0]}x{resized_img.size[1]} pixels",
+                        "mode": resized_img.mode,
+                        "note": "Image attached for visual inspection (resized for efficiency)"
+                    }
+                    step_log.observations = json.dumps(obs_data, indent=2)
+                    print(f"[Callback] ✓ Attached resized image for VLM inspection")
+                except Exception as e:
+                    print(f"[Callback] Error attaching image: {e}")
+        
+            # Pattern 2: Multiple images from refine_segmentation
+            
+            elif obs_data.get("status") == "ready_for_visual_analysis":
+                paths = obs_data.get("image_paths", {})
+                original = paths.get("original")
+                segmented = paths.get("segmented")
+            
+                if original and segmented:
+                    print(f"[Callback] Attaching both original and segmented images")
+                    try:
+                        orig_img = Image.open(original)
+                        seg_img = Image.open(segmented)
+                    
+                        # Resize both images
+                        resized_orig = resize_image(orig_img)
+                        resized_seg = resize_image(seg_img)
+                    
+                        # Attach both resized images as list
+                        step_log.observations_images = [resized_orig, resized_seg]
+                    
+                        obs_data["images_info"] = {
+                            "image_order": ["original", "segmented"],
+                            "original_size": f"{orig_img.size[0]}x{orig_img.size[1]}",
+                            "resized_size": f"{resized_orig.size[0]}x{resized_orig.size[1]}",
+                            "note": "Both images attached for visual comparison (resized for efficiency)"
+                        }
+                        step_log.observations = json.dumps(obs_data, indent=2)
+                        print(f"[Callback] ✓ Attached both resized images for VLM inspection")
+                    except Exception as e:
+                        print(f"[Callback] Error attaching images: {e}")
+    
+        except json.JSONDecodeError:
+            pass
+        except Exception as e:
+            print(f"[Callback] Error in attach_images_callback: {e}")
+
+
+    @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.
+        """
+        if not isinstance(step_log, ActionStep):
+            return
+    
+        current_step = step_log.step_number
+    
+        # Clear images from ALL previous steps (keeping only current)
+        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:
+                    print(f"  [Memory] Clearing images from step {previous_step.step_number}")
+                    previous_step.observations_images = None
+
+                    
+    def __init__(self):
+        self.instructions = """
+        You are an assistant for the cellpose-sam segmentation tool.
+
+        ## PRIMARY WORKFLOW - IMAGE SEGMENTATION 
+
+        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
+        3. carefully analyze the image metadata and matched parameters:
+           - consider cell density based on image dimensions
+           - compare matched parameter values to image characteristics
+           - consider if adjustments would likely improve the segmentation
+        4. Be conservative: if you make changes, assess if they should differ significantly from the original values
+        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 receive metadata about both original and segmented images
+           - Use the provided information to assess segmentation quality
+        8. Based on the metadata and any quality metrics returned:
+           - Identify potential segmentation issues based on reported metrics
+           - If refinement is needed, use knowledge graph and RAG tools to understand parameter effects
+           - Decide which parameters to adjust based on the segmentation analysis
+           - Re-run run_cellpose_sam with adjusted parameters   
+        **CRITICAL: Call refine_cellpose_sam_segmentation AT MOST 2 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
+
+        ## DOCUMENTATION QUERY WORKFLOW ##
+
+        - "What is X": use `search_documentation_vector`
+        - "How does X affect Y": use `search_knowledge_graph`  
+        - Complex analysis: use `hybrid_search`
+        - Parameter relationships: use `get_parameter_relationships`        
+
+        ## RESPONSE STYLE ##
+        - 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 the metadata and metrics provided by the tools
+
+        """
+        
+        self.model = self._initialize_model()
+        self.agent = self._create_agent()
+        
+
+    def _initialize_model(self):
+        """Initializes the TransformersModel for the agent with VLM support."""
+        clear_gpu_cache()
+
+        return InferenceClientModel(
+            model_id=settings.AGENT_MODEL_ID,
+            token = settings.HF_TOKEN
+        )
+
+
+
+    def _create_agent(self):
+        """Creates the ToolCallingAgent with all available tools and memory management."""
+        return ToolCallingAgent(
+            model=self.model,
+            tools=all_tools,
+            instructions=self.instructions,
+            max_steps=10,
+            step_callbacks=[
+                self.attach_images_callback,
+                self.manage_image_memory,
+            ]
+        )
+
+    @observe()
+    def run(self, task: str):
+        """Runs the agent on a given task with Langfuse tracing."""
+        print(f"\n{'='*60}\nTASK: {task}\n{'='*60}")
+        
+        langfuse.update_current_trace(
+            input={"task": task},
+            user_id="user_001",
+            tags=["rag", "cellpose", "knowledge-graph", "vision"],
+            metadata={"agent_type": "ToolCallingAgent", "model_id": settings.AGENT_MODEL_ID}
+        )
+
+        try:
+            final_answer = self.agent.run(task)
+            print("\n--- Final Answer from Agent ---\n", final_answer)
+            langfuse.update_current_trace(output={"final_answer": final_answer})
+            return final_answer
+        except Exception as e:
+            print(f"Agent run failed: {e}")
+            langfuse.update_current_trace(output={"error": str(e)})
+            raise
+        finally:
+            clear_gpu_cache()

agents/agent.py~

@@ -0,0 +1,224 @@
+"""
+CellposeAgent with proper VLM configuration
+"""
+import torch
+import json 
+from datetime import datetime
+from PIL import Image
+from smolagents import ToolCallingAgent, InferenceClientModel
+from smolagents.agents import ActionStep
+from langfuse import get_client, observe
+
+from config import settings
+from utils.gpu import clear_gpu_cache
+from tools import all_tools
+
+
+langfuse = get_client()
+
+
+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 to reduce token consumption.
+        """
+        if not isinstance(step_log, ActionStep):
+            return
+    
+        if not step_log.observations:
+            return
+    
+        def resize_image(img: Image.Image, max_size: int = 1024) -> Image.Image:
+            """Resize image maintaining aspect ratio, max dimension = max_size."""
+            if max(img.size) <= max_size:
+                return img
+        
+            ratio = max_size / max(img.size)
+            new_size = tuple(int(dim * ratio) for dim in img.size)
+            resized = img.resize(new_size, Image.Resampling.LANCZOS)
+            print(f"    Resized {img.size} → {resized.size}")
+            return resized
+    
+        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)
+                    resized_img = resize_image(img)
+                
+                    # Attach resized PIL Image
+                    step_log.observations_images = [resized_img]
+                    
+                    # Keep metadata for context
+                    obs_data["image_info"] = {
+                        "original_dimensions": f"{img.size[0]}x{img.size[1]} pixels",
+                        "resized_dimensions": f"{resized_img.size[0]}x{resized_img.size[1]} pixels",
+                        "mode": resized_img.mode,
+                        "note": "Image attached for visual inspection (resized for efficiency)"
+                    }
+                    step_log.observations = json.dumps(obs_data, indent=2)
+                    print(f"[Callback] ✓ Attached resized image for VLM inspection")
+                except Exception as e:
+                    print(f"[Callback] Error attaching image: {e}")
+        
+            # Pattern 2: Multiple images from refine_segmentation
+            
+elif obs_data.get("status") == "ready_for_visual_analysis":
+            paths = obs_data.get("image_paths", {})
+            original = paths.get("original")
+            segmented = paths.get("segmented")
+            
+            if original and segmented:
+                print(f"[Callback] Attaching both original and segmented images")
+                try:
+                    orig_img = Image.open(original)
+                    seg_img = Image.open(segmented)
+                    
+                    # Resize both images
+                    resized_orig = resize_image(orig_img)
+                    resized_seg = resize_image(seg_img)
+                    
+                    # Attach both resized images as list
+                    step_log.observations_images = [resized_orig, resized_seg]
+                    
+                    obs_data["images_info"] = {
+                        "image_order": ["original", "segmented"],
+                        "original_size": f"{orig_img.size[0]}x{orig_img.size[1]}",
+                        "resized_size": f"{resized_orig.size[0]}x{resized_orig.size[1]}",
+                        "note": "Both images attached for visual comparison (resized for efficiency)"
+                    }
+                    step_log.observations = json.dumps(obs_data, indent=2)
+                    print(f"[Callback] ✓ Attached both resized images for VLM inspection")
+                except Exception as e:
+                    print(f"[Callback] Error attaching images: {e}")
+    
+    except json.JSONDecodeError:
+        pass
+    except Exception as e:
+        print(f"[Callback] Error in attach_images_callback: {e}")
+
+
+@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.
+    """
+    if not isinstance(step_log, ActionStep):
+        return
+    
+    current_step = step_log.step_number
+    
+    # Clear images from ALL previous steps (keeping only current)
+    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:
+                print(f"  [Memory] Clearing images from step {previous_step.step_number}")
+                previous_step.observations_images = None
+                    
+    def __init__(self):
+        self.instructions = """
+        You are an assistant for the cellpose-sam segmentation tool.
+
+        ## PRIMARY WORKFLOW - IMAGE SEGMENTATION 
+
+        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
+        3. carefully analyze the image metadata and matched parameters:
+           - consider cell density based on image dimensions
+           - compare matched parameter values to image characteristics
+           - consider if adjustments would likely improve the segmentation
+        4. Be conservative: if you make changes, assess if they should differ significantly from the original values
+        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 receive metadata about both original and segmented images
+           - Use the provided information to assess segmentation quality
+        8. Based on the metadata and any quality metrics returned:
+           - Identify potential segmentation issues based on reported metrics
+           - If refinement is needed, use knowledge graph and RAG tools to understand parameter effects
+           - Decide which parameters to adjust based on the segmentation analysis
+           - Re-run run_cellpose_sam with adjusted parameters   
+        **CRITICAL: Call refine_cellpose_sam_segmentation AT MOST 2 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
+
+        ## DOCUMENTATION QUERY WORKFLOW ##
+
+        - "What is X": use `search_documentation_vector`
+        - "How does X affect Y": use `search_knowledge_graph`  
+        - Complex analysis: use `hybrid_search`
+        - Parameter relationships: use `get_parameter_relationships`        
+
+        ## RESPONSE STYLE ##
+        - 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 the metadata and metrics provided by the tools
+
+        """
+        
+        self.model = self._initialize_model()
+        self.agent = self._create_agent()
+        
+
+    def _initialize_model(self):
+        """Initializes the TransformersModel for the agent with VLM support."""
+        clear_gpu_cache()
+
+        return InferenceClientModel(
+            model_id=settings.AGENT_MODEL_ID,
+            token = settings.HF_TOKEN
+        )
+
+
+
+    def _create_agent(self):
+        """Creates the ToolCallingAgent with all available tools and memory management."""
+        return ToolCallingAgent(
+            model=self.model,
+            tools=all_tools,
+            instructions=self.instructions,
+            max_steps=10,
+            step_callbacks=[
+                self.attach_images_callback,
+                self.manage_image_memory,
+            ]
+        )
+
+    @observe()
+    def run(self, task: str):
+        """Runs the agent on a given task with Langfuse tracing."""
+        print(f"\n{'='*60}\nTASK: {task}\n{'='*60}")
+        
+        langfuse.update_current_trace(
+            input={"task": task},
+            user_id="user_001",
+            tags=["rag", "cellpose", "knowledge-graph", "vision"],
+            metadata={"agent_type": "ToolCallingAgent", "model_id": settings.AGENT_MODEL_ID}
+        )
+
+        try:
+            final_answer = self.agent.run(task)
+            print("\n--- Final Answer from Agent ---\n", final_answer)
+            langfuse.update_current_trace(output={"final_answer": final_answer})
+            return final_answer
+        except Exception as e:
+            print(f"Agent run failed: {e}")
+            langfuse.update_current_trace(output={"error": str(e)})
+            raise
+        finally:
+            clear_gpu_cache()

app.py

@@ -1,14 +1,188 @@
+"""
+Gradio web interface for CellposeAgent
+"""
 import gradio as gr
-import spaces
-import torch
+from pathlib import Path
+from langfuse import get_client
+from openinference.instrumentation.smolagents import SmolagentsInstrumentor
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+from config import settings
+from agents.agent import CellposeAgent
+from stores import neo4j_store
+from utils.prechecks import check_hf_persistent_storage
 
-@spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+def setup_observability():
+    """Initializes Langfuse and Smolagents instrumentation."""
+    get_client()
+    SmolagentsInstrumentor().instrument()
+    print("✓ Observability and instrumentation initialized.")
+
+
+def initialize_app():
+    """Initialize the application and verify prerequisites."""
+    print("\n--- Initializing Cellpose Agent Application ---")
+    
+    # Setup observability
+    setup_observability()
+    
+    # Configure LlamaIndex
+    settings.configure_llama_index()
+
+    # check for cellpose-db
+    check_hf_persistent_storage(
+        repo_id = "hmgill/Cellpose-DB",
+        target = "cellpose_db",
+        file_or_folder="folder"
+    )
+
+    # check for cellpose sam
+    check_hf_persistent_storage(
+        repo_id = "hmgill/Cellpose-SAM-Checkpoint",
+        target = "sam_vit_h_4b8939.pth",
+        file_or_folder="file"
+    )
+
+    # Verify knowledge graph is ready
+    try:
+        node_count, _ = neo4j_store.check_graph_status()
+        if node_count == 0:
+            print("\n❌ WARNING: The knowledge graph is empty.")
+            print("Please run the setup script to build the knowledge graph:")
+            print("\n    python setup_kg.py\n")
+            return False
+        print(f"✓ Knowledge graph is ready with {node_count} nodes.")
+    except Exception as e:
+        print(f"❌ ERROR: Could not connect to Neo4j: {e}")
+        print("Please ensure Neo4j is running and accessible.")
+        return False
+    
+    return True
+
+
+def process_image_task(image_path: str, task_text: str, agent: CellposeAgent) -> str:
+    """
+    Process a user task with the CellposeAgent.
+    
+    Args:
+        image_path: Path to the uploaded image file
+        task_text: User's text prompt/question
+        agent: Initialized CellposeAgent instance
+    
+    Returns:
+        str: Agent's response
+    """
+    if not image_path:
+        return "⚠️ Please upload an image first."
+    
+    if not task_text:
+        task_text = f"What parameters would work best for my image {image_path}?"
+    
+    try:
+        result = agent.run(task_text)
+        get_client().flush()
+        return result
+    except Exception as e:
+        return f"❌ Error processing task: {str(e)}"
+
+
+def create_gradio_interface():
+    """Creates and configures the Gradio interface."""
+    
+    # Initialize the agent once at startup
+    if not initialize_app():
+        raise RuntimeError("Failed to initialize application. Please check logs.")
+    
+    agent = CellposeAgent()
+    print("✓ CellposeAgent initialized and ready.")
+    
+    with gr.Blocks(title="Cellpose-SAM Agent", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            """
+            # 🔬 Cellpose-SAM Segmentation Agent
+            
+            Upload a microscopy image and ask the AI agent to recommend optimal segmentation parameters,
+            run segmentation, or answer questions about the cellpose-sam pipeline.
+            """
+        )
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Image upload
+                image_input = gr.Image(
+                    label="Upload Microscopy Image",
+                    type="filepath",
+                    height=300
+                )
+                
+                # Task input
+                task_input = gr.Textbox(
+                    label="Task / Question",
+                    placeholder="e.g., 'What parameters would work best for this image?' or leave empty for default",
+                    lines=3
+                )
+                
+                # Submit button
+                submit_btn = gr.Button("Run Agent", variant="primary", size="lg")
+                
+                # Example tasks
+                gr.Markdown("### 💡 Example Tasks")
+                gr.Examples(
+                    examples=[
+                        ["What parameters would work best for this image?"],
+                        ["Analyze this image and run segmentation with optimal parameters."],
+                        ["What is the flow_threshold parameter and how does it affect segmentation?"],
+                        ["Run segmentation with diameter=30, flow_threshold=0.5, cellprob_threshold=0, min_size=20"],
+                    ],
+                    inputs=task_input,
+                    label="Click to use:"
+                )
+            
+            with gr.Column(scale=1):
+                # Output
+                output = gr.Textbox(
+                    label="Agent Response",
+                    lines=20,
+                    max_lines=30,
+                    show_copy_button=True
+                )
+        
+        # Event handler
+        submit_btn.click(
+            fn=lambda img, task: process_image_task(img, task, agent),
+            inputs=[image_input, task_input],
+            outputs=output
+        )
+        
+        gr.Markdown(
+            """
+            ---
+            ### 📚 What can this agent do?
+            
+            - **Parameter Recommendation**: Analyzes your image and suggests optimal segmentation parameters
+            - **Automated Segmentation**: Runs the full cellpose-sam pipeline with parameter refinement
+            - **Visual Analysis**: Uses vision-language models to assess segmentation quality
+            - **Documentation Search**: Answers questions about parameters using RAG and knowledge graphs
+            - **Iterative Refinement**: Automatically adjusts parameters based on visual feedback
+            
+            ### 🔍 How it works
+            
+            1. Upload your microscopy image
+            2. The agent finds similar images and recommends parameters
+            3. Visually analyzes your image to validate recommendations
+            4. Runs segmentation and checks quality
+            5. Refines parameters if needed (up to 2 iterations)
+            """
+        )
+    
+    return demo
+
+
+def main():
+    """Launch the Gradio application."""
+    demo = create_gradio_interface()
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()

config/settings.py

@@ -0,0 +1,60 @@
+"""
+
+"""
+
+import os
+import torch
+from llama_index.core import Settings
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.llms.huggingface import HuggingFaceLLM
+from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+from llama_index.core.prompts import PromptTemplate
+
+
+
+# ---- Model IDs ---
+AGENT_MODEL_ID = "google/gemma-3-12b-it"
+EMBEDDING_MODEL_ID = "clip-ViT-B-32"
+
+
+# --- Environment & Paths ---
+CHROMADB = os.getenv("CHROMADB")
+CELLPOSE_SAM = os.getenv("CELLPOSE_SAM")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+NEO4J_URI = os.getenv("NEO4J_URI")
+NEO4J_USERNAME = os.getenv("NEO4J_USERNAME")
+NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
+NEO4J_DATABASE = os.getenv("NEO4J_DATABASE")
+
+
+# --- LlamaIndex Global Settings ---
+def configure_llama_index():
+
+    """
+    Configures global LlamaIndex settings for the embedding model and the LLM.
+    """
+    
+    print("✓ Configuring LlamaIndex settings...")
+
+    # Gemma 3 Prompt Template 
+    query_wrapper_prompt = PromptTemplate(
+        "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{query_str}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    )
+    
+    llm = HuggingFaceInferenceAPI(
+        model_name=AGENT_MODEL_ID,
+        token = HF_TOKEN,
+        provider = "auto"
+    )
+
+    Settings.llm = llm
+
+    Settings.embed_model = HuggingFaceEmbedding(
+        model_name=f"sentence-transformers/{EMBEDDING_MODEL_ID}"
+    )
+    
+    Settings.chunk_size = 512
+    Settings.chunk_overlap = 50
+    
+    print("✓ LlamaIndex configured to use local Embedding Model and LLM.")

config/settings.py~

@@ -0,0 +1,60 @@
+"""
+
+"""
+
+import os
+import torch
+from llama_index.core import Settings
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.llms.huggingface import HuggingFaceLLM
+from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+from llama_index.core.prompts import PromptTemplate
+
+
+
+# ---- Model IDs ---
+AGENT_MODEL_ID = "google/gemma-3-12b-it"
+EMBEDDING_MODEL_ID = "clip-ViT-B-32"
+
+
+# --- Environment & Paths ---
+CHROMADB = os.getenv("CHROMADB", "./data/cellpose_db/")
+CELLPOSE_SAM = os.getenv("CELLPOSE_SAM", "./data/sam_vit_h_4b8939.pth")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+NEO4J_URI = os.getenv("NEO4J_URI", "neo4j+s://8d0af37b.databases.neo4j.io")
+NEO4J_USERNAME = os.getenv("NEO4J_USERNAME", "neo4j")
+NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "b5zqfnglm_CWHVYpmuXBR8oDyjaOqvT17L8pBUnfUJ0")
+NEO4J_DATABASE = os.getenv("NEO4J_DATABASE", "neo4j")
+
+
+# --- LlamaIndex Global Settings ---
+def configure_llama_index():
+
+    """
+    Configures global LlamaIndex settings for the embedding model and the LLM.
+    """
+    
+    print("✓ Configuring LlamaIndex settings...")
+
+    # Gemma 3 Prompt Template 
+    query_wrapper_prompt = PromptTemplate(
+        "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{query_str}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    )
+    
+    llm = HuggingFaceInferenceAPI(
+        model_name=AGENT_MODEL_ID,
+        token = HF_TOKEN,
+        provider = "auto"
+    )
+
+    Settings.llm = llm
+
+    Settings.embed_model = HuggingFaceEmbedding(
+        model_name=f"sentence-transformers/{EMBEDDING_MODEL_ID}"
+    )
+    
+    Settings.chunk_size = 512
+    Settings.chunk_overlap = 50
+    
+    print("✓ LlamaIndex configured to use local Embedding Model and LLM.")

models/embeddings.py

@@ -0,0 +1,36 @@
+"""
+
+"""
+
+from PIL import Image
+from sentence_transformers import SentenceTransformer
+from config import settings
+
+# --- Global Singleton for Embedding Model ---
+_embedding_model = None
+
+def get_embedding_model():
+    """
+    Initializes and returns the SentenceTransformer model (singleton pattern).
+    """
+    global _embedding_model
+    if _embedding_model is None:
+        print("Initializing embedding model...")
+        _embedding_model = SentenceTransformer(settings.EMBEDDING_MODEL_ID)
+        print(f"✓ Embedding model initialized ({settings.EMBEDDING_MODEL_ID})")
+    return _embedding_model
+
+def get_image_embedding(image_path: str) -> list[float]:
+    """
+    Generates a CLIP embedding for a given image file.
+
+    Args:
+        image_path (str): The path to the image file.
+
+    Returns:
+        list[float]: The image embedding as a list of floats.
+    """
+    model = get_embedding_model()
+    img = Image.open(image_path).convert("RGB")
+    embedding = model.encode(img, convert_to_numpy=True)
+    return embedding.tolist()

models/reranker.py

@@ -0,0 +1,29 @@
+"""
+
+"""
+
+from llama_index.core.postprocessor import SentenceTransformerRerank
+
+# --- Global Singleton for Reranker Model ---
+_reranker_model = None
+
+def get_reranker():
+    """
+    Initializes and returns the SentenceTransformerRerank model (singleton pattern).
+    This model will download on first use.
+    """
+    
+    global _reranker_model
+    
+    if _reranker_model is None:
+        
+        print("Initializing Cross-Encoder Reranker model...")
+        
+        # A popular, lightweight, and effective cross-encoder
+        _reranker_model = SentenceTransformerRerank(
+            model="cross-encoder/ms-marco-MiniLM-L-6-v2", 
+            top_n=3  # The number of documents to return after reranking
+        )
+        print("✓ Reranker model initialized.")
+        
+    return _reranker_model

stores/__init__.py

@@ -0,0 +1,15 @@
+from .chroma_store import get_client as get_chroma_client
+from .neo4j_store import (
+    get_graph_store,
+    check_graph_status,
+    initialize_knowledge_graph,
+    get_kg_index
+)
+
+__all__ = [
+    "get_chroma_client",
+    "get_graph_store",
+    "check_graph_status",
+    "initialize_knowledge_graph",
+    "get_kg_index"
+]

stores/chroma_store.py

@@ -0,0 +1,22 @@
+"""
+
+"""
+
+import chromadb
+from config import settings
+
+# --- Global Singleton for ChromaDB Client ---
+_chroma_client = None
+
+    
+
+def get_client():
+    """
+    Initializes and returns the ChromaDB persistent client (singleton pattern).
+    """
+    global _chroma_client
+    if _chroma_client is None:
+        print("Initializing ChromaDB client...")
+        _chroma_client = chromadb.PersistentClient(path=settings.CHROMADB)
+        print(f"✓ ChromaDB client connected to path: {settings.CHROMADB}")
+    return _chroma_client

stores/chroma_store.py~

@@ -0,0 +1,22 @@
+"""
+
+"""
+
+import chromadb
+from config import settings
+
+# --- Global Singleton for ChromaDB Client ---
+_chroma_client = None
+
+    
+
+def get_client():
+    """
+    Initializes and returns the ChromaDB persistent client (singleton pattern).
+    """
+    global _chroma_client
+    if _chroma_client is None:
+        print("Initializing ChromaDB client...")
+        _chroma_client = chromadb.PersistentClient(path=settings.CHROMADB)
+        print(f"✓ ChromaDB client connected to path: {settings.CHROMA_DB_PATH}")
+    return _chroma_client

stores/neo4j_store.py

@@ -0,0 +1,69 @@
+"""
+
+"""
+
+from llama_index.core import Document, KnowledgeGraphIndex, StorageContext
+from llama_index.graph_stores.neo4j import Neo4jGraphStore
+from neo4j import GraphDatabase
+
+from config import settings
+from stores import chroma_store
+
+# --- Global Singleton for KG Index ---
+_kg_index = None
+
+def get_graph_store():
+    """Initializes and returns the Neo4jGraphStore."""
+    return Neo4jGraphStore(
+        username=settings.NEO4J_USERNAME,
+        password=settings.NEO4J_PASSWORD,
+        url=settings.NEO4J_URI,
+        database=settings.NEO4J_DATABASE,
+    )
+
+def check_graph_status():
+    """Checks if the Neo4j graph contains any nodes or relationships."""
+    driver = GraphDatabase.driver(
+        settings.NEO4J_URI,
+        auth=(settings.NEO4J_USERNAME, settings.NEO4J_PASSWORD)
+    )
+    with driver.session(database=settings.NEO4J_DATABASE) as session:
+        nodes_result = session.run("MATCH (n) RETURN count(n) as count")
+        node_count = nodes_result.single()['count']
+        rels_result = session.run("MATCH ()-[r]->() RETURN count(r) as count")
+        rel_count = rels_result.single()['count']
+    driver.close()
+    return node_count, rel_count
+
+def initialize_knowledge_graph():
+    """Builds the knowledge graph from documents in ChromaDB and stores it in Neo4j."""
+    print("\n--- Building Knowledge Graph in Neo4j ---")
+    chroma_client = chroma_store.get_client()
+    doc_collection = chroma_client.get_collection(name='cellpose_docs')
+    doc_data = doc_collection.get()
+
+    documents = [
+        Document(text=text, metadata=meta)
+        for text, meta in zip(doc_data['documents'], doc_data['metadatas'])
+    ]
+
+    storage_context = StorageContext.from_defaults(graph_store=get_graph_store())
+    
+    KnowledgeGraphIndex.from_documents(
+        documents,
+        storage_context=storage_context,
+        max_triplets_per_chunk=3,
+        include_embeddings=True,
+        show_progress=True
+    )
+    print("✓ Knowledge Graph built and stored in Neo4j successfully.")
+
+def get_kg_index():
+    """Loads the KnowledgeGraphIndex from the existing Neo4j graph store."""
+    global _kg_index
+    if _kg_index is None:
+        print("Loading Knowledge Graph index from Neo4j...")
+        storage_context = StorageContext.from_defaults(graph_store=get_graph_store())
+        _kg_index = KnowledgeGraphIndex(nodes=[], storage_context=storage_context)
+        print("✓ Knowledge Graph index loaded.")
+    return _kg_index

tools/__init__.py

@@ -0,0 +1,35 @@
+from .segmentation import (
+    get_segmentation_parameters,
+    run_cellpose_sam,
+    refine_cellpose_sam_segmentation
+)
+from .search import (
+    list_all_collections,
+    search_documentation_vector,
+    search_knowledge_graph,
+    hybrid_search,
+    get_parameter_relationships,
+)
+
+all_tools = [
+    get_segmentation_parameters,
+    run_cellpose_sam,
+    refine_cellpose_sam_segmentation,
+    list_all_collections,
+    search_documentation_vector,
+    search_knowledge_graph,
+    hybrid_search,
+    get_parameter_relationships,
+]
+
+__all__ = [
+    "all_tools",
+    "get_segmentation_parameters",
+    "run_cellpose_sam",
+    "refine_cellpose_sam_segmentation",
+    "list_all_collections",
+    "search_documentation_vector",
+    "search_knowledge_graph",
+    "hybrid_search",
+    "get_parameter_relationships",
+]

tools/search.py

@@ -0,0 +1,101 @@
+"""
+
+"""
+
+# project/tools/search.py
+from smolagents import tool
+from langfuse import get_client
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.chroma import ChromaVectorStore
+
+from stores import get_chroma_client, get_kg_index
+from models.reranker import get_reranker
+
+langfuse = get_client()
+
+@tool
+def list_all_collections() -> list[str]:
+    """Lists the names of all available collections in the ChromaDB database."""
+    # This is fine because it has no arguments.
+    print("\n--- TOOL CALLED: list_all_collections ---")
+    client = get_chroma_client()
+    collections = client.list_collections()
+    return [c.name for c in collections]
+
+
+@tool
+def search_documentation_vector(query: str) -> str:
+    """
+    Searches cellpose documentation using vector search followed by a reranking step.
+
+    Args:
+        query (str): The question or search term to look up in the documentation.
+    """
+    print(f"\n--- TOOL CALLED: search_documentation_vector (with Reranker) for '{query}' ---")
+    try:
+        client = get_chroma_client()
+        collection = client.get_collection(name='cellpose_docs')
+        vector_store = ChromaVectorStore(chroma_collection=collection)
+        vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
+        
+        query_engine = vector_index.as_query_engine(
+            similarity_top_k=25,
+            node_postprocessors=[get_reranker()]
+        )
+        response = query_engine.query(query)
+        return str(response)
+    except Exception as e:
+        return f"Error searching documentation: {e}"
+
+
+@tool
+def search_knowledge_graph(query: str) -> str:
+    """
+    Searches using knowledge graph relationships (Neo4j). Best for "how" and "why" questions.
+
+    Args:
+        query (str): The question about relationships between concepts (e.g., parameters).
+    """
+    print(f"\n--- TOOL CALLED: search_knowledge_graph for '{query}' ---")
+    try:
+        kg_index = get_kg_index()
+        query_engine = kg_index.as_query_engine(
+            include_text=True, response_mode="tree_summarize"
+        )
+        response = query_engine.query(query)
+        return str(response)
+    except Exception as e:
+        return f"Error querying knowledge graph: {e}."
+
+
+@tool
+def get_parameter_relationships(parameter_name: str) -> str:
+    """
+    Gets information about how a parameter relates to others using the knowledge graph.
+
+    Args:
+        parameter_name (str): The specific parameter name to investigate (e.g., 'flow_threshold').
+    """
+    print(f"\n--- TOOL CALLED: get_parameter_relationships for '{parameter_name}' ---")
+    query = f"What is {parameter_name} and how does it relate to other parameters?"
+    return search_knowledge_graph(query)
+
+
+@tool
+def hybrid_search(query: str) -> str:
+    """
+    Combines reranked vector search and knowledge graph search for complex questions.
+
+    Args:
+        query (str): The complex question that may require both semantic and relational understanding.
+    """
+    print(f"\n--- TOOL CALLED: hybrid_search (with Reranker) for '{query}' ---")
+    try:
+        vector_response_str = search_documentation_vector(query)
+        kg_response = search_knowledge_graph(query)
+
+        return f"Vector Search Results (Reranked):\n{vector_response_str}\n\nKnowledge Graph Insights:\n{kg_response}"
+    
+    except Exception as e:
+        print(f"--- Hybrid search failed, falling back to vector search: {e} ---")
+        return search_documentation_vector(query)

tools/segmentation.py

@@ -0,0 +1,532 @@
+"""
+Segmentation tools for cellpose-sam pipeline with proper smolagents VLM integration.
+"""
+import base64
+import json
+import re
+from typing import Any, Dict, TYPE_CHECKING
+import numpy as np
+import cv2
+import torch
+from PIL import Image
+from skimage.measure import regionprops
+from cellpose import models
+from segment_anything import sam_model_registry, SamPredictor
+
+from smolagents import tool
+from smolagents.agents import ActionStep
+from langfuse import get_client
+
+from stores import chroma_store
+from models.embeddings import get_image_embedding
+from utils.image_utils import resize_and_encode_image
+from config import settings
+
+
+langfuse = get_client()
+
+
+# --- Global State and Caching ---
+_image_cache: Dict[str, tuple[str, str]] = {}
+_cellpose_model = None
+_sam_predictor = None
+
+
+def get_cellpose_model():
+    """Initialize Cellpose model (singleton)"""
+    global _cellpose_model
+    if _cellpose_model is None:
+        _cellpose_model = models.CellposeModel(gpu=torch.cuda.is_available())
+    return _cellpose_model
+
+
+def get_sam_predictor():
+    """Initialize SAM predictor (singleton)"""
+    global _sam_predictor
+    if _sam_predictor is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        sam = sam_model_registry["vit_h"](checkpoint=settings.CELLPOSE_SAM)
+        sam.to(device=device)
+        _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)
+    _image_cache[image_path] = (image_base64, media_type)
+    return image_base64, media_type
+
+
+def parse_parameters_from_text(param_text: str) -> dict:
+    """Extract parameter values from parameter text string."""
+    defaults = {
+        'diameter': 25,
+        'flow_threshold': 0.6,
+        'cellprob_threshold': 0,
+        'min_size': 15
+    }
+    
+    params = defaults.copy()
+    
+    patterns = {
+        'diameter': r'diameter[=:]\s*(\d+)',
+        'flow_threshold': r'flow_threshold[=:]\s*([\d.]+)',
+        'cellprob_threshold': r'cellprob_threshold[=:]\s*([-\d.]+)',
+        'min_size': r'min_size[=:]\s*(\d+)'
+    }
+    
+    for param_name, pattern in patterns.items():
+        match = re.search(pattern, param_text, re.IGNORECASE)
+        if match:
+            value = match.group(1)
+            if param_name in ['diameter', 'min_size']:
+                params[param_name] = int(value)
+            else:
+                params[param_name] = float(value)
+    
+    return params
+
+
+@tool
+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.
+    
+    Args:
+        image_path (str): Path to the image file to segment.
+        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}' ---")
+
+    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}"})
+
+    try:
+        # Get similar parameters from ChromaDB
+        client = chroma_store.get_client()
+        collection = client.get_collection(name='cellpose-sam_parameters_by_image_similarity')
+        query_embedding = get_image_embedding(image_path)
+        
+        results = collection.query(query_embeddings=[query_embedding], n_results=1)
+
+        if not (results['metadatas'] and results['metadatas'][0]):
+            return json.dumps({"error": "No similar images found in the database."})
+
+        matched_parameters = results['metadatas'][0][0].get('parameter_text', 'N/A')
+        matched_image = results['metadatas'][0][0].get('image_name', 'N/A')
+        distance = results['distances'][0][0]
+
+        print(f"Most similar: {matched_image} (distance: {distance:.3f})")
+        print(f"Recommended: {matched_parameters}")
+
+        # Parse parameters
+        params = parse_parameters_from_text(matched_parameters)
+
+        # Analyze image
+        image = np.array(Image.open(image_path).convert("RGB"))
+        image_shape = image.shape
+        stats = {
+            'size': (image_shape[0] * image_shape[1]),
+            'mean_intensity': float(np.mean(image)),
+            'stdev_intensity': float(np.std(image)),
+            'min_intensity': int(np.min(image)),
+            'max_intensity': int(np.max(image)),
+        }
+
+        # Log to Langfuse WITH image (for observability)
+        try:
+            langfuse.update_current_trace(
+                input={
+                    "image_path": image_path,
+                    "query_image": {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:{media_type};base64,{image_base64}"
+                        }
+                    },
+                    "image_stats": stats
+                },
+                metadata={
+                    "matched_image": matched_image,
+                    "similarity_distance": float(distance),
+                    "matched_parameters": matched_parameters,
+                    "parsed_parameters": params
+                }
+            )
+        except Exception as log_error:
+            print(f"Warning: Could not log to Langfuse: {log_error}")
+
+        # Determine confidence level
+        if distance < 0.2:
+            confidence = "high"
+            confidence_note = "Very similar image found. Parameters should work well as-is."
+        elif distance < 0.4:
+            confidence = "medium"
+            confidence_note = "Similar image found. Parameters are a good starting point but may need minor adjustments."
+        else:
+            confidence = "low"
+            confidence_note = "No very similar images found. Parameters may need significant adjustment based on visual inspection."
+
+        # Return WITHOUT base64 (image already attached to ActionStep)
+        response = {
+            "status": "success",
+            "image_path": image_path,
+            "recommended_parameters": params,
+            "matched_image": matched_image,
+            "similarity_distance": float(distance),
+            "confidence": confidence,
+            "image_stats": stats,
+            "raw_parameter_text": matched_parameters,
+            "visual_guidance": "IMAGE NOW VISIBLE: The input image is now attached to this step. "
+                              "Please visually inspect the image to assess cell morphology, density, "
+                              "and boundaries before deciding whether to adjust the recommended parameters.",
+            "recommendation": f"{confidence_note}\n\nRecommended parameters:\n"
+                             f"- diameter: {params['diameter']}\n"
+                             f"- flow_threshold: {params['flow_threshold']}\n"
+                             f"- cellprob_threshold: {params['cellprob_threshold']}\n"
+                             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"diameter={params['diameter']}, flow_threshold={params['flow_threshold']}, "
+                             f"cellprob_threshold={params['cellprob_threshold']}, min_size={params['min_size']})"
+        }
+
+        return json.dumps(response, indent=2)
+
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+@tool
+def run_cellpose_sam(
+    image_path: str,
+    diameter: int = None,
+    flow_threshold: float = None,
+    cellprob_threshold: float = None,
+    min_size: int = None,
+    output_path: str = None,
+    use_recommended_params: bool = True,
+    agent: Any = None
+) -> str:
+    """
+    Runs cellpose-sam segmentation pipeline on an image with specified parameters.
+    Returns results WITHOUT base64 images to prevent GPU memory issues.
+    
+    Args:
+        image_path (str): Path to the image file to segment
+        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)
+        min_size (int): Minimum cell size in pixels
+        output_path (str): Optional path to save the overlay image
+        use_recommended_params (bool): If True and params not provided, get recommendations
+        agent (Any, optional): The agent instance
+    
+    Returns:
+        str: JSON string with segmentation results (paths and stats, NO base64)
+    """
+    print(f"\n--- TOOL CALLED: run_cellpose_sam for '{image_path}' ---")
+    
+    try:
+        # Load and cache input image
+        input_image_base64, input_media_type = _get_cached_image(image_path) or _load_and_cache_image(image_path)
+    except Exception as e:
+        return json.dumps({"error": f"Could not read input image: {e}"})
+    
+    # Auto-fetch recommended parameters if needed
+    if use_recommended_params and all(p is None for p in [diameter, flow_threshold, cellprob_threshold, min_size]):
+        print("No parameters provided. Fetching recommended parameters...")
+        param_response = get_segmentation_parameters(image_path, agent=agent)
+        
+        try:
+            param_data = json.loads(param_response)
+            if param_data.get("status") == "success":
+                rec_params = param_data["recommended_parameters"]
+                diameter = diameter or rec_params.get('diameter', 25)
+                flow_threshold = flow_threshold or rec_params.get('flow_threshold', 0.6)
+                cellprob_threshold = cellprob_threshold or rec_params.get('cellprob_threshold', 0)
+                min_size = min_size or rec_params.get('min_size', 15)
+            else:
+                diameter, flow_threshold, cellprob_threshold, min_size = 25, 0.6, 0, 15
+        except json.JSONDecodeError:
+            diameter, flow_threshold, cellprob_threshold, min_size = 25, 0.6, 0, 15
+    else:
+        diameter = diameter if diameter is not None else 25
+        flow_threshold = flow_threshold if flow_threshold is not None else 0.6
+        cellprob_threshold = cellprob_threshold if cellprob_threshold is not None else 0
+        min_size = min_size if min_size is not None else 15
+    
+    print(f"Final parameters: diameter={diameter}, flow_threshold={flow_threshold}, "
+          f"cellprob_threshold={cellprob_threshold}, min_size={min_size}")
+    
+    try:
+        # Read image
+        img = cv2.imread(image_path)
+        if img is None:
+            return json.dumps({"error": f"Could not read image at {image_path}"})
+        
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        cellpose_model = get_cellpose_model()
+        sam_predictor = get_sam_predictor()
+        
+        # Run Cellpose
+        print("Running Cellpose...")
+        masks_cellpose, flows, styles = cellpose_model.eval(
+            img_rgb,
+            diameter=diameter,
+            flow_threshold=flow_threshold,
+            cellprob_threshold=cellprob_threshold,
+            min_size=min_size
+        )
+        
+        if masks_cellpose.max() == 0:
+            return json.dumps({
+                "status": "no_cells_detected",
+                "message": "No cells detected. Try adjusting parameters.",
+                "parameters": {
+                    "diameter": diameter,
+                    "flow_threshold": flow_threshold,
+                    "cellprob_threshold": cellprob_threshold,
+                    "min_size": min_size
+                }
+            })
+        
+        print(f"Cellpose detected {masks_cellpose.max()} regions")
+        
+        # SAM refinement
+        sam_predictor.set_image(img_rgb)
+        props = regionprops(masks_cellpose)
+        boxes = np.array([prop.bbox for prop in props])
+        boxes = boxes[:, [1,0,3,2]]
+        
+        print(f"Refining {len(boxes)} masks with SAM...")
+        
+        combined_masks = np.zeros(img_rgb.shape[:2], dtype=np.uint16)
+        colored_overlay = img_rgb.copy().astype(np.float32)
+        
+        for i, box in enumerate(boxes):
+            masks, scores, _ = sam_predictor.predict(box=box, multimask_output=True)
+            best_mask = masks[np.argmax(scores)]
+            combined_masks[best_mask] = i + 1
+            color = np.random.randint(0, 255, 3)
+            colored_overlay[best_mask] = colored_overlay[best_mask] * 0.6 + color * 0.4
+        
+        # Generate output path
+        if output_path is None:
+            base_name = image_path.rsplit('.', 1)[0]
+            output_path = f"{base_name}_cellpose_sam_overlay.png"
+        
+        # Save output
+        cv2.imwrite(output_path, cv2.cvtColor(colored_overlay.astype(np.uint8), cv2.COLOR_RGB2BGR))
+        
+        # Load and cache output image
+        output_image_base64, output_media_type = _load_and_cache_image(output_path)
+        
+        # Log to Langfuse WITH both images
+        try:
+            langfuse.update_current_trace(
+                input={
+                    "image_path": image_path,
+                    "input_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{input_media_type};base64,{input_image_base64}"}
+                    }
+                },
+                output={
+                    "cell_count": int(masks_cellpose.max()),
+                    "output_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{output_media_type};base64,{output_image_base64}"}
+                    },
+                    "output_path": output_path
+                },
+                metadata={
+                    "parameters": {
+                        "diameter": diameter,
+                        "flow_threshold": flow_threshold,
+                        "cellprob_threshold": cellprob_threshold,
+                        "min_size": min_size
+                    }
+                }
+            )
+        except Exception as log_error:
+            print(f"Warning: Could not log output to Langfuse: {log_error}")
+        
+        # Return WITHOUT base64
+        result = {
+            "status": "success",
+            "cell_count": int(masks_cellpose.max()),
+            "output_path": output_path,
+            "input_path": image_path,
+            "parameters": {
+                "diameter": diameter,
+                "flow_threshold": flow_threshold,
+                "cellprob_threshold": cellprob_threshold,
+                "min_size": min_size
+            },
+            "summary": f"Detected {masks_cellpose.max()} cells. Output saved to: {output_path}",
+            "next_step": "Call refine_cellpose_sam_segmentation to visually analyze the segmentation quality and decide if parameter adjustments are needed."
+        }
+        
+        return json.dumps(result, indent=2)
+        
+    except Exception as e:
+        return json.dumps({"error": f"Error during segmentation: {e}"})
+
+
+@tool
+def refine_cellpose_sam_segmentation(
+    original_image_path: str,
+    segmentation_output_path: str,
+    current_parameters: dict,
+    agent: Any = None,
+) -> str:
+    """
+    Provides both original and segmented images to the VLM for visual quality assessment.
+    The VLM will be able to see both images and provide informed analysis.
+    
+    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.
+    
+    Before calling, consider using search_knowledge_graph or hybrid_search to refresh
+    your understanding of how cellpose parameters affect segmentation.
+    
+    Common issues and fixes:
+    - Under-segmentation (cells merged): decrease flow_threshold or diameter
+    - Over-segmentation (cells fragmented): increase flow_threshold or min_size
+    - Too few cells: decrease cellprob_threshold or flow_threshold
+    - 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
+        current_parameters: Dict with current diameter, flow_threshold, cellprob_threshold, min_size
+        agent: The agent instance (passed automatically)
+    
+    Returns:
+        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"Current parameters: {current_parameters}")
+    
+    try:
+        # Load both images (for cache)
+        original_b64, original_type = _get_cached_image(original_image_path) or _load_and_cache_image(original_image_path)
+        segmented_b64, segmented_type = _get_cached_image(segmentation_output_path) or _load_and_cache_image(segmentation_output_path)
+        
+        # CRITICAL: Attach BOTH images to ActionStep so VLM can see them
+        if agent is not None and hasattr(agent, 'memory') and hasattr(agent.memory, 'steps'):
+            current_steps = [s for s in agent.memory.steps if isinstance(s, ActionStep)]
+            if current_steps:
+                current_step = current_steps[-1]
+                
+                # Load both as PIL Images
+                original_img = Image.open(original_image_path).convert("RGB")
+                segmented_img = Image.open(segmentation_output_path).convert("RGB")
+                
+                # CRITICAL: Use .copy() for both images
+                current_step.observations_images = [original_img.copy(), segmented_img.copy()]
+                print(f"✓ Attached both images to ActionStep for VLM comparison")
+        
+        # Get image dimensions for context
+        original_img_array = np.array(Image.open(original_image_path).convert("RGB"))
+        img_size = original_img_array.shape[0] * original_img_array.shape[1]
+        
+        # Log to Langfuse WITH both images
+        try:
+            langfuse.update_current_trace(
+                input={
+                    "tool": "refine_cellpose_sam_segmentation",
+                    "original_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{original_type};base64,{original_b64}"}
+                    },
+                    "segmented_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{segmented_type};base64,{segmented_b64}"}
+                    },
+                    "current_parameters": current_parameters
+                },
+                metadata={
+                    "original_path": original_image_path,
+                    "segmented_path": segmentation_output_path
+                }
+            )
+        except Exception as log_error:
+            print(f"Warning: Could not log to Langfuse: {log_error}")
+        
+        # Return analysis guidance WITHOUT base64
+        analysis = {
+            "status": "ready_for_visual_analysis",
+            "images_attached": "BOTH IMAGES NOW VISIBLE: The first image is the original input, "
+                              "the second is the segmented overlay. Compare them visually to assess quality.",
+            "image_paths": {
+                "original": original_image_path,
+                "segmented": segmentation_output_path
+            },
+            "current_parameters": current_parameters,
+            "image_info": {
+                "dimensions": f"{original_img_array.shape[1]}x{original_img_array.shape[0]}",
+                "total_pixels": img_size
+            },
+            "visual_analysis_checklist": [
+                "1. Do the colored masks accurately cover entire cells without extending beyond boundaries?",
+                "2. Are neighboring cells properly separated, or are they merged together?",
+                "3. Are there many small false positive detections (noise)?",
+                "4. Are any large, obvious cells being missed completely?",
+                "5. Overall quality assessment: excellent, good, needs_refinement, or poor?"
+            ],
+            "parameter_adjustment_guide": {
+                "under_segmentation": {
+                    "symptoms": "Masks don't reach cell edges, cells appear merged",
+                    "solution": "Decrease flow_threshold by 0.1-0.2 OR decrease diameter by 10-20%"
+                },
+                "over_segmentation": {
+                    "symptoms": "Masks extend past boundaries, cells fragmented into pieces",
+                    "solution": "Increase flow_threshold by 0.1-0.2 OR increase min_size to 2-3x current value"
+                },
+                "too_few_cells": {
+                    "symptoms": "Obvious cells in image are not being detected",
+                    "solution": "Decrease cellprob_threshold by 1-2 OR decrease flow_threshold by 0.1-0.2"
+                },
+                "too_many_false_positives": {
+                    "symptoms": "Many tiny spurious detections, background noise detected as cells",
+                    "solution": "Increase cellprob_threshold by 1-2 OR increase min_size to 2-3x current value"
+                }
+            },
+            "next_steps": {
+                "if_good": "If segmentation looks accurate, inform the user of success and provide the output_path.",
+                "if_needs_refinement": "Based on your visual analysis, adjust the appropriate parameters and call run_cellpose_sam again with the new values.",
+                "important": "You can only call refine_cellpose_sam_segmentation AT MOST 2 TIMES total. If this is your second call, you must make a final decision."
+            }
+        }
+        
+        return json.dumps(analysis, indent=2)
+        
+    except Exception as e:
+        error_result = {
+            "status": "error",
+            "error": str(e),
+            "message": "Could not load images for refinement. Check that both file paths are valid."
+        }
+        return json.dumps(error_result, indent=2)

tools/segmentation.py~

@@ -0,0 +1,531 @@
+"""
+Segmentation tools for cellpose-sam pipeline with proper smolagents VLM integration.
+"""
+import base64
+import json
+import re
+from typing import Any, Dict, TYPE_CHECKING
+import numpy as np
+import cv2
+import torch
+from PIL import Image
+from skimage.measure import regionprops
+from cellpose import models
+from segment_anything import sam_model_registry, SamPredictor
+
+from smolagents import tool
+from smolagents.agents import ActionStep
+from langfuse import get_client
+
+from stores import chroma_store
+from models.embeddings import get_image_embedding
+from utils.image_utils import resize_and_encode_image
+
+
+langfuse = get_client()
+
+
+# --- Global State and Caching ---
+_image_cache: Dict[str, tuple[str, str]] = {}
+_cellpose_model = None
+_sam_predictor = None
+
+
+def get_cellpose_model():
+    """Initialize Cellpose model (singleton)"""
+    global _cellpose_model
+    if _cellpose_model is None:
+        _cellpose_model = models.CellposeModel(gpu=torch.cuda.is_available())
+    return _cellpose_model
+
+
+def get_sam_predictor():
+    """Initialize SAM predictor (singleton)"""
+    global _sam_predictor
+    if _sam_predictor is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
+        sam.to(device=device)
+        _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)
+    _image_cache[image_path] = (image_base64, media_type)
+    return image_base64, media_type
+
+
+def parse_parameters_from_text(param_text: str) -> dict:
+    """Extract parameter values from parameter text string."""
+    defaults = {
+        'diameter': 25,
+        'flow_threshold': 0.6,
+        'cellprob_threshold': 0,
+        'min_size': 15
+    }
+    
+    params = defaults.copy()
+    
+    patterns = {
+        'diameter': r'diameter[=:]\s*(\d+)',
+        'flow_threshold': r'flow_threshold[=:]\s*([\d.]+)',
+        'cellprob_threshold': r'cellprob_threshold[=:]\s*([-\d.]+)',
+        'min_size': r'min_size[=:]\s*(\d+)'
+    }
+    
+    for param_name, pattern in patterns.items():
+        match = re.search(pattern, param_text, re.IGNORECASE)
+        if match:
+            value = match.group(1)
+            if param_name in ['diameter', 'min_size']:
+                params[param_name] = int(value)
+            else:
+                params[param_name] = float(value)
+    
+    return params
+
+
+@tool
+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.
+    
+    Args:
+        image_path (str): Path to the image file to segment.
+        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}' ---")
+
+    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}"})
+
+    try:
+        # Get similar parameters from ChromaDB
+        client = chroma_store.get_client()
+        collection = client.get_collection(name='cellpose-sam_parameters_by_image_similarity')
+        query_embedding = get_image_embedding(image_path)
+        
+        results = collection.query(query_embeddings=[query_embedding], n_results=1)
+
+        if not (results['metadatas'] and results['metadatas'][0]):
+            return json.dumps({"error": "No similar images found in the database."})
+
+        matched_parameters = results['metadatas'][0][0].get('parameter_text', 'N/A')
+        matched_image = results['metadatas'][0][0].get('image_name', 'N/A')
+        distance = results['distances'][0][0]
+
+        print(f"Most similar: {matched_image} (distance: {distance:.3f})")
+        print(f"Recommended: {matched_parameters}")
+
+        # Parse parameters
+        params = parse_parameters_from_text(matched_parameters)
+
+        # Analyze image
+        image = np.array(Image.open(image_path).convert("RGB"))
+        image_shape = image.shape
+        stats = {
+            'size': (image_shape[0] * image_shape[1]),
+            'mean_intensity': float(np.mean(image)),
+            'stdev_intensity': float(np.std(image)),
+            'min_intensity': int(np.min(image)),
+            'max_intensity': int(np.max(image)),
+        }
+
+        # Log to Langfuse WITH image (for observability)
+        try:
+            langfuse.update_current_trace(
+                input={
+                    "image_path": image_path,
+                    "query_image": {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:{media_type};base64,{image_base64}"
+                        }
+                    },
+                    "image_stats": stats
+                },
+                metadata={
+                    "matched_image": matched_image,
+                    "similarity_distance": float(distance),
+                    "matched_parameters": matched_parameters,
+                    "parsed_parameters": params
+                }
+            )
+        except Exception as log_error:
+            print(f"Warning: Could not log to Langfuse: {log_error}")
+
+        # Determine confidence level
+        if distance < 0.2:
+            confidence = "high"
+            confidence_note = "Very similar image found. Parameters should work well as-is."
+        elif distance < 0.4:
+            confidence = "medium"
+            confidence_note = "Similar image found. Parameters are a good starting point but may need minor adjustments."
+        else:
+            confidence = "low"
+            confidence_note = "No very similar images found. Parameters may need significant adjustment based on visual inspection."
+
+        # Return WITHOUT base64 (image already attached to ActionStep)
+        response = {
+            "status": "success",
+            "image_path": image_path,
+            "recommended_parameters": params,
+            "matched_image": matched_image,
+            "similarity_distance": float(distance),
+            "confidence": confidence,
+            "image_stats": stats,
+            "raw_parameter_text": matched_parameters,
+            "visual_guidance": "IMAGE NOW VISIBLE: The input image is now attached to this step. "
+                              "Please visually inspect the image to assess cell morphology, density, "
+                              "and boundaries before deciding whether to adjust the recommended parameters.",
+            "recommendation": f"{confidence_note}\n\nRecommended parameters:\n"
+                             f"- diameter: {params['diameter']}\n"
+                             f"- flow_threshold: {params['flow_threshold']}\n"
+                             f"- cellprob_threshold: {params['cellprob_threshold']}\n"
+                             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"diameter={params['diameter']}, flow_threshold={params['flow_threshold']}, "
+                             f"cellprob_threshold={params['cellprob_threshold']}, min_size={params['min_size']})"
+        }
+
+        return json.dumps(response, indent=2)
+
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+@tool
+def run_cellpose_sam(
+    image_path: str,
+    diameter: int = None,
+    flow_threshold: float = None,
+    cellprob_threshold: float = None,
+    min_size: int = None,
+    output_path: str = None,
+    use_recommended_params: bool = True,
+    agent: Any = None
+) -> str:
+    """
+    Runs cellpose-sam segmentation pipeline on an image with specified parameters.
+    Returns results WITHOUT base64 images to prevent GPU memory issues.
+    
+    Args:
+        image_path (str): Path to the image file to segment
+        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)
+        min_size (int): Minimum cell size in pixels
+        output_path (str): Optional path to save the overlay image
+        use_recommended_params (bool): If True and params not provided, get recommendations
+        agent (Any, optional): The agent instance
+    
+    Returns:
+        str: JSON string with segmentation results (paths and stats, NO base64)
+    """
+    print(f"\n--- TOOL CALLED: run_cellpose_sam for '{image_path}' ---")
+    
+    try:
+        # Load and cache input image
+        input_image_base64, input_media_type = _get_cached_image(image_path) or _load_and_cache_image(image_path)
+    except Exception as e:
+        return json.dumps({"error": f"Could not read input image: {e}"})
+    
+    # Auto-fetch recommended parameters if needed
+    if use_recommended_params and all(p is None for p in [diameter, flow_threshold, cellprob_threshold, min_size]):
+        print("No parameters provided. Fetching recommended parameters...")
+        param_response = get_segmentation_parameters(image_path, agent=agent)
+        
+        try:
+            param_data = json.loads(param_response)
+            if param_data.get("status") == "success":
+                rec_params = param_data["recommended_parameters"]
+                diameter = diameter or rec_params.get('diameter', 25)
+                flow_threshold = flow_threshold or rec_params.get('flow_threshold', 0.6)
+                cellprob_threshold = cellprob_threshold or rec_params.get('cellprob_threshold', 0)
+                min_size = min_size or rec_params.get('min_size', 15)
+            else:
+                diameter, flow_threshold, cellprob_threshold, min_size = 25, 0.6, 0, 15
+        except json.JSONDecodeError:
+            diameter, flow_threshold, cellprob_threshold, min_size = 25, 0.6, 0, 15
+    else:
+        diameter = diameter if diameter is not None else 25
+        flow_threshold = flow_threshold if flow_threshold is not None else 0.6
+        cellprob_threshold = cellprob_threshold if cellprob_threshold is not None else 0
+        min_size = min_size if min_size is not None else 15
+    
+    print(f"Final parameters: diameter={diameter}, flow_threshold={flow_threshold}, "
+          f"cellprob_threshold={cellprob_threshold}, min_size={min_size}")
+    
+    try:
+        # Read image
+        img = cv2.imread(image_path)
+        if img is None:
+            return json.dumps({"error": f"Could not read image at {image_path}"})
+        
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        cellpose_model = get_cellpose_model()
+        sam_predictor = get_sam_predictor()
+        
+        # Run Cellpose
+        print("Running Cellpose...")
+        masks_cellpose, flows, styles = cellpose_model.eval(
+            img_rgb,
+            diameter=diameter,
+            flow_threshold=flow_threshold,
+            cellprob_threshold=cellprob_threshold,
+            min_size=min_size
+        )
+        
+        if masks_cellpose.max() == 0:
+            return json.dumps({
+                "status": "no_cells_detected",
+                "message": "No cells detected. Try adjusting parameters.",
+                "parameters": {
+                    "diameter": diameter,
+                    "flow_threshold": flow_threshold,
+                    "cellprob_threshold": cellprob_threshold,
+                    "min_size": min_size
+                }
+            })
+        
+        print(f"Cellpose detected {masks_cellpose.max()} regions")
+        
+        # SAM refinement
+        sam_predictor.set_image(img_rgb)
+        props = regionprops(masks_cellpose)
+        boxes = np.array([prop.bbox for prop in props])
+        boxes = boxes[:, [1,0,3,2]]
+        
+        print(f"Refining {len(boxes)} masks with SAM...")
+        
+        combined_masks = np.zeros(img_rgb.shape[:2], dtype=np.uint16)
+        colored_overlay = img_rgb.copy().astype(np.float32)
+        
+        for i, box in enumerate(boxes):
+            masks, scores, _ = sam_predictor.predict(box=box, multimask_output=True)
+            best_mask = masks[np.argmax(scores)]
+            combined_masks[best_mask] = i + 1
+            color = np.random.randint(0, 255, 3)
+            colored_overlay[best_mask] = colored_overlay[best_mask] * 0.6 + color * 0.4
+        
+        # Generate output path
+        if output_path is None:
+            base_name = image_path.rsplit('.', 1)[0]
+            output_path = f"{base_name}_cellpose_sam_overlay.png"
+        
+        # Save output
+        cv2.imwrite(output_path, cv2.cvtColor(colored_overlay.astype(np.uint8), cv2.COLOR_RGB2BGR))
+        
+        # Load and cache output image
+        output_image_base64, output_media_type = _load_and_cache_image(output_path)
+        
+        # Log to Langfuse WITH both images
+        try:
+            langfuse.update_current_trace(
+                input={
+                    "image_path": image_path,
+                    "input_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{input_media_type};base64,{input_image_base64}"}
+                    }
+                },
+                output={
+                    "cell_count": int(masks_cellpose.max()),
+                    "output_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{output_media_type};base64,{output_image_base64}"}
+                    },
+                    "output_path": output_path
+                },
+                metadata={
+                    "parameters": {
+                        "diameter": diameter,
+                        "flow_threshold": flow_threshold,
+                        "cellprob_threshold": cellprob_threshold,
+                        "min_size": min_size
+                    }
+                }
+            )
+        except Exception as log_error:
+            print(f"Warning: Could not log output to Langfuse: {log_error}")
+        
+        # Return WITHOUT base64
+        result = {
+            "status": "success",
+            "cell_count": int(masks_cellpose.max()),
+            "output_path": output_path,
+            "input_path": image_path,
+            "parameters": {
+                "diameter": diameter,
+                "flow_threshold": flow_threshold,
+                "cellprob_threshold": cellprob_threshold,
+                "min_size": min_size
+            },
+            "summary": f"Detected {masks_cellpose.max()} cells. Output saved to: {output_path}",
+            "next_step": "Call refine_cellpose_sam_segmentation to visually analyze the segmentation quality and decide if parameter adjustments are needed."
+        }
+        
+        return json.dumps(result, indent=2)
+        
+    except Exception as e:
+        return json.dumps({"error": f"Error during segmentation: {e}"})
+
+
+@tool
+def refine_cellpose_sam_segmentation(
+    original_image_path: str,
+    segmentation_output_path: str,
+    current_parameters: dict,
+    agent: Any = None,
+) -> str:
+    """
+    Provides both original and segmented images to the VLM for visual quality assessment.
+    The VLM will be able to see both images and provide informed analysis.
+    
+    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.
+    
+    Before calling, consider using search_knowledge_graph or hybrid_search to refresh
+    your understanding of how cellpose parameters affect segmentation.
+    
+    Common issues and fixes:
+    - Under-segmentation (cells merged): decrease flow_threshold or diameter
+    - Over-segmentation (cells fragmented): increase flow_threshold or min_size
+    - Too few cells: decrease cellprob_threshold or flow_threshold
+    - 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
+        current_parameters: Dict with current diameter, flow_threshold, cellprob_threshold, min_size
+        agent: The agent instance (passed automatically)
+    
+    Returns:
+        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"Current parameters: {current_parameters}")
+    
+    try:
+        # Load both images (for cache)
+        original_b64, original_type = _get_cached_image(original_image_path) or _load_and_cache_image(original_image_path)
+        segmented_b64, segmented_type = _get_cached_image(segmentation_output_path) or _load_and_cache_image(segmentation_output_path)
+        
+        # CRITICAL: Attach BOTH images to ActionStep so VLM can see them
+        if agent is not None and hasattr(agent, 'memory') and hasattr(agent.memory, 'steps'):
+            current_steps = [s for s in agent.memory.steps if isinstance(s, ActionStep)]
+            if current_steps:
+                current_step = current_steps[-1]
+                
+                # Load both as PIL Images
+                original_img = Image.open(original_image_path).convert("RGB")
+                segmented_img = Image.open(segmentation_output_path).convert("RGB")
+                
+                # CRITICAL: Use .copy() for both images
+                current_step.observations_images = [original_img.copy(), segmented_img.copy()]
+                print(f"✓ Attached both images to ActionStep for VLM comparison")
+        
+        # Get image dimensions for context
+        original_img_array = np.array(Image.open(original_image_path).convert("RGB"))
+        img_size = original_img_array.shape[0] * original_img_array.shape[1]
+        
+        # Log to Langfuse WITH both images
+        try:
+            langfuse.update_current_trace(
+                input={
+                    "tool": "refine_cellpose_sam_segmentation",
+                    "original_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{original_type};base64,{original_b64}"}
+                    },
+                    "segmented_image": {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{segmented_type};base64,{segmented_b64}"}
+                    },
+                    "current_parameters": current_parameters
+                },
+                metadata={
+                    "original_path": original_image_path,
+                    "segmented_path": segmentation_output_path
+                }
+            )
+        except Exception as log_error:
+            print(f"Warning: Could not log to Langfuse: {log_error}")
+        
+        # Return analysis guidance WITHOUT base64
+        analysis = {
+            "status": "ready_for_visual_analysis",
+            "images_attached": "BOTH IMAGES NOW VISIBLE: The first image is the original input, "
+                              "the second is the segmented overlay. Compare them visually to assess quality.",
+            "image_paths": {
+                "original": original_image_path,
+                "segmented": segmentation_output_path
+            },
+            "current_parameters": current_parameters,
+            "image_info": {
+                "dimensions": f"{original_img_array.shape[1]}x{original_img_array.shape[0]}",
+                "total_pixels": img_size
+            },
+            "visual_analysis_checklist": [
+                "1. Do the colored masks accurately cover entire cells without extending beyond boundaries?",
+                "2. Are neighboring cells properly separated, or are they merged together?",
+                "3. Are there many small false positive detections (noise)?",
+                "4. Are any large, obvious cells being missed completely?",
+                "5. Overall quality assessment: excellent, good, needs_refinement, or poor?"
+            ],
+            "parameter_adjustment_guide": {
+                "under_segmentation": {
+                    "symptoms": "Masks don't reach cell edges, cells appear merged",
+                    "solution": "Decrease flow_threshold by 0.1-0.2 OR decrease diameter by 10-20%"
+                },
+                "over_segmentation": {
+                    "symptoms": "Masks extend past boundaries, cells fragmented into pieces",
+                    "solution": "Increase flow_threshold by 0.1-0.2 OR increase min_size to 2-3x current value"
+                },
+                "too_few_cells": {
+                    "symptoms": "Obvious cells in image are not being detected",
+                    "solution": "Decrease cellprob_threshold by 1-2 OR decrease flow_threshold by 0.1-0.2"
+                },
+                "too_many_false_positives": {
+                    "symptoms": "Many tiny spurious detections, background noise detected as cells",
+                    "solution": "Increase cellprob_threshold by 1-2 OR increase min_size to 2-3x current value"
+                }
+            },
+            "next_steps": {
+                "if_good": "If segmentation looks accurate, inform the user of success and provide the output_path.",
+                "if_needs_refinement": "Based on your visual analysis, adjust the appropriate parameters and call run_cellpose_sam again with the new values.",
+                "important": "You can only call refine_cellpose_sam_segmentation AT MOST 2 TIMES total. If this is your second call, you must make a final decision."
+            }
+        }
+        
+        return json.dumps(analysis, indent=2)
+        
+    except Exception as e:
+        error_result = {
+            "status": "error",
+            "error": str(e),
+            "message": "Could not load images for refinement. Check that both file paths are valid."
+        }
+        return json.dumps(error_result, indent=2)

utils/__init__.py

@@ -0,0 +1,24 @@
+from .gpu import (
+    clear_gpu_cache,
+    get_max_memory,
+    monitor_and_clear_cache
+)
+
+from .image_utils import (
+    resize_and_encode_image
+)
+
+from .prechecks import (
+    check_hf_persistent_storage
+)
+
+__all__ = __all__ = [
+    # GPU utilities
+    "clear_gpu_cache",
+    "get_max_memory",
+    "monitor_and_clear_cache",
+    # Image utilities
+    "resize_and_encode_image",
+    # precheck
+    "check_hf_persistent_storage"
+]

utils/__init__.py~

@@ -0,0 +1,23 @@
+from .gpu import (
+    clear_gpu_cache,
+    get_max_memory,
+    monitor_and_clear_cache
+)
+
+from .image_utils import (
+    resize_and_encode_image
+)
+
+from .precheck import (
+    ""
+)
+
+__all__ = __all__ = [
+    # GPU utilities
+    "clear_gpu_cache",
+    "get_max_memory",
+    "monitor_and_clear_cache",
+    # Image utilities
+    "resize_and_encode_image",
+    
+]

utils/gpu.py

@@ -0,0 +1,80 @@
+"""
+
+"""
+
+import torch
+import gc
+
+def clear_gpu_cache():
+    """Frees up GPU memory by clearing cache and collecting garbage."""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        gc.collect()
+        print("✓ GPU cache cleared.")
+
+
+def get_max_memory(memory_fraction=0.85, cpu_memory="50GB"):
+    """
+    Automatically configure max memory per GPU.
+    
+    When used with device_map="auto", this tells the model loader how much memory
+    it CAN use per GPU during the INITIAL model loading phase. If a model's layers
+    don't fit on one GPU with this limit, the loader will automatically split the
+    model across multiple GPUs.
+    
+    Args:
+        memory_fraction: Fraction of GPU memory to allocate (0.0-1.0). 
+                        Default 0.85 leaves 15% headroom.
+        cpu_memory: Maximum CPU memory to use as offload space.
+    
+    Returns:
+        dict: Memory limits per device, or None if no CUDA available
+    """
+    if not torch.cuda.is_available():
+        print("⚠ No CUDA GPUs available")
+        return None
+    
+    max_memory = {}
+    total_available = 0
+    
+    for i in range(torch.cuda.device_count()):
+        props = torch.cuda.get_device_properties(i)
+        total_memory = props.total_memory
+        usable_memory = int(total_memory * memory_fraction)
+        max_memory[i] = usable_memory
+        total_available += usable_memory
+        
+        print(f"GPU {i} ({props.name}): "
+              f"{usable_memory / 1024**3:.2f}GB / {total_memory / 1024**3:.2f}GB "
+              f"({memory_fraction*100:.0f}% limit)")
+    
+    # CPU memory for offloading if needed
+    max_memory["cpu"] = cpu_memory
+    
+    print(f"✓ Total GPU memory available for models: {total_available / 1024**3:.2f}GB")
+    print(f"✓ CPU offload memory: {cpu_memory}")
+    
+    return max_memory
+
+def monitor_and_clear_cache(threshold=0.90):
+    """
+    Monitor GPU memory and clear cache if usage exceeds threshold.
+    Call this periodically during long-running operations.
+    
+    Args:
+        threshold: Memory usage fraction (0.0-1.0) that triggers cache clearing
+    """
+    if not torch.cuda.is_available():
+        return
+    
+    for i in range(torch.cuda.device_count()):
+        props = torch.cuda.get_device_properties(i)
+        allocated = torch.cuda.memory_allocated(i)
+        total = props.total_memory
+        usage = allocated / total
+        
+        if usage > threshold:
+            print(f"⚠ GPU {i} usage at {usage*100:.1f}%, clearing cache...")
+            torch.cuda.empty_cache()
+            gc.collect()

utils/image_utils.py

@@ -0,0 +1,31 @@
+"""
+Image utilities for encoding and resizing
+"""
+import base64
+from io import BytesIO
+from PIL import Image
+
+
+def resize_and_encode_image(image_path: str, size: tuple = (512, 512)) -> tuple[str, str]:
+    """
+    Resize an image to specified size and encode as base64.
+    
+    Args:
+        image_path (str): Path to the image file
+        size (tuple): Target size as (width, height), default (1024, 1024)
+    
+    Returns:
+        tuple: (base64_string, media_type)
+    """
+    # Open and convert to RGB
+    img = Image.open(image_path).convert("RGB")
+    
+    # Resize with high-quality resampling
+    img_resized = img.resize(size, Image.Resampling.LANCZOS)
+    
+    # Encode to base64
+    buffered = BytesIO()
+    img_resized.save(buffered, format="PNG")
+    img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    
+    return img_base64, "image/png"

utils/prechecks.py

@@ -0,0 +1,44 @@
+"""
+
+"""
+
+from pathlib import Path
+from huggingface_hub import hf_hub_download, snapshot_download
+
+def check_hf_persistent_storage(
+    repo_id: str = None,
+    repo_type: str = "model",
+    file_or_folder="file",
+    target: str = None,
+    destination: str = "/data/"
+):
+
+    file_path = Path(destination) / target
+
+    def _download_file():
+        try:
+            if file_or_folder == "file":
+                hf_hub_download(
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    filename=target,
+                    local_dir=destination
+                )
+            elif file_or_folder == "folder":
+                snapshot_download(
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    allow_patterns=f"{target}/**",
+                    local_dir=destination
+                )
+
+            print(f"Successfully downloaded '{target}' to '{destination}'.")
+        except Exception as e:
+            print(f"An error occurred during the download: {e}")
+
+    # Check if the file exists at the specified path
+    if not file_path.exists():
+        _download_file()
+    else:
+        print(f"File '{file_path}' already exists. No download needed.")
+

utils/prechecks.py~

@@ -0,0 +1,44 @@
+"""
+
+"""
+
+from pathlib import Path
+from huggingface_hub import hf_hub_download, snapshot_download
+
+def check_hf_persistent_storage(
+    repo_id: str = None,
+    repo_type: str = "model",
+    file_or_folder="file",
+    target: str = None,
+    destination: str = "./data/"
+):
+
+    file_path = Path(destination) / target
+
+    def _download_file():
+        try:
+            if file_or_folder == "file":
+                hf_hub_download(
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    filename=target,
+                    local_dir=destination
+                )
+            elif file_or_folder == "folder":
+                snapshot_download(
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    allow_patterns=f"{target}/**",
+                    local_dir=destination
+                )
+
+            print(f"Successfully downloaded '{target}' to '{destination}'.")
+        except Exception as e:
+            print(f"An error occurred during the download: {e}")
+
+    # Check if the file exists at the specified path
+    if not file_path.exists():
+        _download_file()
+    else:
+        print(f"File '{file_path}' already exists. No download needed.")
+