# Architecture Overview The AI Imaging Agent uses a **two-stage pipeline** that combines fast text retrieval with vision-language model selection to recommend imaging tools. ## System Architecture ```mermaid graph TB subgraph "User Interface" UI[Gradio Chat Interface] end subgraph "API Layer" Pipeline[RAGImagingPipeline] Validator[File Validator] MetaExtractor[Metadata Extractor] end subgraph "Stage 1: Retrieval" Embedder[BGE-M3 Text Embedder] FAISS[FAISS Vector Index] Reranker[CrossEncoder Reranker] Catalog[Software Catalog JSONL] end subgraph "Stage 2: Agent Selection" Agent[PydanticAI Agent] VLM[GPT-4o/4o-mini VLM] Tools[Agent Tools] end UI --> Pipeline Pipeline --> Validator Pipeline --> MetaExtractor Pipeline --> Embedder Embedder --> FAISS FAISS --> Reranker Catalog -.-> FAISS Reranker --> Agent Agent --> VLM Agent --> Tools Agent --> UI ``` ## Design Principles ### 1. Two-Stage Pipeline **Why two stages?** - **Speed**: Text retrieval is fast (~100ms), VLM calls are slower (~2-5s) - **Cost**: Only run expensive VLM on top candidates - **Quality**: Combine semantic search (Stage 1) with reasoning (Stage 2) ### 2. No Generation in Retrieval Stage 1 uses **no LLMs**: - Deterministic text search - Reproducible results - Fast iteration - Lower cost ### 3. Single VLM Call in Selection Stage 2 makes **exactly one VLM call**: - Sees all candidates at once - Performs comparative reasoning - Returns complete rankings - Efficient use of context window ### 4. Vision + Text Integration VLM receives: - **Visual**: PNG preview of image - **Textual**: Query, metadata, candidate descriptions - **Structured**: Candidate metadata table Enables image-aware tool selection. ## Data Flow ### Input Processing ``` User uploads: scan.dcm "Segment the lungs" ↓ File Validation - Size check (< 200MB for DICOM) - Format validation - Security checks ↓ Metadata Extraction - Format: DICOM - Modality: CT - Dimensions: 512×512×300 (3D) - Spacing: 0.7×0.7×1.5mm ↓ Preview Generation - Extract middle slice: scan_preview.png - Format: PNG, RGB - Preserve metadata separately ``` ### Stage 1: Retrieval ``` Query: "Segment the lungs" Uploaded: scan.dcm (DICOM, CT, 3D) ↓ Query Enhancement Enhanced: "Segment the lungs format:DICOM format:CT format:3D" ↓ Metadata-Aware Hinting + image metadata summary (modality/anatomy/dims) ↓ Embedding (BGE-M3) Vector: [0.23, -0.15, 0.87, ..., 0.34] # 1024 dims ↓ FAISS Search Top 20 candidates by cosine similarity ↓ Retry Broadening (if low results) retry with a shorter query formulation ↓ CrossEncoder Reranking Re-score with cross-attention Top 8 candidates → Candidates passed to Stage 2 ``` ### Stage 2: Agent Selection ``` Inputs: - User query: "Segment the lungs" - Image preview: scan_preview.png - Candidates: [tool1, tool2, ..., tool8] - Metadata: DICOM, CT, 3D, 512×512×300 ↓ VLM Prompt Construction System: "You are an imaging tool expert..." User text: Query + metadata + candidate table User image: PNG preview ↓ VLM Call (GPT-4o) - Analyzes image content (CT thorax) - Reads candidate descriptions - Considers format compatibility - Reasons about task alignment ↓ Response (Structured) { "status": "complete", "recommendations": [ { "rank": 1, "name": "TotalSegmentator", "accuracy": 95, "explanation": "...", "reason": "task_match" }, ... ] } → Formatted recommendations to user ``` ## Key Components ### api/pipeline.py **RAGImagingPipeline**: Main orchestrator ```python class RAGImagingPipeline: def __init__(self, catalog_path, index_dir): self.retriever = TextRetriever(...) # Stage 2 (selection/ranking) is handled by the PydanticAI agent # configured in generator/prompts.py using models from generator/schema.py def recommend(self, query, files): # Stage 1: Retrieval candidates = self.retriever.retrieve(query) # Stage 2: Selection via PydanticAI agent recommendations = run_selection_agent( query=query, candidates=candidates, files=files, ) return recommendations ``` **Responsibilities**: - File validation - Metadata extraction - Pipeline orchestration - Error handling ### retriever/ **Text-based retrieval, no LLMs** Components: - `text_embedder.py`: BGE-M3 embedding model - `vector_index.py`: FAISS index management - `reranker.py`: CrossEncoder reranking - `software_doc.py`: Catalog schema and loading **Retrieval flow**: 1. Embed query → vector 2. FAISS search → top-N by similarity 3. CrossEncoder → rerank with cross-attention 4. Return top-K candidates ### generator/ **VLM-based tool selection building blocks** Components: - `schema.py`: Pydantic models for agent responses and tool recommendations - `prompts.py`: System and tool-selection prompts used by the PydanticAI agent **Selection logic**: - Implemented in the PydanticAI agent (`agent/agent.py`) using these schemas and prompts - Single VLM call with all candidates - Structured output (Pydantic schemas) with ranked recommendations - Vision + text multimodal input ### agent/ **PydanticAI conversational agent** Components: - `agent.py`: Agent definition and tools - `state.py`: ChatState dataclass - `tools.py`: Agent tools (search, repo_info, demo_exec) **Tools**: - `search_alternative`: Request alternative search - `repo_info`: Fetch GitHub repository details - `run_gradio_demo`: Execute Gradio Space demos ### utils/ **Shared utilities** - `image_meta.py`: DICOM/NIfTI/TIFF metadata extraction - `file_validator.py`: Size and format validation - `previews.py`: Image conversion to PNG - `tags.py`: Control tag parsing (`[EXCLUDE:...]`, etc.) - `config.py`: Configuration management ### ui/ **Gradio interface** Components: - `app.py`: Gradio application - `components.py`: Reusable UI components - `handlers.py`: Message handlers - `formatters.py`: Response formatting - `visualizations.py`: Previews and traces **Key function**: ```python def respond(message: str, files: list, state: dict) -> tuple: """ Main interaction function. Returns: (reply, media, updated_state) """ ``` ## Module Boundaries Clear separation of concerns: | Module | Purpose | Dependencies | |--------|---------|--------------| | `api/` | Pipeline orchestration | `retriever/`, `generator/`, `utils/` | | `retriever/` | Text search only | None (pure retrieval) | | `generator/` | VLM selection only | None (pure generation) | | `agent/` | Conversational logic | `api/`, `utils/` | | `ui/` | Interface only | `agent/`, `api/` | | `utils/` | Shared functionality | None (pure utilities) | **Benefits**: - Independent testing - Clear interfaces - Modular replacement - No circular dependencies ## Data Schemas ### Software Catalog JSONL format, based on schema.org SoftwareSourceCode: ```json { "name": "TotalSegmentator", "description": "Automated multi-organ segmentation...", "url": "https://github.com/wasserth/TotalSegmentator", "codeRepository": "https://github.com/wasserth/TotalSegmentator", "programmingLanguage": "Python", "license": "Apache-2.0", "keywords": ["segmentation", "medical-imaging", "CT"], "applicationCategory": "Medical Imaging", "operatingSystem": ["Linux", "Windows", "macOS"], "softwareRequirements": ["Python 3.9+", "PyTorch"], "supportingData": { "modalities": ["CT", "MRI"], "dimensions": ["3D"], "formats": ["DICOM", "NIfTI"], "tasks": ["segmentation"], "demo_url": "https://huggingface.co/spaces/..." } } ``` ### Agent Response Pydantic models in `generator/schema.py`: ```python class ToolRecommendation(BaseModel): rank: int name: str accuracy_score: int # 0-100 explanation: str reason: ToolReason # Enum supporting_data: dict class AgentResponse(BaseModel): status: ConversationStatus # Enum recommendations: list[ToolRecommendation] message: str | None ``` **Validation**: - Type checking via Pydantic - Enum constraints - Field aliases for LLM compatibility ## Extension Points ### Adding New Models In `config.yaml`: ```yaml available_models: - display_name: "Custom Model" name: "model-name" base_url: "https://api.example.com/v1" api_key_env: "CUSTOM_API_KEY" ``` ### Adding New Tools Add a tool in `agent/agent.py` and route implementation to `agent/tools/` modules: ```python @agent.tool async def new_tool(ctx: RunContext[AgentState], param: str) -> str: """Tool description for the agent.""" # Delegate to ai_agent.agent.tools.* implementation return result ``` ### Custom Metadata Extractors In `utils/image_meta.py`: ```python def extract_custom_format(file_path: str) -> dict: """Extract metadata from custom format.""" # Implementation return metadata ``` ## Security Considerations ### User Data - **Images**: Sent to OpenAI API (preview PNG) if gpt is selected - **Metadata**: Processed locally, sent to VLM as text - **Queries**: Sent to OpenAI API **Privacy**: User data sees OpenAI's VLM API only. ### Catalog Integrity - Software catalog is curated - SHA1 checksums verify integrity - No user-generated catalog entries ### Demo Execution - Calls external Gradio Spaces (user choice) - No credentials shared with demos - User's image uploaded to public spaces (warn users) ## Next Steps - Deep dive into [Retrieval Pipeline](retrieval.md) - Learn about [Agent & VLM Selection](agent.md) - Explore [Software Catalog](catalog.md)