app.py

"""
Gradio web interface for CellposeAgent
"""
import os
import json
import gradio as gr
from pathlib import Path
from langfuse import get_client
from openinference.instrumentation.smolagents import SmolagentsInstrumentor
from smolagents.agents import ActionStep

from config import settings
from agents.agent import CellposeAgent
from stores import neo4j_store
from utils.prechecks import check_hf_persistent_storage


def setup_observability():
    """Initializes Langfuse and Smolagents instrumentation."""
    try:
        if os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"):
            get_client()
            SmolagentsInstrumentor().instrument()
            print("✓ Observability and instrumentation initialized.")
        else:
            print("⚠️  Langfuse keys not found. Observability disabled.")
    except Exception as e:
        print(f"⚠️  Could not initialize observability: {e}")


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) -> tuple[str, str | None]:
    """
    Process a user task with the CellposeAgent.
    
    The image_path is now passed SEPARATELY to agent.run() so it gets stored
    in global context. This prevents the LLM from needing to reproduce the 
    exact path string (which it often corrupts).
    
    Args:
        image_path: Path to the uploaded image file
        task_text: User's text prompt/question
        agent: Initialized CellposeAgent instance
    
    Returns:
        tuple: (agent's text response, path to segmented image or None)
    """
    if not image_path:
        return "⚠️ Please upload an image first.", None
    
    # Build task WITHOUT embedding the full path
    # The agent instructions tell it to use image_path="" and the tool will resolve from context
    if not task_text or task_text.strip() == "":
        task = "Analyze the uploaded image and recommend optimal segmentation parameters. Then run segmentation."
    else:
        # Don't include the path in the task - it will be available via context
        task = task_text
    
    try:
        # Pass image_path SEPARATELY - this stores it in global context
        # Tools will retrieve the correct path from context instead of relying on LLM
        result = agent.run(task=task, image_path=image_path)
        get_client().flush()
        
        # Extract output image path from agent's memory (more reliable than parsing text)
        output_image_path = None
        try:
            # Search through agent's memory steps in reverse order (most recent first)
            for step in reversed(agent.agent.memory.steps):
                if isinstance(step, ActionStep) and step.observations:
                    try:
                        obs_data = json.loads(step.observations)
                        
                        # Check if this step contains segmentation output
                        if obs_data.get("status") == "success" and "output_path" in obs_data:
                            candidate_path = obs_data["output_path"]
                            
                            # Verify the file exists
                            if Path(candidate_path).exists():
                                output_image_path = candidate_path
                                print(f"✓ Found segmentation output: {output_image_path}")
                                break
                                
                    except (json.JSONDecodeError, Exception) as parse_error:
                        continue
            
            if output_image_path is None:
                print("ℹ️ No segmentation output found in agent memory")
                    
        except Exception as e:
            print(f"Warning: Could not extract output image from agent memory: {e}")
        
        return result, output_image_path
        
    except Exception as e:
        return f"❌ Error processing task: {str(e)}", None


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?' or 'Run segmentation' (image path handled automatically)",
                    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):
                # Text output
                output = gr.Textbox(
                    label="Agent Response",
                    lines=12,
                    max_lines=20,
                    show_copy_button=True
                )
                
                # Image output for segmentation results
                output_image = gr.Image(
                    label="Segmentation Result",
                    type="filepath",
                    height=400
                )
        
        # Event handler
        submit_btn.click(
            fn=lambda img, task: process_image_task(img, task, agent),
            inputs=[image_input, task_input],
            outputs=[output, output_image]
        )
        
        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)
            6. **Displays the segmented overlay image with colored cell masks**
            
            ### ⚙️ Technical Note
            
            Image paths are handled automatically by the system - the agent doesn't need to 
            reproduce exact file paths, preventing common path corruption issues with LLMs.
            """
        )
    
    return demo


def main():
    """Launch the Gradio application."""
    demo = create_gradio_interface()
    demo.launch(server_port=7860)


if __name__ == "__main__":
    main()