diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index 938f8f5596c44c52db4638b8f8a7991fe43a38c9..363ffa89cc8cbaab95cef83bf2dda3940130f5ec 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -1,17 +1,17 @@ -FROM ghcr.io/astral-sh/uv:python3.12-bookworm - -# Install just and other system dependencies -RUN apt-get update && apt-get install -y \ - sudo \ - curl \ - && curl --proto '=https' --tlsv1.2 -sSf https://just.systems/install.sh | bash -s -- --to /usr/local/bin \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* - -# Create non-root user with UID/GID typically used by VS Code (1000:1000) -RUN useradd -ms /bin/bash -u 1000 vscode \ - && apt-get update && apt-get install -y sudo \ - && echo "vscode ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers - -USER vscode +FROM ghcr.io/astral-sh/uv:python3.12-bookworm + +# Install just and other system dependencies +RUN apt-get update && apt-get install -y \ + sudo \ + curl \ + && curl --proto '=https' --tlsv1.2 -sSf https://just.systems/install.sh | bash -s -- --to /usr/local/bin \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + +# Create non-root user with UID/GID typically used by VS Code (1000:1000) +RUN useradd -ms /bin/bash -u 1000 vscode \ + && apt-get update && apt-get install -y sudo \ + && echo "vscode ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers + +USER vscode WORKDIR /workspaces \ No newline at end of file diff --git a/.dockerignore b/.dockerignore index 4dc6a61d4f3983f76275a971b8ec33248a7a0cb0..cfd71789eb104a00f79fa9cc54159d701bee166e 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,6 +1,12 @@ -.git -__pycache__ -*.pyc -.env -.env.* -tests \ No newline at end of file +.git +__pycache__ +*.pyc +.env +.env.* +tests + +# Frontend artifacts: copied via the node stage in the multi-stage build, +# so the python stage never needs the host's node_modules / dist tree. +src/frontend/node_modules +src/frontend/dist +src/frontend/.vite \ No newline at end of file diff --git a/.env.dist b/.env.dist index 769b1a5db589a383cc537f94f1139403fb3f8afc..b284f47fb9648a66fe7891aa7b35f922ee714879 100644 --- a/.env.dist +++ b/.env.dist @@ -26,6 +26,12 @@ LOG_PROMPTS=0 # write selector prompt snapshots # Path to config.yaml CONFIG_PATH=path/to/custom/config.yaml +# Shared password for the new FastAPI frontend (leave empty to disable auth). +APP_PASSWORD= + +# Session TTL in seconds for the in-process session store (defaults to 6h). +SESSION_TTL_SECONDS=21600 + # GraphDB GRAPHDB_GRAPH= GRAPHDB_URL= diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..7ff3b29e7f55dfd29ae5594ced5316a0ace86a80 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,9 @@ +*.png filter=lfs diff=lfs merge=lfs -text +*.tif filter=lfs diff=lfs merge=lfs -text +*.tiff filter=lfs diff=lfs merge=lfs -text +*.jpg filter=lfs diff=lfs merge=lfs -text +*.jpeg filter=lfs diff=lfs merge=lfs -text +*.gif filter=lfs diff=lfs merge=lfs -text +*.ico filter=lfs diff=lfs merge=lfs -text +*.webp filter=lfs diff=lfs merge=lfs -text +*.bmp filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/deploy-to-hf.yml b/.github/workflows/deploy-to-hf.yml index 45037852038f76a9ef3ed8cb15411752290b1057..d0dfd1fddbc16bd8f45d7b69f46e89a2c363d157 100644 --- a/.github/workflows/deploy-to-hf.yml +++ b/.github/workflows/deploy-to-hf.yml @@ -1,54 +1,54 @@ -name: Deploy to Hugging Face Space - -on: - push: - branches: - - main - workflow_dispatch: - -jobs: - deploy: - runs-on: ubuntu-latest - - steps: - - name: Checkout code - uses: actions/checkout@v4 - with: - fetch-depth: 1 - lfs: false - - - name: Prepare clean HF deploy branch without assets - run: | - git config user.email "ci@github.actions" - git config user.name "github-actions[bot]" - - original_readme="$(cat README.md)" - - cat > README.md < README.md < - -

- ---- - -## ๐Ÿš€ Quick Start - -### Prerequisites - -- Python 3.10โ€“3.12 -- OpenAI API key (or compatible API endpoint) -- Internet connection for model calls - -### Installation - -```bash -# Clone the repository -git clone -cd ai-agent - -# Create virtual environment -python -m venv .venv - -# Activate virtual environment -# On Linux/macOS: -source .venv/bin/activate -# On Windows: -.venv\Scripts\activate - -# Install the package -pip install --upgrade pip -pip install -e . - -# For development (includes test dependencies) -pip install -e ".[dev]" -``` - -### Configuration - -Create a `.env` file at the repository root: - -```dotenv -# Required: OpenAI API key -OPENAI_API_KEY=sk-xxxx - -# Optional: GitHub token for repo info tool -GITHUB_TOKEN=ghp_xxxx - -# Optional: Alternative model providers (EPFL, etc.) -EPFL_API_KEY=sk-xxxx -EPFL_API_KEY_EMBEDDER=sk-xxxx - -# Software catalog path -SOFTWARE_CATALOG=dataset/catalog.jsonl - -# Pipeline configuration -TOP_K=8 # Number of candidates to retrieve -NUM_CHOICES=3 # Number of tools to recommend -AGENT_OUTPUT_RETRIES=3 # Structured output validation retries -EMBED_CATALOG_ON_START=1 # Pre-embed catalog if FAISS is empty - -# Logging configuration -LOGLEVEL_CONSOLE=WARNING -LOGLEVEL_FILE=INFO -FILE_LOG=1 -LOG_DIR=logs -LOG_PROMPTS=0 # Set to 1 to save prompt snapshots for debugging - -# Custom config path -CONFIG_PATH=config.yaml -``` - -### Model Configuration - -The agent model can be configured via `config.yaml`: - -```yaml -# AI Agent Model Configuration - -# Default/fallback model (used for CLI and initial startup) -agent_model: - name: "gpt-5.1" - base_url: null # null for default OpenAI endpoint - api_key_env: "OPENAI_API_KEY" - -# Available models for UI dropdown -available_models: - - display_name: "gpt-4o-mini" - name: "gpt-4o-mini" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-4o" - name: "gpt-4o" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-5-mini" - name: "gpt-5-mini" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-5.1" - name: "gpt-5.1" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - -retrieval: - embedder: - backend: "remote" # "remote" or "local" - model_name: "Qwen/Qwen3-Embedding-8B" - base_url: "https://inference-rcp.epfl.ch/v1" - api_key_env: "EPFL_API_KEY_EMBEDDER" - timeout_s: 20 - # local example: - # backend: "local" - # model_name: "BAAI/bge-m3" - # device: "cpu" # optional - - reranker: - backend: "remote" # "remote" or "local" - model_name: "BAAI/bge-reranker-v2-m3" - base_url: "https://inference-rcp.epfl.ch/v1" - api_key_env: "EPFL_API_KEY_EMBEDDER" - timeout_s: 20 - # local example: - # backend: "local" - # model_name: "BAAI/bge-reranker-v2-m3" - # device: "cpu" # optional -``` - -### Running the App - -```bash -# Start the chat interface -ai_agent chat - -# Open your browser to: -# http://127.0.0.1:7860 -``` - -Try uploading a cat image and asking: -> "I want to segment the cat from this image" - ---- - -## ๐Ÿ’ฌ Usage - -### Chat Interface - -The chat interface provides a natural conversation flow: - -1. **Upload Files**: Drop images (PNG, JPG, TIFF, DICOM, NIfTI, etc.) or other supported files -2. **Describe Your Task**: Use natural language like "segment the lungs" or "register brain MRI" -3. **Review Recommendations**: Get ranked tool suggestions with accuracy scores and explanations -4. **Run Demos**: Click "Run demo" to execute tools directly on your uploaded images -5. **Iterate**: Ask for alternatives, refine your query, or upload different files - -### Supported File Formats - -**Images:** -- Standard: PNG, JPG, JPEG, WEBP, BMP, GIF -- Medical: DICOM (.dcm), NIfTI (.nii, .nii.gz), TIFF stacks -- Scientific: Multi-page TIFF, TIFF with metadata - -**Other Files:** -- Data: CSV, JSON, XML -- Media: MP3, MP4 - -### Example Queries - -- "Segment the lungs from this CT scan" -- "Register these two brain MRI images" -- "Extract text from this medical report image" -- "Classify what organ is shown in this ultrasound" -- "Detect tumors in this MRI scan" -- "I need to analyze DICOM files, what tools are available?" - -### Understanding Results - -Each recommendation includes: -- **Rank**: Priority order (1 = best match) -- **Accuracy Score**: Confidence level (0-100%) -- **Explanation**: Why this tool matches your request -- **Metadata**: Supported modalities, dimensions, formats, license -- **Demo Link**: Direct link to runnable example - ---- - -## ๐Ÿ—๏ธ Architecture - -### Pipeline Overview - -The system follows a two-stage architecture: - -``` -User Input (Image + Text Query) - โ†“ -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ RETRIEVAL STAGE โ”‚ -โ”‚ - BGE-M3 Embeddings โ”‚ -โ”‚ - FAISS Vector Search โ”‚ -โ”‚ - CrossEncoder Reranking โ”‚ -โ”‚ - Format Token Matching โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ - โ†“ Top-K Candidates -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ AGENT SELECTION โ”‚ -โ”‚ - Pydantic AI Agent โ”‚ -โ”‚ - OpenAI VLM โ”‚ -โ”‚ - Image + Metadata Analysis โ”‚ -โ”‚ - Multi-Tool Reasoning โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ - โ†“ - Ranked Recommendations -``` - -### Retrieval Stage - -**No LLM calls** - purely text-based search: - -1. **Query Construction**: User task + format tokens from uploaded files -2. **Embedding**: BGE-M3 model generates query embedding -3. **Vector Search**: FAISS retrieves top candidates -4. **Reranking**: CrossEncoder refines results for precision -5. **Retry Broadening**: If too few hits, retry with a shorter/broader query - -### Agent Selection Stage - -**Single VLM call** - multimodal reasoning: - -1. **Input Preparation**: - - Text: User query + candidate table + file metadata - - Image: PNG preview (converted from any format) - - Context: Original file format, dimensions, modality - -2. **Agent Tools**: - - `search_tools`: Search catalog with query - - `search_alternative`: Find alternatives (iterative) - - `repo_info`: Fetch GitHub documentation via DeepWiki MCP - -3. **Output**: Ranked tool selections with accuracy scores and explanations - -### Key Components - -- **`api/pipeline.py`**: RAG retrieval orchestrator -- **`agent/agent.py`**: Pydantic AI agent with tool definitions -- **`retriever/`**: Embedding, FAISS indexing, reranking -- **`generator/`**: Prompts and schema for tool selection -- **`ui/`**: Gradio chat interface components -- **`utils/`**: Image processing, metadata extraction, file validation -- **`catalog/`**: Catalog syncing from GraphDB (optional) - ---- - -## โš™๏ธ Configuration - -### Environment Variables - -| Variable | Description | Default | Required | -|----------|-------------|---------|----------| -| `OPENAI_API_KEY` | OpenAI API key | - | โœ… | -| `EPFL_API_KEY_EMBEDDER` | API key for remote embedder and reranker endpoints | - | โœ… (when `retrieval.embedder.backend: remote` and/or `retrieval.reranker.backend: remote`) | -| `GITHUB_TOKEN` | GitHub token for repo info | - | โŒ | -| `SOFTWARE_CATALOG` | Path to catalog JSONL | `dataset/catalog.jsonl` | โœ… | -| `TOP_K` | Retrieval candidates count | `8` | โŒ | -| `NUM_CHOICES` | Tools to recommend | `3` | โŒ | -| `AGENT_OUTPUT_RETRIES` | Structured output validation retries | `3` | โŒ | -| `EMBED_CATALOG_ON_START` | Pre-embed catalog on startup when FAISS is empty | `1` | โŒ | -| `LOGLEVEL_CONSOLE` | Console log level | `WARNING` | โŒ | -| `LOGLEVEL_FILE` | File log level | `INFO` | โŒ | -| `FILE_LOG` | Enable file logging | `1` | โŒ | -| `LOG_DIR` | Log directory | `logs` | โŒ | -| `LOG_PROMPTS` | Save prompt snapshots | `0` | โŒ | -| `CONFIG_PATH` | Model config file | `config.yaml` | โœ… | - -### GraphDB Catalog Sync (Optional) - -For automatic catalog syncing from a GraphDB instance: - -```dotenv -GRAPHDB_URL=https://your-graphdb.example.com -GRAPHDB_GRAPH=your-graph-name -GRAPHDB_USER=username -GRAPHDB_PASSWORD=password -GRAPHDB_QUERY_FILE=/path/to/query.rq -SYNC_EVERY_HOURS=24 # Auto-refresh interval (0 to disable) -OUTPUT_JSONLD=dataset/catalog.jsonld -OUTPUT_JSONL=dataset/catalog.jsonl -``` - -Run manual sync: -```bash -ai_agent sync -``` - ---- - -## ๐Ÿ“‹ Catalog Format - -The catalog is a JSONL file where each line is a `SoftwareDoc` following schema.org SoftwareSourceCode structure. - -### Minimal Example - -```json -{ - "name": "3d-lungs-segmentation", - "description": "3D lung segmentation from CT; returns a mask/overlay.", - - "applicationCategory": ["Medical Imaging"], - "featureList": ["segmentation"], - "imagingModality": ["CT"], - "dims": [3], - "anatomy": ["lung"], - "keywords": ["mask", "overlay", "lung segmentation", "CT"], - - "programmingLanguage": "Python", - "requiresGPU": false, - "isAccessibleForFree": true, - "license": "Apache-2.0", - - "supportingData": [ - { - "datasetFormat": "TIFF", - "bodySite": "lung", - "imagingModality": "CT", - "hasDimensionality": 3 - }, - { - "datasetFormat": "DICOM", - "bodySite": "lung", - "imagingModality": "CT", - "hasDimensionality": 3 - } - ], - - "runnableExample": [ - { - "hostType": "gradio", - "url": "https://huggingface.co/spaces/qchapp/3d-lungs-segmentation", - "name": "HF Space" - } - ] -} -``` - -### Key Fields - -- **name**: Unique identifier for the tool -- **description**: Clear explanation of what the tool does -- **featureList**: Operations (e.g., segmentation, registration, classification) -- **imagingModality**: Medical imaging types (CT, MRI, XR, US, PET) -- **dims**: Supported dimensions (2D, 3D, 4D) -- **anatomy**: Body parts/organs -- **supportingData**: Format compatibility information (critical for matching) -- **runnableExample**: Links to live demos (HuggingFace Spaces, notebooks, web apps) - ---- - -## ๐Ÿ”ง Development - -### Project Structure - -``` -ai-agent/ -โ”œโ”€โ”€ src/ai_agent/ -โ”‚ โ”œโ”€โ”€ agent/ # Pydantic AI agent and tools -โ”‚ โ”‚ โ”œโ”€โ”€ agent.py # Agent definition -โ”‚ โ”‚ โ”œโ”€โ”€ models.py # Agent state models -โ”‚ โ”‚ โ”œโ”€โ”€ tools/ # Agent tool implementations -โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ search_tool.py -โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ search_alternative_tool.py -โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ gradio_space_tool.py -โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ repo_info_tool.py -โ”‚ โ”‚ โ”‚ โ””โ”€โ”€ deepwiki_tool.py -โ”‚ โ”‚ โ””โ”€โ”€ utils.py -โ”‚ โ”œโ”€โ”€ api/ # Pipeline orchestration -โ”‚ โ”‚ โ””โ”€โ”€ pipeline.py # RAGImagingPipeline -โ”‚ โ”œโ”€โ”€ retriever/ # Retrieval components -โ”‚ โ”‚ โ”œโ”€โ”€ text_embedder.py -โ”‚ โ”‚ โ”œโ”€โ”€ vector_index.py -โ”‚ โ”‚ โ”œโ”€โ”€ reranker.py -โ”‚ โ”‚ โ””โ”€โ”€ software_doc.py -โ”‚ โ”œโ”€โ”€ generator/ # Agent prompts and schemas -โ”‚ โ”‚ โ”œโ”€โ”€ prompts.py -โ”‚ โ”‚ โ””โ”€โ”€ schema.py -โ”‚ โ”œโ”€โ”€ ui/ # Gradio interface -โ”‚ โ”‚ โ”œโ”€โ”€ app.py -โ”‚ โ”‚ โ”œโ”€โ”€ handlers.py -โ”‚ โ”‚ โ”œโ”€โ”€ components.py -โ”‚ โ”‚ โ”œโ”€โ”€ formatters.py -โ”‚ โ”‚ โ”œโ”€โ”€ state.py -โ”‚ โ”‚ โ””โ”€โ”€ visualizations.py -โ”‚ โ”œโ”€โ”€ utils/ # Shared utilities -โ”‚ โ”‚ โ”œโ”€โ”€ config.py # Configuration management -โ”‚ โ”‚ โ”œโ”€โ”€ file_validator.py -โ”‚ โ”‚ โ”œโ”€โ”€ image_meta.py # Metadata extraction -โ”‚ โ”‚ โ”œโ”€โ”€ image_io.py -โ”‚ โ”‚ โ”œโ”€โ”€ previews.py -โ”‚ โ”‚ โ””โ”€โ”€ tags.py -โ”‚ โ”œโ”€โ”€ catalog/ # Catalog syncing -โ”‚ โ”‚ โ””โ”€โ”€ sync.py -โ”‚ โ””โ”€โ”€ cli.py # CLI entry point -โ”œโ”€โ”€ tests/ # Test suite -โ”‚ โ”œโ”€โ”€ test_retrieval_pipeline.py -โ”‚ โ”œโ”€โ”€ test_repo_summary.py -โ”‚ โ””โ”€โ”€ data/ -โ”œโ”€โ”€ artifacts/ # Generated artifacts -โ”‚ โ””โ”€โ”€ rag_index/ # FAISS index -โ”œโ”€โ”€ dataset/ # Catalog data -โ”‚ โ””โ”€โ”€ catalog.jsonl -โ”œโ”€โ”€ logs/ # Application logs -โ”œโ”€โ”€ config.yaml # Model configuration -โ”œโ”€โ”€ pyproject.toml # Project metadata & dependencies -โ”œโ”€โ”€ Dockerfile # Production Docker image -โ”œโ”€โ”€ tools/image/Dockerfile # Development Docker image -โ””โ”€โ”€ justfile # Task runner commands -``` - -### Local Development - -```bash -# Install in development mode -pip install -e ".[dev]" - -# Run tests -pytest tests/ -``` - -### Testing - -Run the test suite: - -```bash -# All tests -pytest tests/ - -# Specific test file -pytest tests/test_retrieval_pipeline.py - -# With verbose output -pytest -v tests/ - -# With coverage -pytest --cov=ai_agent tests/ -``` - -### Logging & Debugging - -**Console Logs**: Set `LOGLEVEL_CONSOLE=DEBUG` for verbose output - -**File Logs**: Automatically saved to `logs/app_YYYYMMDD.log` (rotates daily) - -**Prompt Snapshots**: Enable `LOG_PROMPTS=1` to save: -- `logs/vlm_selector_YYYYMMDD_HHMMSS.txt` - System/user prompts - ---- - -## ๐Ÿ“š API & CLI Reference - -### CLI Commands - -```bash -# Launch chat interface -ai_agent chat - -# Sync catalog from GraphDB -ai_agent sync -``` - -## ๐Ÿ—บ๏ธ Maintainer Guide - -For full project documentation with detailed folder responsibilities, environment defaults, and improvement guidelines, see [docs/guide.md](docs/guide.md). - ---- - -## ๐Ÿ“ Changelog - -See [CHANGELOG.md](CHANGELOG.md) for detailed version history. - -### Recent Highlights - -**[1.0.0]** -- โœจ New chat-based interface (`ai_agent chat`) with rich media and tool integration -- ๐Ÿ› ๏ธ Fully agent-based architecture replacing legacy pipelines -- ๐Ÿ” Smarter retrieval with automatic retry -- ๐Ÿ”— DeepWiki MCP integration for fast GitHub repository documentation access -- ๐Ÿ”ง YAML configuration (`config.yaml`) for flexible model and backend setup -- ๐ŸŽจ Redesigned UI with Imaging Plaza branding and improved UX -- โšก Performance improvements (pre-embedding, caching, faster startup) -- ๐Ÿงน Major cleanup: removed deprecated code paths, legacy UI, and outdated tests - -**[0.1.3] - 2025-10-22** -- Gradio space runner tool -- Repository info tool -- UI fixes and polish - ---- - -## ๐Ÿ“„ License - -This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details. - ---- - -## ๐Ÿ™ Credits & Acknowledgments - -**Developed by**: Imaging Plaza Team - -**Technologies:** -- [Pydantic AI](https://github.com/pydantic/pydantic-ai) - AI agent framework -- [OpenAI](https://openai.com) - GPT vision model -- [FAISS](https://github.com/facebookresearch/faiss) - Vector search -- [BGE-M3](https://huggingface.co/BAAI/bge-m3) - Multilingual embeddings -- [Gradio](https://gradio.app) - Interactive web UI -- [DeepWiki](https://deepwiki.com) - GitHub repository documentation - -**Medical Imaging Formats:** -- [pydicom](https://github.com/pydicom/pydicom) - DICOM support -- [nibabel](https://nipy.org/nibabel/) - NIfTI support - ---- - -## ๐Ÿ“ฎ Support - -For issues, questions, or contributions, please contact the Imaging Plaza team. - ---- - -**๐Ÿฅ Medical Disclaimer**: This software is a tool recommendation system, not a diagnostic tool. Always consult qualified medical professionals for clinical decisions. +--- +title: AI Agent +emoji: ๐Ÿค– +colorFrom: blue +colorTo: green +sdk: docker +app_port: 7860 +pinned: false +--- + +# AI Imaging Agent (Imaging Plaza) + +[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) +[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE) + +An intelligent RAG + AI agent system that helps users discover the right imaging software for their images and tasks. Upload an image, describe what you want to do, and get ranked tool recommendations with links to runnable demos. + +## โœจ Key Features + +- **๐Ÿค– Conversational AI Agent**: Natural language interaction with multi-turn context +- **๐Ÿ” Smart Retrieval**: BGE-M3 embeddings + FAISS + CrossEncoder reranking +- **๐Ÿ‘๏ธ Vision-Aware Selection**: VLM-based tool selection considering both image content and metadata +- **๐Ÿฅ Medical Imaging Focus**: Specialized support for CT, MRI, DICOM, NIfTI, and other medical formats +- **๐ŸŽฏ Format-Aware Matching**: IO compatibility scoring based on file formats and dimensions +- **๐Ÿš€ Demo Integration**: Direct execution of Gradio Space demos on your images +- **๐Ÿ“Š Rich UI**: Chat interface with image previews, file management, and execution traces + +

+ +

+ +--- + +## ๐Ÿš€ Quick Start + +### Prerequisites + +- Python 3.10โ€“3.12 +- OpenAI API key (or compatible API endpoint) +- Internet connection for model calls + +### Installation + +```bash +# Clone the repository +git clone +cd ai-agent + +# Create virtual environment +python -m venv .venv + +# Activate virtual environment +# On Linux/macOS: +source .venv/bin/activate +# On Windows: +.venv\Scripts\activate + +# Install the package +pip install --upgrade pip +pip install -e . + +# For development (includes test dependencies) +pip install -e ".[dev]" +``` + +### Configuration + +Create a `.env` file at the repository root: + +```dotenv +# Required: OpenAI API key +OPENAI_API_KEY=sk-xxxx + +# Optional: GitHub token for repo info tool +GITHUB_TOKEN=ghp_xxxx + +# Optional: Alternative model providers (EPFL, etc.) +EPFL_API_KEY=sk-xxxx +EPFL_API_KEY_EMBEDDER=sk-xxxx + +# Software catalog path +SOFTWARE_CATALOG=dataset/catalog.jsonl + +# Pipeline configuration +TOP_K=8 # Number of candidates to retrieve +NUM_CHOICES=3 # Number of tools to recommend +AGENT_OUTPUT_RETRIES=3 # Structured output validation retries +EMBED_CATALOG_ON_START=1 # Pre-embed catalog if FAISS is empty + +# Logging configuration +LOGLEVEL_CONSOLE=WARNING +LOGLEVEL_FILE=INFO +FILE_LOG=1 +LOG_DIR=logs +LOG_PROMPTS=0 # Set to 1 to save prompt snapshots for debugging + +# Custom config path +CONFIG_PATH=config.yaml +``` + +### Model Configuration + +The agent model can be configured via `config.yaml`: + +```yaml +# AI Agent Model Configuration + +# Default/fallback model (used for CLI and initial startup) +agent_model: + name: "gpt-5.1" + base_url: null # null for default OpenAI endpoint + api_key_env: "OPENAI_API_KEY" + +# Available models for UI dropdown +available_models: + - display_name: "gpt-4o-mini" + name: "gpt-4o-mini" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-4o" + name: "gpt-4o" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-5-mini" + name: "gpt-5-mini" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-5.1" + name: "gpt-5.1" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + +retrieval: + embedder: + backend: "remote" # "remote" or "local" + model_name: "Qwen/Qwen3-Embedding-8B" + base_url: "https://inference-rcp.epfl.ch/v1" + api_key_env: "EPFL_API_KEY_EMBEDDER" + timeout_s: 20 + # local example: + # backend: "local" + # model_name: "BAAI/bge-m3" + # device: "cpu" # optional + + reranker: + backend: "remote" # "remote" or "local" + model_name: "BAAI/bge-reranker-v2-m3" + base_url: "https://inference-rcp.epfl.ch/v1" + api_key_env: "EPFL_API_KEY_EMBEDDER" + timeout_s: 20 + # local example: + # backend: "local" + # model_name: "BAAI/bge-reranker-v2-m3" + # device: "cpu" # optional +``` + +### Running the App + +```bash +# Start the chat interface +ai_agent chat + +# Open your browser to: +# http://127.0.0.1:7860 +``` + +Try uploading a cat image and asking: +> "I want to segment the cat from this image" + +--- + +## ๐Ÿ’ฌ Usage + +### Chat Interface + +The chat interface provides a natural conversation flow: + +1. **Upload Files**: Drop images (PNG, JPG, TIFF, DICOM, NIfTI, etc.) or other supported files +2. **Describe Your Task**: Use natural language like "segment the lungs" or "register brain MRI" +3. **Review Recommendations**: Get ranked tool suggestions with accuracy scores and explanations +4. **Run Demos**: Click "Run demo" to execute tools directly on your uploaded images +5. **Iterate**: Ask for alternatives, refine your query, or upload different files + +### Supported File Formats + +**Images:** +- Standard: PNG, JPG, JPEG, WEBP, BMP, GIF +- Medical: DICOM (.dcm), NIfTI (.nii, .nii.gz), TIFF stacks +- Scientific: Multi-page TIFF, TIFF with metadata + +**Other Files:** +- Data: CSV, JSON, XML +- Media: MP3, MP4 + +### Example Queries + +- "Segment the lungs from this CT scan" +- "Register these two brain MRI images" +- "Extract text from this medical report image" +- "Classify what organ is shown in this ultrasound" +- "Detect tumors in this MRI scan" +- "I need to analyze DICOM files, what tools are available?" + +### Understanding Results + +Each recommendation includes: +- **Rank**: Priority order (1 = best match) +- **Accuracy Score**: Confidence level (0-100%) +- **Explanation**: Why this tool matches your request +- **Metadata**: Supported modalities, dimensions, formats, license +- **Demo Link**: Direct link to runnable example + +--- + +## ๐Ÿ—๏ธ Architecture + +### Pipeline Overview + +The system follows a two-stage architecture: + +``` +User Input (Image + Text Query) + โ†“ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ RETRIEVAL STAGE โ”‚ +โ”‚ - BGE-M3 Embeddings โ”‚ +โ”‚ - FAISS Vector Search โ”‚ +โ”‚ - CrossEncoder Reranking โ”‚ +โ”‚ - Format Token Matching โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ†“ Top-K Candidates +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ AGENT SELECTION โ”‚ +โ”‚ - Pydantic AI Agent โ”‚ +โ”‚ - OpenAI VLM โ”‚ +โ”‚ - Image + Metadata Analysis โ”‚ +โ”‚ - Multi-Tool Reasoning โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ†“ + Ranked Recommendations +``` + +### Retrieval Stage + +**No LLM calls** - purely text-based search: + +1. **Query Construction**: User task + format tokens from uploaded files +2. **Embedding**: BGE-M3 model generates query embedding +3. **Vector Search**: FAISS retrieves top candidates +4. **Reranking**: CrossEncoder refines results for precision +5. **Retry Broadening**: If too few hits, retry with a shorter/broader query + +### Agent Selection Stage + +**Single VLM call** - multimodal reasoning: + +1. **Input Preparation**: + - Text: User query + candidate table + file metadata + - Image: PNG preview (converted from any format) + - Context: Original file format, dimensions, modality + +2. **Agent Tools**: + - `search_tools`: Search catalog with query + - `search_alternative`: Find alternatives (iterative) + - `repo_info`: Fetch GitHub documentation via DeepWiki MCP + +3. **Output**: Ranked tool selections with accuracy scores and explanations + +### Key Components + +- **`api/pipeline.py`**: RAG retrieval orchestrator +- **`agent/agent.py`**: Pydantic AI agent with tool definitions +- **`retriever/`**: Embedding, FAISS indexing, reranking +- **`generator/`**: Prompts and schema for tool selection +- **`ui/`**: Gradio chat interface components +- **`utils/`**: Image processing, metadata extraction, file validation +- **`catalog/`**: Catalog syncing from GraphDB (optional) + +--- + +## โš™๏ธ Configuration + +### Environment Variables + +| Variable | Description | Default | Required | +|----------|-------------|---------|----------| +| `OPENAI_API_KEY` | OpenAI API key | - | โœ… | +| `EPFL_API_KEY_EMBEDDER` | API key for remote embedder and reranker endpoints | - | โœ… (when `retrieval.embedder.backend: remote` and/or `retrieval.reranker.backend: remote`) | +| `GITHUB_TOKEN` | GitHub token for repo info | - | โŒ | +| `SOFTWARE_CATALOG` | Path to catalog JSONL | `dataset/catalog.jsonl` | โœ… | +| `TOP_K` | Retrieval candidates count | `8` | โŒ | +| `NUM_CHOICES` | Tools to recommend | `3` | โŒ | +| `AGENT_OUTPUT_RETRIES` | Structured output validation retries | `3` | โŒ | +| `EMBED_CATALOG_ON_START` | Pre-embed catalog on startup when FAISS is empty | `1` | โŒ | +| `LOGLEVEL_CONSOLE` | Console log level | `WARNING` | โŒ | +| `LOGLEVEL_FILE` | File log level | `INFO` | โŒ | +| `FILE_LOG` | Enable file logging | `1` | โŒ | +| `LOG_DIR` | Log directory | `logs` | โŒ | +| `LOG_PROMPTS` | Save prompt snapshots | `0` | โŒ | +| `CONFIG_PATH` | Model config file | `config.yaml` | โœ… | + +### GraphDB Catalog Sync (Optional) + +For automatic catalog syncing from a GraphDB instance: + +```dotenv +GRAPHDB_URL=https://your-graphdb.example.com +GRAPHDB_GRAPH=your-graph-name +GRAPHDB_USER=username +GRAPHDB_PASSWORD=password +GRAPHDB_QUERY_FILE=/path/to/query.rq +SYNC_EVERY_HOURS=24 # Auto-refresh interval (0 to disable) +OUTPUT_JSONLD=dataset/catalog.jsonld +OUTPUT_JSONL=dataset/catalog.jsonl +``` + +Run manual sync: +```bash +ai_agent sync +``` + +--- + +## ๐Ÿ“‹ Catalog Format + +The catalog is a JSONL file where each line is a `SoftwareDoc` following schema.org SoftwareSourceCode structure. + +### Minimal Example + +```json +{ + "name": "3d-lungs-segmentation", + "description": "3D lung segmentation from CT; returns a mask/overlay.", + + "applicationCategory": ["Medical Imaging"], + "featureList": ["segmentation"], + "imagingModality": ["CT"], + "dims": [3], + "anatomy": ["lung"], + "keywords": ["mask", "overlay", "lung segmentation", "CT"], + + "programmingLanguage": "Python", + "requiresGPU": false, + "isAccessibleForFree": true, + "license": "Apache-2.0", + + "supportingData": [ + { + "datasetFormat": "TIFF", + "bodySite": "lung", + "imagingModality": "CT", + "hasDimensionality": 3 + }, + { + "datasetFormat": "DICOM", + "bodySite": "lung", + "imagingModality": "CT", + "hasDimensionality": 3 + } + ], + + "runnableExample": [ + { + "hostType": "gradio", + "url": "https://huggingface.co/spaces/qchapp/3d-lungs-segmentation", + "name": "HF Space" + } + ] +} +``` + +### Key Fields + +- **name**: Unique identifier for the tool +- **description**: Clear explanation of what the tool does +- **featureList**: Operations (e.g., segmentation, registration, classification) +- **imagingModality**: Medical imaging types (CT, MRI, XR, US, PET) +- **dims**: Supported dimensions (2D, 3D, 4D) +- **anatomy**: Body parts/organs +- **supportingData**: Format compatibility information (critical for matching) +- **runnableExample**: Links to live demos (HuggingFace Spaces, notebooks, web apps) + +--- + +## ๐Ÿ”ง Development + +### Project Structure + +``` +ai-agent/ +โ”œโ”€โ”€ src/ai_agent/ +โ”‚ โ”œโ”€โ”€ agent/ # Pydantic AI agent and tools +โ”‚ โ”‚ โ”œโ”€โ”€ agent.py # Agent definition +โ”‚ โ”‚ โ”œโ”€โ”€ models.py # Agent state models +โ”‚ โ”‚ โ”œโ”€โ”€ tools/ # Agent tool implementations +โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ search_tool.py +โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ search_alternative_tool.py +โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ gradio_space_tool.py +โ”‚ โ”‚ โ”‚ โ”œโ”€โ”€ repo_info_tool.py +โ”‚ โ”‚ โ”‚ โ””โ”€โ”€ deepwiki_tool.py +โ”‚ โ”‚ โ””โ”€โ”€ utils.py +โ”‚ โ”œโ”€โ”€ api/ # Pipeline orchestration +โ”‚ โ”‚ โ””โ”€โ”€ pipeline.py # RAGImagingPipeline +โ”‚ โ”œโ”€โ”€ retriever/ # Retrieval components +โ”‚ โ”‚ โ”œโ”€โ”€ text_embedder.py +โ”‚ โ”‚ โ”œโ”€โ”€ vector_index.py +โ”‚ โ”‚ โ”œโ”€โ”€ reranker.py +โ”‚ โ”‚ โ””โ”€โ”€ software_doc.py +โ”‚ โ”œโ”€โ”€ generator/ # Agent prompts and schemas +โ”‚ โ”‚ โ”œโ”€โ”€ prompts.py +โ”‚ โ”‚ โ””โ”€โ”€ schema.py +โ”‚ โ”œโ”€โ”€ ui/ # Gradio interface +โ”‚ โ”‚ โ”œโ”€โ”€ app.py +โ”‚ โ”‚ โ”œโ”€โ”€ handlers.py +โ”‚ โ”‚ โ”œโ”€โ”€ components.py +โ”‚ โ”‚ โ”œโ”€โ”€ formatters.py +โ”‚ โ”‚ โ”œโ”€โ”€ state.py +โ”‚ โ”‚ โ””โ”€โ”€ visualizations.py +โ”‚ โ”œโ”€โ”€ utils/ # Shared utilities +โ”‚ โ”‚ โ”œโ”€โ”€ config.py # Configuration management +โ”‚ โ”‚ โ”œโ”€โ”€ file_validator.py +โ”‚ โ”‚ โ”œโ”€โ”€ image_meta.py # Metadata extraction +โ”‚ โ”‚ โ”œโ”€โ”€ image_io.py +โ”‚ โ”‚ โ”œโ”€โ”€ previews.py +โ”‚ โ”‚ โ””โ”€โ”€ tags.py +โ”‚ โ”œโ”€โ”€ catalog/ # Catalog syncing +โ”‚ โ”‚ โ””โ”€โ”€ sync.py +โ”‚ โ””โ”€โ”€ cli.py # CLI entry point +โ”œโ”€โ”€ tests/ # Test suite +โ”‚ โ”œโ”€โ”€ test_retrieval_pipeline.py +โ”‚ โ”œโ”€โ”€ test_repo_summary.py +โ”‚ โ””โ”€โ”€ data/ +โ”œโ”€โ”€ artifacts/ # Generated artifacts +โ”‚ โ””โ”€โ”€ rag_index/ # FAISS index +โ”œโ”€โ”€ dataset/ # Catalog data +โ”‚ โ””โ”€โ”€ catalog.jsonl +โ”œโ”€โ”€ logs/ # Application logs +โ”œโ”€โ”€ config.yaml # Model configuration +โ”œโ”€โ”€ pyproject.toml # Project metadata & dependencies +โ”œโ”€โ”€ Dockerfile # Production Docker image +โ”œโ”€โ”€ tools/image/Dockerfile # Development Docker image +โ””โ”€โ”€ justfile # Task runner commands +``` + +### Local Development + +```bash +# Install in development mode +pip install -e ".[dev]" + +# Run tests +pytest tests/ +``` + +### Testing + +Run the test suite: + +```bash +# All tests +pytest tests/ + +# Specific test file +pytest tests/test_retrieval_pipeline.py + +# With verbose output +pytest -v tests/ + +# With coverage +pytest --cov=ai_agent tests/ +``` + +### Logging & Debugging + +**Console Logs**: Set `LOGLEVEL_CONSOLE=DEBUG` for verbose output + +**File Logs**: Automatically saved to `logs/app_YYYYMMDD.log` (rotates daily) + +**Prompt Snapshots**: Enable `LOG_PROMPTS=1` to save: +- `logs/vlm_selector_YYYYMMDD_HHMMSS.txt` - System/user prompts + +--- + +## ๐Ÿ“š API & CLI Reference + +### CLI Commands + +```bash +# Launch chat interface +ai_agent chat + +# Sync catalog from GraphDB +ai_agent sync +``` + +## ๐Ÿ—บ๏ธ Maintainer Guide + +For full project documentation with detailed folder responsibilities, environment defaults, and improvement guidelines, see [docs/guide.md](docs/guide.md). + +--- + +## ๐Ÿ“ Changelog + +See [CHANGELOG.md](CHANGELOG.md) for detailed version history. + +### Recent Highlights + +**[1.0.0]** +- โœจ New chat-based interface (`ai_agent chat`) with rich media and tool integration +- ๐Ÿ› ๏ธ Fully agent-based architecture replacing legacy pipelines +- ๐Ÿ” Smarter retrieval with automatic retry +- ๐Ÿ”— DeepWiki MCP integration for fast GitHub repository documentation access +- ๐Ÿ”ง YAML configuration (`config.yaml`) for flexible model and backend setup +- ๐ŸŽจ Redesigned UI with Imaging Plaza branding and improved UX +- โšก Performance improvements (pre-embedding, caching, faster startup) +- ๐Ÿงน Major cleanup: removed deprecated code paths, legacy UI, and outdated tests + +**[0.1.3] - 2025-10-22** +- Gradio space runner tool +- Repository info tool +- UI fixes and polish + +--- + +## ๐Ÿ“„ License + +This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details. + +--- + +## ๐Ÿ™ Credits & Acknowledgments + +**Developed by**: Imaging Plaza Team + +**Technologies:** +- [Pydantic AI](https://github.com/pydantic/pydantic-ai) - AI agent framework +- [OpenAI](https://openai.com) - GPT vision model +- [FAISS](https://github.com/facebookresearch/faiss) - Vector search +- [BGE-M3](https://huggingface.co/BAAI/bge-m3) - Multilingual embeddings +- [Gradio](https://gradio.app) - Interactive web UI +- [DeepWiki](https://deepwiki.com) - GitHub repository documentation + +**Medical Imaging Formats:** +- [pydicom](https://github.com/pydicom/pydicom) - DICOM support +- [nibabel](https://nipy.org/nibabel/) - NIfTI support + +--- + +## ๐Ÿ“ฎ Support + +For issues, questions, or contributions, please contact the Imaging Plaza team. + +--- + +**๐Ÿฅ Medical Disclaimer**: This software is a tool recommendation system, not a diagnostic tool. Always consult qualified medical professionals for clinical decisions. \ No newline at end of file diff --git a/config.yaml b/config.yaml index 6c4b8021ba8bdf304dcd691cb6e774b20c36bf0d..19fd26426a05ceeef8bd961f8c09040b193f67ea 100644 --- a/config.yaml +++ b/config.yaml @@ -1,64 +1,64 @@ -# AI Agent Model Configuration - -# Default/fallback model (used for CLI and initial startup) -agent_model: - # name: "gpt-5.1" - # base_url: null # null for default OpenAI endpoint - # api_key_env: "OPENAI_API_KEY" - name: "openai/gpt-oss-120b" - base_url: "https://inference-rcp.epfl.ch/v1" - api_key_env: "EPFL_API_KEY" - -# Available models for UI dropdown -available_models: - - display_name: "gpt-4o-mini" - name: "gpt-4o-mini" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-4o" - name: "gpt-4o" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-5-mini" - name: "gpt-5-mini" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-5.1" - name: "gpt-5.1" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "openai/gpt-oss-120b [EPFL]" - name: "openai/gpt-oss-120b" - base_url: "https://inference-rcp.epfl.ch/v1" - provider: "EPFL" - api_key_env: "EPFL_API_KEY" - -# Retrieval stack (embedder + reranker) -retrieval: - embedder: - backend: "remote" # "remote" or "local" - model_name: "Qwen/Qwen3-Embedding-8B" - base_url: "https://inference-rcp.epfl.ch/v1" - api_key_env: "EPFL_API_KEY_EMBEDDER" - timeout_s: 20 - # local example: - # backend: "local" - # model_name: "BAAI/bge-m3" - - reranker: - backend: "remote" # "remote" or "local" - model_name: "BAAI/bge-reranker-v2-m3" - base_url: "https://inference-rcp.epfl.ch/v1" - api_key_env: "EPFL_API_KEY_EMBEDDER" - timeout_s: 20 - # local example: - # backend: "local" - # model_name: "BAAI/bge-reranker-v2-m3" +# AI Agent Model Configuration + +# Default/fallback model (used for CLI and initial startup) +agent_model: + # name: "gpt-5.1" + # base_url: null # null for default OpenAI endpoint + # api_key_env: "OPENAI_API_KEY" + name: "openai/gpt-oss-120b" + base_url: "https://inference-rcp.epfl.ch/v1" + api_key_env: "EPFL_API_KEY" + +# Available models for UI dropdown +available_models: + - display_name: "gpt-4o-mini" + name: "gpt-4o-mini" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-4o" + name: "gpt-4o" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-5-mini" + name: "gpt-5-mini" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-5.1" + name: "gpt-5.1" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "openai/gpt-oss-120b [EPFL]" + name: "openai/gpt-oss-120b" + base_url: "https://inference-rcp.epfl.ch/v1" + provider: "EPFL" + api_key_env: "EPFL_API_KEY" + +# Retrieval stack (embedder + reranker) +retrieval: + embedder: + backend: "remote" # "remote" or "local" + model_name: "Qwen/Qwen3-Embedding-8B" + base_url: "https://inference-rcp.epfl.ch/v1" + api_key_env: "EPFL_API_KEY_EMBEDDER" + timeout_s: 20 + # local example: + # backend: "local" + # model_name: "BAAI/bge-m3" + + reranker: + backend: "remote" # "remote" or "local" + model_name: "BAAI/bge-reranker-v2-m3" + base_url: "https://inference-rcp.epfl.ch/v1" + api_key_env: "EPFL_API_KEY_EMBEDDER" + timeout_s: 20 + # local example: + # backend: "local" + # model_name: "BAAI/bge-reranker-v2-m3" diff --git a/data/query.rq b/data/query.rq new file mode 100644 index 0000000000000000000000000000000000000000..de5f162699fe810d3b62b58608ea81f0a8a55b39 --- /dev/null +++ b/data/query.rq @@ -0,0 +1,8 @@ +CONSTRUCT {{ + ?s ?p ?o +}} +WHERE {{ + GRAPH <{graph}> {{ + ?s ?p ?o . + }} +}} diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000000000000000000000000000000000000..f182b1d04247521711bbd2ff371551871ff26ac6 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,25 @@ +services: + ai-agent: + build: . + image: ai-agent:dev + container_name: ai-agent + restart: unless-stopped + env_file: .env + environment: + HOST: 0.0.0.0 + PORT: "7860" + expose: + - "7860" + volumes: + - ./data:/home/user/app/data + - ./dataset:/home/user/app/dataset + - ./artifacts:/home/user/app/artifacts + - ./logs:/home/user/app/logs + + cloudflared: + image: cloudflare/cloudflared:latest + container_name: ai-agent-cloudflared + restart: unless-stopped + command: tunnel --no-autoupdate --url http://ai-agent:7860 + depends_on: + - ai-agent diff --git a/docs/architecture/agent.md b/docs/architecture/agent.md index 2739d39a7e54bb45c1bd6bb1fe6880bffbc7d08e..fee0d42417b2fb38972df5bd502affc68d3b5b31 100644 --- a/docs/architecture/agent.md +++ b/docs/architecture/agent.md @@ -1,619 +1,619 @@ -# Agent & VLM Selection - -The second stage of the pipeline uses a vision-language model (VLM) with PydanticAI agents to select and rank the best tools from candidates. - -## Overview - -**Goal**: Select the most relevant tools using vision + text understanding - -**Characteristics**: - -- ๐Ÿง  Intelligent reasoning with explanations -- ๐Ÿ‘๏ธ Vision-aware (analyzes image content) -- ๐ŸŽฏ Comparative ranking of candidates -- ๐Ÿ’ฌ Conversational with context -- ๐Ÿ“Š Structured output (Pydantic schemas) - -## Architecture - -```mermaid -graph TB - A[User Message + Files] --> B[PydanticAI Agent] - B --> C{Agent Router} - C -->|Tool Call| D[Agent Tools] - C -->|LLM Reasoning| E[GPT-4o/4o-mini] - D --> B - E --> F[ToolSelection Schema] - F --> G[Structured Response] - G --> B - B --> H[Formatted Reply] -``` - -## PydanticAI Agent - -### Agent Framework - -**Framework**: [PydanticAI](https://ai.pydantic.dev/) - -**Benefits**: - -- Type-safe with Pydantic models -- Structured output validation -- Built-in tool support -- Async/await support -- Easy testing with dependency injection - -### Agent Definition - -```python -from pydantic_ai import Agent -from pydantic_ai.models.openai import OpenAIResponsesModel -from pydantic_ai.providers.openai import OpenAIProvider -from ai_agent.generator.prompts import get_agent_system_prompt -from ai_agent.generator.schema import ToolSelection -from ai_agent.agent.utils import AgentState - -provider = OpenAIProvider(api_key=os.getenv("OPENAI_API_KEY")) -openai_model = OpenAIResponsesModel(model_name="gpt-4o-mini", provider=provider) - -agent = Agent( - model=openai_model, - system_prompt=get_agent_system_prompt(num_choices=3), - deps_type=AgentState, -) -``` - -**Key parameters**: - -- `model`: VLM model to use (configurable via `config.yaml`) -- `system_prompt`: Agent role, scoring rules, and output format (from `generator/prompts.py`) -- `deps_type`: `AgentState` โ€” tracks tool calls, quotas, and session overrides -- `output_type`: `ToolSelection` โ€” passed to `agent.run_sync()` to enforce structured JSON output (not set on the `Agent` constructor) - -### Conversation State - -```python -from pydantic import BaseModel, Field -from typing import List, Optional, Dict, Any, Set - -class AgentState(BaseModel): - """Holds incremental tool call logs and runtime overrides.""" - - tool_calls: List[Dict[str, Any]] = Field(default_factory=list) - tool_counts: Dict[str, int] = Field(default_factory=dict) - disabled_tools: Set[str] = Field(default_factory=set) - excluded_tools: List[str] = Field(default_factory=list) - - # Runtime overrides (session-only) - override_model: Optional[str] = None - override_base_url: Optional[str] = None - override_top_k: Optional[int] = None - override_num_choices: Optional[int] = None - - image_paths: List[str] = Field(default_factory=list) - original_formats: List[str] = Field(default_factory=list) -``` - -**Passed to every tool call** via dependency injection. Also carries per-tool call counts for quota enforcement. - -## Agent Tools - -Tools extend agent capabilities beyond chat: - -### search_alternative - -Request alternative search with different query formulation: - -```python -@agent.tool(retries=2, prepare=cap_prepare) -@limit_tool_calls("search_alternative", cap=3) -async def search_alternative( - ctx: RunContext[AgentState], - alternative_query: str, - excluded: List[str] | None = None, - top_k: int = 12, -) -> List[dict]: - """Search for tools using an alternative query formulation.""" - - inp = SearchAlternativeInput( - alternative_query=alternative_query, - excluded=excluded or [], - top_k=top_k, - original_formats=ctx.deps.original_formats, - image_paths=ctx.deps.image_paths, - ) - out = tool_search_alternative(inp) - return [c.model_dump(mode="python") for c in out.candidates] -``` - -**Usage**: - -- Agent invokes when user asks for alternatives -- Up to 3 calls per conversation -- Formulates semantically different queries - -**Example**: -``` -User: Show me alternatives -Agent: [Calls search_alternative with "pulmonary segmentation CT"] -``` - -### repo_info - -Fetch GitHub repository details: - -```python -@agent.tool(retries=2, prepare=cap_prepare) -@limit_tool_calls("repo_info", cap=12) -async def repo_info(ctx: RunContext[AgentState], url: str, tool_name: str | None = None) -> dict: - """Fetch a short summary of a GitHub repository.""" - - # Normalize to canonical GitHub URL - norm_url = coerce_github_url_or_none(url) - - # Call tool_repo_summary (tries DeepWiki MCP first, falls back to repocards) - out = await tool_repo_summary(RepoSummaryInput(url=norm_url, tool_name=tool_name)) - return out.model_dump(mode="python") -``` - -**Data sources**: - -1. **DeepWiki MCP**: Pre-indexed, fast, no rate limits -2. **Repocards**: Direct fetch, fallback for new repos - -**Returns**: - -- Repository description -- Stars, language, topics -- Last update date -- License information - -**Example**: -``` -User: Tell me about TotalSegmentator -Agent: [Calls repo_info("https://github.com/wasserth/TotalSegmentator")] - - TotalSegmentator is an automated multi-organ segmentation tool... - โญ 1.2k stars | Python | Apache-2.0 license - Topics: segmentation, medical-imaging, deep-learning -``` - -### run_example - -Execute Gradio Space demos (optional, experimental): - -```python -@agent.tool(retries=0, prepare=cap_prepare) -@limit_tool_calls("run_example", cap=1) -async def run_example( - ctx: RunContext[AgentState], - tool_name: str, - endpoint_url: str | None = None, - extra_text: str | None = None, -) -> dict: - """Run an example / demo for a given tool via its Gradio space.""" - - out = tool_run_example(RunExampleInput( - tool_name=tool_name, - endpoint_url=endpoint_url, - extra_text=extra_text, - )) - return out.model_dump(mode="python") -``` - -**Status**: Partially implemented, limited to specific demo formats. - -## Selection and Ranking - -The PydanticAI agent performs tool selection and ranking directly as part of its LLM reasoning step. There is no separate `VLMToolSelector` class โ€” the agent's system prompt (defined in `generator/prompts.py`) encodes the scoring rules, and the `ToolSelection` Pydantic schema (defined in `generator/schema.py`) enforces structured output. - -### System Prompt - -The agent system prompt is assembled by `get_agent_system_prompt()` in `generator/prompts.py` and covers: - -- **Image analysis**: Instructions to analyze the attached preview image and reference visual observations in explanations -- **Tool call sequence**: When to call `search_tools`, `search_alternative`, `repo_info`, and `run_example` -- **Scoring rules**: Accuracy (0โ€“100) = Task match (40) + Format compatibility (30) + Features (30) -- **Output format**: Single JSON object matching the `ToolSelection` schema - -```python -from ai_agent.generator.prompts import get_agent_system_prompt - -# Generates a prompt that instructs the agent to return up to N ranked choices -system_prompt = get_agent_system_prompt(num_choices=3) -``` - -### Selection Process - -#### Step 1: Tool Calls (Retrieval) - -The agent calls `search_tools` once (and optionally `search_alternative` up to 3 times) to retrieve candidate tools from the vector index: - -``` -Agent โ†’ search_tools(query="segment lungs", top_k=12) - โ† [TotalSegmentator, MedSAM, nnU-Net, ...] -``` - -#### Step 2: Verification - -For each finalist the agent plans to recommend, it calls `repo_info` to fetch up-to-date GitHub metadata: - -``` -Agent โ†’ repo_info(url="https://github.com/wasserth/TotalSegmentator") - โ† {stars: 1200, language: "Python", topics: [...], description: "..."} -``` - -#### Step 3: Structured Output - -The agent returns one JSON object (no prose) that is validated against the `ToolSelection` schema: - -```python -run_result = agent_instance.run_sync( - user_prompt, # text + optional BinaryContent image - deps=deps, # AgentState with image_paths, excluded_tools, etc. - output_type=ToolSelection, - usage_limits=UsageLimits(tool_calls_limit=20), -) -result = run_result.output # ToolSelection instance -``` - -**Multimodal input**: - -- Text: User task + hidden metadata (format hints, image dimensions) -- Image: PNG preview bytes passed as `BinaryContent(data=image_bytes, media_type="image/png")` -- Context: Conversation history prepended to the prompt - -### Structured Response Schema - -The `ToolSelection` Pydantic model (in `generator/schema.py`) validates the agent output: - -```python -from ai_agent.generator.schema import ( - ToolSelection, ToolChoice, Conversation, - ConversationStatus, NoToolReason -) - -class ToolChoice(BaseModel): - name: str - rank: int - accuracy: float # 0-100 - why: str - demo_link: Optional[str] = None - -class Conversation(BaseModel): - status: ConversationStatus - question: Optional[str] = None # required if status=needs_clarification - context: Optional[str] = None # required if status=needs_clarification - options: Optional[List[str]] = None - -class ToolSelection(BaseModel): - conversation: Conversation - choices: List[ToolChoice] = [] - explanation: Optional[str] = None - reason: Optional[NoToolReason] = None -``` - -**Example response** (`ToolSelection`): -```json -{ - "conversation": {"status": "complete", "question": null, "context": null}, - "choices": [ - { - "rank": 1, - "name": "TotalSegmentator", - "accuracy": 95.0, - "why": "Specifically designed for automated multi-organ CT segmentation...", - "demo_link": "https://huggingface.co/spaces/..." - }, - { - "rank": 2, - "name": "MedSAM", - "accuracy": 85.0, - "why": "Flexible SAM-based segmentation supporting DICOM input...", - "demo_link": "https://huggingface.co/spaces/..." - } - ], - "explanation": null, - "reason": null -} -``` - -### Validation - -Pydantic validates: - -- All required fields present -- Types correct (`int`, `float`, `str`, enum) -- `accuracy` within 0โ€“100 range -- `ConversationStatus` is one of the allowed enum values -- `NoToolReason` is a valid enum value when `choices` is empty - -`ToolSelection.normalize()` also enforces consistency rules automatically (e.g. setting `status=complete` when choices are returned, `status=needs_clarification` when a question is present). - -## Conversation States - -State machine for conversation flow: - -```python -from ai_agent.generator.schema import ConversationStatus - -class ConversationStatus(str, Enum): - COMPLETE = "complete" # Recommendations provided (or no tool found) - NEEDS_CLARIFICATION = "needs_clarification" # Agent needs more info -``` - -### Complete - -Normal successful response: - -```python -{ - "conversation": {"status": "complete", "question": null, "context": null}, - "choices": [...], - "explanation": null, - "reason": null -} -``` - -**Triggers**: - -- Query is clear -- Candidates found -- Image/metadata sufficient - -### Needs Clarification - -Agent requests more information: - -```python -{ - "conversation": { - "status": "needs_clarification", - "question": "Which specific organ would you like to segment?", - "context": "Several segmentation tools available; target organ narrows choices.", - "options": ["Lungs", "Brain", "Liver", "Other (briefly specify)"] - }, - "choices": [], - "explanation": null, - "reason": null -} -``` - -**Triggers**: -- Ambiguous query -- Multiple valid interpretations -- Missing critical information - -**Example flow**: -``` -User: Segment this MRI -Agent: [STATUS: needs_clarification] Which organ would you like to segment? -User: The brain -Agent: [STATUS: complete] Here are brain segmentation tools... -``` - -### No Tool Terminal - -No suitable tools in catalog โ€” `status` is still `complete`, but `choices` is empty and a `reason` + `explanation` are provided: - -```python -{ - "conversation": {"status": "complete", "question": null, "context": null}, - "choices": [], - "reason": "no_task_match", - "explanation": "No tools in the catalog handle audio processing. This catalog covers imaging analysis software." -} -``` - -Available `NoToolReason` values: `no_suitable_tool`, `no_modality_match`, `no_task_match`, `no_dimension_match`, `invalid_files`. - -## Ranking Logic - -### Scoring Factors - -The agent considers: - -#### High Priority -1. **Task Match**: Tool designed for this specific task -2. **Format Compatibility**: Supports user's file format -3. **Visual Analysis**: Image content matches tool's domain - -#### Medium Priority -4. **Modality Alignment**: CT tool for CT image, MRI for MRI -5. **Dimension Match**: 3D tool for 3D volume -6. **Feature Coverage**: Specific capabilities mentioned - -#### Low Priority -7. **License**: Open-source preference (if no preference stated) -8. **Demo Availability**: Has runnable demo -9. **Popularity**: Community adoption - -### Explanation Generation - -Each recommendation includes explanation: - -**Good explanation template**: -``` -{Tool} is {specifically designed / well-suited} for {task} -on {modality} images. It supports {format} input {with/without} -preprocessing and provides {key features}. {Caveats if any}. -``` - -**Example**: -``` -TotalSegmentator is specifically designed for automated multi-organ -segmentation on CT scans. It supports DICOM input without preprocessing -and can segment 104 anatomical structures including lungs, air airways, -and vessels. It works best on whole-body CT but also performs well on -thoracic scans. -``` - -### Rank Assignment - -- **Rank 1**: Best overall match (highest accuracy score) -- **Rank 2**: Strong alternative or different approach -- **Rank 3**: Fallback option or specialized capability - -**Important**: Ranks are relative to **this specific query**, not absolute tool quality. - -## Model Configuration - -### Model Selection - -Available via `config.yaml`: - -```yaml -agent_model: - name: "gpt-4o-mini" - base_url: null - api_key_env: "OPENAI_API_KEY" -``` - -### Model Comparison - -| Model | Vision | Speed | Cost | Best For | -|-------|--------|-------|------|----------| -| gpt-4o-mini | โœ… | โšกโšกโšก | $ | Most queries, fast iteration | -| gpt-4o | โœ…โœ… | โšกโšก | $$ | Complex visual analysis | -| gpt-5.1 | โœ…โœ…โœ… | โšก | $$$ | Maximum accuracy needed | - -### Custom Endpoints - -Support for OpenAI-compatible APIs: - -```yaml -agent_model: - name: "llama-3.2-vision" - base_url: "https://inference.epfl.ch/v1" - api_key_env: "EPFL_API_KEY" -``` - -## Error Handling - -### Agent Errors - -**Tool quota exceeded** (handled gracefully in `run_agent`): -```python -except UsageLimitExceeded: - # Returns a ToolSelection with empty choices and an explanation - result = ToolSelection( - conversation=Conversation(status=ConversationStatus.COMPLETE, ...), - choices=[], - explanation="Tool call limit reached. Try a more specific query.", - ) -``` - -**Invalid structured output**: - -PydanticAI automatically retries the LLM call (up to `retries=2` per tool) if the model returns output that fails `ToolSelection` validation. The `ToolSelection.normalize()` model validator also auto-corrects minor inconsistencies. - -**API Errors**: -```python -except Exception as e: - log.warning(f"Agent execution encountered an error: {e}") - raise # propagated to the UI layer -``` - -### Graceful Degradation - -If the agent fails after all retries: - -1. Return empty `choices` with an `explanation` describing what was searched -2. UI surfaces the explanation so users can refine their query -3. Suggest manual exploration of the catalog - - - -## Testing - -### Unit Tests - -Test agent selection with PydanticAI's built-in test model (your catalog should contain the choice provided below, i.e. the `TotalSegmentator` tool): - -```python -from pydantic_ai import Agent -from pydantic_ai.models.test import TestModel -from ai_agent.generator.schema import ToolSelection, Conversation, ConversationStatus, ToolChoice -from ai_agent.agent.utils import AgentState - -def test_agent_selection(): - test_model = TestModel() - test_agent = Agent(model=test_model, deps_type=AgentState) - - mock_output = ToolSelection( - conversation=Conversation(status=ConversationStatus.COMPLETE), - choices=[ - ToolChoice(name="TotalSegmentator", rank=1, accuracy=95.0, why="Best CT segmenter") - ] - ) - - with test_agent.override(model=test_model): - result = test_agent.run_sync("segment lungs", deps=AgentState(), output_type=ToolSelection) - - assert result.output.conversation.status == ConversationStatus.COMPLETE - assert len(result.output.choices) == 1 - assert result.output.choices[0].rank == 1 -``` - -### Integration Tests - -Test with real VLM (expensive, slow): - -```python -@pytest.mark.integration -def test_real_agent(): - from ai_agent.agent.agent import run_agent - - with open("tests/data/sample.tif", "rb") as f: - image_bytes = f.read() - - result = run_agent( - task="I want to segment the lungs of this CT scan", - image_paths=["tests/data/sample.tif"], - image_bytes=image_bytes, - ) - - assert result.conversation.status == ConversationStatus.COMPLETE - assert len(result.choices) > 0 - assert all(0 <= c.accuracy <= 100 for c in result.choices) -``` - -## Next Steps - -- Learn about [Software Catalog](catalog.md) -- Return to [Architecture Overview](overview.md) -- Explore [Retrieval Pipeline](retrieval.md) +# Agent & VLM Selection + +The second stage of the pipeline uses a vision-language model (VLM) with PydanticAI agents to select and rank the best tools from candidates. + +## Overview + +**Goal**: Select the most relevant tools using vision + text understanding + +**Characteristics**: + +- ๐Ÿง  Intelligent reasoning with explanations +- ๐Ÿ‘๏ธ Vision-aware (analyzes image content) +- ๐ŸŽฏ Comparative ranking of candidates +- ๐Ÿ’ฌ Conversational with context +- ๐Ÿ“Š Structured output (Pydantic schemas) + +## Architecture + +```mermaid +graph TB + A[User Message + Files] --> B[PydanticAI Agent] + B --> C{Agent Router} + C -->|Tool Call| D[Agent Tools] + C -->|LLM Reasoning| E[GPT-4o/4o-mini] + D --> B + E --> F[ToolSelection Schema] + F --> G[Structured Response] + G --> B + B --> H[Formatted Reply] +``` + +## PydanticAI Agent + +### Agent Framework + +**Framework**: [PydanticAI](https://ai.pydantic.dev/) + +**Benefits**: + +- Type-safe with Pydantic models +- Structured output validation +- Built-in tool support +- Async/await support +- Easy testing with dependency injection + +### Agent Definition + +```python +from pydantic_ai import Agent +from pydantic_ai.models.openai import OpenAIResponsesModel +from pydantic_ai.providers.openai import OpenAIProvider +from ai_agent.generator.prompts import get_agent_system_prompt +from ai_agent.generator.schema import ToolSelection +from ai_agent.agent.utils import AgentState + +provider = OpenAIProvider(api_key=os.getenv("OPENAI_API_KEY")) +openai_model = OpenAIResponsesModel(model_name="gpt-4o-mini", provider=provider) + +agent = Agent( + model=openai_model, + system_prompt=get_agent_system_prompt(num_choices=3), + deps_type=AgentState, +) +``` + +**Key parameters**: + +- `model`: VLM model to use (configurable via `config.yaml`) +- `system_prompt`: Agent role, scoring rules, and output format (from `generator/prompts.py`) +- `deps_type`: `AgentState` โ€” tracks tool calls, quotas, and session overrides +- `output_type`: `ToolSelection` โ€” passed to `agent.run_sync()` to enforce structured JSON output (not set on the `Agent` constructor) + +### Conversation State + +```python +from pydantic import BaseModel, Field +from typing import List, Optional, Dict, Any, Set + +class AgentState(BaseModel): + """Holds incremental tool call logs and runtime overrides.""" + + tool_calls: List[Dict[str, Any]] = Field(default_factory=list) + tool_counts: Dict[str, int] = Field(default_factory=dict) + disabled_tools: Set[str] = Field(default_factory=set) + excluded_tools: List[str] = Field(default_factory=list) + + # Runtime overrides (session-only) + override_model: Optional[str] = None + override_base_url: Optional[str] = None + override_top_k: Optional[int] = None + override_num_choices: Optional[int] = None + + image_paths: List[str] = Field(default_factory=list) + original_formats: List[str] = Field(default_factory=list) +``` + +**Passed to every tool call** via dependency injection. Also carries per-tool call counts for quota enforcement. + +## Agent Tools + +Tools extend agent capabilities beyond chat: + +### search_alternative + +Request alternative search with different query formulation: + +```python +@agent.tool(retries=2, prepare=cap_prepare) +@limit_tool_calls("search_alternative", cap=3) +async def search_alternative( + ctx: RunContext[AgentState], + alternative_query: str, + excluded: List[str] | None = None, + top_k: int = 12, +) -> List[dict]: + """Search for tools using an alternative query formulation.""" + + inp = SearchAlternativeInput( + alternative_query=alternative_query, + excluded=excluded or [], + top_k=top_k, + original_formats=ctx.deps.original_formats, + image_paths=ctx.deps.image_paths, + ) + out = tool_search_alternative(inp) + return [c.model_dump(mode="python") for c in out.candidates] +``` + +**Usage**: + +- Agent invokes when user asks for alternatives +- Up to 3 calls per conversation +- Formulates semantically different queries + +**Example**: +``` +User: Show me alternatives +Agent: [Calls search_alternative with "pulmonary segmentation CT"] +``` + +### repo_info + +Fetch GitHub repository details: + +```python +@agent.tool(retries=2, prepare=cap_prepare) +@limit_tool_calls("repo_info", cap=12) +async def repo_info(ctx: RunContext[AgentState], url: str, tool_name: str | None = None) -> dict: + """Fetch a short summary of a GitHub repository.""" + + # Normalize to canonical GitHub URL + norm_url = coerce_github_url_or_none(url) + + # Call tool_repo_summary (tries DeepWiki MCP first, falls back to repocards) + out = await tool_repo_summary(RepoSummaryInput(url=norm_url, tool_name=tool_name)) + return out.model_dump(mode="python") +``` + +**Data sources**: + +1. **DeepWiki MCP**: Pre-indexed, fast, no rate limits +2. **Repocards**: Direct fetch, fallback for new repos + +**Returns**: + +- Repository description +- Stars, language, topics +- Last update date +- License information + +**Example**: +``` +User: Tell me about TotalSegmentator +Agent: [Calls repo_info("https://github.com/wasserth/TotalSegmentator")] + + TotalSegmentator is an automated multi-organ segmentation tool... + โญ 1.2k stars | Python | Apache-2.0 license + Topics: segmentation, medical-imaging, deep-learning +``` + +### run_example + +Execute Gradio Space demos (optional, experimental): + +```python +@agent.tool(retries=0, prepare=cap_prepare) +@limit_tool_calls("run_example", cap=1) +async def run_example( + ctx: RunContext[AgentState], + tool_name: str, + endpoint_url: str | None = None, + extra_text: str | None = None, +) -> dict: + """Run an example / demo for a given tool via its Gradio space.""" + + out = tool_run_example(RunExampleInput( + tool_name=tool_name, + endpoint_url=endpoint_url, + extra_text=extra_text, + )) + return out.model_dump(mode="python") +``` + +**Status**: Partially implemented, limited to specific demo formats. + +## Selection and Ranking + +The PydanticAI agent performs tool selection and ranking directly as part of its LLM reasoning step. There is no separate `VLMToolSelector` class โ€” the agent's system prompt (defined in `generator/prompts.py`) encodes the scoring rules, and the `ToolSelection` Pydantic schema (defined in `generator/schema.py`) enforces structured output. + +### System Prompt + +The agent system prompt is assembled by `get_agent_system_prompt()` in `generator/prompts.py` and covers: + +- **Image analysis**: Instructions to analyze the attached preview image and reference visual observations in explanations +- **Tool call sequence**: When to call `search_tools`, `search_alternative`, `repo_info`, and `run_example` +- **Scoring rules**: Accuracy (0โ€“100) = Task match (40) + Format compatibility (30) + Features (30) +- **Output format**: Single JSON object matching the `ToolSelection` schema + +```python +from ai_agent.generator.prompts import get_agent_system_prompt + +# Generates a prompt that instructs the agent to return up to N ranked choices +system_prompt = get_agent_system_prompt(num_choices=3) +``` + +### Selection Process + +#### Step 1: Tool Calls (Retrieval) + +The agent calls `search_tools` once (and optionally `search_alternative` up to 3 times) to retrieve candidate tools from the vector index: + +``` +Agent โ†’ search_tools(query="segment lungs", top_k=12) + โ† [TotalSegmentator, MedSAM, nnU-Net, ...] +``` + +#### Step 2: Verification + +For each finalist the agent plans to recommend, it calls `repo_info` to fetch up-to-date GitHub metadata: + +``` +Agent โ†’ repo_info(url="https://github.com/wasserth/TotalSegmentator") + โ† {stars: 1200, language: "Python", topics: [...], description: "..."} +``` + +#### Step 3: Structured Output + +The agent returns one JSON object (no prose) that is validated against the `ToolSelection` schema: + +```python +run_result = agent_instance.run_sync( + user_prompt, # text + optional BinaryContent image + deps=deps, # AgentState with image_paths, excluded_tools, etc. + output_type=ToolSelection, + usage_limits=UsageLimits(tool_calls_limit=20), +) +result = run_result.output # ToolSelection instance +``` + +**Multimodal input**: + +- Text: User task + hidden metadata (format hints, image dimensions) +- Image: PNG preview bytes passed as `BinaryContent(data=image_bytes, media_type="image/png")` +- Context: Conversation history prepended to the prompt + +### Structured Response Schema + +The `ToolSelection` Pydantic model (in `generator/schema.py`) validates the agent output: + +```python +from ai_agent.generator.schema import ( + ToolSelection, ToolChoice, Conversation, + ConversationStatus, NoToolReason +) + +class ToolChoice(BaseModel): + name: str + rank: int + accuracy: float # 0-100 + why: str + demo_link: Optional[str] = None + +class Conversation(BaseModel): + status: ConversationStatus + question: Optional[str] = None # required if status=needs_clarification + context: Optional[str] = None # required if status=needs_clarification + options: Optional[List[str]] = None + +class ToolSelection(BaseModel): + conversation: Conversation + choices: List[ToolChoice] = [] + explanation: Optional[str] = None + reason: Optional[NoToolReason] = None +``` + +**Example response** (`ToolSelection`): +```json +{ + "conversation": {"status": "complete", "question": null, "context": null}, + "choices": [ + { + "rank": 1, + "name": "TotalSegmentator", + "accuracy": 95.0, + "why": "Specifically designed for automated multi-organ CT segmentation...", + "demo_link": "https://huggingface.co/spaces/..." + }, + { + "rank": 2, + "name": "MedSAM", + "accuracy": 85.0, + "why": "Flexible SAM-based segmentation supporting DICOM input...", + "demo_link": "https://huggingface.co/spaces/..." + } + ], + "explanation": null, + "reason": null +} +``` + +### Validation + +Pydantic validates: + +- All required fields present +- Types correct (`int`, `float`, `str`, enum) +- `accuracy` within 0โ€“100 range +- `ConversationStatus` is one of the allowed enum values +- `NoToolReason` is a valid enum value when `choices` is empty + +`ToolSelection.normalize()` also enforces consistency rules automatically (e.g. setting `status=complete` when choices are returned, `status=needs_clarification` when a question is present). + +## Conversation States + +State machine for conversation flow: + +```python +from ai_agent.generator.schema import ConversationStatus + +class ConversationStatus(str, Enum): + COMPLETE = "complete" # Recommendations provided (or no tool found) + NEEDS_CLARIFICATION = "needs_clarification" # Agent needs more info +``` + +### Complete + +Normal successful response: + +```python +{ + "conversation": {"status": "complete", "question": null, "context": null}, + "choices": [...], + "explanation": null, + "reason": null +} +``` + +**Triggers**: + +- Query is clear +- Candidates found +- Image/metadata sufficient + +### Needs Clarification + +Agent requests more information: + +```python +{ + "conversation": { + "status": "needs_clarification", + "question": "Which specific organ would you like to segment?", + "context": "Several segmentation tools available; target organ narrows choices.", + "options": ["Lungs", "Brain", "Liver", "Other (briefly specify)"] + }, + "choices": [], + "explanation": null, + "reason": null +} +``` + +**Triggers**: +- Ambiguous query +- Multiple valid interpretations +- Missing critical information + +**Example flow**: +``` +User: Segment this MRI +Agent: [STATUS: needs_clarification] Which organ would you like to segment? +User: The brain +Agent: [STATUS: complete] Here are brain segmentation tools... +``` + +### No Tool Terminal + +No suitable tools in catalog โ€” `status` is still `complete`, but `choices` is empty and a `reason` + `explanation` are provided: + +```python +{ + "conversation": {"status": "complete", "question": null, "context": null}, + "choices": [], + "reason": "no_task_match", + "explanation": "No tools in the catalog handle audio processing. This catalog covers imaging analysis software." +} +``` + +Available `NoToolReason` values: `no_suitable_tool`, `no_modality_match`, `no_task_match`, `no_dimension_match`, `invalid_files`. + +## Ranking Logic + +### Scoring Factors + +The agent considers: + +#### High Priority +1. **Task Match**: Tool designed for this specific task +2. **Format Compatibility**: Supports user's file format +3. **Visual Analysis**: Image content matches tool's domain + +#### Medium Priority +4. **Modality Alignment**: CT tool for CT image, MRI for MRI +5. **Dimension Match**: 3D tool for 3D volume +6. **Feature Coverage**: Specific capabilities mentioned + +#### Low Priority +7. **License**: Open-source preference (if no preference stated) +8. **Demo Availability**: Has runnable demo +9. **Popularity**: Community adoption + +### Explanation Generation + +Each recommendation includes explanation: + +**Good explanation template**: +``` +{Tool} is {specifically designed / well-suited} for {task} +on {modality} images. It supports {format} input {with/without} +preprocessing and provides {key features}. {Caveats if any}. +``` + +**Example**: +``` +TotalSegmentator is specifically designed for automated multi-organ +segmentation on CT scans. It supports DICOM input without preprocessing +and can segment 104 anatomical structures including lungs, air airways, +and vessels. It works best on whole-body CT but also performs well on +thoracic scans. +``` + +### Rank Assignment + +- **Rank 1**: Best overall match (highest accuracy score) +- **Rank 2**: Strong alternative or different approach +- **Rank 3**: Fallback option or specialized capability + +**Important**: Ranks are relative to **this specific query**, not absolute tool quality. + +## Model Configuration + +### Model Selection + +Available via `config.yaml`: + +```yaml +agent_model: + name: "gpt-4o-mini" + base_url: null + api_key_env: "OPENAI_API_KEY" +``` + +### Model Comparison + +| Model | Vision | Speed | Cost | Best For | +|-------|--------|-------|------|----------| +| gpt-4o-mini | โœ… | โšกโšกโšก | $ | Most queries, fast iteration | +| gpt-4o | โœ…โœ… | โšกโšก | $$ | Complex visual analysis | +| gpt-5.1 | โœ…โœ…โœ… | โšก | $$$ | Maximum accuracy needed | + +### Custom Endpoints + +Support for OpenAI-compatible APIs: + +```yaml +agent_model: + name: "llama-3.2-vision" + base_url: "https://inference.epfl.ch/v1" + api_key_env: "EPFL_API_KEY" +``` + +## Error Handling + +### Agent Errors + +**Tool quota exceeded** (handled gracefully in `run_agent`): +```python +except UsageLimitExceeded: + # Returns a ToolSelection with empty choices and an explanation + result = ToolSelection( + conversation=Conversation(status=ConversationStatus.COMPLETE, ...), + choices=[], + explanation="Tool call limit reached. Try a more specific query.", + ) +``` + +**Invalid structured output**: + +PydanticAI automatically retries the LLM call (up to `retries=2` per tool) if the model returns output that fails `ToolSelection` validation. The `ToolSelection.normalize()` model validator also auto-corrects minor inconsistencies. + +**API Errors**: +```python +except Exception as e: + log.warning(f"Agent execution encountered an error: {e}") + raise # propagated to the UI layer +``` + +### Graceful Degradation + +If the agent fails after all retries: + +1. Return empty `choices` with an `explanation` describing what was searched +2. UI surfaces the explanation so users can refine their query +3. Suggest manual exploration of the catalog + + + +## Testing + +### Unit Tests + +Test agent selection with PydanticAI's built-in test model (your catalog should contain the choice provided below, i.e. the `TotalSegmentator` tool): + +```python +from pydantic_ai import Agent +from pydantic_ai.models.test import TestModel +from ai_agent.generator.schema import ToolSelection, Conversation, ConversationStatus, ToolChoice +from ai_agent.agent.utils import AgentState + +def test_agent_selection(): + test_model = TestModel() + test_agent = Agent(model=test_model, deps_type=AgentState) + + mock_output = ToolSelection( + conversation=Conversation(status=ConversationStatus.COMPLETE), + choices=[ + ToolChoice(name="TotalSegmentator", rank=1, accuracy=95.0, why="Best CT segmenter") + ] + ) + + with test_agent.override(model=test_model): + result = test_agent.run_sync("segment lungs", deps=AgentState(), output_type=ToolSelection) + + assert result.output.conversation.status == ConversationStatus.COMPLETE + assert len(result.output.choices) == 1 + assert result.output.choices[0].rank == 1 +``` + +### Integration Tests + +Test with real VLM (expensive, slow): + +```python +@pytest.mark.integration +def test_real_agent(): + from ai_agent.agent.agent import run_agent + + with open("tests/data/sample.tif", "rb") as f: + image_bytes = f.read() + + result = run_agent( + task="I want to segment the lungs of this CT scan", + image_paths=["tests/data/sample.tif"], + image_bytes=image_bytes, + ) + + assert result.conversation.status == ConversationStatus.COMPLETE + assert len(result.choices) > 0 + assert all(0 <= c.accuracy <= 100 for c in result.choices) +``` + +## Next Steps + +- Learn about [Software Catalog](catalog.md) +- Return to [Architecture Overview](overview.md) +- Explore [Retrieval Pipeline](retrieval.md) diff --git a/docs/architecture/catalog.md b/docs/architecture/catalog.md index 71427a68a0532c313da73f6d8357983de14e1bd8..101d50d741a132d478a8aab73c24d2083e64922f 100644 --- a/docs/architecture/catalog.md +++ b/docs/architecture/catalog.md @@ -1,455 +1,455 @@ -# Software Catalog - -The software catalog is the foundation of the AI Imaging Agent, containing curated information about imaging analysis tools. - -## Overview - -**Format**: JSON Lines (JSONL) -**Location**: `dataset/catalog.jsonl` -**Schema**: Based on schema.org SoftwareSourceCode -**Size**: ~150 tools currently - -## Catalog Schema - -### Core Fields - -Based on [schema.org/SoftwareSourceCode](https://schema.org/SoftwareSourceCode): - -```json -{ - "@type": "SoftwareSourceCode", - "name": "TotalSegmentator", - "description": "Tool for automated segmentation of 104 anatomical structures", - "url": "https://github.com/wasserth/TotalSegmentator", - "codeRepository": "https://github.com/wasserth/TotalSegmentator", - "programmingLanguage": "Python", - "runtimePlatform": "PyTorch", - "license": "Apache-2.0", - "keywords": ["segmentation", "CT", "MRI", "medical-imaging"], - "applicationCategory": "Medical Imaging", - "operatingSystem": ["Linux", "Windows", "macOS"], - "softwareVersion": "2.0.0", - "datePublished": "2022-09-01", - "dateModified": "2024-01-15", - "author": { - "@type": "Person", - "name": "Jakob Wasserthal" - } -} -``` - -### Extended Fields - -Custom fields in `supportingData`: - -```json -{ - "supportingData": { - "modalities": ["CT", "MRI"], - "dimensions": ["3D"], - "formats": ["DICOM", "NIfTI", "PNG"], - "tasks": ["segmentation", "organ-segmentation"], - "demo_url": "https://huggingface.co/spaces/username/totalsegmentator", - "paper_url": "https://doi.org/10.1000/example", - "citations": 150, - "github_stars": 1200 - } -} -``` - -### Field Descriptions - -#### name -Canonical tool name (matches repository or published name) - -**Example**: `"TotalSegmentator"`, `"nnU-Net"`, `"MedSAM"` - -#### description -Brief description of tool's purpose and capabilities - -**Guidelines**: - -- 1-2 sentences -- Mention key features -- Include domain/modality if specific - -#### url -Primary landing page (usually GitHub repo) - -#### codeRepository -Source code repository URL (GitHub, GitLab, etc.) - -#### programmingLanguage -Primary language(s) - -**Common values**: `"Python"`, `"C++"`, `"JavaScript"`, `"Jupyter Notebook"` - -#### license -Software license identifier (SPDX format) - -**Common values**: - -- `"Apache-2.0"`: Permissive, commercial OK -- `"MIT"`: Very permissive -- `"GPL-3.0"`: Copyleft -- `"BSD-3-Clause"`: Permissive -- `"Proprietary"`: Restricted - -#### keywords -Array of relevant tags/keywords - -**Categories**: - -- **Tasks**: segmentation, classification, registration, detection -- **Modalities**: CT, MRI, X-ray, ultrasound, microscopy -- **Techniques**: deep-learning, traditional-cv, machine-learning -- **Domains**: medical-imaging, scientific-imaging, neuroscience - -#### supportingData.modalities -Medical imaging modalities supported - -**Standard values**: - -- `"CT"`: Computed Tomography -- `"MRI"`: Magnetic Resonance Imaging -- `"XR"`: X-ray radiography -- `"US"`: Ultrasound -- `"PET"`: Positron Emission Tomography -- `"SPECT"`: Single-Photon Emission CT -- `"OCT"`: Optical Coherence Tomography -- `"Microscopy"`: Various microscopy types - -#### supportingData.dimensions -Spatial dimensions supported - -**Values**: `["2D"]`, `["3D"]`, `["2D", "3D"]`, `["4D"]` - -- **2D**: Single slice images -- **3D**: Volumetric data -- **4D**: Time-series volumes (3D + time) - -#### supportingData.formats -File formats supported for input/output - -**Common values**: - -- Medical: `"DICOM"`, `"NIfTI"`, `"NRRD"`, `"Analyze"` -- Standard: `"PNG"`, `"JPEG"`, `"TIFF"`, `"BMP"` -- Scientific: `"HDF5"`, `"Zarr"`, `"OME-TIFF"` -- Other: `"NumPy"`, `"MAT"` - -#### supportingData.tasks -Analysis tasks the tool performs - -**Common values**: - -- `"segmentation"`: Image segmentation -- `"classification"`: Image classification -- `"detection"`: Object detection -- `"registration"`: Image registration/alignment -- `"reconstruction"`: 3D reconstruction -- `"enhancement"`: Image enhancement -- `"analysis"`: General analysis - -#### supportingData.demo_url -Link to runnable demo (HuggingFace Space, Colab, web app) - -**Preferred**: HuggingFace Gradio Spaces (best integration) - -**Example**: `"https://huggingface.co/spaces/username/toolname"` - -## Catalog Structure - -### File Format - -JSON Lines (JSONL): Each line is a complete JSON object - -```jsonl -{"@type": "SoftwareSourceCode", "name": "Tool1", ...} -{"@type": "SoftwareSourceCode", "name": "Tool2", ...} -{"@type": "SoftwareSourceCode", "name": "Tool3", ...} -``` - -**Benefits**: - -- Easy to append new tools -- Stream processing for large catalogs -- Each line independently parseable -- Git-friendly (line-based diffs) - -### Catalog Loading - -```python -import json - -def load_catalog(path: str) -> list[dict]: - tools = [] - with open(path) as f: - for line in f: - if line.strip(): - tools.append(json.loads(line)) - return tools -``` - -### Validation - -Tools are validated on load: - -```python -from pydantic import BaseModel, HttpUrl - -class SoftwareSourceCode(BaseModel): - name: str - description: str - url: HttpUrl - license: str - keywords: list[str] - supportingData: dict - - class Config: - extra = "allow" # Allow additional schema.org fields -``` - -## Catalog Management - -### Adding New Tools - -1. **Create entry** following schema: - -```json -{ - "@type": "SoftwareSourceCode", - "name": "NewTool", - "description": "Brief description of the tool", - "url": "https://github.com/user/newtool", - "codeRepository": "https://github.com/user/newtool", - "programmingLanguage": "Python", - "license": "MIT", - "keywords": ["segmentation", "CT"], - "supportingData": { - "modalities": ["CT"], - "dimensions": ["3D"], - "formats": ["DICOM", "NIfTI"], - "tasks": ["segmentation"], - "demo_url": "https://huggingface.co/spaces/user/newtool" - } -} -``` - -2. **Append to catalog.jsonl** (as single line, no pretty printing) - -3. **Update checksum**: - -```bash -shasum dataset/catalog.jsonl > dataset/catalog.jsonl.sha1 -``` - -4. **Sync catalog**: - -```bash -ai_agent sync -``` - -This rebuilds the embeddings and FAISS index. - -### Updating Existing Tools - -1. **Find tool** in `catalog.jsonl` -2. **Edit JSON** (update fields) -3. **Validate JSON** syntax -4. **Update checksum** and **sync** - -### Removing Tools - -1. **Delete line** from `catalog.jsonl` -2. **Update checksum** and **sync** - - - -## Synchronization - -### Auto-Sync - -Configured via `.env`: - -```dotenv -SYNC_EVERY_HOURS=24 -``` - -**Process**: -1. Background thread checks catalog every 24h -2. Compares SHA1 checksum -3. If changed: - - Reload catalog - - Re-embed all tools - - Rebuild FAISS index - -### Manual Sync - -```bash -ai_agent sync -``` - -**Output**: -``` -[sync] 150 โ†’ dataset/catalog.jsonl -[sync] Rebuilding embeddings... -[sync] Embedding 150 tools... (5.2s) -[sync] Building FAISS index... -[sync] Saved to artifacts/rag_index/ -[sync] Sync complete. -``` - -## Embeddings and Index - -### Embedding Process - -For each tool, create text representation: - -```python -tool_text = f"{tool['name']} {tool['description']} {' '.join(tool['keywords'])}" - -# Optional: Include supportingData -if 'supportingData' in tool: - sd = tool['supportingData'] - tool_text += f" {' '.join(sd.get('modalities', []))}" - tool_text += f" {' '.join(sd.get('tasks', []))}" - -# Embed -embedding = embedder.encode(tool_text, normalize_embeddings=True) -``` - -### Index Storage - -``` -artifacts/rag_index/ -โ”œโ”€โ”€ index.faiss # FAISS IndexFlatIP -โ””โ”€โ”€ meta.json # Tool IDs, config, timestamps -``` - -**meta.json** structure: - -```json -{ - "tool_ids": ["tool1", "tool2", ...], - "version": "1.0", - "embedding_model": "BAAI/bge-m3", - "embedding_dim": 1024, - "num_tools": 150, - "created_at": "2024-03-01T12:00:00Z", - "catalog_sha1": "abc123..." -} -``` - -## Quality Assurance - -### Validation Rules - -1. **Required fields**: name, description, url, license -2. **Valid URLs**: Well-formed HTTP/HTTPS URLs -3. **Standard licenses**: SPDX identifiers preferred -4. **Consistent keywords**: Use standard terminology -5. **Demo URLs**: Verify demos are live and accessible - -### Automated Checks - -```python -def validate_catalog(catalog_path): - errors = [] - - with open(catalog_path) as f: - for i, line in enumerate(f, 1): - try: - tool = json.loads(line) - - # Required fields - for field in ['name', 'description', 'url']: - if field not in tool: - errors.append(f"Line {i}: Missing {field}") - - # URL validation - if not tool['url'].startswith('http'): - errors.append(f"Line {i}: Invalid URL") - - # supportingData structure - if 'supportingData' in tool: - sd = tool['supportingData'] - if 'demo_url' in sd and sd['demo_url']: - if not sd['demo_url'].startswith('http'): - errors.append(f"Line {i}: Invalid demo_url") - - except json.JSONDecodeError as e: - errors.append(f"Line {i}: JSON syntax error - {e}") - - return errors -``` - -## Best Practices - -### Tool Descriptions - -โœ… **Good**: -``` -"Automated multi-organ segmentation for CT and MRI supporting 104 anatomical structures" -``` - -โŒ **Bad**: -``` -"A tool" # Too vague -"The best segmentation tool ever created with amazing accuracy..." # Too marketing-y -``` - -### Keywords - -โœ… **Good**: -``` -["segmentation", "CT", "MRI", "medical-imaging", "deep-learning", "organ-segmentation"] -``` - -โŒ **Bad**: -``` -["cool", "awesome", "the best"] # Not searchable terms -``` - -### Demo URLs - -โœ… **Preferred**: -- HuggingFace Gradio Spaces -- Google Colab notebooks -- Live web demos - -โŒ **Avoid**: -- Dead links -- Paywalled demos -- Demos requiring registration - -## Next Steps - -- Return to [Architecture Overview](overview.md) -- Learn about [Retrieval Pipeline](retrieval.md) -- Explore [Agent & VLM Selection](agent.md) +# Software Catalog + +The software catalog is the foundation of the AI Imaging Agent, containing curated information about imaging analysis tools. + +## Overview + +**Format**: JSON Lines (JSONL) +**Location**: `dataset/catalog.jsonl` +**Schema**: Based on schema.org SoftwareSourceCode +**Size**: ~150 tools currently + +## Catalog Schema + +### Core Fields + +Based on [schema.org/SoftwareSourceCode](https://schema.org/SoftwareSourceCode): + +```json +{ + "@type": "SoftwareSourceCode", + "name": "TotalSegmentator", + "description": "Tool for automated segmentation of 104 anatomical structures", + "url": "https://github.com/wasserth/TotalSegmentator", + "codeRepository": "https://github.com/wasserth/TotalSegmentator", + "programmingLanguage": "Python", + "runtimePlatform": "PyTorch", + "license": "Apache-2.0", + "keywords": ["segmentation", "CT", "MRI", "medical-imaging"], + "applicationCategory": "Medical Imaging", + "operatingSystem": ["Linux", "Windows", "macOS"], + "softwareVersion": "2.0.0", + "datePublished": "2022-09-01", + "dateModified": "2024-01-15", + "author": { + "@type": "Person", + "name": "Jakob Wasserthal" + } +} +``` + +### Extended Fields + +Custom fields in `supportingData`: + +```json +{ + "supportingData": { + "modalities": ["CT", "MRI"], + "dimensions": ["3D"], + "formats": ["DICOM", "NIfTI", "PNG"], + "tasks": ["segmentation", "organ-segmentation"], + "demo_url": "https://huggingface.co/spaces/username/totalsegmentator", + "paper_url": "https://doi.org/10.1000/example", + "citations": 150, + "github_stars": 1200 + } +} +``` + +### Field Descriptions + +#### name +Canonical tool name (matches repository or published name) + +**Example**: `"TotalSegmentator"`, `"nnU-Net"`, `"MedSAM"` + +#### description +Brief description of tool's purpose and capabilities + +**Guidelines**: + +- 1-2 sentences +- Mention key features +- Include domain/modality if specific + +#### url +Primary landing page (usually GitHub repo) + +#### codeRepository +Source code repository URL (GitHub, GitLab, etc.) + +#### programmingLanguage +Primary language(s) + +**Common values**: `"Python"`, `"C++"`, `"JavaScript"`, `"Jupyter Notebook"` + +#### license +Software license identifier (SPDX format) + +**Common values**: + +- `"Apache-2.0"`: Permissive, commercial OK +- `"MIT"`: Very permissive +- `"GPL-3.0"`: Copyleft +- `"BSD-3-Clause"`: Permissive +- `"Proprietary"`: Restricted + +#### keywords +Array of relevant tags/keywords + +**Categories**: + +- **Tasks**: segmentation, classification, registration, detection +- **Modalities**: CT, MRI, X-ray, ultrasound, microscopy +- **Techniques**: deep-learning, traditional-cv, machine-learning +- **Domains**: medical-imaging, scientific-imaging, neuroscience + +#### supportingData.modalities +Medical imaging modalities supported + +**Standard values**: + +- `"CT"`: Computed Tomography +- `"MRI"`: Magnetic Resonance Imaging +- `"XR"`: X-ray radiography +- `"US"`: Ultrasound +- `"PET"`: Positron Emission Tomography +- `"SPECT"`: Single-Photon Emission CT +- `"OCT"`: Optical Coherence Tomography +- `"Microscopy"`: Various microscopy types + +#### supportingData.dimensions +Spatial dimensions supported + +**Values**: `["2D"]`, `["3D"]`, `["2D", "3D"]`, `["4D"]` + +- **2D**: Single slice images +- **3D**: Volumetric data +- **4D**: Time-series volumes (3D + time) + +#### supportingData.formats +File formats supported for input/output + +**Common values**: + +- Medical: `"DICOM"`, `"NIfTI"`, `"NRRD"`, `"Analyze"` +- Standard: `"PNG"`, `"JPEG"`, `"TIFF"`, `"BMP"` +- Scientific: `"HDF5"`, `"Zarr"`, `"OME-TIFF"` +- Other: `"NumPy"`, `"MAT"` + +#### supportingData.tasks +Analysis tasks the tool performs + +**Common values**: + +- `"segmentation"`: Image segmentation +- `"classification"`: Image classification +- `"detection"`: Object detection +- `"registration"`: Image registration/alignment +- `"reconstruction"`: 3D reconstruction +- `"enhancement"`: Image enhancement +- `"analysis"`: General analysis + +#### supportingData.demo_url +Link to runnable demo (HuggingFace Space, Colab, web app) + +**Preferred**: HuggingFace Gradio Spaces (best integration) + +**Example**: `"https://huggingface.co/spaces/username/toolname"` + +## Catalog Structure + +### File Format + +JSON Lines (JSONL): Each line is a complete JSON object + +```jsonl +{"@type": "SoftwareSourceCode", "name": "Tool1", ...} +{"@type": "SoftwareSourceCode", "name": "Tool2", ...} +{"@type": "SoftwareSourceCode", "name": "Tool3", ...} +``` + +**Benefits**: + +- Easy to append new tools +- Stream processing for large catalogs +- Each line independently parseable +- Git-friendly (line-based diffs) + +### Catalog Loading + +```python +import json + +def load_catalog(path: str) -> list[dict]: + tools = [] + with open(path) as f: + for line in f: + if line.strip(): + tools.append(json.loads(line)) + return tools +``` + +### Validation + +Tools are validated on load: + +```python +from pydantic import BaseModel, HttpUrl + +class SoftwareSourceCode(BaseModel): + name: str + description: str + url: HttpUrl + license: str + keywords: list[str] + supportingData: dict + + class Config: + extra = "allow" # Allow additional schema.org fields +``` + +## Catalog Management + +### Adding New Tools + +1. **Create entry** following schema: + +```json +{ + "@type": "SoftwareSourceCode", + "name": "NewTool", + "description": "Brief description of the tool", + "url": "https://github.com/user/newtool", + "codeRepository": "https://github.com/user/newtool", + "programmingLanguage": "Python", + "license": "MIT", + "keywords": ["segmentation", "CT"], + "supportingData": { + "modalities": ["CT"], + "dimensions": ["3D"], + "formats": ["DICOM", "NIfTI"], + "tasks": ["segmentation"], + "demo_url": "https://huggingface.co/spaces/user/newtool" + } +} +``` + +2. **Append to catalog.jsonl** (as single line, no pretty printing) + +3. **Update checksum**: + +```bash +shasum dataset/catalog.jsonl > dataset/catalog.jsonl.sha1 +``` + +4. **Sync catalog**: + +```bash +ai_agent sync +``` + +This rebuilds the embeddings and FAISS index. + +### Updating Existing Tools + +1. **Find tool** in `catalog.jsonl` +2. **Edit JSON** (update fields) +3. **Validate JSON** syntax +4. **Update checksum** and **sync** + +### Removing Tools + +1. **Delete line** from `catalog.jsonl` +2. **Update checksum** and **sync** + + + +## Synchronization + +### Auto-Sync + +Configured via `.env`: + +```dotenv +SYNC_EVERY_HOURS=24 +``` + +**Process**: +1. Background thread checks catalog every 24h +2. Compares SHA1 checksum +3. If changed: + - Reload catalog + - Re-embed all tools + - Rebuild FAISS index + +### Manual Sync + +```bash +ai_agent sync +``` + +**Output**: +``` +[sync] 150 โ†’ dataset/catalog.jsonl +[sync] Rebuilding embeddings... +[sync] Embedding 150 tools... (5.2s) +[sync] Building FAISS index... +[sync] Saved to artifacts/rag_index/ +[sync] Sync complete. +``` + +## Embeddings and Index + +### Embedding Process + +For each tool, create text representation: + +```python +tool_text = f"{tool['name']} {tool['description']} {' '.join(tool['keywords'])}" + +# Optional: Include supportingData +if 'supportingData' in tool: + sd = tool['supportingData'] + tool_text += f" {' '.join(sd.get('modalities', []))}" + tool_text += f" {' '.join(sd.get('tasks', []))}" + +# Embed +embedding = embedder.encode(tool_text, normalize_embeddings=True) +``` + +### Index Storage + +``` +artifacts/rag_index/ +โ”œโ”€โ”€ index.faiss # FAISS IndexFlatIP +โ””โ”€โ”€ meta.json # Tool IDs, config, timestamps +``` + +**meta.json** structure: + +```json +{ + "tool_ids": ["tool1", "tool2", ...], + "version": "1.0", + "embedding_model": "BAAI/bge-m3", + "embedding_dim": 1024, + "num_tools": 150, + "created_at": "2024-03-01T12:00:00Z", + "catalog_sha1": "abc123..." +} +``` + +## Quality Assurance + +### Validation Rules + +1. **Required fields**: name, description, url, license +2. **Valid URLs**: Well-formed HTTP/HTTPS URLs +3. **Standard licenses**: SPDX identifiers preferred +4. **Consistent keywords**: Use standard terminology +5. **Demo URLs**: Verify demos are live and accessible + +### Automated Checks + +```python +def validate_catalog(catalog_path): + errors = [] + + with open(catalog_path) as f: + for i, line in enumerate(f, 1): + try: + tool = json.loads(line) + + # Required fields + for field in ['name', 'description', 'url']: + if field not in tool: + errors.append(f"Line {i}: Missing {field}") + + # URL validation + if not tool['url'].startswith('http'): + errors.append(f"Line {i}: Invalid URL") + + # supportingData structure + if 'supportingData' in tool: + sd = tool['supportingData'] + if 'demo_url' in sd and sd['demo_url']: + if not sd['demo_url'].startswith('http'): + errors.append(f"Line {i}: Invalid demo_url") + + except json.JSONDecodeError as e: + errors.append(f"Line {i}: JSON syntax error - {e}") + + return errors +``` + +## Best Practices + +### Tool Descriptions + +โœ… **Good**: +``` +"Automated multi-organ segmentation for CT and MRI supporting 104 anatomical structures" +``` + +โŒ **Bad**: +``` +"A tool" # Too vague +"The best segmentation tool ever created with amazing accuracy..." # Too marketing-y +``` + +### Keywords + +โœ… **Good**: +``` +["segmentation", "CT", "MRI", "medical-imaging", "deep-learning", "organ-segmentation"] +``` + +โŒ **Bad**: +``` +["cool", "awesome", "the best"] # Not searchable terms +``` + +### Demo URLs + +โœ… **Preferred**: +- HuggingFace Gradio Spaces +- Google Colab notebooks +- Live web demos + +โŒ **Avoid**: +- Dead links +- Paywalled demos +- Demos requiring registration + +## Next Steps + +- Return to [Architecture Overview](overview.md) +- Learn about [Retrieval Pipeline](retrieval.md) +- Explore [Agent & VLM Selection](agent.md) diff --git a/docs/architecture/overview.md b/docs/architecture/overview.md index 6f59b39c176ce1e943905c6672f6cf61514f469e..023d7aac7dc2abf1c1c93c17feaa02b47d4a9358 100644 --- a/docs/architecture/overview.md +++ b/docs/architecture/overview.md @@ -1,459 +1,459 @@ -# 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) +# 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) diff --git a/docs/architecture/retrieval.md b/docs/architecture/retrieval.md index 0b38aeb5f7974764cd00d98052db61f087dbf727..3e87de9d8e162d1f624e272d54d188358e11addd 100644 --- a/docs/architecture/retrieval.md +++ b/docs/architecture/retrieval.md @@ -1,387 +1,387 @@ -# Retrieval Pipeline - -The retrieval stage is the first phase of the AI Imaging Agent's two-stage pipeline. It performs fast text-based search to find candidate tools from the software catalog. - -## Overview - -**Goal**: Quickly narrow down the software catalog to most relevant candidates - -**Characteristics**: - -- โšก Fast (~100-300ms) -- ๐Ÿ”ข Deterministic and reproducible -- ๐Ÿšซ No LLM calls -- ๐Ÿ’ฐ Low cost (no API fees) - -## Pipeline Stages - -```mermaid -graph LR - A[User Query] --> B[Query Enhancement] - B --> C[Embedding] - C --> D[FAISS Search] - D --> E[CrossEncoder Rerank] - E --> F[Top-K Candidates] -``` - -## Step 1: Query Enhancement - -### Format Token Injection - -When users upload files, format tokens are added to the query: - -```python -# User uploads: scan.dcm -# User query: "segment lungs" - -# Enhanced query: -"segment lungs format:DICOM format:CT format:3D" -``` - -**Format tokens added**: - -- File extension (`format:DICOM`, `format:NIfTI`) -- Image modality from metadata (`format:CT`, `format:MRI`) -- Dimensionality (`format:2D`, `format:3D`, `format:4D`) - -**Why this helps**: - -- Matches tools that support specific formats -- Boosts DICOM-compatible tools for DICOM input -- Ensures dimension compatibility (3D tools for volumes) - -### Control Tag Processing - -Special tags are extracted and processed: - -```python -query = "segment lungs [EXCLUDE:tool1|tool2]" - -# Extracted: -clean_query = "segment lungs" -excluded_tools = ["tool1", "tool2"] -``` - -**Supported tags**: -- `[EXCLUDE:tool1|tool2]`: Filter tools from results - -## Step 2: Metadata-Aware Querying - -The pipeline does not perform semantic vocabulary expansion. Instead, retrieval combines: - -- cleaned task text -- format tokens inferred from uploaded files (for example `format:DICOM`) -- compact image metadata hints (modality/anatomy/dimensionality when available) - -This keeps retrieval deterministic and closely tied to the user's data. - -### Alternative Query Generation - -On retry (when initial results < 5 tools): - -```python -# Initial query -query1 = "segment rare pulmonary structure" -results = 2 tools # Too few! - -# Retry 1: Broader formulation (keep first 2-3 words) -query2 = "segment rare pulmonary" -results = 7 tools # Better! - -# Retry 2: If still insufficient, repeat with same broadening strategy -``` - -**Max retries**: 2 - -## Step 3: Embedding - -### BGE-M3 Model - -**Model**: `BAAI/bge-m3` - -**Characteristics**: - -- Multilingual (but used for English) -- 1024-dimensional embeddings -- Trained for retrieval tasks -- Fast inference (~10ms per query) - -**Embedding process**: - -```python -from sentence_transformers import SentenceTransformer - -model = SentenceTransformer("BAAI/bge-m3") -query_vector = model.encode( - query, - normalize_embeddings=True # L2 normalization for cosine similarity -) -# Returns: np.array of shape (1024,) -``` - -### Catalog Embedding - -Software tools are pre-embedded during indexing: - -```python -# For each tool in catalog: -tool_text = f"{tool.name} {tool.description} {' '.join(tool.keywords)}" -tool_vector = model.encode(tool_text, normalize_embeddings=True) - -# Store in FAISS index -faiss_index.add(tool_vector) -``` - -**Index structure**: - -- FAISS IndexFlatIP (inner product = cosine similarity for normalized vectors) -- ~150 tools in current catalog -- Index size: ~600KB - -## Step 4: FAISS Search - -### Vector Search - -FAISS performs fast similarity search: - -```python -import faiss - -# Search for top 20 most similar tools -scores, indices = faiss_index.search( - query_vector.reshape(1, -1), - k=20 -) - -# Returns: -# scores: [0.85, 0.82, 0.79, ...] # Cosine similarities -# indices: [42, 17, 89, ...] # Tool IDs in catalog -``` - -**Search algorithm**: - -- IndexFlatIP: Exact search (brute force) -- Fast for catalog size (~150 tools) -- Could use IVF for larger catalogs (>10k tools) - -**Why top-20**: - -- More candidates than needed (default final: 8) -- Provides options for reranking -- Balances recall vs. later stage cost - -### Candidate Retrieval - -```python -candidates = [catalog[idx] for idx in indices[:20]] -candidate_scores = scores[:20].tolist() - -# Example candidates: -[ - { - "name": "TotalSegmentator", - "score": 0.85, - "description": "Automated multi-organ segmentation...", - ... - }, - ... -] -``` - -## Step 5: CrossEncoder Reranking - -### Why Rerank? - -**BiEncoder (BGE-M3)** limitations: - -- Encodes query and documents independently -- No query-document interaction -- Misses subtle relevance signals - -**CrossEncoder** benefits: - -- Jointly encodes query + document -- Cross-attention between query and doc -- More accurate relevance scoring -- Slower (not suitable for entire catalog) - -### Reranking Model - -**Model**: `BAAI/bge-reranker-v2-m3` - -**Characteristics**: - -- Trained on MS-MARCO passage ranking -- 6 layers, fast inference (~50ms per pair) -- Direct relevance score (no embedding) - -**Reranking process**: - -```python -from sentence_transformers import CrossEncoder - -reranker = CrossEncoder("BAAI/bge-reranker-v2-m3") - -# Score each (query, candidate) pair -pairs = [(query, candidate.description) for candidate in candidates] -rerank_scores = reranker.predict(pairs) - -# Re-sort by rerank scores -sorted_indices = np.argsort(rerank_scores)[::-1] -reranked_candidates = [candidates[i] for i in sorted_indices][:8] -``` - -**Output**: Top-8 candidates with refined ranking - -## Output Format - -### Candidate Schema - -Each candidate passed to Stage 2: - -```python -{ - "name": "TotalSegmentator", - "description": "Automated multi-organ segmentation for CT and MRI", - "url": "https://github.com/wasserth/TotalSegmentator", - "keywords": ["segmentation", "medical-imaging", "CT", "MRI"], - "license": "Apache-2.0", - "supporting_data": { - "modalities": ["CT", "MRI"], - "dimensions": ["3D"], - "formats": ["DICOM", "NIfTI"], - "demo_url": "https://huggingface.co/spaces/..." - }, - "retrieval_score": 0.85, # FAISS or rerank score -} -``` - -**Fields used by VLM**: - -- Essential for understanding tool capability -- Formatted as table in VLM prompt -- Enables comparative reasoning - -## Index Management - -### Building the Index - -Done during catalog sync: - -```bash -ai_agent sync -``` - -**Process**: - -1. Load catalog JSONL -2. Embed each tool description -3. Build FAISS index -4. Save to disk: `artifacts/rag_index/` - -**Files**: -``` -artifacts/rag_index/ -โ”œโ”€โ”€ index.faiss # FAISS binary index -โ””โ”€โ”€ meta.json # Metadata (tool IDs, config) -``` - -### Loading the Index - -At startup: - -```python -from retriever.vector_index import VectorIndex - -index = VectorIndex() -index.load("artifacts/rag_index") - -# Ready for queries -results = index.search(query, k=20) -``` - -### Updating the Index - -When catalog changes: -1. Sync detects new/modified tools -2. Re-embed entire catalog (fast, ~2 seconds) -3. Rebuild FAISS index -4. Reload in pipeline (no restart needed) - -## Performance Optimization - -### Caching - -**Model loading**: - -- BGE-M3 and CrossEncoder loaded once at startup -- Kept in memory for entire session - -**Index loading**: - -- FAISS index loaded once -- Small enough to fit in memory (~MB) - -### Batch Processing - -For multiple queries (testing, batch mode): - -```python -# Batch embed multiple queries -query_vectors = model.encode(queries, batch_size=32) - -# Batch FAISS search -scores, indices = index.search(query_vectors, k=20) -``` - -### GPU Acceleration - -Models can use GPU if available: - -```python -model = SentenceTransformer("BAAI/bge-m3", device="cuda") -reranker = CrossEncoder("...", device="cuda") -``` - -## Retrieval Metrics - -### Monitored Metrics - -During retrieval: - -- **Number of candidates found**: Should be โ‰ฅ8 for good coverage -- **Average similarity score**: Higher = better match -- **Reranking impact**: Score change after reranking -- **Retry usage**: Whether broadening retry was triggered - -### Logging - -Retrieval events logged: - -``` -INFO retriever.vector_index: FAISS search: query="segment lungs", results=20, top_score=0.85 -INFO retriever.reranker: Reranking 20 candidates, top_score_change=+0.12 -INFO retriever.pipeline: Final candidates: 8, avg_score=0.78 -``` - -## Limitations - -### Current Limitations - -1. **English only**: No multilingual support (though model is capable) -2. **Small catalog**: ~150 tools (FAISS overkill, but scales) -3. **No filtering**: Can't filter by license, modality in retrieval (done in Stage 2) -4. **Heuristic retries**: Broadening strategy is simple prefix-based shortening - -### Future Enhancements - -- **Hybrid search**: Combine semantic + keyword (BM25) -- **Metadata filters**: Pre-filter by modality, license, format -- **Personalization**: User history, preferences -- **Adaptive retries**: Learn better broadening formulations from query logs - -## Next Steps - -- Learn about [Agent & VLM Selection](agent.md) -- Explore [Software Catalog](catalog.md) -- Return to [Architecture Overview](overview.md) +# Retrieval Pipeline + +The retrieval stage is the first phase of the AI Imaging Agent's two-stage pipeline. It performs fast text-based search to find candidate tools from the software catalog. + +## Overview + +**Goal**: Quickly narrow down the software catalog to most relevant candidates + +**Characteristics**: + +- โšก Fast (~100-300ms) +- ๐Ÿ”ข Deterministic and reproducible +- ๐Ÿšซ No LLM calls +- ๐Ÿ’ฐ Low cost (no API fees) + +## Pipeline Stages + +```mermaid +graph LR + A[User Query] --> B[Query Enhancement] + B --> C[Embedding] + C --> D[FAISS Search] + D --> E[CrossEncoder Rerank] + E --> F[Top-K Candidates] +``` + +## Step 1: Query Enhancement + +### Format Token Injection + +When users upload files, format tokens are added to the query: + +```python +# User uploads: scan.dcm +# User query: "segment lungs" + +# Enhanced query: +"segment lungs format:DICOM format:CT format:3D" +``` + +**Format tokens added**: + +- File extension (`format:DICOM`, `format:NIfTI`) +- Image modality from metadata (`format:CT`, `format:MRI`) +- Dimensionality (`format:2D`, `format:3D`, `format:4D`) + +**Why this helps**: + +- Matches tools that support specific formats +- Boosts DICOM-compatible tools for DICOM input +- Ensures dimension compatibility (3D tools for volumes) + +### Control Tag Processing + +Special tags are extracted and processed: + +```python +query = "segment lungs [EXCLUDE:tool1|tool2]" + +# Extracted: +clean_query = "segment lungs" +excluded_tools = ["tool1", "tool2"] +``` + +**Supported tags**: +- `[EXCLUDE:tool1|tool2]`: Filter tools from results + +## Step 2: Metadata-Aware Querying + +The pipeline does not perform semantic vocabulary expansion. Instead, retrieval combines: + +- cleaned task text +- format tokens inferred from uploaded files (for example `format:DICOM`) +- compact image metadata hints (modality/anatomy/dimensionality when available) + +This keeps retrieval deterministic and closely tied to the user's data. + +### Alternative Query Generation + +On retry (when initial results < 5 tools): + +```python +# Initial query +query1 = "segment rare pulmonary structure" +results = 2 tools # Too few! + +# Retry 1: Broader formulation (keep first 2-3 words) +query2 = "segment rare pulmonary" +results = 7 tools # Better! + +# Retry 2: If still insufficient, repeat with same broadening strategy +``` + +**Max retries**: 2 + +## Step 3: Embedding + +### BGE-M3 Model + +**Model**: `BAAI/bge-m3` + +**Characteristics**: + +- Multilingual (but used for English) +- 1024-dimensional embeddings +- Trained for retrieval tasks +- Fast inference (~10ms per query) + +**Embedding process**: + +```python +from sentence_transformers import SentenceTransformer + +model = SentenceTransformer("BAAI/bge-m3") +query_vector = model.encode( + query, + normalize_embeddings=True # L2 normalization for cosine similarity +) +# Returns: np.array of shape (1024,) +``` + +### Catalog Embedding + +Software tools are pre-embedded during indexing: + +```python +# For each tool in catalog: +tool_text = f"{tool.name} {tool.description} {' '.join(tool.keywords)}" +tool_vector = model.encode(tool_text, normalize_embeddings=True) + +# Store in FAISS index +faiss_index.add(tool_vector) +``` + +**Index structure**: + +- FAISS IndexFlatIP (inner product = cosine similarity for normalized vectors) +- ~150 tools in current catalog +- Index size: ~600KB + +## Step 4: FAISS Search + +### Vector Search + +FAISS performs fast similarity search: + +```python +import faiss + +# Search for top 20 most similar tools +scores, indices = faiss_index.search( + query_vector.reshape(1, -1), + k=20 +) + +# Returns: +# scores: [0.85, 0.82, 0.79, ...] # Cosine similarities +# indices: [42, 17, 89, ...] # Tool IDs in catalog +``` + +**Search algorithm**: + +- IndexFlatIP: Exact search (brute force) +- Fast for catalog size (~150 tools) +- Could use IVF for larger catalogs (>10k tools) + +**Why top-20**: + +- More candidates than needed (default final: 8) +- Provides options for reranking +- Balances recall vs. later stage cost + +### Candidate Retrieval + +```python +candidates = [catalog[idx] for idx in indices[:20]] +candidate_scores = scores[:20].tolist() + +# Example candidates: +[ + { + "name": "TotalSegmentator", + "score": 0.85, + "description": "Automated multi-organ segmentation...", + ... + }, + ... +] +``` + +## Step 5: CrossEncoder Reranking + +### Why Rerank? + +**BiEncoder (BGE-M3)** limitations: + +- Encodes query and documents independently +- No query-document interaction +- Misses subtle relevance signals + +**CrossEncoder** benefits: + +- Jointly encodes query + document +- Cross-attention between query and doc +- More accurate relevance scoring +- Slower (not suitable for entire catalog) + +### Reranking Model + +**Model**: `BAAI/bge-reranker-v2-m3` + +**Characteristics**: + +- Trained on MS-MARCO passage ranking +- 6 layers, fast inference (~50ms per pair) +- Direct relevance score (no embedding) + +**Reranking process**: + +```python +from sentence_transformers import CrossEncoder + +reranker = CrossEncoder("BAAI/bge-reranker-v2-m3") + +# Score each (query, candidate) pair +pairs = [(query, candidate.description) for candidate in candidates] +rerank_scores = reranker.predict(pairs) + +# Re-sort by rerank scores +sorted_indices = np.argsort(rerank_scores)[::-1] +reranked_candidates = [candidates[i] for i in sorted_indices][:8] +``` + +**Output**: Top-8 candidates with refined ranking + +## Output Format + +### Candidate Schema + +Each candidate passed to Stage 2: + +```python +{ + "name": "TotalSegmentator", + "description": "Automated multi-organ segmentation for CT and MRI", + "url": "https://github.com/wasserth/TotalSegmentator", + "keywords": ["segmentation", "medical-imaging", "CT", "MRI"], + "license": "Apache-2.0", + "supporting_data": { + "modalities": ["CT", "MRI"], + "dimensions": ["3D"], + "formats": ["DICOM", "NIfTI"], + "demo_url": "https://huggingface.co/spaces/..." + }, + "retrieval_score": 0.85, # FAISS or rerank score +} +``` + +**Fields used by VLM**: + +- Essential for understanding tool capability +- Formatted as table in VLM prompt +- Enables comparative reasoning + +## Index Management + +### Building the Index + +Done during catalog sync: + +```bash +ai_agent sync +``` + +**Process**: + +1. Load catalog JSONL +2. Embed each tool description +3. Build FAISS index +4. Save to disk: `artifacts/rag_index/` + +**Files**: +``` +artifacts/rag_index/ +โ”œโ”€โ”€ index.faiss # FAISS binary index +โ””โ”€โ”€ meta.json # Metadata (tool IDs, config) +``` + +### Loading the Index + +At startup: + +```python +from retriever.vector_index import VectorIndex + +index = VectorIndex() +index.load("artifacts/rag_index") + +# Ready for queries +results = index.search(query, k=20) +``` + +### Updating the Index + +When catalog changes: +1. Sync detects new/modified tools +2. Re-embed entire catalog (fast, ~2 seconds) +3. Rebuild FAISS index +4. Reload in pipeline (no restart needed) + +## Performance Optimization + +### Caching + +**Model loading**: + +- BGE-M3 and CrossEncoder loaded once at startup +- Kept in memory for entire session + +**Index loading**: + +- FAISS index loaded once +- Small enough to fit in memory (~MB) + +### Batch Processing + +For multiple queries (testing, batch mode): + +```python +# Batch embed multiple queries +query_vectors = model.encode(queries, batch_size=32) + +# Batch FAISS search +scores, indices = index.search(query_vectors, k=20) +``` + +### GPU Acceleration + +Models can use GPU if available: + +```python +model = SentenceTransformer("BAAI/bge-m3", device="cuda") +reranker = CrossEncoder("...", device="cuda") +``` + +## Retrieval Metrics + +### Monitored Metrics + +During retrieval: + +- **Number of candidates found**: Should be โ‰ฅ8 for good coverage +- **Average similarity score**: Higher = better match +- **Reranking impact**: Score change after reranking +- **Retry usage**: Whether broadening retry was triggered + +### Logging + +Retrieval events logged: + +``` +INFO retriever.vector_index: FAISS search: query="segment lungs", results=20, top_score=0.85 +INFO retriever.reranker: Reranking 20 candidates, top_score_change=+0.12 +INFO retriever.pipeline: Final candidates: 8, avg_score=0.78 +``` + +## Limitations + +### Current Limitations + +1. **English only**: No multilingual support (though model is capable) +2. **Small catalog**: ~150 tools (FAISS overkill, but scales) +3. **No filtering**: Can't filter by license, modality in retrieval (done in Stage 2) +4. **Heuristic retries**: Broadening strategy is simple prefix-based shortening + +### Future Enhancements + +- **Hybrid search**: Combine semantic + keyword (BM25) +- **Metadata filters**: Pre-filter by modality, license, format +- **Personalization**: User history, preferences +- **Adaptive retries**: Learn better broadening formulations from query logs + +## Next Steps + +- Learn about [Agent & VLM Selection](agent.md) +- Explore [Software Catalog](catalog.md) +- Return to [Architecture Overview](overview.md) diff --git a/docs/development/contributing.md b/docs/development/contributing.md index 43ecb114b8a99d1d30958d9b7a9ac232cfa71a24..bf39e2e3a77aeb77cac40650496dc28af59fe8f6 100644 --- a/docs/development/contributing.md +++ b/docs/development/contributing.md @@ -1,315 +1,315 @@ -# Contributing - -Thank you for your interest in contributing to the AI Imaging Agent! This guide will help you get started. - -## Getting Started - -### 1. Fork and Clone - -```bash -# Fork the repository on GitHub -# Then clone your fork -git clone https://github.com/YOUR_USERNAME/ai-agent.git -cd ai-agent -``` - -### 2. Set Up Development Environment - -```bash -# Create virtual environment -python -m venv .venv -source .venv/bin/activate # On Windows: .venv\Scripts\activate - -# Install in development mode -pip install -e ".[dev]" -``` - -### 3. Create a Branch - -```bash -git checkout -b feature/your-feature-name -# or -git checkout -b fix/issue-description -``` - -## Development Workflow - -### Making Changes - -1. **Make your changes** in the appropriate module -2. **Test your changes** (see [Testing](testing.md)) -3. **Update documentation** if needed -4. **Update CHANGELOG.md** following [Keep a Changelog](https://keepachangelog.com/) format - - - -#### Type Checking - -MyPy for type checking: - -```bash -mypy src/ -``` - -### Running Tests - -```bash -# Run all tests -pytest tests/ - -# Run specific test file -pytest tests/test_retrieval_pipeline.py - -# Run with coverage -pytest --cov=ai_agent tests/ -``` - -## Contribution Guidelines - -### Code Quality - -- **Follow PEP 8**: Use Black for formatting -- **Type hints**: Add type annotations to new functions -- **Docstrings**: Document all public functions and classes -- **Tests**: Add tests for new functionality -- **No warnings**: Fix any linter warnings before submitting - -### Commit Messages - -Use clear, descriptive commit messages: - -``` -feat: Add alternative search tool for agent - -- Implement search_alternative tool -- Add retry logic for insufficient results -- Update agent prompt with tool description -``` - -**Format**: - -- `feat:` New features -- `fix:` Bug fixes -- `docs:` Documentation changes -- `test:` Test additions/changes -- `refactor:` Code refactoring -- `chore:` Maintenance tasks - -### Pull Requests - -1. **Update CHANGELOG.md** under `[Unreleased]` section -2. **Write clear PR description** explaining changes -3. **Link related issues** if applicable -4. **Request review** from maintainers - -**PR description template**: - -```markdown -## Description -Brief description of changes - -## Type of Change -- [ ] Bug fix -- [ ] New feature -- [ ] Documentation update -- [ ] Refactoring - -## Testing -How to test these changes - -## Checklist -- [ ] Code follows project style -- [ ] Tests added/updated -- [ ] Documentation updated -- [ ] CHANGELOG.md updated -``` - -## Areas to Contribute - -### High Priority - -- **Catalog expansion**: Add more imaging tools -- **Demo integration**: Improve Gradio Space execution and add new spaces for tools -- **Format support**: Add new image format handlers - - - -### Feature Requests - -Before implementing new features: - -1. **Check existing issues** for similar requests -2. **Open an issue** to discuss the feature -3. **Get feedback** from maintainers -4. **Implement** after discussion - -## Documentation - -### Writing Documentation - -Documentation lives in `docs/` and uses MkDocs Material. - -```bash -# Install MkDocs -pip install mkdocs-material - -# Serve locally -mkdocs serve - -# Open http://127.0.0.1:8000 -``` - -### Documentation Style - -- Use **clear headings** and structure -- Include **code examples** where relevant -- Add **warnings** and **tips** for important information -- Keep **language simple** and accessible - -## Adding Tools to Catalog - -### Process - -1. **Create tool entry** in `dataset/catalog.jsonl`: - -```json -{ - "@type": "SoftwareSourceCode", - "name": "ToolName", - "description": "Tool description", - "url": "https://github.com/user/tool", - "license": "Apache-2.0", - "keywords": ["segmentation", "CT"], - "supportingData": { - "modalities": ["CT"], - "dimensions": ["3D"], - "formats": ["DICOM"], - "demo_url": "https://huggingface.co/spaces/user/tool" - } -} -``` - -2. **Validate entry**: - -```bash -# Check JSON syntax -python -c "import json; print(json.loads('YOUR_JSON_HERE'))" -``` - -3. **Update checksum**: - -```bash -shasum dataset/catalog.jsonl > dataset/catalog.jsonl.sha1 -``` - -4. **Sync catalog**: - -```bash -ai_agent sync -``` - -5. **Test retrieval**: - -```bash -ai_agent chat -# Try queries that should return your new tool -``` - -### Tool Criteria - -Tools should: - -- โœ… Be relevant to imaging analysis -- โœ… Be actively maintained -- โœ… Have clear documentation -- โœ… Preferably have a runnable demo -- โœ… Be open-source or have free tier - -## Reporting Issues - -### Bug Reports - -Include: - -- **Description** of the bug -- **Steps to reproduce** -- **Expected behavior** -- **Actual behavior** -- **Environment** (OS, Python version, etc.) -- **Logs** if available - -### Feature Requests - -Include: - -- **Use case** for the feature -- **Proposed solution** (if you have one) -- **Alternatives considered** -- **Examples** of similar features - -## Code Review Process - -1. **Automated checks** run on PR (tests, linting) -2. **Maintainer review** provides feedback -3. **Address feedback** and update PR -4. **Approval** from maintainer(s) -5. **Merge** into main branch - -## Release Process - -Releases follow semantic versioning: - -1. Update version in `pyproject.toml` -2. Update `CHANGELOG.md` with release date -3. Create git tag: `git tag v0.2.0` -4. Push tag: `git push origin v0.2.0` -5. GitHub Actions deploys documentation - -## Getting Help - -- **GitHub Discussions**: Ask questions -- **GitHub Issues**: Report bugs - -## Code of Conduct - -Be respectful and professional: - -- Use welcoming and inclusive language -- Respect differing viewpoints -- Accept constructive criticism -- Focus on what's best for the community - -## License - -By contributing, you agree that your contributions will be licensed under the Apache 2.0 License. - -## Next Steps - -- Review [Project Structure](structure.md) -- Learn about [Testing](testing.md) -- Read [Architecture Overview](../architecture/overview.md) +# Contributing + +Thank you for your interest in contributing to the AI Imaging Agent! This guide will help you get started. + +## Getting Started + +### 1. Fork and Clone + +```bash +# Fork the repository on GitHub +# Then clone your fork +git clone https://github.com/YOUR_USERNAME/ai-agent.git +cd ai-agent +``` + +### 2. Set Up Development Environment + +```bash +# Create virtual environment +python -m venv .venv +source .venv/bin/activate # On Windows: .venv\Scripts\activate + +# Install in development mode +pip install -e ".[dev]" +``` + +### 3. Create a Branch + +```bash +git checkout -b feature/your-feature-name +# or +git checkout -b fix/issue-description +``` + +## Development Workflow + +### Making Changes + +1. **Make your changes** in the appropriate module +2. **Test your changes** (see [Testing](testing.md)) +3. **Update documentation** if needed +4. **Update CHANGELOG.md** following [Keep a Changelog](https://keepachangelog.com/) format + + + +#### Type Checking + +MyPy for type checking: + +```bash +mypy src/ +``` + +### Running Tests + +```bash +# Run all tests +pytest tests/ + +# Run specific test file +pytest tests/test_retrieval_pipeline.py + +# Run with coverage +pytest --cov=ai_agent tests/ +``` + +## Contribution Guidelines + +### Code Quality + +- **Follow PEP 8**: Use Black for formatting +- **Type hints**: Add type annotations to new functions +- **Docstrings**: Document all public functions and classes +- **Tests**: Add tests for new functionality +- **No warnings**: Fix any linter warnings before submitting + +### Commit Messages + +Use clear, descriptive commit messages: + +``` +feat: Add alternative search tool for agent + +- Implement search_alternative tool +- Add retry logic for insufficient results +- Update agent prompt with tool description +``` + +**Format**: + +- `feat:` New features +- `fix:` Bug fixes +- `docs:` Documentation changes +- `test:` Test additions/changes +- `refactor:` Code refactoring +- `chore:` Maintenance tasks + +### Pull Requests + +1. **Update CHANGELOG.md** under `[Unreleased]` section +2. **Write clear PR description** explaining changes +3. **Link related issues** if applicable +4. **Request review** from maintainers + +**PR description template**: + +```markdown +## Description +Brief description of changes + +## Type of Change +- [ ] Bug fix +- [ ] New feature +- [ ] Documentation update +- [ ] Refactoring + +## Testing +How to test these changes + +## Checklist +- [ ] Code follows project style +- [ ] Tests added/updated +- [ ] Documentation updated +- [ ] CHANGELOG.md updated +``` + +## Areas to Contribute + +### High Priority + +- **Catalog expansion**: Add more imaging tools +- **Demo integration**: Improve Gradio Space execution and add new spaces for tools +- **Format support**: Add new image format handlers + + + +### Feature Requests + +Before implementing new features: + +1. **Check existing issues** for similar requests +2. **Open an issue** to discuss the feature +3. **Get feedback** from maintainers +4. **Implement** after discussion + +## Documentation + +### Writing Documentation + +Documentation lives in `docs/` and uses MkDocs Material. + +```bash +# Install MkDocs +pip install mkdocs-material + +# Serve locally +mkdocs serve + +# Open http://127.0.0.1:8000 +``` + +### Documentation Style + +- Use **clear headings** and structure +- Include **code examples** where relevant +- Add **warnings** and **tips** for important information +- Keep **language simple** and accessible + +## Adding Tools to Catalog + +### Process + +1. **Create tool entry** in `dataset/catalog.jsonl`: + +```json +{ + "@type": "SoftwareSourceCode", + "name": "ToolName", + "description": "Tool description", + "url": "https://github.com/user/tool", + "license": "Apache-2.0", + "keywords": ["segmentation", "CT"], + "supportingData": { + "modalities": ["CT"], + "dimensions": ["3D"], + "formats": ["DICOM"], + "demo_url": "https://huggingface.co/spaces/user/tool" + } +} +``` + +2. **Validate entry**: + +```bash +# Check JSON syntax +python -c "import json; print(json.loads('YOUR_JSON_HERE'))" +``` + +3. **Update checksum**: + +```bash +shasum dataset/catalog.jsonl > dataset/catalog.jsonl.sha1 +``` + +4. **Sync catalog**: + +```bash +ai_agent sync +``` + +5. **Test retrieval**: + +```bash +ai_agent chat +# Try queries that should return your new tool +``` + +### Tool Criteria + +Tools should: + +- โœ… Be relevant to imaging analysis +- โœ… Be actively maintained +- โœ… Have clear documentation +- โœ… Preferably have a runnable demo +- โœ… Be open-source or have free tier + +## Reporting Issues + +### Bug Reports + +Include: + +- **Description** of the bug +- **Steps to reproduce** +- **Expected behavior** +- **Actual behavior** +- **Environment** (OS, Python version, etc.) +- **Logs** if available + +### Feature Requests + +Include: + +- **Use case** for the feature +- **Proposed solution** (if you have one) +- **Alternatives considered** +- **Examples** of similar features + +## Code Review Process + +1. **Automated checks** run on PR (tests, linting) +2. **Maintainer review** provides feedback +3. **Address feedback** and update PR +4. **Approval** from maintainer(s) +5. **Merge** into main branch + +## Release Process + +Releases follow semantic versioning: + +1. Update version in `pyproject.toml` +2. Update `CHANGELOG.md` with release date +3. Create git tag: `git tag v0.2.0` +4. Push tag: `git push origin v0.2.0` +5. GitHub Actions deploys documentation + +## Getting Help + +- **GitHub Discussions**: Ask questions +- **GitHub Issues**: Report bugs + +## Code of Conduct + +Be respectful and professional: + +- Use welcoming and inclusive language +- Respect differing viewpoints +- Accept constructive criticism +- Focus on what's best for the community + +## License + +By contributing, you agree that your contributions will be licensed under the Apache 2.0 License. + +## Next Steps + +- Review [Project Structure](structure.md) +- Learn about [Testing](testing.md) +- Read [Architecture Overview](../architecture/overview.md) diff --git a/docs/development/structure.md b/docs/development/structure.md index 87e07740a046a25f20196cba0c5c1933390e35ad..09f5fef4691dc33191a797fbb17de75465f619ec 100644 --- a/docs/development/structure.md +++ b/docs/development/structure.md @@ -1,362 +1,362 @@ -# Project Structure - -The AI Imaging Agent is organized into modular components with clear separation of concerns. - -## Directory Layout - -``` -ai-agent/ -โ”œโ”€โ”€ .github/ -โ”‚ โ””โ”€โ”€ workflows/ # CI/CD workflows -โ”‚ โ””โ”€โ”€ deploy_docs.yml # Documentation deployment -โ”œโ”€โ”€ artifacts/ -โ”‚ โ””โ”€โ”€ rag_index/ # FAISS index and embeddings -โ”œโ”€โ”€ dataset/ -โ”‚ โ””โ”€โ”€ catalog.jsonl # Software catalog -โ”œโ”€โ”€ docs/ # MkDocs documentation -โ”œโ”€โ”€ logs/ # Application logs -โ”œโ”€โ”€ src/ -โ”‚ โ””โ”€โ”€ ai_agent/ # Main package -โ”‚ โ”œโ”€โ”€ agent/ # PydanticAI agent -โ”‚ โ”œโ”€โ”€ api/ # Pipeline orchestration -โ”‚ โ”œโ”€โ”€ catalog/ # Catalog management -โ”‚ โ”œโ”€โ”€ generator/ # VLM selection (schemas) -โ”‚ โ”œโ”€โ”€ retriever/ # Text retrieval -โ”‚ โ”œโ”€โ”€ ui/ # Gradio interface -โ”‚ โ””โ”€โ”€ utils/ # Shared utilities -โ”œโ”€โ”€ tests/ # Test suite -โ”œโ”€โ”€ config.yaml # Model configuration -โ”œโ”€โ”€ mkdocs.yml # Documentation config -โ”œโ”€โ”€ pyproject.toml # Package metadata -โ””โ”€โ”€ README.md # Project readme -``` - -## Core Modules - -### src/ai_agent/ - -Main package containing all application code. - -#### agent/ - -PydanticAI conversational agent implementation. - -``` -agent/ -โ”œโ”€โ”€ __init__.py -โ”œโ”€โ”€ agent.py # Agent definition, tool adapters -โ”œโ”€โ”€ models.py # Agent output/log models -โ”œโ”€โ”€ utils.py # Agent state and tool quota helpers -โ””โ”€โ”€ tools/ # Tool implementations (search, repo_info, mcp) -``` - -**Key components**: - -- `agent.py`: Agent instance, system prompt, tool definitions -- `models.py`: Agent output and tool usage schemas -- `utils.py`: `AgentState` plus call caps/prepare hooks -- `tools/`: Tool implementations (search, alternatives, repo info, mcp tools) - -**Dependencies**: `api/`, `utils/` - -#### api/ - -Pipeline orchestration and core logic. - -``` -api/ -โ”œโ”€โ”€ __init__.py -โ””โ”€โ”€ pipeline.py # RAGImagingPipeline main class -``` - -**Responsibilities**: - -- File validation and metadata extraction -- Retrieval + VLM selection orchestration -- Error handling and logging -- Index management - -**Dependencies**: `retriever/`, `generator/`, `utils/` - -#### catalog/ - -Software catalog synchronization. - -``` -catalog/ -โ”œโ”€โ”€ __init__.py -โ””โ”€โ”€ sync.py # Catalog sync logic -``` - -**Functions**: - -- Load catalog from JSONL -- Check for changes (SHA1) -- Trigger index rebuild - -**Dependencies**: `retriever/` - -#### generator/ - -VLM selection schemas and types. - -``` -generator/ -โ”œโ”€โ”€ __init__.py -โ””โ”€โ”€ schema.py # Pydantic models for responses -``` - -**Models**: - -- `ToolRecommendation`: Individual tool recommendation -- `AgentResponse`: Complete response with status -- `ConversationStatus`: Enum for conversation states -- `ToolReason`: Enum for recommendation reasons - -**Dependencies**: None (pure schemas) - -#### retriever/ - -Text-based retrieval pipeline. - -``` -retriever/ -โ”œโ”€โ”€ __init__.py -โ”œโ”€โ”€ text_embedder.py # BGE-M3 embedding model -โ”œโ”€โ”€ vector_index.py # FAISS index management -โ”œโ”€โ”€ reranker.py # CrossEncoder reranking -โ””โ”€โ”€ software_doc.py # Catalog schema and loading -``` - -**Pipeline flow**: - -1. `text_embedder.py`: Embed query -2. `vector_index.py`: FAISS search -3. `reranker.py`: CrossEncoder reranking -4. Output: Top-K candidates - -**Dependencies**: None (pure retrieval) - -#### ui/ - -Gradio web interface. - -``` -ui/ -โ”œโ”€โ”€ __init__.py -โ”œโ”€โ”€ app.py # Gradio app definition -โ”œโ”€โ”€ components.py # Reusable UI components -โ”œโ”€โ”€ formatters.py # Response formatting -โ”œโ”€โ”€ handlers.py # Message handlers -โ”œโ”€โ”€ state.py # UI state management -โ””โ”€โ”€ visualizations.py # Preview and trace rendering -``` - -**Key files**: - -- `app.py`: Main Gradio interface -- `handlers.py`: `respond()` function - core interaction logic -- `formatters.py`: Format recommendations as markdown/cards -- `components.py`: Reusable Gradio components - -**Dependencies**: `agent/`, `api/` - -#### utils/ - -Shared utilities. - -``` -utils/ -โ”œโ”€โ”€ __init__.py -โ”œโ”€โ”€ config.py # Configuration loading -โ”œโ”€โ”€ file_validator.py # File validation -โ”œโ”€โ”€ image_meta.py # Metadata extraction (DICOM, NIfTI, TIFF) -โ”œโ”€โ”€ previews.py # Image preview generation -โ””โ”€โ”€ tags.py # Control tag parsing -``` - -**Common utilities**: - -- `config.py`: Load `config.yaml` with Pydantic validation -- `file_validator.py`: Size limits, format checks -- `image_meta.py`: Extract DICOM/NIfTI/TIFF metadata -- `previews.py`: Convert medical images to PNG -- `tags.py`: Parse exclusion tags and strip control tags from queries - -**Dependencies**: None (pure utilities) - -#### cli.py - -Command-line interface entry point. - -```python -def main(): - # Parse arguments - # Route to chat or sync -``` - -**Commands**: - -- `ai_agent chat`: Launch UI -- `ai_agent sync`: Sync catalog - -### tests/ - -Test suite. - -``` -tests/ -โ”œโ”€โ”€ data/ -โ”‚ โ””โ”€โ”€ test_data.json # Test cases -โ”œโ”€โ”€ test_retrieval_pipeline.py -โ”œโ”€โ”€ test_deepwiki_repo_info.py -โ””โ”€โ”€ ... -``` - -**Test categories**: - -- Unit tests: Individual components -- Integration tests: Full pipeline -- End-to-end tests: Real API calls (optional) - -## Configuration Files - -### pyproject.toml - -Python package metadata and dependencies. - -```toml -[project] -name = "ai_agent" -version = "1.0.0" -dependencies = [...] - -[project.scripts] -ai_agent = "ai_agent.cli:main" -``` - -### config.yaml - -Model configuration. - -```yaml -agent_model: - name: "gpt-4o-mini" - base_url: null - api_key_env: "OPENAI_API_KEY" - -available_models: - - display_name: "gpt-4o-mini" - name: "gpt-4o-mini" - ... -``` - -### mkdocs.yml - -Documentation configuration. - -```yaml -site_name: AI Imaging Agent -theme: - name: material -nav: [...] -``` - -### .env - -Environment variables (not committed). - -```dotenv -OPENAI_API_KEY=sk-xxxx -SOFTWARE_CATALOG=dataset/catalog.jsonl -``` - -## Data Files - -### dataset/catalog.jsonl - -Software catalog in JSON Lines format. - -Each line is a complete JSON object following schema.org SoftwareSourceCode. - -### artifacts/rag_index/ - -Pre-built FAISS index and metadata. - -``` -artifacts/rag_index/ -โ”œโ”€โ”€ index.faiss # FAISS binary index -โ””โ”€โ”€ meta.json # Tool IDs, config, timestamps -``` - -## Module Boundaries - -Clear separation prevents circular dependencies: - -``` -ui/ โ†’ agent/ โ†’ api/ โ†’ retriever/ - โ†’ generator/ - โ†’ utils/ -``` - -**Rules**: - -- `utils/`: No dependencies on other modules -- `retriever/`: Pure retrieval, no generation -- `generator/`: Pure schemas, no retrieval -- `api/`: Orchestrates retriever + generator -- `agent/`: Uses api for tool calls -- `ui/`: Top-level, depends on agent + api - -## Import Patterns - -All imports use absolute paths from `ai_agent`: - -```python -from ai_agent.retriever.vector_index import VectorIndex -from ai_agent.utils.config import load_config -from ai_agent.agent.utils import AgentState -``` - -**Never use** relative imports like `from ..utils import ...` - -## Extension Points - -### Adding New Tools - -Add tool adapters to `agent/agent.py` and implement logic in `agent/tools/`: - -```python -@agent.tool -async def new_tool(ctx: RunContext[AgentState], param: str) -> str: - """Tool description.""" - # Implementation - return result -``` - -### Adding New Metadata Extractors - -Add to `utils/image_meta.py`: - -```python -def extract_custom_format(file_path: str) -> dict: - """Extract metadata from custom format.""" - # Implementation - return metadata -``` - -### Adding New Retrieval Models - -Replace in `retriever/text_embedder.py`: - -```python -class TextEmbedder: - def __init__(self, model_name="new-embedding-model"): - self.model = SentenceTransformer(model_name) -``` - -## Next Steps - -- Learn about [Contributing](contributing.md) -- Explore [Testing](testing.md) -- Return to [Architecture Overview](../architecture/overview.md) +# Project Structure + +The AI Imaging Agent is organized into modular components with clear separation of concerns. + +## Directory Layout + +``` +ai-agent/ +โ”œโ”€โ”€ .github/ +โ”‚ โ””โ”€โ”€ workflows/ # CI/CD workflows +โ”‚ โ””โ”€โ”€ deploy_docs.yml # Documentation deployment +โ”œโ”€โ”€ artifacts/ +โ”‚ โ””โ”€โ”€ rag_index/ # FAISS index and embeddings +โ”œโ”€โ”€ dataset/ +โ”‚ โ””โ”€โ”€ catalog.jsonl # Software catalog +โ”œโ”€โ”€ docs/ # MkDocs documentation +โ”œโ”€โ”€ logs/ # Application logs +โ”œโ”€โ”€ src/ +โ”‚ โ””โ”€โ”€ ai_agent/ # Main package +โ”‚ โ”œโ”€โ”€ agent/ # PydanticAI agent +โ”‚ โ”œโ”€โ”€ api/ # Pipeline orchestration +โ”‚ โ”œโ”€โ”€ catalog/ # Catalog management +โ”‚ โ”œโ”€โ”€ generator/ # VLM selection (schemas) +โ”‚ โ”œโ”€โ”€ retriever/ # Text retrieval +โ”‚ โ”œโ”€โ”€ ui/ # Gradio interface +โ”‚ โ””โ”€โ”€ utils/ # Shared utilities +โ”œโ”€โ”€ tests/ # Test suite +โ”œโ”€โ”€ config.yaml # Model configuration +โ”œโ”€โ”€ mkdocs.yml # Documentation config +โ”œโ”€โ”€ pyproject.toml # Package metadata +โ””โ”€โ”€ README.md # Project readme +``` + +## Core Modules + +### src/ai_agent/ + +Main package containing all application code. + +#### agent/ + +PydanticAI conversational agent implementation. + +``` +agent/ +โ”œโ”€โ”€ __init__.py +โ”œโ”€โ”€ agent.py # Agent definition, tool adapters +โ”œโ”€โ”€ models.py # Agent output/log models +โ”œโ”€โ”€ utils.py # Agent state and tool quota helpers +โ””โ”€โ”€ tools/ # Tool implementations (search, repo_info, mcp) +``` + +**Key components**: + +- `agent.py`: Agent instance, system prompt, tool definitions +- `models.py`: Agent output and tool usage schemas +- `utils.py`: `AgentState` plus call caps/prepare hooks +- `tools/`: Tool implementations (search, alternatives, repo info, mcp tools) + +**Dependencies**: `api/`, `utils/` + +#### api/ + +Pipeline orchestration and core logic. + +``` +api/ +โ”œโ”€โ”€ __init__.py +โ””โ”€โ”€ pipeline.py # RAGImagingPipeline main class +``` + +**Responsibilities**: + +- File validation and metadata extraction +- Retrieval + VLM selection orchestration +- Error handling and logging +- Index management + +**Dependencies**: `retriever/`, `generator/`, `utils/` + +#### catalog/ + +Software catalog synchronization. + +``` +catalog/ +โ”œโ”€โ”€ __init__.py +โ””โ”€โ”€ sync.py # Catalog sync logic +``` + +**Functions**: + +- Load catalog from JSONL +- Check for changes (SHA1) +- Trigger index rebuild + +**Dependencies**: `retriever/` + +#### generator/ + +VLM selection schemas and types. + +``` +generator/ +โ”œโ”€โ”€ __init__.py +โ””โ”€โ”€ schema.py # Pydantic models for responses +``` + +**Models**: + +- `ToolRecommendation`: Individual tool recommendation +- `AgentResponse`: Complete response with status +- `ConversationStatus`: Enum for conversation states +- `ToolReason`: Enum for recommendation reasons + +**Dependencies**: None (pure schemas) + +#### retriever/ + +Text-based retrieval pipeline. + +``` +retriever/ +โ”œโ”€โ”€ __init__.py +โ”œโ”€โ”€ text_embedder.py # BGE-M3 embedding model +โ”œโ”€โ”€ vector_index.py # FAISS index management +โ”œโ”€โ”€ reranker.py # CrossEncoder reranking +โ””โ”€โ”€ software_doc.py # Catalog schema and loading +``` + +**Pipeline flow**: + +1. `text_embedder.py`: Embed query +2. `vector_index.py`: FAISS search +3. `reranker.py`: CrossEncoder reranking +4. Output: Top-K candidates + +**Dependencies**: None (pure retrieval) + +#### ui/ + +Gradio web interface. + +``` +ui/ +โ”œโ”€โ”€ __init__.py +โ”œโ”€โ”€ app.py # Gradio app definition +โ”œโ”€โ”€ components.py # Reusable UI components +โ”œโ”€โ”€ formatters.py # Response formatting +โ”œโ”€โ”€ handlers.py # Message handlers +โ”œโ”€โ”€ state.py # UI state management +โ””โ”€โ”€ visualizations.py # Preview and trace rendering +``` + +**Key files**: + +- `app.py`: Main Gradio interface +- `handlers.py`: `respond()` function - core interaction logic +- `formatters.py`: Format recommendations as markdown/cards +- `components.py`: Reusable Gradio components + +**Dependencies**: `agent/`, `api/` + +#### utils/ + +Shared utilities. + +``` +utils/ +โ”œโ”€โ”€ __init__.py +โ”œโ”€โ”€ config.py # Configuration loading +โ”œโ”€โ”€ file_validator.py # File validation +โ”œโ”€โ”€ image_meta.py # Metadata extraction (DICOM, NIfTI, TIFF) +โ”œโ”€โ”€ previews.py # Image preview generation +โ””โ”€โ”€ tags.py # Control tag parsing +``` + +**Common utilities**: + +- `config.py`: Load `config.yaml` with Pydantic validation +- `file_validator.py`: Size limits, format checks +- `image_meta.py`: Extract DICOM/NIfTI/TIFF metadata +- `previews.py`: Convert medical images to PNG +- `tags.py`: Parse exclusion tags and strip control tags from queries + +**Dependencies**: None (pure utilities) + +#### cli.py + +Command-line interface entry point. + +```python +def main(): + # Parse arguments + # Route to chat or sync +``` + +**Commands**: + +- `ai_agent chat`: Launch UI +- `ai_agent sync`: Sync catalog + +### tests/ + +Test suite. + +``` +tests/ +โ”œโ”€โ”€ data/ +โ”‚ โ””โ”€โ”€ test_data.json # Test cases +โ”œโ”€โ”€ test_retrieval_pipeline.py +โ”œโ”€โ”€ test_deepwiki_repo_info.py +โ””โ”€โ”€ ... +``` + +**Test categories**: + +- Unit tests: Individual components +- Integration tests: Full pipeline +- End-to-end tests: Real API calls (optional) + +## Configuration Files + +### pyproject.toml + +Python package metadata and dependencies. + +```toml +[project] +name = "ai_agent" +version = "1.0.0" +dependencies = [...] + +[project.scripts] +ai_agent = "ai_agent.cli:main" +``` + +### config.yaml + +Model configuration. + +```yaml +agent_model: + name: "gpt-4o-mini" + base_url: null + api_key_env: "OPENAI_API_KEY" + +available_models: + - display_name: "gpt-4o-mini" + name: "gpt-4o-mini" + ... +``` + +### mkdocs.yml + +Documentation configuration. + +```yaml +site_name: AI Imaging Agent +theme: + name: material +nav: [...] +``` + +### .env + +Environment variables (not committed). + +```dotenv +OPENAI_API_KEY=sk-xxxx +SOFTWARE_CATALOG=dataset/catalog.jsonl +``` + +## Data Files + +### dataset/catalog.jsonl + +Software catalog in JSON Lines format. + +Each line is a complete JSON object following schema.org SoftwareSourceCode. + +### artifacts/rag_index/ + +Pre-built FAISS index and metadata. + +``` +artifacts/rag_index/ +โ”œโ”€โ”€ index.faiss # FAISS binary index +โ””โ”€โ”€ meta.json # Tool IDs, config, timestamps +``` + +## Module Boundaries + +Clear separation prevents circular dependencies: + +``` +ui/ โ†’ agent/ โ†’ api/ โ†’ retriever/ + โ†’ generator/ + โ†’ utils/ +``` + +**Rules**: + +- `utils/`: No dependencies on other modules +- `retriever/`: Pure retrieval, no generation +- `generator/`: Pure schemas, no retrieval +- `api/`: Orchestrates retriever + generator +- `agent/`: Uses api for tool calls +- `ui/`: Top-level, depends on agent + api + +## Import Patterns + +All imports use absolute paths from `ai_agent`: + +```python +from ai_agent.retriever.vector_index import VectorIndex +from ai_agent.utils.config import load_config +from ai_agent.agent.utils import AgentState +``` + +**Never use** relative imports like `from ..utils import ...` + +## Extension Points + +### Adding New Tools + +Add tool adapters to `agent/agent.py` and implement logic in `agent/tools/`: + +```python +@agent.tool +async def new_tool(ctx: RunContext[AgentState], param: str) -> str: + """Tool description.""" + # Implementation + return result +``` + +### Adding New Metadata Extractors + +Add to `utils/image_meta.py`: + +```python +def extract_custom_format(file_path: str) -> dict: + """Extract metadata from custom format.""" + # Implementation + return metadata +``` + +### Adding New Retrieval Models + +Replace in `retriever/text_embedder.py`: + +```python +class TextEmbedder: + def __init__(self, model_name="new-embedding-model"): + self.model = SentenceTransformer(model_name) +``` + +## Next Steps + +- Learn about [Contributing](contributing.md) +- Explore [Testing](testing.md) +- Return to [Architecture Overview](../architecture/overview.md) diff --git a/docs/development/testing.md b/docs/development/testing.md index 4162e4fdf41ace7c3534da54e48ab414b7568138..7787dabb963f8b742e8a094f988f0245f7609f76 100644 --- a/docs/development/testing.md +++ b/docs/development/testing.md @@ -1,410 +1,410 @@ -# Testing (still under development) - -The AI Imaging Agent uses pytest for testing. This guide covers running tests and writing new ones. - -**Note:** We are still developing some tests for the agent, hence this part is not relevant for now. - -## Running Tests - -### Basic Usage - -```bash -# Run all tests -pytest - -# Run specific test file -pytest tests/test_retrieval_pipeline.py - -# Run specific test -pytest tests/test_retrieval_pipeline.py::test_basic_retrieval - -# Run with verbose output -pytest -v - -# Run with coverage -pytest --cov=ai_agent --cov-report=html -``` - -### Test Categories - -Tests are marked by category: - -```bash -# Run only unit tests -pytest -m unit - -# Run only integration tests -pytest -m integration - -# Skip slow tests -pytest -m "not slow" -``` - -## Test Organization - -### Directory Structure - -``` -tests/ -โ”œโ”€โ”€ data/ -โ”‚ โ”œโ”€โ”€ test_data.json # Test cases -โ”‚ โ””โ”€โ”€ 0002.DCM # Sample DICOM file -โ”œโ”€โ”€ test_retrieval_pipeline.py # Retrieval tests -โ”œโ”€โ”€ test_deepwiki_repo_info.py # Repo info tests -โ”œโ”€โ”€ test_gpt4o_vision.py # VLM tests (integration) -โ””โ”€โ”€ __pycache__/ -``` - -### Test File Naming - -- `test_*.py`: Test files -- `*_test.py`: Alternative naming (less common) - -### Test Function Naming - -```python -def test_basic_retrieval(): - """Test basic retrieval functionality.""" - pass - -def test_edge_case_empty_query(): - """Test handling of empty query.""" - pass - -def test_integration_full_pipeline(): - """Integration test for complete pipeline.""" - pass -``` - -## Writing Tests - -### Unit Test Example - -```python -import pytest -from ai_agent.retriever.vector_index import VectorIndex - -def test_vector_index_search(): - """Test FAISS vector search.""" - # Arrange - index = VectorIndex() - index.load("artifacts/rag_index") - - query = "segment lungs CT" - - # Act - results = index.search(query, k=5) - - # Assert - assert len(results) == 5 - assert all(r['score'] > 0 for r in results) - assert 'TotalSegmentator' in [r['name'] for r in results] -``` - -### Integration Test Example - -```python -import pytest -from ai_agent.api.pipeline import RAGImagingPipeline - -@pytest.mark.integration -def test_full_pipeline_with_image(): - """Integration test with real image and VLM call.""" - # Arrange - pipeline = RAGImagingPipeline( - catalog_path="dataset/catalog.jsonl", - index_dir="artifacts/rag_index" - ) - - # Act - result = pipeline.recommend( - query="segment lungs", - files=["tests/data/chest_ct.dcm"] - ) - - # Assert - assert result.status == "complete" - assert len(result.recommendations) > 0 - assert result.recommendations[0].accuracy_score > 70 -``` - -### Parametrized Tests - -```python -@pytest.mark.parametrize("query,expected_tool", [ - ("segment brain MRI", "FreeSurfer"), - ("segment lungs CT", "TotalSegmentator"), - ("classify chest X-ray", "CheXNet"), -]) -def test_retrieval_for_queries(query, expected_tool): - """Test retrieval returns expected tools for various queries.""" - index = VectorIndex() - index.load("artifacts/rag_index") - - results = index.search(query, k=10) - tool_names = [r['name'] for r in results] - - assert expected_tool in tool_names -``` - -### Fixtures - -```python -import pytest - -@pytest.fixture -def pipeline(): - """Provide initialized pipeline for tests.""" - return RAGImagingPipeline( - catalog_path="dataset/catalog.jsonl", - index_dir="artifacts/rag_index" - ) - -@pytest.fixture -def sample_dicom(): - """Provide path to sample DICOM file.""" - return "tests/data/0002.DCM" - -def test_with_fixtures(pipeline, sample_dicom): - """Test using fixtures.""" - result = pipeline.recommend( - query="analyze DICOM", - files=[sample_dicom] - ) - assert result is not None -``` - - - -## Test Data - -### Using Test Cases - -Load test cases from JSON: - -```python -import json - -def load_test_cases(): - """Load test cases from data file.""" - with open("tests/data/test_data.json") as f: - return json.load(f) - -@pytest.mark.parametrize("test_case", load_test_cases()) -def test_from_json(test_case): - """Test using cases from JSON file.""" - query = test_case["query"] - expected = test_case["expected_tool"] - - # Test logic here - assert expected in results -``` - -### Sample Data Files - -Keep sample files small: - -- **DICOM**: Single slice, low resolution -- **NIfTI**: Small volume (e.g., 64ร—64ร—64) -- **Images**: PNG/JPG under 1 MB - - - -## Continuous Integration - -### GitHub Actions - -Tests run automatically on: - -- Pull requests -- Pushes to main - -### CI Configuration - -```yaml -# .github/workflows/test.yml -name: Tests - -on: [push, pull_request] - -jobs: - test: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: '3.10' - - run: pip install -e ".[dev]" - - run: pytest --cov=ai_agent -``` - -## Best Practices - -### Do's - -โœ… **Test edge cases**: Empty inputs, invalid data, etc. -โœ… **Test error handling**: Verify exceptions are caught -โœ… **Use descriptive names**: `test_retrieval_with_empty_query` not `test1` -โœ… **Keep tests isolated**: Each test should be independent -โœ… **Use fixtures**: Avoid repeating setup code -โœ… **Mock expensive operations**: VLM calls, network requests - -### Don'ts - -โŒ **Don't test implementation details**: Test behavior, not internal state -โŒ **Don't make tests depend on each other**: Each should run independently -โŒ **Don't commit large test files**: Keep test data small -โŒ **Don't skip error checking**: Test both success and failure paths - -## Performance Testing - -### Benchmarking - -Use pytest-benchmark: - -```python -def test_retrieval_performance(benchmark): - """Benchmark retrieval speed.""" - index = VectorIndex() - index.load("artifacts/rag_index") - - result = benchmark(index.search, "segment lungs", k=10) - - assert len(result) == 10 -``` - -### Profiling - -```bash -# Profile tests -pytest --profile - -# Generate SVG profile -pytest --profile-svg -``` - -## Debugging Tests - -### Running in Debug Mode - -```python -# Add to test -import pdb; pdb.set_trace() - -# Run pytest -pytest tests/test_file.py -``` - -### Verbose Output - -```bash -# Show print statements -pytest -s - -# Very verbose -pytest -vv - -# Show local variables on failure -pytest -l -``` - -### Running Single Test - -```bash -# Run one test function -pytest tests/test_file.py::test_function_name -v -``` - -## Next Steps - -- Review [Project Structure](structure.md) -- Read [Contributing Guide](contributing.md) -- Explore [Architecture](../architecture/overview.md) +# Testing (still under development) + +The AI Imaging Agent uses pytest for testing. This guide covers running tests and writing new ones. + +**Note:** We are still developing some tests for the agent, hence this part is not relevant for now. + +## Running Tests + +### Basic Usage + +```bash +# Run all tests +pytest + +# Run specific test file +pytest tests/test_retrieval_pipeline.py + +# Run specific test +pytest tests/test_retrieval_pipeline.py::test_basic_retrieval + +# Run with verbose output +pytest -v + +# Run with coverage +pytest --cov=ai_agent --cov-report=html +``` + +### Test Categories + +Tests are marked by category: + +```bash +# Run only unit tests +pytest -m unit + +# Run only integration tests +pytest -m integration + +# Skip slow tests +pytest -m "not slow" +``` + +## Test Organization + +### Directory Structure + +``` +tests/ +โ”œโ”€โ”€ data/ +โ”‚ โ”œโ”€โ”€ test_data.json # Test cases +โ”‚ โ””โ”€โ”€ 0002.DCM # Sample DICOM file +โ”œโ”€โ”€ test_retrieval_pipeline.py # Retrieval tests +โ”œโ”€โ”€ test_deepwiki_repo_info.py # Repo info tests +โ”œโ”€โ”€ test_gpt4o_vision.py # VLM tests (integration) +โ””โ”€โ”€ __pycache__/ +``` + +### Test File Naming + +- `test_*.py`: Test files +- `*_test.py`: Alternative naming (less common) + +### Test Function Naming + +```python +def test_basic_retrieval(): + """Test basic retrieval functionality.""" + pass + +def test_edge_case_empty_query(): + """Test handling of empty query.""" + pass + +def test_integration_full_pipeline(): + """Integration test for complete pipeline.""" + pass +``` + +## Writing Tests + +### Unit Test Example + +```python +import pytest +from ai_agent.retriever.vector_index import VectorIndex + +def test_vector_index_search(): + """Test FAISS vector search.""" + # Arrange + index = VectorIndex() + index.load("artifacts/rag_index") + + query = "segment lungs CT" + + # Act + results = index.search(query, k=5) + + # Assert + assert len(results) == 5 + assert all(r['score'] > 0 for r in results) + assert 'TotalSegmentator' in [r['name'] for r in results] +``` + +### Integration Test Example + +```python +import pytest +from ai_agent.api.pipeline import RAGImagingPipeline + +@pytest.mark.integration +def test_full_pipeline_with_image(): + """Integration test with real image and VLM call.""" + # Arrange + pipeline = RAGImagingPipeline( + catalog_path="dataset/catalog.jsonl", + index_dir="artifacts/rag_index" + ) + + # Act + result = pipeline.recommend( + query="segment lungs", + files=["tests/data/chest_ct.dcm"] + ) + + # Assert + assert result.status == "complete" + assert len(result.recommendations) > 0 + assert result.recommendations[0].accuracy_score > 70 +``` + +### Parametrized Tests + +```python +@pytest.mark.parametrize("query,expected_tool", [ + ("segment brain MRI", "FreeSurfer"), + ("segment lungs CT", "TotalSegmentator"), + ("classify chest X-ray", "CheXNet"), +]) +def test_retrieval_for_queries(query, expected_tool): + """Test retrieval returns expected tools for various queries.""" + index = VectorIndex() + index.load("artifacts/rag_index") + + results = index.search(query, k=10) + tool_names = [r['name'] for r in results] + + assert expected_tool in tool_names +``` + +### Fixtures + +```python +import pytest + +@pytest.fixture +def pipeline(): + """Provide initialized pipeline for tests.""" + return RAGImagingPipeline( + catalog_path="dataset/catalog.jsonl", + index_dir="artifacts/rag_index" + ) + +@pytest.fixture +def sample_dicom(): + """Provide path to sample DICOM file.""" + return "tests/data/0002.DCM" + +def test_with_fixtures(pipeline, sample_dicom): + """Test using fixtures.""" + result = pipeline.recommend( + query="analyze DICOM", + files=[sample_dicom] + ) + assert result is not None +``` + + + +## Test Data + +### Using Test Cases + +Load test cases from JSON: + +```python +import json + +def load_test_cases(): + """Load test cases from data file.""" + with open("tests/data/test_data.json") as f: + return json.load(f) + +@pytest.mark.parametrize("test_case", load_test_cases()) +def test_from_json(test_case): + """Test using cases from JSON file.""" + query = test_case["query"] + expected = test_case["expected_tool"] + + # Test logic here + assert expected in results +``` + +### Sample Data Files + +Keep sample files small: + +- **DICOM**: Single slice, low resolution +- **NIfTI**: Small volume (e.g., 64ร—64ร—64) +- **Images**: PNG/JPG under 1 MB + + + +## Continuous Integration + +### GitHub Actions + +Tests run automatically on: + +- Pull requests +- Pushes to main + +### CI Configuration + +```yaml +# .github/workflows/test.yml +name: Tests + +on: [push, pull_request] + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + - run: pip install -e ".[dev]" + - run: pytest --cov=ai_agent +``` + +## Best Practices + +### Do's + +โœ… **Test edge cases**: Empty inputs, invalid data, etc. +โœ… **Test error handling**: Verify exceptions are caught +โœ… **Use descriptive names**: `test_retrieval_with_empty_query` not `test1` +โœ… **Keep tests isolated**: Each test should be independent +โœ… **Use fixtures**: Avoid repeating setup code +โœ… **Mock expensive operations**: VLM calls, network requests + +### Don'ts + +โŒ **Don't test implementation details**: Test behavior, not internal state +โŒ **Don't make tests depend on each other**: Each should run independently +โŒ **Don't commit large test files**: Keep test data small +โŒ **Don't skip error checking**: Test both success and failure paths + +## Performance Testing + +### Benchmarking + +Use pytest-benchmark: + +```python +def test_retrieval_performance(benchmark): + """Benchmark retrieval speed.""" + index = VectorIndex() + index.load("artifacts/rag_index") + + result = benchmark(index.search, "segment lungs", k=10) + + assert len(result) == 10 +``` + +### Profiling + +```bash +# Profile tests +pytest --profile + +# Generate SVG profile +pytest --profile-svg +``` + +## Debugging Tests + +### Running in Debug Mode + +```python +# Add to test +import pdb; pdb.set_trace() + +# Run pytest +pytest tests/test_file.py +``` + +### Verbose Output + +```bash +# Show print statements +pytest -s + +# Very verbose +pytest -vv + +# Show local variables on failure +pytest -l +``` + +### Running Single Test + +```bash +# Run one test function +pytest tests/test_file.py::test_function_name -v +``` + +## Next Steps + +- Review [Project Structure](structure.md) +- Read [Contributing Guide](contributing.md) +- Explore [Architecture](../architecture/overview.md) diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md index 9059569a65dbe31caf56e77750904a8405f8e6e5..dcd4270f64c77d6721c1dd486e8ef0019d4a19d5 100644 --- a/docs/getting-started/configuration.md +++ b/docs/getting-started/configuration.md @@ -1,165 +1,165 @@ -# Configuration - -Before running the AI Imaging Agent, you need to configure it with your API keys and preferences. - -## Environment Variables - -Create a `.env` file in the repository root with the following configuration: - -```dotenv -# Required: OpenAI API key -OPENAI_API_KEY=sk-xxxx - -# Optional: GitHub token for repository info tool -GITHUB_TOKEN=ghp_xxxx - -# Optional: Alternative model providers -EPFL_API_KEY=sk-xxxx - -# Software catalog path -SOFTWARE_CATALOG=dataset/catalog.jsonl - -# Logging configuration -LOGLEVEL_CONSOLE=WARNING -LOGLEVEL_FILE=INFO -FILE_LOG=1 -LOG_DIR=logs -LOG_PROMPTS=0 # Set to 1 to save prompt snapshots for debugging - -# Custom config path -CONFIG_PATH=config.yaml -``` - -## Required Configuration - -### OpenAI API Key - -The AI Imaging Agent requires an OpenAI API key for the vision-language model: - -1. Sign up for an account at [OpenAI](https://platform.openai.com/) -2. Navigate to [API Keys](https://platform.openai.com/api-keys) -3. Create a new API key -4. Add it to your `.env` file: - -```dotenv -OPENAI_API_KEY=sk-your-actual-key-here -``` - -## Model Configuration - -The agent model can be configured via `config.yaml`: - -```yaml -# AI Agent Model Configuration - -# Default/fallback model (used for CLI and initial startup) -agent_model: - name: "gpt-4o-mini" - base_url: null # null for default OpenAI endpoint - api_key_env: "OPENAI_API_KEY" - -# Available models for UI dropdown -available_models: - - display_name: "gpt-4o-mini" - name: "gpt-4o-mini" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-4o" - name: "gpt-4o" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" - - - display_name: "gpt-5.1" - name: "gpt-5.1" - base_url: null - provider: "OpenAI" - api_key_env: "OPENAI_API_KEY" -``` - -### Using Alternative Model Providers - -You can configure custom OpenAI-compatible endpoints: - -```yaml -available_models: - - display_name: "EPFL Inference" - name: "gpt-4o-mini" - base_url: "https://inference.epfl.ch/v1" - provider: "EPFL" - api_key_env: "EPFL_API_KEY" -``` - -Then add the corresponding API key to your `.env`: - -```dotenv -EPFL_API_KEY=your-epfl-key -``` - -## Optional Configuration - -### GitHub Token - -For the repository info tool (optional): - -```dotenv -GITHUB_TOKEN=ghp_your_github_personal_access_token -``` - -This enables the agent to fetch detailed information about GitHub repositories. - -### Pipeline Parameters - -Adjust retrieval and recommendation settings directly in the app setting. You can change the `TOP_K` and `NUM_CHOICES` parameters. - -### Logging - -Configure logging behavior: - -```dotenv -# Console log level (DEBUG, INFO, WARNING, ERROR) -LOGLEVEL_CONSOLE=WARNING - -# File log level -LOGLEVEL_FILE=INFO - -# Enable file logging (0 or 1) -FILE_LOG=1 - -# Log directory -LOG_DIR=logs - -# Save VLM prompts and images for debugging (0 or 1) -LOG_PROMPTS=0 -``` - -!!! tip "Debug Mode" - Set `LOG_PROMPTS=1` to save VLM prompts and images to the `logs/` directory. This is useful for debugging but will increase disk usage. - -### Software Catalog - -Specify the path to your software catalog: - -```dotenv -SOFTWARE_CATALOG=dataset/catalog.jsonl -``` - -The catalog should be in JSONL format following the schema.org SoftwareSourceCode structure. - -## Verification - -After configuring, verify your setup: - -```bash -# Check that environment variables are loaded -python -c "from dotenv import load_dotenv; import os; load_dotenv(); print('API Key:', 'SET' if os.getenv('OPENAI_API_KEY') else 'NOT SET')" -``` - -## Next Steps - -With configuration complete, you're ready to: - -- [Run the Quick Start](quickstart.md) -- Learn about [Using the Chat Interface](../user-guide/chat-interface.md) +# Configuration + +Before running the AI Imaging Agent, you need to configure it with your API keys and preferences. + +## Environment Variables + +Create a `.env` file in the repository root with the following configuration: + +```dotenv +# Required: OpenAI API key +OPENAI_API_KEY=sk-xxxx + +# Optional: GitHub token for repository info tool +GITHUB_TOKEN=ghp_xxxx + +# Optional: Alternative model providers +EPFL_API_KEY=sk-xxxx + +# Software catalog path +SOFTWARE_CATALOG=dataset/catalog.jsonl + +# Logging configuration +LOGLEVEL_CONSOLE=WARNING +LOGLEVEL_FILE=INFO +FILE_LOG=1 +LOG_DIR=logs +LOG_PROMPTS=0 # Set to 1 to save prompt snapshots for debugging + +# Custom config path +CONFIG_PATH=config.yaml +``` + +## Required Configuration + +### OpenAI API Key + +The AI Imaging Agent requires an OpenAI API key for the vision-language model: + +1. Sign up for an account at [OpenAI](https://platform.openai.com/) +2. Navigate to [API Keys](https://platform.openai.com/api-keys) +3. Create a new API key +4. Add it to your `.env` file: + +```dotenv +OPENAI_API_KEY=sk-your-actual-key-here +``` + +## Model Configuration + +The agent model can be configured via `config.yaml`: + +```yaml +# AI Agent Model Configuration + +# Default/fallback model (used for CLI and initial startup) +agent_model: + name: "gpt-4o-mini" + base_url: null # null for default OpenAI endpoint + api_key_env: "OPENAI_API_KEY" + +# Available models for UI dropdown +available_models: + - display_name: "gpt-4o-mini" + name: "gpt-4o-mini" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-4o" + name: "gpt-4o" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" + + - display_name: "gpt-5.1" + name: "gpt-5.1" + base_url: null + provider: "OpenAI" + api_key_env: "OPENAI_API_KEY" +``` + +### Using Alternative Model Providers + +You can configure custom OpenAI-compatible endpoints: + +```yaml +available_models: + - display_name: "EPFL Inference" + name: "gpt-4o-mini" + base_url: "https://inference.epfl.ch/v1" + provider: "EPFL" + api_key_env: "EPFL_API_KEY" +``` + +Then add the corresponding API key to your `.env`: + +```dotenv +EPFL_API_KEY=your-epfl-key +``` + +## Optional Configuration + +### GitHub Token + +For the repository info tool (optional): + +```dotenv +GITHUB_TOKEN=ghp_your_github_personal_access_token +``` + +This enables the agent to fetch detailed information about GitHub repositories. + +### Pipeline Parameters + +Adjust retrieval and recommendation settings directly in the app setting. You can change the `TOP_K` and `NUM_CHOICES` parameters. + +### Logging + +Configure logging behavior: + +```dotenv +# Console log level (DEBUG, INFO, WARNING, ERROR) +LOGLEVEL_CONSOLE=WARNING + +# File log level +LOGLEVEL_FILE=INFO + +# Enable file logging (0 or 1) +FILE_LOG=1 + +# Log directory +LOG_DIR=logs + +# Save VLM prompts and images for debugging (0 or 1) +LOG_PROMPTS=0 +``` + +!!! tip "Debug Mode" + Set `LOG_PROMPTS=1` to save VLM prompts and images to the `logs/` directory. This is useful for debugging but will increase disk usage. + +### Software Catalog + +Specify the path to your software catalog: + +```dotenv +SOFTWARE_CATALOG=dataset/catalog.jsonl +``` + +The catalog should be in JSONL format following the schema.org SoftwareSourceCode structure. + +## Verification + +After configuring, verify your setup: + +```bash +# Check that environment variables are loaded +python -c "from dotenv import load_dotenv; import os; load_dotenv(); print('API Key:', 'SET' if os.getenv('OPENAI_API_KEY') else 'NOT SET')" +``` + +## Next Steps + +With configuration complete, you're ready to: + +- [Run the Quick Start](quickstart.md) +- Learn about [Using the Chat Interface](../user-guide/chat-interface.md) diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 135e3f5cf0bac866e4dfb9e5cdac01ae41199eca..74d91859ee1cc7acf87cdc21535f97d066616c2b 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -1,143 +1,143 @@ -# Installation - -This guide will help you install and set up the AI Imaging Agent on your system. - -## Prerequisites - -Before installing, ensure you have: - -- **Python 3.10โ€“3.12** installed -- **pip** (Python package manager) -- **OpenAI API key** (or compatible API endpoint) -- Internet connection for model calls - -## Installation Steps - -### 1. Clone the Repository - -```bash -git clone https://github.com/imaging-plaza/ai-agent.git -cd ai-agent -``` - -### 2. Create Virtual Environment - -It's recommended to use a virtual environment to isolate dependencies: - -=== "Linux/macOS" - - ```bash - python -m venv .venv - source .venv/bin/activate - ``` - -=== "Windows" - - ```bash - python -m venv .venv - .venv\Scripts\activate - ``` - -### 3. Install the Package - -For regular use: - -```bash -pip install --upgrade pip -pip install -e . -``` - -For development (includes test dependencies): - -```bash -pip install -e ".[dev]" -``` - -## Verify Installation - -Verify that the installation was successful: - -```bash -ai_agent --help -``` - -You should see the available commands: - -``` -usage: ai_agent [-h] {chat,sync} - -AI Agent CLI - -positional arguments: - {chat,sync} 'chat' launches the chat UI; 'sync' runs one catalog refresh. -``` - -## Next Steps - -Now that you have installed the AI Imaging Agent, proceed to: - -- [Configuration](configuration.md) - Set up your environment and API keys -- [Quick Start](quickstart.md) - Run your first query - -## Troubleshooting - -### Python Version Issues - -If you encounter issues with Python version compatibility: - -```bash -# Check your Python version -python --version - -# Use a specific Python version -python3.10 -m venv .venv -``` - -### Installation Errors - -If you encounter dependency conflicts: - -```bash -# Upgrade pip first -pip install --upgrade pip setuptools wheel - -# Try installing again -pip install -e . -``` - -### Missing System Dependencies - -Some packages may require system libraries: - -=== "Ubuntu/Debian" - - ```bash - sudo apt-get update - sudo apt-get install python3-dev build-essential - ``` - -=== "macOS" - - ```bash - # Using Homebrew - brew install python@3.10 - ``` - -=== "Windows" - - Ensure you have [Microsoft C++ Build Tools](https://visualstudio.microsoft.com/visual-cpp-build-tools/) installed. - -## Docker Installation (Alternative) - -A Dockerfile is available for containerized deployment: - -```bash -# Build the Docker image -docker build -t ai-agent -f Dockerfile . - -# Run the container -docker run -p 7860:7860 --env-file .env ai-agent -``` - -!!! note - Make sure to create a `.env` file with your configuration before running the Docker container. +# Installation + +This guide will help you install and set up the AI Imaging Agent on your system. + +## Prerequisites + +Before installing, ensure you have: + +- **Python 3.10โ€“3.12** installed +- **pip** (Python package manager) +- **OpenAI API key** (or compatible API endpoint) +- Internet connection for model calls + +## Installation Steps + +### 1. Clone the Repository + +```bash +git clone https://github.com/imaging-plaza/ai-agent.git +cd ai-agent +``` + +### 2. Create Virtual Environment + +It's recommended to use a virtual environment to isolate dependencies: + +=== "Linux/macOS" + + ```bash + python -m venv .venv + source .venv/bin/activate + ``` + +=== "Windows" + + ```bash + python -m venv .venv + .venv\Scripts\activate + ``` + +### 3. Install the Package + +For regular use: + +```bash +pip install --upgrade pip +pip install -e . +``` + +For development (includes test dependencies): + +```bash +pip install -e ".[dev]" +``` + +## Verify Installation + +Verify that the installation was successful: + +```bash +ai_agent --help +``` + +You should see the available commands: + +``` +usage: ai_agent [-h] {chat,sync} + +AI Agent CLI + +positional arguments: + {chat,sync} 'chat' launches the chat UI; 'sync' runs one catalog refresh. +``` + +## Next Steps + +Now that you have installed the AI Imaging Agent, proceed to: + +- [Configuration](configuration.md) - Set up your environment and API keys +- [Quick Start](quickstart.md) - Run your first query + +## Troubleshooting + +### Python Version Issues + +If you encounter issues with Python version compatibility: + +```bash +# Check your Python version +python --version + +# Use a specific Python version +python3.10 -m venv .venv +``` + +### Installation Errors + +If you encounter dependency conflicts: + +```bash +# Upgrade pip first +pip install --upgrade pip setuptools wheel + +# Try installing again +pip install -e . +``` + +### Missing System Dependencies + +Some packages may require system libraries: + +=== "Ubuntu/Debian" + + ```bash + sudo apt-get update + sudo apt-get install python3-dev build-essential + ``` + +=== "macOS" + + ```bash + # Using Homebrew + brew install python@3.10 + ``` + +=== "Windows" + + Ensure you have [Microsoft C++ Build Tools](https://visualstudio.microsoft.com/visual-cpp-build-tools/) installed. + +## Docker Installation (Alternative) + +A Dockerfile is available for containerized deployment: + +```bash +# Build the Docker image +docker build -t ai-agent -f Dockerfile . + +# Run the container +docker run -p 7860:7860 --env-file .env ai-agent +``` + +!!! note + Make sure to create a `.env` file with your configuration before running the Docker container. diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index d863b73361b97751cea1e27199294e9557a752a1..21c01c84c7a8c6b98fb0724e92bd59fd4d06c434 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -1,190 +1,190 @@ -# Quick Start - -This guide will walk you through your first interaction with the AI Imaging Agent. - -## Starting the Application - -Once you've [installed](installation.md) and [configured](configuration.md) the agent, start the chat interface: - -```bash -ai_agent chat -``` - -You should see output like: - -``` -Running on local URL: http://127.0.0.1:7860 -``` - -Open your web browser and navigate to **http://127.0.0.1:7860** - -## Your First Query - -### Example 1: Object Segmentation - -Let's try a simple segmentation task: - -1. **Upload an Image**: Click the upload area or drag and drop an image (e.g., a photo of a cat) - -2. **Type Your Request**: In the chat input, type: - ``` - I want to segment the cat from this image - ``` - -3. **Review Recommendations**: The agent will return ranked tool recommendations with: - - Tool names and descriptions - - Accuracy scores - - Explanations for why each tool matches your task - - Links to runnable demos - -4. **Run a Demo** (optional): Click the "Run Demo" button to see the result of the tool on your uploaded image - -### Example 2: Medical Image Analysis - -For medical imaging tasks: - -1. **Upload a Medical Image**: Upload a DICOM file, NIfTI volume, or medical image - -2. **Describe Your Task**: - ``` - Segment the lungs from this CT scan - ``` - -3. **Get Format-Aware Results**: The agent considers: - - Your image format (DICOM, NIfTI, etc.) - - Image dimensions (2D, 3D, 4D) - - Medical imaging modality (CT, MRI, etc.) - -### Example 3: General Computer Vision - -For general tasks: - -``` -Detect all objects in this image -``` - -``` -Extract text from this document image -``` - -``` -Classify what type of animal is in this picture -``` - -## Understanding the Interface - -### Chat Panel - -- **Message History**: Scroll to see previous interactions -- **Rich Media**: Images, files, and tool cards are rendered inline -- **Code Blocks**: Formatted code and JSON responses - -### Sidebar - -- **Uploaded Files**: View all files you've uploaded in the session -- **Preview Images**: See converted image previews -- **Debug Info**: View conversation state and excluded tools (if in debug mode) - -### Tool Recommendation Cards - -Each recommended tool shows: - -- **Rank**: Priority order (1 = best match) -- **Name**: Tool/software name -- **Accuracy Score**: Confidence level (0-100%) -- **Description**: What the tool does -- **Explanation**: Why it matches your request -- **Metadata**: - - Supported modalities (CT, MRI, etc.) - - Dimensions (2D, 3D, etc.) - - File formats (DICOM, NIfTI, PNG, etc.) - - License information - - Tags and categories -- **Demo Link**: Direct link to runnable example - -## Advanced Usage - -### Multi-Turn Conversations - -The agent maintains conversation context: - -``` -You: I have a lung CT scan -Agent: [Provides general information about lung CT analysis tools] - -You: I want to segment the airways -Agent: [Provides specific airway segmentation tools] - -You: Show me alternatives -Agent: [Provides different tool options] -``` - -### Excluding Tools - -Exclude specific tools from results: - -``` -Find lung segmentation tools [EXCLUDE:totalsegmentator|medicalsam] -``` - -### Requesting Alternatives - -If initial results don't match your needs: - -``` -Show me alternative tools - -Can you search for other options? - -What else is available? -``` - -## CLI Commands - -The agent provides two main commands: - -### Launch Chat Interface - -```bash -ai_agent chat -``` - -Starts the Gradio web interface with automatic catalog synchronization. - -### Sync Catalog - -```bash -ai_agent sync -``` - -Manually synchronize the software catalog without launching the UI. - -## Tips for Best Results - -!!! tip "Be Specific" - The more specific your request, the better the recommendations: - - - โŒ "Process this image" - - โœ… "Segment the liver from this abdominal CT scan" - -!!! tip "Upload First" - Upload your image before describing the task. The agent can see image content and metadata. - -!!! tip "Mention Formats" - If you need specific format support, mention it: - - "I need a tool that works with DICOM files" - -!!! tip "Use Natural Language" - No need to use technical jargon - conversational language works fine: - - "Help me find tumors in this MRI" works just as well as "Tumor detection in MRI volumes" - -## Next Steps - -Now that you've run your first queries: - -- Learn more about [Using the Chat Interface](../user-guide/chat-interface.md) -- Explore [Supported File Formats](../user-guide/file-formats.md) -- Understand [How Recommendations Work](../user-guide/recommendations.md) -- Dive into the [Architecture Overview](../architecture/overview.md) +# Quick Start + +This guide will walk you through your first interaction with the AI Imaging Agent. + +## Starting the Application + +Once you've [installed](installation.md) and [configured](configuration.md) the agent, start the chat interface: + +```bash +ai_agent chat +``` + +You should see output like: + +``` +Running on local URL: http://127.0.0.1:7860 +``` + +Open your web browser and navigate to **http://127.0.0.1:7860** + +## Your First Query + +### Example 1: Object Segmentation + +Let's try a simple segmentation task: + +1. **Upload an Image**: Click the upload area or drag and drop an image (e.g., a photo of a cat) + +2. **Type Your Request**: In the chat input, type: + ``` + I want to segment the cat from this image + ``` + +3. **Review Recommendations**: The agent will return ranked tool recommendations with: + - Tool names and descriptions + - Accuracy scores + - Explanations for why each tool matches your task + - Links to runnable demos + +4. **Run a Demo** (optional): Click the "Run Demo" button to see the result of the tool on your uploaded image + +### Example 2: Medical Image Analysis + +For medical imaging tasks: + +1. **Upload a Medical Image**: Upload a DICOM file, NIfTI volume, or medical image + +2. **Describe Your Task**: + ``` + Segment the lungs from this CT scan + ``` + +3. **Get Format-Aware Results**: The agent considers: + - Your image format (DICOM, NIfTI, etc.) + - Image dimensions (2D, 3D, 4D) + - Medical imaging modality (CT, MRI, etc.) + +### Example 3: General Computer Vision + +For general tasks: + +``` +Detect all objects in this image +``` + +``` +Extract text from this document image +``` + +``` +Classify what type of animal is in this picture +``` + +## Understanding the Interface + +### Chat Panel + +- **Message History**: Scroll to see previous interactions +- **Rich Media**: Images, files, and tool cards are rendered inline +- **Code Blocks**: Formatted code and JSON responses + +### Sidebar + +- **Uploaded Files**: View all files you've uploaded in the session +- **Preview Images**: See converted image previews +- **Debug Info**: View conversation state and excluded tools (if in debug mode) + +### Tool Recommendation Cards + +Each recommended tool shows: + +- **Rank**: Priority order (1 = best match) +- **Name**: Tool/software name +- **Accuracy Score**: Confidence level (0-100%) +- **Description**: What the tool does +- **Explanation**: Why it matches your request +- **Metadata**: + - Supported modalities (CT, MRI, etc.) + - Dimensions (2D, 3D, etc.) + - File formats (DICOM, NIfTI, PNG, etc.) + - License information + - Tags and categories +- **Demo Link**: Direct link to runnable example + +## Advanced Usage + +### Multi-Turn Conversations + +The agent maintains conversation context: + +``` +You: I have a lung CT scan +Agent: [Provides general information about lung CT analysis tools] + +You: I want to segment the airways +Agent: [Provides specific airway segmentation tools] + +You: Show me alternatives +Agent: [Provides different tool options] +``` + +### Excluding Tools + +Exclude specific tools from results: + +``` +Find lung segmentation tools [EXCLUDE:totalsegmentator|medicalsam] +``` + +### Requesting Alternatives + +If initial results don't match your needs: + +``` +Show me alternative tools + +Can you search for other options? + +What else is available? +``` + +## CLI Commands + +The agent provides two main commands: + +### Launch Chat Interface + +```bash +ai_agent chat +``` + +Starts the Gradio web interface with automatic catalog synchronization. + +### Sync Catalog + +```bash +ai_agent sync +``` + +Manually synchronize the software catalog without launching the UI. + +## Tips for Best Results + +!!! tip "Be Specific" + The more specific your request, the better the recommendations: + + - โŒ "Process this image" + - โœ… "Segment the liver from this abdominal CT scan" + +!!! tip "Upload First" + Upload your image before describing the task. The agent can see image content and metadata. + +!!! tip "Mention Formats" + If you need specific format support, mention it: + + "I need a tool that works with DICOM files" + +!!! tip "Use Natural Language" + No need to use technical jargon - conversational language works fine: + + "Help me find tumors in this MRI" works just as well as "Tumor detection in MRI volumes" + +## Next Steps + +Now that you've run your first queries: + +- Learn more about [Using the Chat Interface](../user-guide/chat-interface.md) +- Explore [Supported File Formats](../user-guide/file-formats.md) +- Understand [How Recommendations Work](../user-guide/recommendations.md) +- Dive into the [Architecture Overview](../architecture/overview.md) diff --git a/docs/guide.md b/docs/guide.md index 13f1f2098852c8e6567db74c99256d0018fa5a21..1c5e56557edb33310b9ac781f6e8d123c6372f28 100644 --- a/docs/guide.md +++ b/docs/guide.md @@ -1,277 +1,277 @@ -# Project Guide - -This guide is a practical map of the entire repository for contributors and maintainers. - -It focuses on: -- What each folder is responsible for -- Which Python environment and package workflow are the defaults -- Which commands are currently valid -- What to improve next in architecture, testing, performance, and developer experience - -## 1) System Summary - -AI Imaging Agent is a RAG plus VLM recommender for imaging software. - -High-level flow: -1. User uploads file(s) and enters a task. -2. Retrieval stage finds candidate tools (BGE-M3 + FAISS + reranker). -3. Agent/VLM stage ranks candidates with image-aware reasoning. -4. UI renders ranked recommendations and optional demo links. - -Primary orchestrator: [src/ai_agent/api/pipeline.py](src/ai_agent/api/pipeline.py) - -## 2) Default Python Environment And Packages (Dev Container Canonical) - -Assume development is done inside the dev container. - -Source of truth: -- Dev container: [.devcontainer/devcontainer.json](../.devcontainer/devcontainer.json) -- Package metadata and pinned dependencies: [pyproject.toml](../pyproject.toml) -- Secondary dependency list: [requirements.txt](../requirements.txt) - -Default environment: -- OS: Debian Bookworm (dev container) -- Python: 3.12 -- Environment manager: uv -- Virtual environment path: .venv - -Recommended commands: - -```bash -uv venv -uv pip install -e . -uv pip install -e ".[dev]" -``` - -Run and test: - -```bash -ai_agent chat -ai_agent sync -pytest tests/ -``` - -Important note on command drift: -- CLI officially supports `chat` and `sync` in [src/ai_agent/cli.py](../src/ai_agent/cli.py). -- [justfile](../justfile) currently references `ai_agent ui`, which does not match current CLI modes. -- Documentation in this guide follows the actual CLI implementation. - -## 3) Repository Top-Level Map - -- [.github/](../.github/): automation and agent instructions -- [.devcontainer/](../.devcontainer/): dev container build and editor defaults -- [docs/](.): MkDocs source pages -- [src/](../src/): application source code -- [tests/](../tests/): test suite -- [data/](../data/): sample data assets -- [tools/](../tools/): container/tooling helpers -- [CHANGELOG.md](../CHANGELOG.md): release history -- [config.yaml](../config.yaml): model/provider configuration -- [mkdocs.yml](../mkdocs.yml): docs site navigation and theme -- [pyproject.toml](../pyproject.toml): package metadata, dependencies, entrypoints - -## 4) Detailed Source Folder Responsibilities - -Package root: [src/ai_agent/](../src/ai_agent) - -### 4.1 [src/ai_agent/agent/](../src/ai_agent/agent) - -Purpose: conversational orchestration using PydanticAI. - -Key files: -- [src/ai_agent/agent/agent.py](../src/ai_agent/agent/agent.py): agent setup, tool wiring, response flow -- [src/ai_agent/agent/models.py](../src/ai_agent/agent/models.py): state/output models -- [src/ai_agent/agent/utils.py](../src/ai_agent/agent/utils.py): helper utilities and guardrails -- [src/ai_agent/agent/tools/](../src/ai_agent/agent/tools): concrete tool implementations -- [src/ai_agent/agent/tools/mcp/](../src/ai_agent/agent/tools/mcp): MCP adapters - -Boundary: -- Should orchestrate tools and policy, not own retrieval internals. - -### 4.2 [src/ai_agent/api/](../src/ai_agent/api) - -Purpose: pipeline orchestration between inputs, retrieval, and selection. - -Key file: -- [src/ai_agent/api/pipeline.py](../src/ai_agent/api/pipeline.py) - -Responsibilities: -- validate files -- extract metadata -- build retrieval query -- call retrieval and selection stages -- manage index refresh/reload behavior - -Boundary: -- Keep UI concerns out of this module. - -### 4.3 [src/ai_agent/retriever/](../src/ai_agent/retriever) - -Purpose: deterministic retrieval stack (no LLM calls). - -Key files: -- [src/ai_agent/retriever/text_embedder.py](../src/ai_agent/retriever/text_embedder.py) -- [src/ai_agent/retriever/vector_index.py](../src/ai_agent/retriever/vector_index.py) -- [src/ai_agent/retriever/reranker.py](../src/ai_agent/retriever/reranker.py) -- [src/ai_agent/retriever/software_doc.py](../src/ai_agent/retriever/software_doc.py) - -Boundary: -- Retrieval quality logic should stay here. - -### 4.4 [src/ai_agent/generator/](../src/ai_agent/generator) - -Purpose: selection schema and prompting primitives. - -Key files: -- [src/ai_agent/generator/prompts.py](../src/ai_agent/generator/prompts.py) -- [src/ai_agent/generator/schema.py](../src/ai_agent/generator/schema.py) - -Boundary: -- Keep this layer focused on schema and prompt contracts, not transport/UI concerns. - -### 4.5 [src/ai_agent/ui/](../src/ai_agent/ui) - -Purpose: Gradio app and interaction handling. - -Key files: -- [src/ai_agent/ui/app.py](../src/ai_agent/ui/app.py) -- [src/ai_agent/ui/handlers.py](../src/ai_agent/ui/handlers.py) -- [src/ai_agent/ui/components.py](../src/ai_agent/ui/components.py) -- [src/ai_agent/ui/formatters.py](../src/ai_agent/ui/formatters.py) -- [src/ai_agent/ui/state.py](../src/ai_agent/ui/state.py) -- [src/ai_agent/ui/visualizations.py](../src/ai_agent/ui/visualizations.py) - -Boundary: -- UI should call orchestrators, not reimplement retrieval/selection decisions. - -### 4.6 [src/ai_agent/utils/](../src/ai_agent/utils) - -Purpose: cross-cutting utility functions. - -Key files: -- [src/ai_agent/utils/config.py](../src/ai_agent/utils/config.py) -- [src/ai_agent/utils/file_validator.py](../src/ai_agent/utils/file_validator.py) -- [src/ai_agent/utils/image_meta.py](../src/ai_agent/utils/image_meta.py) -- [src/ai_agent/utils/image_io.py](../src/ai_agent/utils/image_io.py) -- [src/ai_agent/utils/previews.py](../src/ai_agent/utils/previews.py) -- [src/ai_agent/utils/tags.py](../src/ai_agent/utils/tags.py) -- [src/ai_agent/utils/temp_file_manager.py](../src/ai_agent/utils/temp_file_manager.py) - -Boundary: -- Keep utilities reusable and independent from UI-specific logic. - -### 4.7 [src/ai_agent/catalog/](../src/ai_agent/catalog) - -Purpose: catalog synchronization and refresh helpers. - -Key file: -- [src/ai_agent/catalog/sync.py](../src/ai_agent/catalog/sync.py) - -Boundary: -- Catalog IO and sync logic should stay isolated from ranking logic. - -### 4.8 [src/ai_agent/core/](../src/ai_agent/core) - -Purpose: shared core coordination such as pipeline registry. - -Key file: -- [src/ai_agent/core/pipeline_registry.py](../src/ai_agent/core/pipeline_registry.py) - -Boundary: -- Keep core primitives minimal and dependency-light. - -### 4.9 [src/ai_agent/queries/](../src/ai_agent/queries) - -Purpose: query assets used by catalog sync/retrieval support. - -Key file: -- [src/ai_agent/queries/get_relevant_software.rq](../src/ai_agent/queries/get_relevant_software.rq) - -Boundary: -- Keep query definitions versioned and testable. - -### 4.10 [src/ai_agent/cli.py](../src/ai_agent/cli.py) - -Purpose: command entry point and mode dispatch. - -Current modes: -- `chat` -- `sync` - -This is the command contract docs should follow. - -## 5) Supporting Folders - -### 5.1 [tests/](../tests) - -Contains unit/integration tests and test fixtures under [tests/data/](../tests/data). - -Improvement target: -- add more focused tests for UI handler edge cases and tool failure handling. - -### 5.2 [tools/](../tools) - -Container and deployment support assets. - -Notable file: -- [tools/image/Dockerfile](../tools/image/Dockerfile) (uv + Python 3.12 baseline) - -### 5.3 [docs/](.) - -Documentation source for MkDocs. - -Add new pages to [mkdocs.yml](../mkdocs.yml) nav to keep docs discoverable. - -## 6) Known Inconsistencies To Track - -1. [justfile](../justfile) uses `ai_agent ui`, while [src/ai_agent/cli.py](../src/ai_agent/cli.py) defines `chat` and `sync`. -2. Installation docs often show pip-first flow, while dev container bootstrap is uv-first. -3. [requirements.txt](../requirements.txt) is looser than [pyproject.toml](../pyproject.toml), which contains current pinned/runtime dependencies. - -## 7) Codebase Improvement Guidelines - -### 7.1 Architecture And Modularity - -1. Keep strict stage boundaries: retrieval logic in `retriever`, selection contracts in `generator`, orchestration in `api`. -2. Minimize cross-layer imports from `ui` to low-level modules. -3. Introduce lightweight interface contracts for tool adapters to reduce coupling in `agent/tools`. -4. Centralize shared constants/env defaults to reduce duplicated configuration behavior. - -### 7.2 Testing And Quality Gates - -1. Add regression tests for format-token query construction and retry broadening behavior. -2. Add failure-path tests for image preview generation and graceful degradation. -3. Add contract tests for agent tool outputs (search, alternative search, repo info). -4. Enforce formatting/lint/type checks in CI (`ruff`, `black --check`, `mypy`, `pytest`). - -### 7.3 Performance And Retrieval Quality - -1. Add benchmark fixtures for retrieval latency and reranker throughput. -2. Track retrieval quality with a small fixed evaluation set (top-k recall, MRR). -3. Cache expensive metadata extraction where safe for repeated files in a session. -4. Make index reload behavior observable with structured counters in logs. - -### 7.4 Developer Experience And CI - -1. Align `just` tasks with real CLI contract (`chat`/`sync`). -2. Add a docs link checker in CI to prevent markdown drift. -3. Document one canonical local workflow (dev container first, optional local pip fallback). -4. Add a short maintainer checklist for release prep and changelog updates. - -## 8) Practical Contributor Checklist - -Before opening a PR: -1. Install/update in editable mode in the active environment. -2. Run tests relevant to changed modules. -3. Validate docs links if docs were touched. -4. Update [CHANGELOG.md](../CHANGELOG.md) for user-visible changes. -5. Confirm command and environment docs still match real behavior. - -## 9) Related References - -- [README.md](../README.md) -- [docs/index.md](index.md) -- [docs/architecture/overview.md](architecture/overview.md) -- [docs/development/structure.md](development/structure.md) -- [AGENTS.md](../AGENTS.md) -- [.github/copilot-instructions.md](../.github/copilot-instructions.md) +# Project Guide + +This guide is a practical map of the entire repository for contributors and maintainers. + +It focuses on: +- What each folder is responsible for +- Which Python environment and package workflow are the defaults +- Which commands are currently valid +- What to improve next in architecture, testing, performance, and developer experience + +## 1) System Summary + +AI Imaging Agent is a RAG plus VLM recommender for imaging software. + +High-level flow: +1. User uploads file(s) and enters a task. +2. Retrieval stage finds candidate tools (BGE-M3 + FAISS + reranker). +3. Agent/VLM stage ranks candidates with image-aware reasoning. +4. UI renders ranked recommendations and optional demo links. + +Primary orchestrator: [src/ai_agent/api/pipeline.py](src/ai_agent/api/pipeline.py) + +## 2) Default Python Environment And Packages (Dev Container Canonical) + +Assume development is done inside the dev container. + +Source of truth: +- Dev container: [.devcontainer/devcontainer.json](../.devcontainer/devcontainer.json) +- Package metadata and pinned dependencies: [pyproject.toml](../pyproject.toml) +- Secondary dependency list: [requirements.txt](../requirements.txt) + +Default environment: +- OS: Debian Bookworm (dev container) +- Python: 3.12 +- Environment manager: uv +- Virtual environment path: .venv + +Recommended commands: + +```bash +uv venv +uv pip install -e . +uv pip install -e ".[dev]" +``` + +Run and test: + +```bash +ai_agent chat +ai_agent sync +pytest tests/ +``` + +Important note on command drift: +- CLI officially supports `chat` and `sync` in [src/ai_agent/cli.py](../src/ai_agent/cli.py). +- [justfile](../justfile) currently references `ai_agent ui`, which does not match current CLI modes. +- Documentation in this guide follows the actual CLI implementation. + +## 3) Repository Top-Level Map + +- [.github/](../.github/): automation and agent instructions +- [.devcontainer/](../.devcontainer/): dev container build and editor defaults +- [docs/](.): MkDocs source pages +- [src/](../src/): application source code +- [tests/](../tests/): test suite +- [data/](../data/): sample data assets +- [tools/](../tools/): container/tooling helpers +- [CHANGELOG.md](../CHANGELOG.md): release history +- [config.yaml](../config.yaml): model/provider configuration +- [mkdocs.yml](../mkdocs.yml): docs site navigation and theme +- [pyproject.toml](../pyproject.toml): package metadata, dependencies, entrypoints + +## 4) Detailed Source Folder Responsibilities + +Package root: [src/ai_agent/](../src/ai_agent) + +### 4.1 [src/ai_agent/agent/](../src/ai_agent/agent) + +Purpose: conversational orchestration using PydanticAI. + +Key files: +- [src/ai_agent/agent/agent.py](../src/ai_agent/agent/agent.py): agent setup, tool wiring, response flow +- [src/ai_agent/agent/models.py](../src/ai_agent/agent/models.py): state/output models +- [src/ai_agent/agent/utils.py](../src/ai_agent/agent/utils.py): helper utilities and guardrails +- [src/ai_agent/agent/tools/](../src/ai_agent/agent/tools): concrete tool implementations +- [src/ai_agent/agent/tools/mcp/](../src/ai_agent/agent/tools/mcp): MCP adapters + +Boundary: +- Should orchestrate tools and policy, not own retrieval internals. + +### 4.2 [src/ai_agent/api/](../src/ai_agent/api) + +Purpose: pipeline orchestration between inputs, retrieval, and selection. + +Key file: +- [src/ai_agent/api/pipeline.py](../src/ai_agent/api/pipeline.py) + +Responsibilities: +- validate files +- extract metadata +- build retrieval query +- call retrieval and selection stages +- manage index refresh/reload behavior + +Boundary: +- Keep UI concerns out of this module. + +### 4.3 [src/ai_agent/retriever/](../src/ai_agent/retriever) + +Purpose: deterministic retrieval stack (no LLM calls). + +Key files: +- [src/ai_agent/retriever/text_embedder.py](../src/ai_agent/retriever/text_embedder.py) +- [src/ai_agent/retriever/vector_index.py](../src/ai_agent/retriever/vector_index.py) +- [src/ai_agent/retriever/reranker.py](../src/ai_agent/retriever/reranker.py) +- [src/ai_agent/retriever/software_doc.py](../src/ai_agent/retriever/software_doc.py) + +Boundary: +- Retrieval quality logic should stay here. + +### 4.4 [src/ai_agent/generator/](../src/ai_agent/generator) + +Purpose: selection schema and prompting primitives. + +Key files: +- [src/ai_agent/generator/prompts.py](../src/ai_agent/generator/prompts.py) +- [src/ai_agent/generator/schema.py](../src/ai_agent/generator/schema.py) + +Boundary: +- Keep this layer focused on schema and prompt contracts, not transport/UI concerns. + +### 4.5 [src/ai_agent/ui/](../src/ai_agent/ui) + +Purpose: Gradio app and interaction handling. + +Key files: +- [src/ai_agent/ui/app.py](../src/ai_agent/ui/app.py) +- [src/ai_agent/ui/handlers.py](../src/ai_agent/ui/handlers.py) +- [src/ai_agent/ui/components.py](../src/ai_agent/ui/components.py) +- [src/ai_agent/ui/formatters.py](../src/ai_agent/ui/formatters.py) +- [src/ai_agent/ui/state.py](../src/ai_agent/ui/state.py) +- [src/ai_agent/ui/visualizations.py](../src/ai_agent/ui/visualizations.py) + +Boundary: +- UI should call orchestrators, not reimplement retrieval/selection decisions. + +### 4.6 [src/ai_agent/utils/](../src/ai_agent/utils) + +Purpose: cross-cutting utility functions. + +Key files: +- [src/ai_agent/utils/config.py](../src/ai_agent/utils/config.py) +- [src/ai_agent/utils/file_validator.py](../src/ai_agent/utils/file_validator.py) +- [src/ai_agent/utils/image_meta.py](../src/ai_agent/utils/image_meta.py) +- [src/ai_agent/utils/image_io.py](../src/ai_agent/utils/image_io.py) +- [src/ai_agent/utils/previews.py](../src/ai_agent/utils/previews.py) +- [src/ai_agent/utils/tags.py](../src/ai_agent/utils/tags.py) +- [src/ai_agent/utils/temp_file_manager.py](../src/ai_agent/utils/temp_file_manager.py) + +Boundary: +- Keep utilities reusable and independent from UI-specific logic. + +### 4.7 [src/ai_agent/catalog/](../src/ai_agent/catalog) + +Purpose: catalog synchronization and refresh helpers. + +Key file: +- [src/ai_agent/catalog/sync.py](../src/ai_agent/catalog/sync.py) + +Boundary: +- Catalog IO and sync logic should stay isolated from ranking logic. + +### 4.8 [src/ai_agent/core/](../src/ai_agent/core) + +Purpose: shared core coordination such as pipeline registry. + +Key file: +- [src/ai_agent/core/pipeline_registry.py](../src/ai_agent/core/pipeline_registry.py) + +Boundary: +- Keep core primitives minimal and dependency-light. + +### 4.9 [src/ai_agent/queries/](../src/ai_agent/queries) + +Purpose: query assets used by catalog sync/retrieval support. + +Key file: +- [src/ai_agent/queries/get_relevant_software.rq](../src/ai_agent/queries/get_relevant_software.rq) + +Boundary: +- Keep query definitions versioned and testable. + +### 4.10 [src/ai_agent/cli.py](../src/ai_agent/cli.py) + +Purpose: command entry point and mode dispatch. + +Current modes: +- `chat` +- `sync` + +This is the command contract docs should follow. + +## 5) Supporting Folders + +### 5.1 [tests/](../tests) + +Contains unit/integration tests and test fixtures under [tests/data/](../tests/data). + +Improvement target: +- add more focused tests for UI handler edge cases and tool failure handling. + +### 5.2 [tools/](../tools) + +Container and deployment support assets. + +Notable file: +- [tools/image/Dockerfile](../tools/image/Dockerfile) (uv + Python 3.12 baseline) + +### 5.3 [docs/](.) + +Documentation source for MkDocs. + +Add new pages to [mkdocs.yml](../mkdocs.yml) nav to keep docs discoverable. + +## 6) Known Inconsistencies To Track + +1. [justfile](../justfile) uses `ai_agent ui`, while [src/ai_agent/cli.py](../src/ai_agent/cli.py) defines `chat` and `sync`. +2. Installation docs often show pip-first flow, while dev container bootstrap is uv-first. +3. [requirements.txt](../requirements.txt) is looser than [pyproject.toml](../pyproject.toml), which contains current pinned/runtime dependencies. + +## 7) Codebase Improvement Guidelines + +### 7.1 Architecture And Modularity + +1. Keep strict stage boundaries: retrieval logic in `retriever`, selection contracts in `generator`, orchestration in `api`. +2. Minimize cross-layer imports from `ui` to low-level modules. +3. Introduce lightweight interface contracts for tool adapters to reduce coupling in `agent/tools`. +4. Centralize shared constants/env defaults to reduce duplicated configuration behavior. + +### 7.2 Testing And Quality Gates + +1. Add regression tests for format-token query construction and retry broadening behavior. +2. Add failure-path tests for image preview generation and graceful degradation. +3. Add contract tests for agent tool outputs (search, alternative search, repo info). +4. Enforce formatting/lint/type checks in CI (`ruff`, `black --check`, `mypy`, `pytest`). + +### 7.3 Performance And Retrieval Quality + +1. Add benchmark fixtures for retrieval latency and reranker throughput. +2. Track retrieval quality with a small fixed evaluation set (top-k recall, MRR). +3. Cache expensive metadata extraction where safe for repeated files in a session. +4. Make index reload behavior observable with structured counters in logs. + +### 7.4 Developer Experience And CI + +1. Align `just` tasks with real CLI contract (`chat`/`sync`). +2. Add a docs link checker in CI to prevent markdown drift. +3. Document one canonical local workflow (dev container first, optional local pip fallback). +4. Add a short maintainer checklist for release prep and changelog updates. + +## 8) Practical Contributor Checklist + +Before opening a PR: +1. Install/update in editable mode in the active environment. +2. Run tests relevant to changed modules. +3. Validate docs links if docs were touched. +4. Update [CHANGELOG.md](../CHANGELOG.md) for user-visible changes. +5. Confirm command and environment docs still match real behavior. + +## 9) Related References + +- [README.md](../README.md) +- [docs/index.md](index.md) +- [docs/architecture/overview.md](architecture/overview.md) +- [docs/development/structure.md](development/structure.md) +- [AGENTS.md](../AGENTS.md) +- [.github/copilot-instructions.md](../.github/copilot-instructions.md) diff --git a/docs/index.md b/docs/index.md index 9dd25c74d062d76e0c97e37043882972400b419e..c42730b1fd57ea25f298a8b503773d2eb371261c 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,81 +1,81 @@ -# AI Imaging Agent - -**An intelligent RAG + AI agent system that helps users discover the right imaging software for their images and tasks.** - -[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) -[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/imaging-plaza/ai-agent/blob/main/LICENSE) - ---- - -## What is AI Imaging Agent? - -AI Imaging Agent (also known as **Imaging Plaza**) is a conversational AI assistant that helps researchers and practitioners find the right imaging analysis tools for their specific needs. Simply upload an image, describe what you want to do, and get ranked software recommendations with links to runnable demos. - -## โœจ Key Features - -- **๐Ÿค– Conversational AI Agent**: Natural language interaction with multi-turn context -- **๐Ÿ” Smart Retrieval**: BGE-M3 embeddings + FAISS + CrossEncoder reranking -- **๐Ÿ‘๏ธ Vision-Aware Selection**: VLM-based tool selection considering both image content and metadata -- **๐Ÿฅ Medical Imaging Focus**: Specialized support for CT, MRI, DICOM, NIfTI, and other medical formats -- **๐ŸŽฏ Format-Aware Matching**: IO compatibility scoring based on file formats and dimensions -- **๐Ÿš€ Demo Integration**: Direct execution of Gradio Space demos on your images -- **๐Ÿ“Š Rich UI**: Chat interface with image previews, file management, and execution traces - -## Quick Example - -```bash -# Install and run -pip install -e . -ai_agent chat -``` - -Then in the web interface: - -1. Upload an image (e.g., a CT scan, or a PNG) -2. Type your request (e.g., _"I want to segment the lungs from this image"_ or _"I want to deblur this image"_) -3. Get ranked tool recommendations with accuracy scores -4. Click "Run demo" to execute tools directly - -## Use Cases - -### Medical Imaging -- Segment organs from CT/MRI scans -- Register brain images -- Detect tumors and anomalies -- Analyze DICOM files - -### Scientific Imaging -- Process microscopy images -- Analyze multidimensional TIFF stacks -- Extract features from scientific images - -### General Computer Vision -- Object detection and segmentation -- Image classification -- OCR and text extraction -- Image enhancement - -## How It Works - -The system uses a **two-stage pipeline**: - -1. **Retrieval Stage**: Fast text search using BGE-M3 embeddings and FAISS to find candidate tools from a curated catalog -2. **Agent Selection**: Vision-language model (GPT-4o) analyzes your image and task to rank the best tools with explanations - -Learn more in the [Architecture Overview](architecture/overview.md). - -## Getting Started - -Ready to try it out? Head over to the [Installation Guide](getting-started/installation.md) to get started! - -## Project Guide - -For maintainers and contributors, see the [Project Guide](guide.md) for a detailed repository map, dev-container environment defaults, and practical codebase improvement guidelines. - -## Project Status - -This project is actively developed and maintained by the Imaging Plaza team. Check the [Changelog](reference/changelog.md) for recent updates. - -## License - -This project is licensed under the Apache 2.0 License - see the [LICENSE](https://github.com/imaging-plaza/ai-agent/blob/main/LICENSE) file for details. +# AI Imaging Agent + +**An intelligent RAG + AI agent system that helps users discover the right imaging software for their images and tasks.** + +[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) +[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/imaging-plaza/ai-agent/blob/main/LICENSE) + +--- + +## What is AI Imaging Agent? + +AI Imaging Agent (also known as **Imaging Plaza**) is a conversational AI assistant that helps researchers and practitioners find the right imaging analysis tools for their specific needs. Simply upload an image, describe what you want to do, and get ranked software recommendations with links to runnable demos. + +## โœจ Key Features + +- **๐Ÿค– Conversational AI Agent**: Natural language interaction with multi-turn context +- **๐Ÿ” Smart Retrieval**: BGE-M3 embeddings + FAISS + CrossEncoder reranking +- **๐Ÿ‘๏ธ Vision-Aware Selection**: VLM-based tool selection considering both image content and metadata +- **๐Ÿฅ Medical Imaging Focus**: Specialized support for CT, MRI, DICOM, NIfTI, and other medical formats +- **๐ŸŽฏ Format-Aware Matching**: IO compatibility scoring based on file formats and dimensions +- **๐Ÿš€ Demo Integration**: Direct execution of Gradio Space demos on your images +- **๐Ÿ“Š Rich UI**: Chat interface with image previews, file management, and execution traces + +## Quick Example + +```bash +# Install and run +pip install -e . +ai_agent chat +``` + +Then in the web interface: + +1. Upload an image (e.g., a CT scan, or a PNG) +2. Type your request (e.g., _"I want to segment the lungs from this image"_ or _"I want to deblur this image"_) +3. Get ranked tool recommendations with accuracy scores +4. Click "Run demo" to execute tools directly + +## Use Cases + +### Medical Imaging +- Segment organs from CT/MRI scans +- Register brain images +- Detect tumors and anomalies +- Analyze DICOM files + +### Scientific Imaging +- Process microscopy images +- Analyze multidimensional TIFF stacks +- Extract features from scientific images + +### General Computer Vision +- Object detection and segmentation +- Image classification +- OCR and text extraction +- Image enhancement + +## How It Works + +The system uses a **two-stage pipeline**: + +1. **Retrieval Stage**: Fast text search using BGE-M3 embeddings and FAISS to find candidate tools from a curated catalog +2. **Agent Selection**: Vision-language model (GPT-4o) analyzes your image and task to rank the best tools with explanations + +Learn more in the [Architecture Overview](architecture/overview.md). + +## Getting Started + +Ready to try it out? Head over to the [Installation Guide](getting-started/installation.md) to get started! + +## Project Guide + +For maintainers and contributors, see the [Project Guide](guide.md) for a detailed repository map, dev-container environment defaults, and practical codebase improvement guidelines. + +## Project Status + +This project is actively developed and maintained by the Imaging Plaza team. Check the [Changelog](reference/changelog.md) for recent updates. + +## License + +This project is licensed under the Apache 2.0 License - see the [LICENSE](https://github.com/imaging-plaza/ai-agent/blob/main/LICENSE) file for details. diff --git a/docs/reference/changelog.md b/docs/reference/changelog.md index 455b8d0d33a38d21fbee4d0ad1d6c5c4d0733c19..b33ca663cb9d0bc463acc27b5eeb63275efbd9e5 100644 --- a/docs/reference/changelog.md +++ b/docs/reference/changelog.md @@ -1,79 +1,79 @@ -# Changelog - -All notable changes to the AI Imaging Agent are documented here. - -For the complete, detailed changelog, see [CHANGELOG.md](https://github.com/imaging-plaza/ai-agent/blob/main/CHANGELOG.md) in the repository. - -## Recent Releases - -### [1.0.0] - -#### ๐Ÿš€ Added - -- **Chat-based interface** (`ai_agent chat`) with conversational AI assistant and tool integration -- **Imaging Plaza UI**: Custom branding, theme, and improved layout -- **Iterative retrieval with automatic retry** for low-result queries -- **Alternative search tool** for agent-driven query refinement -- **YAML configuration** (`config.yaml`) for flexible model and backend setup -- **DeepWiki MCP integration** for fast GitHub repository documentation access - ---- - -#### ๐Ÿ”„ Changed - -- CLI updated: `ai_agent chat` replaces deprecated `ai_agent ui` -- Retrieval pipeline enhanced with smarter expansion and retry logic -- UI state management simplified -- Agent-based architecture fully replaces legacy pipelines - ---- - -#### ๐Ÿงน Removed - -- `VLMToolSelector` (replaced by agent-based tool selection) -- Legacy refine intent detection system -- Deprecated UI command (`ai_agent ui`) -- Outdated tests and unused code paths - -### [0.1.3] - 2025-10-22 - -#### Added -- Gradio Space runner tool -- Repository info tool - -#### Fixed -- Gradio UI context binding -- Chatbot message format migration -- Cache cleaning -- PNG preview handling - -### [0.1.2] - 2025-10-07 - -#### Added -- Pydantic AI pipeline with tools -- Better runnable example handling - -### [0.1.1] - 2025-10-02 - -#### Added -- Experimental Pydantic AI agent skeleton -- Multimodal agent pathway - -### [0.1.0] - 2025-09-30 - -#### Added -- Initial chat functionality - -## Versioning - -This project follows [Semantic Versioning](https://semver.org/): - -- **MAJOR**: Incompatible API changes -- **MINOR**: New functionality (backwards-compatible) -- **PATCH**: Bug fixes (backwards-compatible) - -## Contributing - -All notable changes should be documented in [CHANGELOG.md](https://github.com/imaging-plaza/ai-agent/blob/main/CHANGELOG.md) following the [Keep a Changelog](https://keepachangelog.com/) format. - -See [Contributing Guide](../development/contributing.md) for details. +# Changelog + +All notable changes to the AI Imaging Agent are documented here. + +For the complete, detailed changelog, see [CHANGELOG.md](https://github.com/imaging-plaza/ai-agent/blob/main/CHANGELOG.md) in the repository. + +## Recent Releases + +### [1.0.0] + +#### ๐Ÿš€ Added + +- **Chat-based interface** (`ai_agent chat`) with conversational AI assistant and tool integration +- **Imaging Plaza UI**: Custom branding, theme, and improved layout +- **Iterative retrieval with automatic retry** for low-result queries +- **Alternative search tool** for agent-driven query refinement +- **YAML configuration** (`config.yaml`) for flexible model and backend setup +- **DeepWiki MCP integration** for fast GitHub repository documentation access + +--- + +#### ๐Ÿ”„ Changed + +- CLI updated: `ai_agent chat` replaces deprecated `ai_agent ui` +- Retrieval pipeline enhanced with smarter expansion and retry logic +- UI state management simplified +- Agent-based architecture fully replaces legacy pipelines + +--- + +#### ๐Ÿงน Removed + +- `VLMToolSelector` (replaced by agent-based tool selection) +- Legacy refine intent detection system +- Deprecated UI command (`ai_agent ui`) +- Outdated tests and unused code paths + +### [0.1.3] - 2025-10-22 + +#### Added +- Gradio Space runner tool +- Repository info tool + +#### Fixed +- Gradio UI context binding +- Chatbot message format migration +- Cache cleaning +- PNG preview handling + +### [0.1.2] - 2025-10-07 + +#### Added +- Pydantic AI pipeline with tools +- Better runnable example handling + +### [0.1.1] - 2025-10-02 + +#### Added +- Experimental Pydantic AI agent skeleton +- Multimodal agent pathway + +### [0.1.0] - 2025-09-30 + +#### Added +- Initial chat functionality + +## Versioning + +This project follows [Semantic Versioning](https://semver.org/): + +- **MAJOR**: Incompatible API changes +- **MINOR**: New functionality (backwards-compatible) +- **PATCH**: Bug fixes (backwards-compatible) + +## Contributing + +All notable changes should be documented in [CHANGELOG.md](https://github.com/imaging-plaza/ai-agent/blob/main/CHANGELOG.md) following the [Keep a Changelog](https://keepachangelog.com/) format. + +See [Contributing Guide](../development/contributing.md) for details. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 2461ac9ce0d8f41501062b8538334a47e9e71c6a..1e4e699fa51e2d6899e9119375248a259721d8b9 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -1,197 +1,197 @@ -# CLI Commands - -The AI Imaging Agent provides a command-line interface for starting the application and managing the software catalog. - -## Available Commands - -### ai_agent chat - -Launch the chat-based user interface. - -```bash -ai_agent chat -``` - -**What it does**: - -1. Performs startup catalog synchronization -2. Loads the FAISS index -3. Initializes the retrieval and agent pipelines -4. Launches the Gradio web interface on `http://127.0.0.1:7860` -5. Starts background catalog refresh (if configured) - -**Options**: None (all configuration via `.env` and `config.yaml`) - -**Example**: - -```bash -$ ai_agent chat -[startup-sync] 150 โ†’ dataset/catalog.jsonl -[startup-refresh] catalog unchanged; keeping existing FAISS index -Running on local URL: http://127.0.0.1:7860 - -To create a public link, set `share=True` in `launch()`. -``` - -**Background Refresh**: - -If `SYNC_EVERY_HOURS` is set in `.env`, the catalog will auto-refresh in the background: - -```dotenv -SYNC_EVERY_HOURS=24 # Check every 24 hours -``` - -### ai_agent sync - -Manually synchronize the software catalog and rebuild the index. - -```bash -ai_agent sync -``` - -**What it does**: - -1. Loads the software catalog from `SOFTWARE_CATALOG` path -2. Embeds all tool descriptions using BGE-M3 -3. Builds FAISS vector index -4. Saves artifacts to `artifacts/rag_index/` - -**When to use**: - -- After editing `catalog.jsonl` -- After adding new tools -- To force index rebuild -- For testing catalog changes - -**Example**: - -```bash -$ ai_agent sync -[sync] 150 โ†’ dataset/catalog.jsonl -[sync] Embedding 150 tools... (5.2s) -[sync] Building FAISS index... -[sync] Saved to artifacts/rag_index/ -[sync] Sync complete. -``` - -## Command Aliases - -Both commands are available with either `ai_agent` or `ai-agent`: - -```bash -ai_agent chat # Works -ai-agent chat # Also works - -ai_agent sync # Works -ai-agent sync # Also works -``` - -## Common Usage Patterns - -### Development Workflow - -```bash -# Edit catalog -vim dataset/catalog.jsonl - -# Sync catalog -ai_agent sync - -# Test changes -ai_agent chat -``` - - - -### Testing & Development - -```bash -# Enable debug logging -export LOGLEVEL_CONSOLE=DEBUG -export LOG_PROMPTS=1 -ai_agent chat -``` - -## Environment Variables - -All configuration is via environment variables (see [Environment Variables Reference](environment.md)). - -## Exit Codes - -- **0**: Success -- **1**: General error (see logs) - -## Troubleshooting - -### Command Not Found - -If you see `command not found: ai_agent`: - -```bash -# Ensure package is installed -pip install -e . - -# Check installation -pip list | grep ai-agent - -# Try with python -m -python -m ai_agent.cli chat -``` - -### Port Already in Use - -If port 7860 is occupied: - -```bash -# Find and kill process -lsof -ti:7860 | xargs kill -9 - -# Or change port in code (ui/app.py) -``` - -### Catalog Load Error - -If catalog fails to load: - -```bash -# Verify catalog exists -ls -lh dataset/catalog.jsonl - -# Verify JSONL syntax -python -c "import json; [json.loads(l) for l in open('dataset/catalog.jsonl')]" - -# Check environment variable -echo $SOFTWARE_CATALOG -``` - -### Index Build Error - -If FAISS index building fails: - -```bash -# Check artifacts directory -ls -lh artifacts/rag_index/ - -# Rebuild manually -rm -rf artifacts/rag_index/ -ai_agent sync -``` - -## Next Steps - -- Configure [Environment Variables](environment.md) -- Review the [Changelog](changelog.md) -- Return to [Getting Started](../getting-started/quickstart.md) +# CLI Commands + +The AI Imaging Agent provides a command-line interface for starting the application and managing the software catalog. + +## Available Commands + +### ai_agent chat + +Launch the chat-based user interface. + +```bash +ai_agent chat +``` + +**What it does**: + +1. Performs startup catalog synchronization +2. Loads the FAISS index +3. Initializes the retrieval and agent pipelines +4. Launches the Gradio web interface on `http://127.0.0.1:7860` +5. Starts background catalog refresh (if configured) + +**Options**: None (all configuration via `.env` and `config.yaml`) + +**Example**: + +```bash +$ ai_agent chat +[startup-sync] 150 โ†’ dataset/catalog.jsonl +[startup-refresh] catalog unchanged; keeping existing FAISS index +Running on local URL: http://127.0.0.1:7860 + +To create a public link, set `share=True` in `launch()`. +``` + +**Background Refresh**: + +If `SYNC_EVERY_HOURS` is set in `.env`, the catalog will auto-refresh in the background: + +```dotenv +SYNC_EVERY_HOURS=24 # Check every 24 hours +``` + +### ai_agent sync + +Manually synchronize the software catalog and rebuild the index. + +```bash +ai_agent sync +``` + +**What it does**: + +1. Loads the software catalog from `SOFTWARE_CATALOG` path +2. Embeds all tool descriptions using BGE-M3 +3. Builds FAISS vector index +4. Saves artifacts to `artifacts/rag_index/` + +**When to use**: + +- After editing `catalog.jsonl` +- After adding new tools +- To force index rebuild +- For testing catalog changes + +**Example**: + +```bash +$ ai_agent sync +[sync] 150 โ†’ dataset/catalog.jsonl +[sync] Embedding 150 tools... (5.2s) +[sync] Building FAISS index... +[sync] Saved to artifacts/rag_index/ +[sync] Sync complete. +``` + +## Command Aliases + +Both commands are available with either `ai_agent` or `ai-agent`: + +```bash +ai_agent chat # Works +ai-agent chat # Also works + +ai_agent sync # Works +ai-agent sync # Also works +``` + +## Common Usage Patterns + +### Development Workflow + +```bash +# Edit catalog +vim dataset/catalog.jsonl + +# Sync catalog +ai_agent sync + +# Test changes +ai_agent chat +``` + + + +### Testing & Development + +```bash +# Enable debug logging +export LOGLEVEL_CONSOLE=DEBUG +export LOG_PROMPTS=1 +ai_agent chat +``` + +## Environment Variables + +All configuration is via environment variables (see [Environment Variables Reference](environment.md)). + +## Exit Codes + +- **0**: Success +- **1**: General error (see logs) + +## Troubleshooting + +### Command Not Found + +If you see `command not found: ai_agent`: + +```bash +# Ensure package is installed +pip install -e . + +# Check installation +pip list | grep ai-agent + +# Try with python -m +python -m ai_agent.cli chat +``` + +### Port Already in Use + +If port 7860 is occupied: + +```bash +# Find and kill process +lsof -ti:7860 | xargs kill -9 + +# Or change port in code (ui/app.py) +``` + +### Catalog Load Error + +If catalog fails to load: + +```bash +# Verify catalog exists +ls -lh dataset/catalog.jsonl + +# Verify JSONL syntax +python -c "import json; [json.loads(l) for l in open('dataset/catalog.jsonl')]" + +# Check environment variable +echo $SOFTWARE_CATALOG +``` + +### Index Build Error + +If FAISS index building fails: + +```bash +# Check artifacts directory +ls -lh artifacts/rag_index/ + +# Rebuild manually +rm -rf artifacts/rag_index/ +ai_agent sync +``` + +## Next Steps + +- Configure [Environment Variables](environment.md) +- Review the [Changelog](changelog.md) +- Return to [Getting Started](../getting-started/quickstart.md) diff --git a/docs/reference/environment.md b/docs/reference/environment.md index 9f9a45c22d1df2905082032b486c505766838975..02ad82204fe7e04bfdf9db2490fa482ecd8b0f6f 100644 --- a/docs/reference/environment.md +++ b/docs/reference/environment.md @@ -1,309 +1,309 @@ -# Environment Variables - -Configuration for the AI Imaging Agent is managed via environment variables, typically defined in a `.env` file. - -## Required Variables - -### OPENAI_API_KEY - -OpenAI API key for vision-language model calls. - -```dotenv -OPENAI_API_KEY=sk-xxxx -``` - -**Where to get it**: [OpenAI API Keys](https://platform.openai.com/api-keys) - -**Required**: Yes (unless using alternative model provider) - -**Used by**: Agent VLM calls, tool selection - -## Optional Variables - -### SOFTWARE_CATALOG - -Path to the software catalog JSONL file. - -```dotenv -SOFTWARE_CATALOG=dataset/catalog.jsonl -``` - -**Default**: `dataset/catalog.jsonl` - -**Required**: No (uses default) - -### TOP_K - -Number of candidate tools to retrieve from FAISS search. - -```dotenv -TOP_K=8 -``` - -**Default**: `8` - -**Range**: 1-50 (recommended: 5-10) - -**Impact**: More candidates = better recall but slower VLM calls - -### NUM_CHOICES - -Number of final tool recommendations to return to user. - -```dotenv -NUM_CHOICES=3 -``` - -**Default**: `3` - -**Range**: 1-10 (recommended: 3-5) - -**Impact**: Too many recommendations can overwhelm users - -### GITHUB_TOKEN - -GitHub personal access token for repository info tool. - -```dotenv -GITHUB_TOKEN=ghp_xxxx -``` - -**Where to get it**: [GitHub Tokens](https://github.com/settings/tokens) - -**Permissions needed**: `public_repo` (read access) - -**Required**: No (tool gracefully degrades without it) - -**Benefits**: Higher API rate limits, access to private repos - -### SYNC_EVERY_HOURS - -Auto-refresh catalog interval in hours. - -```dotenv -SYNC_EVERY_HOURS=24 -``` - -**Default**: `0` (disabled) - -**Range**: 0 (disabled) or โ‰ฅ1 - -**Behavior**: Background thread checks catalog every N hours and rebuilds index if changed - -## Logging Configuration - -### LOGLEVEL_CONSOLE - -Console logging level. - -```dotenv -LOGLEVEL_CONSOLE=WARNING -``` - -**Options**: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL` - -**Default**: `WARNING` - -**Recommendation**: -- Development: `DEBUG` or `INFO` -- Production: `WARNING` - -### LOGLEVEL_FILE - -File logging level. - -```dotenv -LOGLEVEL_FILE=INFO -``` - -**Options**: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL` - -**Default**: `INFO` - -**Files written to**: `logs/app_YYYYMMDD.log` - -### FILE_LOG - -Enable file logging. - -```dotenv -FILE_LOG=1 -``` - -**Options**: `0` (disabled), `1` (enabled) - -**Default**: `1` - -### LOG_DIR - -Directory for log files. - -```dotenv -LOG_DIR=logs -``` - -**Default**: `logs` - -**Created automatically** if it doesn't exist - -### LOG_PROMPTS - -Save VLM prompts and images for debugging. - -```dotenv -LOG_PROMPTS=1 -``` - -**Options**: `0` (disabled), `1` (enabled) - -**Default**: `0` - -**Writes to**: `logs/prompts/YYYYMMDD_HHMMSS/` - -**Contents**: -- `prompt.txt`: Text prompt sent to VLM -- `image_*.png`: Images included in prompt -- `response.json`: VLM response -- `metadata.json`: Request metadata - -**Warning**: Can consume significant disk space over time - -## Model Configuration - -### CONFIG_PATH - -Path to YAML model configuration file. - -```dotenv -CONFIG_PATH=config.yaml -``` - -**Default**: `config.yaml` - -**See**: `config.yaml` for model configuration details - -### Alternative Model Providers - -For custom OpenAI-compatible endpoints, configure in `config.yaml`: - -```yaml -agent_model: - name: "model-name" - base_url: "https://api.example.com/v1" - api_key_env: "CUSTOM_API_KEY" -``` - -Then in `.env`: - -```dotenv -CUSTOM_API_KEY=your-key-here -``` - - - -## .env File Example - -Complete example `.env` file: - -```dotenv -# Required -OPENAI_API_KEY=sk-xxxx - -# Optional: Alternative providers -EPFL_API_KEY=sk-xxxx -GITHUB_TOKEN=ghp_xxxx - -# Catalog -SOFTWARE_CATALOG=dataset/catalog.jsonl -SYNC_EVERY_HOURS=24 - -# Pipeline -TOP_K=8 -NUM_CHOICES=3 - -# Logging -LOGLEVEL_CONSOLE=WARNING -LOGLEVEL_FILE=INFO -FILE_LOG=1 -LOG_DIR=logs -LOG_PROMPTS=0 # Set to 1 for debugging - -# Model configuration -CONFIG_PATH=config.yaml -``` - -## Loading Environment Variables - -### Automatic Loading - -The application automatically loads `.env` from the repository root: - -```python -from dotenv import load_dotenv -load_dotenv() # Loads .env automatically -``` - -### Manual Loading - -```bash -# Export manually -export OPENAI_API_KEY=sk-xxxx -export TOP_K=8 - -# Or source .env -set -a -source .env -set +a -``` - -### Docker - -Pass environment variables to Docker: - -```bash -docker run --env-file .env ai-agent -``` - -## Security Best Practices - -!!! warning "Never commit .env files" - Add `.env` to `.gitignore` to prevent accidental commits - -!!! warning "Protect API keys" - Treat API keys as sensitive credentials: - - Never share in public repositories - - Rotate keys if exposed - - Use environment-specific keys (dev/prod) - -!!! tip "Use .env.example" - Create `.env.example` with dummy values for documentation: - - ```dotenv - OPENAI_API_KEY=sk-your-key-here - GITHUB_TOKEN=ghp-your-token-here - ``` - -## Next Steps - -- Review [CLI Commands](cli.md) -- Check [Configuration Guide](../getting-started/configuration.md) -- See [Changelog](changelog.md) +# Environment Variables + +Configuration for the AI Imaging Agent is managed via environment variables, typically defined in a `.env` file. + +## Required Variables + +### OPENAI_API_KEY + +OpenAI API key for vision-language model calls. + +```dotenv +OPENAI_API_KEY=sk-xxxx +``` + +**Where to get it**: [OpenAI API Keys](https://platform.openai.com/api-keys) + +**Required**: Yes (unless using alternative model provider) + +**Used by**: Agent VLM calls, tool selection + +## Optional Variables + +### SOFTWARE_CATALOG + +Path to the software catalog JSONL file. + +```dotenv +SOFTWARE_CATALOG=dataset/catalog.jsonl +``` + +**Default**: `dataset/catalog.jsonl` + +**Required**: No (uses default) + +### TOP_K + +Number of candidate tools to retrieve from FAISS search. + +```dotenv +TOP_K=8 +``` + +**Default**: `8` + +**Range**: 1-50 (recommended: 5-10) + +**Impact**: More candidates = better recall but slower VLM calls + +### NUM_CHOICES + +Number of final tool recommendations to return to user. + +```dotenv +NUM_CHOICES=3 +``` + +**Default**: `3` + +**Range**: 1-10 (recommended: 3-5) + +**Impact**: Too many recommendations can overwhelm users + +### GITHUB_TOKEN + +GitHub personal access token for repository info tool. + +```dotenv +GITHUB_TOKEN=ghp_xxxx +``` + +**Where to get it**: [GitHub Tokens](https://github.com/settings/tokens) + +**Permissions needed**: `public_repo` (read access) + +**Required**: No (tool gracefully degrades without it) + +**Benefits**: Higher API rate limits, access to private repos + +### SYNC_EVERY_HOURS + +Auto-refresh catalog interval in hours. + +```dotenv +SYNC_EVERY_HOURS=24 +``` + +**Default**: `0` (disabled) + +**Range**: 0 (disabled) or โ‰ฅ1 + +**Behavior**: Background thread checks catalog every N hours and rebuilds index if changed + +## Logging Configuration + +### LOGLEVEL_CONSOLE + +Console logging level. + +```dotenv +LOGLEVEL_CONSOLE=WARNING +``` + +**Options**: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL` + +**Default**: `WARNING` + +**Recommendation**: +- Development: `DEBUG` or `INFO` +- Production: `WARNING` + +### LOGLEVEL_FILE + +File logging level. + +```dotenv +LOGLEVEL_FILE=INFO +``` + +**Options**: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL` + +**Default**: `INFO` + +**Files written to**: `logs/app_YYYYMMDD.log` + +### FILE_LOG + +Enable file logging. + +```dotenv +FILE_LOG=1 +``` + +**Options**: `0` (disabled), `1` (enabled) + +**Default**: `1` + +### LOG_DIR + +Directory for log files. + +```dotenv +LOG_DIR=logs +``` + +**Default**: `logs` + +**Created automatically** if it doesn't exist + +### LOG_PROMPTS + +Save VLM prompts and images for debugging. + +```dotenv +LOG_PROMPTS=1 +``` + +**Options**: `0` (disabled), `1` (enabled) + +**Default**: `0` + +**Writes to**: `logs/prompts/YYYYMMDD_HHMMSS/` + +**Contents**: +- `prompt.txt`: Text prompt sent to VLM +- `image_*.png`: Images included in prompt +- `response.json`: VLM response +- `metadata.json`: Request metadata + +**Warning**: Can consume significant disk space over time + +## Model Configuration + +### CONFIG_PATH + +Path to YAML model configuration file. + +```dotenv +CONFIG_PATH=config.yaml +``` + +**Default**: `config.yaml` + +**See**: `config.yaml` for model configuration details + +### Alternative Model Providers + +For custom OpenAI-compatible endpoints, configure in `config.yaml`: + +```yaml +agent_model: + name: "model-name" + base_url: "https://api.example.com/v1" + api_key_env: "CUSTOM_API_KEY" +``` + +Then in `.env`: + +```dotenv +CUSTOM_API_KEY=your-key-here +``` + + + +## .env File Example + +Complete example `.env` file: + +```dotenv +# Required +OPENAI_API_KEY=sk-xxxx + +# Optional: Alternative providers +EPFL_API_KEY=sk-xxxx +GITHUB_TOKEN=ghp_xxxx + +# Catalog +SOFTWARE_CATALOG=dataset/catalog.jsonl +SYNC_EVERY_HOURS=24 + +# Pipeline +TOP_K=8 +NUM_CHOICES=3 + +# Logging +LOGLEVEL_CONSOLE=WARNING +LOGLEVEL_FILE=INFO +FILE_LOG=1 +LOG_DIR=logs +LOG_PROMPTS=0 # Set to 1 for debugging + +# Model configuration +CONFIG_PATH=config.yaml +``` + +## Loading Environment Variables + +### Automatic Loading + +The application automatically loads `.env` from the repository root: + +```python +from dotenv import load_dotenv +load_dotenv() # Loads .env automatically +``` + +### Manual Loading + +```bash +# Export manually +export OPENAI_API_KEY=sk-xxxx +export TOP_K=8 + +# Or source .env +set -a +source .env +set +a +``` + +### Docker + +Pass environment variables to Docker: + +```bash +docker run --env-file .env ai-agent +``` + +## Security Best Practices + +!!! warning "Never commit .env files" + Add `.env` to `.gitignore` to prevent accidental commits + +!!! warning "Protect API keys" + Treat API keys as sensitive credentials: + - Never share in public repositories + - Rotate keys if exposed + - Use environment-specific keys (dev/prod) + +!!! tip "Use .env.example" + Create `.env.example` with dummy values for documentation: + + ```dotenv + OPENAI_API_KEY=sk-your-key-here + GITHUB_TOKEN=ghp-your-token-here + ``` + +## Next Steps + +- Review [CLI Commands](cli.md) +- Check [Configuration Guide](../getting-started/configuration.md) +- See [Changelog](changelog.md) diff --git a/docs/user-guide/advanced-features.md b/docs/user-guide/advanced-features.md index 084a13fef7cf823662b239e232e41af686e06fa5..78dad8f38dbaca1b552834584696a34a3ae5fc07 100644 --- a/docs/user-guide/advanced-features.md +++ b/docs/user-guide/advanced-features.md @@ -1,387 +1,387 @@ -# Advanced Features (not tested for now..) - -The AI Imaging Agent includes several advanced features for power users and specialized use cases. - -## Control Tags - -Control tags modify agent behavior using special syntax in your queries. - -### Exclude Tools - -Filter out specific tools from results: - -``` -Find lung segmentation tools [EXCLUDE:totalsegmentator|medicalsam] -``` - -**Syntax**: `[EXCLUDE:tool1|tool2|tool3]` - -**Use cases**: - -- You've already tried certain tools -- Exclude tools you don't have access to -- Filter by licensing (exclude proprietary tools) -- Remove tools with specific limitations - -**Example**: -``` -You: Segment kidneys [EXCLUDE:totalsegmentator] -Agent: [Returns kidney segmentation tools except TotalSegmentator] - -You: Find open-source options [EXCLUDE:proprietarytool1|proprietarytool2] -Agent: [Returns only open-source tools] -``` - -### Notes - -Only `[EXCLUDE:...]` is currently interpreted as a control tag for retrieval filtering. - -## Alternative Searches - -Request the agent to search with different strategies. - -### Requesting Alternatives - -Use natural language: - -``` -Can you search for alternatives? -Show me other options -Find different tools -What else is available? -``` - -**What happens**: - -- Agent formulates alternative query -- Uses different phrasing/keywords for broader coverage -- Searches with different emphasis -- Returns new set of recommendations - -**Limit**: Up to 3 alternative searches per conversation - -### When to Use - -- Initial results don't quite match -- Want to see different approaches -- Exploring the catalog -- Looking for specialized tools - -**Example conversation**: -``` -You: Segment lungs from this CT -Agent: [Provides general lung segmentation tools] - -You: Can you search for alternatives? -Agent: [Searches with emphasis on "airway segmentation", "pulmonary analysis"] - -You: Show me other options -Agent: [Searches with emphasis on "CT thorax processing", "respiratory imaging"] -``` - -## Multi-Model Support - -### Selecting Different Models - -The UI provides a model selector dropdown: - -Available models (configurable in `config.yaml`): - -- **gpt-4o-mini**: Faster, lower cost -- **gpt-4o**: Higher accuracy, multimodal -- **gpt-5.1**: Latest capabilities (if available) -- **Custom endpoints**: EPFL, local servers, etc. - -### Model Trade-offs - -| Model | Speed | Cost | Accuracy | Vision | -|-------|-------|------|----------|--------| -| gpt-4o-mini | โšกโšกโšก | ๐Ÿ’ฐ | โญโญโญ | โœ… | -| gpt-4o | โšกโšก | ๐Ÿ’ฐ๐Ÿ’ฐ | โญโญโญโญ | โœ…โœ… | -| gpt-5.1 | โšก | ๐Ÿ’ฐ๐Ÿ’ฐ๐Ÿ’ฐ | โญโญโญโญโญ | โœ…โœ…โœ… | - -### When to Switch Models - -**Use gpt-4o-mini when**: - -- Doing quick explorations -- Cost is a concern -- Tasks are straightforward -- Query is well-specified - -**Use gpt-4o when**: - -- Complex visual analysis needed -- Accuracy is critical -- Ambiguous queries -- Multi-step reasoning required - -**Use gpt-5.1 when**: - -- Maximum accuracy needed -- Complex multi-modal tasks -- Research/publication work - -## Repository Info Tool - -### What It Does - -The agent can fetch detailed information about GitHub repositories: - -``` -You: Tell me about TotalSegmentator -Agent: [Fetches repo info from GitHub via DeepWiki or repocards] - -Repository: wasserth/TotalSegmentator -Description: Automated multi-organ segmentation in CT and MR images -Stars: 1.2k -Language: Python -Topics: segmentation, medical-imaging, deep-learning -Last Updated: 2024-03-15 -License: Apache-2.0 -``` - -### Data Sources - -1. **DeepWiki MCP** (primary): Fast, pre-indexed repository documentation -2. **Repocards** (fallback): Direct library-based fetch - -### Usage - -Ask about tools naturally: - -``` -What is [tool name]? -Tell me more about [repository] -Show me details for [tool] -``` - -## Conversation State Management - -### State Tracking - -The agent maintains state across conversation: - -- **Uploaded files**: All files in session -- **Preview images**: Converted images for VLM -- **Excluded tools**: Tools filtered via `[EXCLUDE:]` -- **Conversation history**: Previous messages and context -- **Turn counter**: Current conversation turn - -### Viewing State - -In the sidebar (debug mode): - -```json -{ - "conversation_turn": 3, - "uploaded_files": ["scan.dcm", "brain.nii"], - "excluded_tools": ["tool1", "tool2"], - "preview_images": ["/tmp/scan_preview.png"] -} -``` - -### Resetting State - -To start fresh: -- Refresh the page -- Clear uploaded files -- Start new conversation - -## Retrieval Query Behavior - -### How It Works - -The retrieval pipeline currently does not add semantic neighbor terms. -Instead, it builds retrieval queries from: - -``` -User text: "segment brain" -+ format hints from files: format:DICOM / format:NIfTI -+ compact image metadata: modality, anatomy, dimensions (when available) -``` - -**Based on**: - -- BGE-M3 embeddings -- Format-aware hinting from uploaded files -- Metadata-aware context from image inspection - -### Benefits - -- โœ… Stronger format compatibility matching -- โœ… Better ranking for modality/dimension-specific tasks -- โœ… More predictable retrieval behavior - -### Customization - -If initial results are too sparse, the pipeline retries with a broader query -formulation automatically. - -## Format-Aware Matching - -### Input Format Tokens - -File uploads add format tokens to queries: - -``` -Uploaded: scan.dcm (DICOM) -Query enhancement: "segment lungs format:DICOM format:CT format:3D" -``` - -### How It Helps - -- **Narrows results**: Shows compatible tools first -- **Boosts relevance**: DICOM tools rank higher for DICOM -- **Compatibility check**: Agent verifies format support - -### Supported Formats - -Tokens added for: -- File extension (`.dcm`, `.nii`, `.png`) -- Detected format (DICOM, NIfTI, TIFF) -- Modality for medical images (CT, MRI, XR) -- Dimensions (2D, 3D, 4D) - -## Iterative Retrieval - -### Auto-Retry on Low Results - -If initial search returns <5 candidates: - -1. **Retry #1**: Alternative query with semantic expansion -2. **Retry #2**: Further expansion with broader terms -3. **Max 2 retries**: Then return best available - -### Why It Matters - -- Handles rare/specialized queries -- Finds tools even with limited matches -- Automatic - no user action needed - -### Example - -``` -Query: "segment rare anatomical structure" -Initial: 2 candidates found -Retry 1: Expanded to "segment anatomy structure region organ" -Result: 7 candidates found โœ“ -``` - -## Debug Features - -### Prompt Logging - -Enable in `.env`: - -```dotenv -LOG_PROMPTS=1 -``` - -**Saves**: - -- VLM prompts sent to API -- Images included in prompts -- Response JSON -- Timestamp and metadata - -**Location**: `logs/prompts/YYYYMMDD_HHMMSS/` - -**Contents**: -``` -logs/prompts/20240315_143022/ -โ”œโ”€โ”€ prompt.txt # Text prompt -โ”œโ”€โ”€ image_0.png # Uploaded image -โ”œโ”€โ”€ response.json # API response -โ””โ”€โ”€ metadata.json # Request metadata -``` - -### Execution Traces - -Always shown in chat (expandable): - -```html -
-๐Ÿ”ง Execution Trace -...detailed logs... -
-``` - -Shows: - -- Tool calls made -- Parameters used -- API responses -- Timing information - -## Catalog Synchronization - -### Auto-Refresh - -Configured via `.env`: - -```dotenv -SYNC_EVERY_HOURS=24 -``` - -**Behavior**: - -- Background thread checks for catalog updates -- Reloads FAISS index if changed -- No UI interruption -- Logs refresh activity - -### Manual Sync - -Force synchronization: - -```bash -ai_agent sync -``` - -Updates: - -- Software catalog -- Embeddings -- FAISS index -- Vocabulary for expansion - -## Advanced Configuration - -### Custom Catalog - -Use your own tool catalog: - -```dotenv -SOFTWARE_CATALOG=/path/to/custom_catalog.jsonl -``` - -**Format**: JSONL with schema.org SoftwareSourceCode - -### API Endpoints - -Configure custom OpenAI-compatible endpoints in `config.yaml`: - -```yaml -available_models: - - display_name: "Local LLM" - name: "llama-3.1" - base_url: "http://localhost:8000/v1" - api_key_env: "LOCAL_API_KEY" -``` - -### Pipeline Parameters - -Fine-tune retrieval: - -```dotenv -TOP_K=8 # Candidates to retrieve -NUM_CHOICES=3 # Final recommendations -RERANK_TOP_N=20 # Candidates before reranking -``` - -## Next Steps - -- Dive into [Architecture Overview](../architecture/overview.md) -- Learn about [Development and Contributing](../development/contributing.md) -- Check [Environment Variables Reference](../reference/environment.md) +# Advanced Features (not tested for now..) + +The AI Imaging Agent includes several advanced features for power users and specialized use cases. + +## Control Tags + +Control tags modify agent behavior using special syntax in your queries. + +### Exclude Tools + +Filter out specific tools from results: + +``` +Find lung segmentation tools [EXCLUDE:totalsegmentator|medicalsam] +``` + +**Syntax**: `[EXCLUDE:tool1|tool2|tool3]` + +**Use cases**: + +- You've already tried certain tools +- Exclude tools you don't have access to +- Filter by licensing (exclude proprietary tools) +- Remove tools with specific limitations + +**Example**: +``` +You: Segment kidneys [EXCLUDE:totalsegmentator] +Agent: [Returns kidney segmentation tools except TotalSegmentator] + +You: Find open-source options [EXCLUDE:proprietarytool1|proprietarytool2] +Agent: [Returns only open-source tools] +``` + +### Notes + +Only `[EXCLUDE:...]` is currently interpreted as a control tag for retrieval filtering. + +## Alternative Searches + +Request the agent to search with different strategies. + +### Requesting Alternatives + +Use natural language: + +``` +Can you search for alternatives? +Show me other options +Find different tools +What else is available? +``` + +**What happens**: + +- Agent formulates alternative query +- Uses different phrasing/keywords for broader coverage +- Searches with different emphasis +- Returns new set of recommendations + +**Limit**: Up to 3 alternative searches per conversation + +### When to Use + +- Initial results don't quite match +- Want to see different approaches +- Exploring the catalog +- Looking for specialized tools + +**Example conversation**: +``` +You: Segment lungs from this CT +Agent: [Provides general lung segmentation tools] + +You: Can you search for alternatives? +Agent: [Searches with emphasis on "airway segmentation", "pulmonary analysis"] + +You: Show me other options +Agent: [Searches with emphasis on "CT thorax processing", "respiratory imaging"] +``` + +## Multi-Model Support + +### Selecting Different Models + +The UI provides a model selector dropdown: + +Available models (configurable in `config.yaml`): + +- **gpt-4o-mini**: Faster, lower cost +- **gpt-4o**: Higher accuracy, multimodal +- **gpt-5.1**: Latest capabilities (if available) +- **Custom endpoints**: EPFL, local servers, etc. + +### Model Trade-offs + +| Model | Speed | Cost | Accuracy | Vision | +|-------|-------|------|----------|--------| +| gpt-4o-mini | โšกโšกโšก | ๐Ÿ’ฐ | โญโญโญ | โœ… | +| gpt-4o | โšกโšก | ๐Ÿ’ฐ๐Ÿ’ฐ | โญโญโญโญ | โœ…โœ… | +| gpt-5.1 | โšก | ๐Ÿ’ฐ๐Ÿ’ฐ๐Ÿ’ฐ | โญโญโญโญโญ | โœ…โœ…โœ… | + +### When to Switch Models + +**Use gpt-4o-mini when**: + +- Doing quick explorations +- Cost is a concern +- Tasks are straightforward +- Query is well-specified + +**Use gpt-4o when**: + +- Complex visual analysis needed +- Accuracy is critical +- Ambiguous queries +- Multi-step reasoning required + +**Use gpt-5.1 when**: + +- Maximum accuracy needed +- Complex multi-modal tasks +- Research/publication work + +## Repository Info Tool + +### What It Does + +The agent can fetch detailed information about GitHub repositories: + +``` +You: Tell me about TotalSegmentator +Agent: [Fetches repo info from GitHub via DeepWiki or repocards] + +Repository: wasserth/TotalSegmentator +Description: Automated multi-organ segmentation in CT and MR images +Stars: 1.2k +Language: Python +Topics: segmentation, medical-imaging, deep-learning +Last Updated: 2024-03-15 +License: Apache-2.0 +``` + +### Data Sources + +1. **DeepWiki MCP** (primary): Fast, pre-indexed repository documentation +2. **Repocards** (fallback): Direct library-based fetch + +### Usage + +Ask about tools naturally: + +``` +What is [tool name]? +Tell me more about [repository] +Show me details for [tool] +``` + +## Conversation State Management + +### State Tracking + +The agent maintains state across conversation: + +- **Uploaded files**: All files in session +- **Preview images**: Converted images for VLM +- **Excluded tools**: Tools filtered via `[EXCLUDE:]` +- **Conversation history**: Previous messages and context +- **Turn counter**: Current conversation turn + +### Viewing State + +In the sidebar (debug mode): + +```json +{ + "conversation_turn": 3, + "uploaded_files": ["scan.dcm", "brain.nii"], + "excluded_tools": ["tool1", "tool2"], + "preview_images": ["/tmp/scan_preview.png"] +} +``` + +### Resetting State + +To start fresh: +- Refresh the page +- Clear uploaded files +- Start new conversation + +## Retrieval Query Behavior + +### How It Works + +The retrieval pipeline currently does not add semantic neighbor terms. +Instead, it builds retrieval queries from: + +``` +User text: "segment brain" ++ format hints from files: format:DICOM / format:NIfTI ++ compact image metadata: modality, anatomy, dimensions (when available) +``` + +**Based on**: + +- BGE-M3 embeddings +- Format-aware hinting from uploaded files +- Metadata-aware context from image inspection + +### Benefits + +- โœ… Stronger format compatibility matching +- โœ… Better ranking for modality/dimension-specific tasks +- โœ… More predictable retrieval behavior + +### Customization + +If initial results are too sparse, the pipeline retries with a broader query +formulation automatically. + +## Format-Aware Matching + +### Input Format Tokens + +File uploads add format tokens to queries: + +``` +Uploaded: scan.dcm (DICOM) +Query enhancement: "segment lungs format:DICOM format:CT format:3D" +``` + +### How It Helps + +- **Narrows results**: Shows compatible tools first +- **Boosts relevance**: DICOM tools rank higher for DICOM +- **Compatibility check**: Agent verifies format support + +### Supported Formats + +Tokens added for: +- File extension (`.dcm`, `.nii`, `.png`) +- Detected format (DICOM, NIfTI, TIFF) +- Modality for medical images (CT, MRI, XR) +- Dimensions (2D, 3D, 4D) + +## Iterative Retrieval + +### Auto-Retry on Low Results + +If initial search returns <5 candidates: + +1. **Retry #1**: Alternative query with semantic expansion +2. **Retry #2**: Further expansion with broader terms +3. **Max 2 retries**: Then return best available + +### Why It Matters + +- Handles rare/specialized queries +- Finds tools even with limited matches +- Automatic - no user action needed + +### Example + +``` +Query: "segment rare anatomical structure" +Initial: 2 candidates found +Retry 1: Expanded to "segment anatomy structure region organ" +Result: 7 candidates found โœ“ +``` + +## Debug Features + +### Prompt Logging + +Enable in `.env`: + +```dotenv +LOG_PROMPTS=1 +``` + +**Saves**: + +- VLM prompts sent to API +- Images included in prompts +- Response JSON +- Timestamp and metadata + +**Location**: `logs/prompts/YYYYMMDD_HHMMSS/` + +**Contents**: +``` +logs/prompts/20240315_143022/ +โ”œโ”€โ”€ prompt.txt # Text prompt +โ”œโ”€โ”€ image_0.png # Uploaded image +โ”œโ”€โ”€ response.json # API response +โ””โ”€โ”€ metadata.json # Request metadata +``` + +### Execution Traces + +Always shown in chat (expandable): + +```html +
+๐Ÿ”ง Execution Trace +...detailed logs... +
+``` + +Shows: + +- Tool calls made +- Parameters used +- API responses +- Timing information + +## Catalog Synchronization + +### Auto-Refresh + +Configured via `.env`: + +```dotenv +SYNC_EVERY_HOURS=24 +``` + +**Behavior**: + +- Background thread checks for catalog updates +- Reloads FAISS index if changed +- No UI interruption +- Logs refresh activity + +### Manual Sync + +Force synchronization: + +```bash +ai_agent sync +``` + +Updates: + +- Software catalog +- Embeddings +- FAISS index +- Vocabulary for expansion + +## Advanced Configuration + +### Custom Catalog + +Use your own tool catalog: + +```dotenv +SOFTWARE_CATALOG=/path/to/custom_catalog.jsonl +``` + +**Format**: JSONL with schema.org SoftwareSourceCode + +### API Endpoints + +Configure custom OpenAI-compatible endpoints in `config.yaml`: + +```yaml +available_models: + - display_name: "Local LLM" + name: "llama-3.1" + base_url: "http://localhost:8000/v1" + api_key_env: "LOCAL_API_KEY" +``` + +### Pipeline Parameters + +Fine-tune retrieval: + +```dotenv +TOP_K=8 # Candidates to retrieve +NUM_CHOICES=3 # Final recommendations +RERANK_TOP_N=20 # Candidates before reranking +``` + +## Next Steps + +- Dive into [Architecture Overview](../architecture/overview.md) +- Learn about [Development and Contributing](../development/contributing.md) +- Check [Environment Variables Reference](../reference/environment.md) diff --git a/docs/user-guide/chat-interface.md b/docs/user-guide/chat-interface.md index 540b3fe1bdd318f19c6923baf2f4e8cd308ce50c..d669bf1fecc033daf2313d254566c25430177e2e 100644 --- a/docs/user-guide/chat-interface.md +++ b/docs/user-guide/chat-interface.md @@ -1,275 +1,275 @@ -# Using the Chat Interface - -The AI Imaging Agent provides a conversational interface for discovering and using imaging software. This guide explains how to interact with the chat interface effectively. - -## Interface Layout - -The interface consists of three main areas: - -### Left Panel: Chat Conversation -- **Message History**: Your conversation with the agent -- **Rich Media Rendering**: Images, tool cards, and files are displayed inline -- **Input Box**: Type your messages at the bottom -- **File Upload**: Attach files via the paperclip icon or drag-and-drop - -### Right Panel: Sidebar -- **Files Tab**: View uploaded files with format information -- **State Tab**: Debug information showing conversation state - -### Header -- **Model Selector**: Choose which AI model to use -- **Settings**: Access configuration options - -## Basic Workflow - -### 1. Upload Files - -Upload images or other files in several ways: - -- **Drag and Drop**: Drag files directly onto the upload area -- **Click to Browse**: Click the upload area to select files -- **Attach to Message**: Use the paperclip icon in the input box - -Files are automatically processed and metadata is extracted. - -### 2. Describe Your Task - -Use natural language to describe what you want to do: - -!!! example "Good Task Descriptions" - - "I want to segment the lungs from this CT scan" - - "Help me detect tumors in this MRI" - - "I need to register these two brain images" - - "Extract text from this medical report" - - "Classify the organ shown in this ultrasound" - -### 3. Review Recommendations - -The agent returns ranked tool recommendations with: - -- **Tool Cards**: Each tool is presented in a card format -- **Accuracy Scores**: Confidence levels for each recommendation -- **Explanations**: Why each tool matches your request -- **Metadata**: Technical details about compatibility - -### 4. Run Demos (Optional) - -The agent may offer to run demos: - -``` -Agent: Would you like me to run the demo with your image? -``` - -Respond with affirmative language: -- "yes" -- "sure" -- "ok" -- "please" -- "go ahead" - -The agent will execute the tool and show results. - -## Multi-Turn Conversations - -The agent maintains context across multiple messages: - -!!! example "Multi-Turn Example" - ``` - You: I have a lung CT scan [uploads file] - - Agent: I can see you have a DICOM CT image. What would you like to do with it? - - You: Segment the airways - - Agent: [Provides airway segmentation tool recommendations] - - You: What about segmenting the whole lung? - - Agent: [Provides lung segmentation tools, remembering you're working with CT] - - You: Show me alternatives - - Agent: [Provides additional options] - ``` - -## Advanced Features - -### Excluding Tools - -Exclude specific tools using the `[EXCLUDE:...]` tag: - -``` -Find segmentation tools [EXCLUDE:totalsegmentator|medicalsam] -``` - -You can exclude multiple tools separated by `|`. - -### Requesting Alternatives - -Ask the agent to search with different strategies: - -``` -Can you search for alternatives? - -Show me other options - -Find different tools for this task -``` - -The agent can perform up to 3 alternative searches per conversation. - -## Understanding Agent Responses - -### Recommendation Cards - -Each recommendation includes: - -#### Header -- **Rank Number**: 1, 2, 3 (1 = best match) -- **Tool Name**: Software/tool identifier -- **Accuracy Score**: 0-100% confidence - -#### Body -- **Description**: What the tool does -- **Explanation**: Why it matches your task -- **Demo Link**: Click to visit runnable example - -#### Footer Metadata -- **Modalities**: CT, MRI, X-ray, etc. -- **Dimensions**: 2D, 3D, 4D -- **Formats**: Supported file formats (DICOM, NIfTI, etc.) -- **License**: Software license information -- **Tags**: Categorization and keywords - -### Execution Traces - -When demos run, you'll see execution details: - -``` -
-Tool Execution Trace - -Image uploaded to Gradio Space -Processing started... -Result: Success -Output saved to: result.png -
-``` - -Click to expand and see full execution logs. - -### Clarification Questions - -Sometimes the agent needs more information: - -``` -Agent: I found several segmentation tools. Which organ are you trying to segment? - -You: The liver - -Agent: [Provides liver-specific segmentation tools] -``` - -## File Management - -### Uploaded Files List - -The sidebar shows all uploaded files with: - -- **Filename**: Original file name -- **Format**: File type/extension -- **Size**: File size -- **Preview**: Thumbnail (for images) - -### Image Previews - -Medical images are automatically converted: - -- **DICOM**: PNG previews; 3D series use orthogonal composite views (MIPs + central slices) rather than a single slice -- **NIfTI**: PNG previews built from orthogonal composite views of the volume -- **TIFF Stacks**: PNG previews built from orthogonal composite views of the stack -- **Standard 2D Images**: Resized PNG preview of the original image - -Previews are used for VLM analysis while preserving original format metadata. - -### Removing Files - -Click the 'X' button next to a file to remove it from the current session. - -## Conversation State - -The debug sidebar shows: - -### Current State -- **Status**: idle, processing, waiting -- **Conversation Turn**: Current turn number -- **Excluded Tools**: Tools filtered from results - -### Preview Images -- Images prepared for VLM analysis -- Format conversions applied - -## Tips for Effective Interaction - -!!! tip "Be Specific About Requirements" - Mention specific needs: - - - "I need a tool that works with NIfTI files" - - "Must support 3D volumes" - - "Looking for open-source options" - -!!! tip "Use Conversational Language" - Natural language works best: - - - โœ… "Help me find tool that segments kidneys" - - โŒ "kidney_segmentation_tool filter:3D" - -!!! tip "Iterate Based on Results" - If initial results aren't perfect, refine: - - - "Can you find tools with higher accuracy?" - - "Show me open-source alternatives" - - "What about tools that support DICOM?" - -!!! tip "Ask Follow-Up Questions" - The agent maintains context: - - - "What about the second recommendation?" - - "Can you compare these two tools?" - - "Which one is fastest?" - -## Troubleshooting - -### No Recommendations - -If the agent can't find suitable tools: - -- Try rephrasing your query -- Be more specific about the task -- Check that your file uploaded successfully -- Ensure your task matches the catalog domain (imaging/medical) - -### Wrong Recommendations - -If recommendations don't match: - -- Provide more context about your specific needs -- Mention required file format support -- Specify modality or domain -- Use the exclude feature to filter out irrelevant tools - -### Demo Execution Fails - -If a demo doesn't run: - -- Check your internet connection -- Verify the demo link is still active -- Try a different recommended tool -- Check file format compatibility - -## Next Steps - -- Learn about [Supported File Formats](file-formats.md) -- Understand [How Recommendations Work](recommendations.md) -- Explore [Running Demos](running-demos.md) -- Check out [Advanced Features](advanced-features.md) +# Using the Chat Interface + +The AI Imaging Agent provides a conversational interface for discovering and using imaging software. This guide explains how to interact with the chat interface effectively. + +## Interface Layout + +The interface consists of three main areas: + +### Left Panel: Chat Conversation +- **Message History**: Your conversation with the agent +- **Rich Media Rendering**: Images, tool cards, and files are displayed inline +- **Input Box**: Type your messages at the bottom +- **File Upload**: Attach files via the paperclip icon or drag-and-drop + +### Right Panel: Sidebar +- **Files Tab**: View uploaded files with format information +- **State Tab**: Debug information showing conversation state + +### Header +- **Model Selector**: Choose which AI model to use +- **Settings**: Access configuration options + +## Basic Workflow + +### 1. Upload Files + +Upload images or other files in several ways: + +- **Drag and Drop**: Drag files directly onto the upload area +- **Click to Browse**: Click the upload area to select files +- **Attach to Message**: Use the paperclip icon in the input box + +Files are automatically processed and metadata is extracted. + +### 2. Describe Your Task + +Use natural language to describe what you want to do: + +!!! example "Good Task Descriptions" + - "I want to segment the lungs from this CT scan" + - "Help me detect tumors in this MRI" + - "I need to register these two brain images" + - "Extract text from this medical report" + - "Classify the organ shown in this ultrasound" + +### 3. Review Recommendations + +The agent returns ranked tool recommendations with: + +- **Tool Cards**: Each tool is presented in a card format +- **Accuracy Scores**: Confidence levels for each recommendation +- **Explanations**: Why each tool matches your request +- **Metadata**: Technical details about compatibility + +### 4. Run Demos (Optional) + +The agent may offer to run demos: + +``` +Agent: Would you like me to run the demo with your image? +``` + +Respond with affirmative language: +- "yes" +- "sure" +- "ok" +- "please" +- "go ahead" + +The agent will execute the tool and show results. + +## Multi-Turn Conversations + +The agent maintains context across multiple messages: + +!!! example "Multi-Turn Example" + ``` + You: I have a lung CT scan [uploads file] + + Agent: I can see you have a DICOM CT image. What would you like to do with it? + + You: Segment the airways + + Agent: [Provides airway segmentation tool recommendations] + + You: What about segmenting the whole lung? + + Agent: [Provides lung segmentation tools, remembering you're working with CT] + + You: Show me alternatives + + Agent: [Provides additional options] + ``` + +## Advanced Features + +### Excluding Tools + +Exclude specific tools using the `[EXCLUDE:...]` tag: + +``` +Find segmentation tools [EXCLUDE:totalsegmentator|medicalsam] +``` + +You can exclude multiple tools separated by `|`. + +### Requesting Alternatives + +Ask the agent to search with different strategies: + +``` +Can you search for alternatives? + +Show me other options + +Find different tools for this task +``` + +The agent can perform up to 3 alternative searches per conversation. + +## Understanding Agent Responses + +### Recommendation Cards + +Each recommendation includes: + +#### Header +- **Rank Number**: 1, 2, 3 (1 = best match) +- **Tool Name**: Software/tool identifier +- **Accuracy Score**: 0-100% confidence + +#### Body +- **Description**: What the tool does +- **Explanation**: Why it matches your task +- **Demo Link**: Click to visit runnable example + +#### Footer Metadata +- **Modalities**: CT, MRI, X-ray, etc. +- **Dimensions**: 2D, 3D, 4D +- **Formats**: Supported file formats (DICOM, NIfTI, etc.) +- **License**: Software license information +- **Tags**: Categorization and keywords + +### Execution Traces + +When demos run, you'll see execution details: + +``` +
+Tool Execution Trace + +Image uploaded to Gradio Space +Processing started... +Result: Success +Output saved to: result.png +
+``` + +Click to expand and see full execution logs. + +### Clarification Questions + +Sometimes the agent needs more information: + +``` +Agent: I found several segmentation tools. Which organ are you trying to segment? + +You: The liver + +Agent: [Provides liver-specific segmentation tools] +``` + +## File Management + +### Uploaded Files List + +The sidebar shows all uploaded files with: + +- **Filename**: Original file name +- **Format**: File type/extension +- **Size**: File size +- **Preview**: Thumbnail (for images) + +### Image Previews + +Medical images are automatically converted: + +- **DICOM**: PNG previews; 3D series use orthogonal composite views (MIPs + central slices) rather than a single slice +- **NIfTI**: PNG previews built from orthogonal composite views of the volume +- **TIFF Stacks**: PNG previews built from orthogonal composite views of the stack +- **Standard 2D Images**: Resized PNG preview of the original image + +Previews are used for VLM analysis while preserving original format metadata. + +### Removing Files + +Click the 'X' button next to a file to remove it from the current session. + +## Conversation State + +The debug sidebar shows: + +### Current State +- **Status**: idle, processing, waiting +- **Conversation Turn**: Current turn number +- **Excluded Tools**: Tools filtered from results + +### Preview Images +- Images prepared for VLM analysis +- Format conversions applied + +## Tips for Effective Interaction + +!!! tip "Be Specific About Requirements" + Mention specific needs: + + - "I need a tool that works with NIfTI files" + - "Must support 3D volumes" + - "Looking for open-source options" + +!!! tip "Use Conversational Language" + Natural language works best: + + - โœ… "Help me find tool that segments kidneys" + - โŒ "kidney_segmentation_tool filter:3D" + +!!! tip "Iterate Based on Results" + If initial results aren't perfect, refine: + + - "Can you find tools with higher accuracy?" + - "Show me open-source alternatives" + - "What about tools that support DICOM?" + +!!! tip "Ask Follow-Up Questions" + The agent maintains context: + + - "What about the second recommendation?" + - "Can you compare these two tools?" + - "Which one is fastest?" + +## Troubleshooting + +### No Recommendations + +If the agent can't find suitable tools: + +- Try rephrasing your query +- Be more specific about the task +- Check that your file uploaded successfully +- Ensure your task matches the catalog domain (imaging/medical) + +### Wrong Recommendations + +If recommendations don't match: + +- Provide more context about your specific needs +- Mention required file format support +- Specify modality or domain +- Use the exclude feature to filter out irrelevant tools + +### Demo Execution Fails + +If a demo doesn't run: + +- Check your internet connection +- Verify the demo link is still active +- Try a different recommended tool +- Check file format compatibility + +## Next Steps + +- Learn about [Supported File Formats](file-formats.md) +- Understand [How Recommendations Work](recommendations.md) +- Explore [Running Demos](running-demos.md) +- Check out [Advanced Features](advanced-features.md) diff --git a/docs/user-guide/file-formats.md b/docs/user-guide/file-formats.md index 40ed7df590930fa9d43c324a9608c00b92e04aaa..c80896b0f811ebe84759d1edf16389156f449432 100644 --- a/docs/user-guide/file-formats.md +++ b/docs/user-guide/file-formats.md @@ -1,282 +1,282 @@ -# Supported File Formats - -The AI Imaging Agent supports a wide range of file formats for medical and scientific imaging, as well as general data files. - -## Image Formats - -### Standard Images - -| Format | Extensions | Description | -|--------|-----------|-------------| -| PNG | `.png` | Portable Network Graphics - lossless compression | -| JPEG | `.jpg`, `.jpeg` | Joint Photographic Experts Group - lossy compression | - -Currently, only PNG and JPEG are accepted for standard images. Other web formats (e.g. WebP, BMP, GIF) should be converted to PNG or JPEG before upload. -**Best for**: General photographs, screenshots, web images - -### Medical Imaging Formats - -#### DICOM - -| Format | Extensions | Description | -|--------|-----------|-------------| -| DICOM | `.dcm`, `.dicom` | Digital Imaging and Communications in Medicine | - -**Features**: - -- Industry standard for medical imaging -- Contains rich metadata (patient info, acquisition parameters) -- Supports multiple modalities (CT, MRI, X-ray, etc.) -- Can store 2D images or 3D volumes - -**Metadata Extracted**: - -- Patient ID, Study Instance UID -- Modality (CT, MR, CR, DX, etc.) -- Image dimensions and spacing -- Acquisition date/time -- Manufacturer and model - -**Example Usage**: -``` -Upload a CT DICOM file and ask: -"Segment the lungs from this scan" -``` - -#### NIfTI - -| Format | Extensions | Description | -|--------|-----------|-------------| -| NIfTI | `.nii`, `.nii.gz` | Neuroimaging Informatics Technology Initiative | - -**Features**: - -- Standard for neuroimaging research -- Supports 3D and 4D (time-series) volumes -- Compact storage with optional gzip compression -- Contains spatial orientation information - -**Metadata Extracted**: - -- Volume dimensions (x, y, z, time) -- Voxel spacing -- Data type and bit depth -- Orientation matrix - -**Example Usage**: -``` -Upload a brain MRI NIfTI file: -"Register this brain scan to MNI space" -``` - -### Scientific Imaging Formats - -#### TIFF/TIFF Stacks - -| Format | Extensions | Description | -|--------|-----------|-------------| -| TIFF | `.tif`, `.tiff` | Tagged Image File Format | - -**Features**: - -- Supports multi-page/multi-frame images -- Common in microscopy and scientific imaging -- Can store extensive metadata -- Lossless compression options - -**Metadata Extracted**: - -- Number of pages/frames (for stacks) -- Dimensions (width, height, channels) -- Color mode (RGB, grayscale, etc.) -- Compression method -- DPI/resolution information - -**Example Usage**: -``` -Upload a microscopy TIFF stack: -"Analyze cell structures in this z-stack" -``` - -## Data Formats - - - -## Format Detection - -The agent automatically detects file formats using: - -1. **File Extension**: Primary detection method -2. **Magic Bytes**: Header inspection for validation -3. **Content Analysis**: Fallback for ambiguous cases - -## Metadata Extraction - -### What Gets Extracted - -For each uploaded file, the agent extracts: - -#### Image Metadata -- **Dimensions**: Width, height, depth (for volumes) -- **Channels**: Grayscale, RGB, RGBA -- **Data Type**: uint8, int16, float32, etc. -- **File Size**: Storage size - -#### Medical Image Metadata -- **Modality**: CT, MRI, X-ray, Ultrasound, PET, etc. -- **Patient Info**: Anonymized IDs -- **Study Info**: Study UID, dates -- **Acquisition Parameters**: Slice thickness, spacing, orientation -- **Equipment**: Manufacturer, model, software version - -#### Format-Specific Metadata -- **DICOM Tags**: Full DICOM header information -- **NIfTI Header**: Spatial orientation, timing information -- **TIFF Tags**: IFD entries, compression, photometric interpretation - -### Why Metadata Matters - -Metadata is used for: - -1. **Format Matching**: Recommend tools that support your file format -2. **Compatibility Scoring**: Prioritize tools that work with your specific format -3. **Context Understanding**: Help VLM understand image characteristics -4. **Demo Execution**: Ensure tools can process your data - -## Preview Generation - -### Automatic Conversion - -Medical and scientific images are converted to PNG previews for VLM analysis: - -| Original Format | Preview Generation | -|----------------|-------------------| -| DICOM (2D) | Single-frame converted to PNG | -| DICOM / NIfTI 3D volumes | Orthogonal 3โ€‘view composite PNG (axial, sagittal, coronal) using middle slices and/or maximum intensity projections (MIPs) | -| NIfTI 4D (time series) | Middle timepoint volume rendered as an orthogonal 3โ€‘view composite (middle slices and/or MIPs) | -| TIFF Stack | Orthogonal 3โ€‘view composite for 3D stacks; otherwise contact sheet or animated GIF preview when appropriate | -| Standard Images | Single-view PNG (content preserved; may be resized/normalized) | - -**Important**: Preview generation is for visual analysis only. Original format metadata is preserved and used for compatibility matching. - -### Multi-Slice Handling - -For 3D volumes, the agent typically builds an orthogonal 3โ€‘view composite preview: - -- **Axial**: Horizontal slices (z-axis) -- **Sagittal**: Side view (x-axis) -- **Coronal**: Front view (y-axis) - -Each view may combine the middle slice with a maximum intensity projection (MIP) to capture both anatomical context and bright structures. When a 3โ€‘view composite cannot be generated (e.g., unusual stack layout), the agent may fall back to a contact sheet or an animated GIF preview of multiple slices. - -## Format Compatibility Matching - -### How It Works - -The retrieval system adds format tokens to your query: - -``` -Original query: "segment lungs" -Enhanced query: "segment lungs format:DICOM format:3D" -``` - -Tools are matched based on: - -1. **Direct Format Support**: Tool explicitly supports your format -2. **Format Category**: Tool supports format family (e.g., medical imaging) -3. **Conversion Capability**: Tool can convert from your format - -### IO Compatibility Scoring - -The VLM considers: - -- **Input Format Match**: Can the tool read your file? -- **Output Format**: What format does the tool produce? -- **Dimension Compatibility**: 2D tool for 2D images, 3D for volumes -- **Modality Specificity**: CT tools for CT images, MRI for MRI - - - -## Unsupported Formats - -Currently not supported: - -- **Proprietary Formats**: Manufacturer-specific formats (e.g., .PAR/.REC) -- **Video Processing**: Limited video analysis capability -- **Raw Data**: Unformatted binary dumps without headers - -## Format Best Practices - -!!! tip "Use Standard Formats" - Stick to standard formats (DICOM, NIfTI, PNG, TIFF) for best tool compatibility. - -!!! tip "Include Metadata" - Use formats that preserve metadata (DICOM, NIfTI) rather than exporting to PNG/JPEG. - -!!! tip "Check Compatibility" - If a tool doesn't work, check the format compatibility in the recommendation metadata. - -!!! tip "Convert When Needed" - Some tools prefer specific formats. Convert using standard tools (ITK-SNAP, 3D Slicer) before upload. - -## Example Workflows by Format - -### DICOM Workflow -``` -1. Upload: chest_ct.dcm -2. Query: "Segment lungs" -3. Agent detects: DICOM, CT modality, 3D volume -4. Results: CT-compatible lung segmentation tools -``` - -### NIfTI Workflow -``` -1. Upload: brain_mri.nii.gz -2. Query: "Skull stripping" -3. Agent detects: NIfTI, 3D volume, likely MRI -4. Results: Brain extraction tools supporting NIfTI -``` - -### TIFF Stack Workflow -``` -1. Upload: microscopy_stack.tif -2. Query: "Cell counting" -3. Agent detects: Multi-frame TIFF, 3D stack -4. Results: Microscopy analysis tools -``` - -## Next Steps - -- Learn about [Understanding Recommendations](recommendations.md) -- Explore [Running Demos](running-demos.md) -- Check [Advanced Features](advanced-features.md) +# Supported File Formats + +The AI Imaging Agent supports a wide range of file formats for medical and scientific imaging, as well as general data files. + +## Image Formats + +### Standard Images + +| Format | Extensions | Description | +|--------|-----------|-------------| +| PNG | `.png` | Portable Network Graphics - lossless compression | +| JPEG | `.jpg`, `.jpeg` | Joint Photographic Experts Group - lossy compression | + +Currently, only PNG and JPEG are accepted for standard images. Other web formats (e.g. WebP, BMP, GIF) should be converted to PNG or JPEG before upload. +**Best for**: General photographs, screenshots, web images + +### Medical Imaging Formats + +#### DICOM + +| Format | Extensions | Description | +|--------|-----------|-------------| +| DICOM | `.dcm`, `.dicom` | Digital Imaging and Communications in Medicine | + +**Features**: + +- Industry standard for medical imaging +- Contains rich metadata (patient info, acquisition parameters) +- Supports multiple modalities (CT, MRI, X-ray, etc.) +- Can store 2D images or 3D volumes + +**Metadata Extracted**: + +- Patient ID, Study Instance UID +- Modality (CT, MR, CR, DX, etc.) +- Image dimensions and spacing +- Acquisition date/time +- Manufacturer and model + +**Example Usage**: +``` +Upload a CT DICOM file and ask: +"Segment the lungs from this scan" +``` + +#### NIfTI + +| Format | Extensions | Description | +|--------|-----------|-------------| +| NIfTI | `.nii`, `.nii.gz` | Neuroimaging Informatics Technology Initiative | + +**Features**: + +- Standard for neuroimaging research +- Supports 3D and 4D (time-series) volumes +- Compact storage with optional gzip compression +- Contains spatial orientation information + +**Metadata Extracted**: + +- Volume dimensions (x, y, z, time) +- Voxel spacing +- Data type and bit depth +- Orientation matrix + +**Example Usage**: +``` +Upload a brain MRI NIfTI file: +"Register this brain scan to MNI space" +``` + +### Scientific Imaging Formats + +#### TIFF/TIFF Stacks + +| Format | Extensions | Description | +|--------|-----------|-------------| +| TIFF | `.tif`, `.tiff` | Tagged Image File Format | + +**Features**: + +- Supports multi-page/multi-frame images +- Common in microscopy and scientific imaging +- Can store extensive metadata +- Lossless compression options + +**Metadata Extracted**: + +- Number of pages/frames (for stacks) +- Dimensions (width, height, channels) +- Color mode (RGB, grayscale, etc.) +- Compression method +- DPI/resolution information + +**Example Usage**: +``` +Upload a microscopy TIFF stack: +"Analyze cell structures in this z-stack" +``` + +## Data Formats + + + +## Format Detection + +The agent automatically detects file formats using: + +1. **File Extension**: Primary detection method +2. **Magic Bytes**: Header inspection for validation +3. **Content Analysis**: Fallback for ambiguous cases + +## Metadata Extraction + +### What Gets Extracted + +For each uploaded file, the agent extracts: + +#### Image Metadata +- **Dimensions**: Width, height, depth (for volumes) +- **Channels**: Grayscale, RGB, RGBA +- **Data Type**: uint8, int16, float32, etc. +- **File Size**: Storage size + +#### Medical Image Metadata +- **Modality**: CT, MRI, X-ray, Ultrasound, PET, etc. +- **Patient Info**: Anonymized IDs +- **Study Info**: Study UID, dates +- **Acquisition Parameters**: Slice thickness, spacing, orientation +- **Equipment**: Manufacturer, model, software version + +#### Format-Specific Metadata +- **DICOM Tags**: Full DICOM header information +- **NIfTI Header**: Spatial orientation, timing information +- **TIFF Tags**: IFD entries, compression, photometric interpretation + +### Why Metadata Matters + +Metadata is used for: + +1. **Format Matching**: Recommend tools that support your file format +2. **Compatibility Scoring**: Prioritize tools that work with your specific format +3. **Context Understanding**: Help VLM understand image characteristics +4. **Demo Execution**: Ensure tools can process your data + +## Preview Generation + +### Automatic Conversion + +Medical and scientific images are converted to PNG previews for VLM analysis: + +| Original Format | Preview Generation | +|----------------|-------------------| +| DICOM (2D) | Single-frame converted to PNG | +| DICOM / NIfTI 3D volumes | Orthogonal 3โ€‘view composite PNG (axial, sagittal, coronal) using middle slices and/or maximum intensity projections (MIPs) | +| NIfTI 4D (time series) | Middle timepoint volume rendered as an orthogonal 3โ€‘view composite (middle slices and/or MIPs) | +| TIFF Stack | Orthogonal 3โ€‘view composite for 3D stacks; otherwise contact sheet or animated GIF preview when appropriate | +| Standard Images | Single-view PNG (content preserved; may be resized/normalized) | + +**Important**: Preview generation is for visual analysis only. Original format metadata is preserved and used for compatibility matching. + +### Multi-Slice Handling + +For 3D volumes, the agent typically builds an orthogonal 3โ€‘view composite preview: + +- **Axial**: Horizontal slices (z-axis) +- **Sagittal**: Side view (x-axis) +- **Coronal**: Front view (y-axis) + +Each view may combine the middle slice with a maximum intensity projection (MIP) to capture both anatomical context and bright structures. When a 3โ€‘view composite cannot be generated (e.g., unusual stack layout), the agent may fall back to a contact sheet or an animated GIF preview of multiple slices. + +## Format Compatibility Matching + +### How It Works + +The retrieval system adds format tokens to your query: + +``` +Original query: "segment lungs" +Enhanced query: "segment lungs format:DICOM format:3D" +``` + +Tools are matched based on: + +1. **Direct Format Support**: Tool explicitly supports your format +2. **Format Category**: Tool supports format family (e.g., medical imaging) +3. **Conversion Capability**: Tool can convert from your format + +### IO Compatibility Scoring + +The VLM considers: + +- **Input Format Match**: Can the tool read your file? +- **Output Format**: What format does the tool produce? +- **Dimension Compatibility**: 2D tool for 2D images, 3D for volumes +- **Modality Specificity**: CT tools for CT images, MRI for MRI + + + +## Unsupported Formats + +Currently not supported: + +- **Proprietary Formats**: Manufacturer-specific formats (e.g., .PAR/.REC) +- **Video Processing**: Limited video analysis capability +- **Raw Data**: Unformatted binary dumps without headers + +## Format Best Practices + +!!! tip "Use Standard Formats" + Stick to standard formats (DICOM, NIfTI, PNG, TIFF) for best tool compatibility. + +!!! tip "Include Metadata" + Use formats that preserve metadata (DICOM, NIfTI) rather than exporting to PNG/JPEG. + +!!! tip "Check Compatibility" + If a tool doesn't work, check the format compatibility in the recommendation metadata. + +!!! tip "Convert When Needed" + Some tools prefer specific formats. Convert using standard tools (ITK-SNAP, 3D Slicer) before upload. + +## Example Workflows by Format + +### DICOM Workflow +``` +1. Upload: chest_ct.dcm +2. Query: "Segment lungs" +3. Agent detects: DICOM, CT modality, 3D volume +4. Results: CT-compatible lung segmentation tools +``` + +### NIfTI Workflow +``` +1. Upload: brain_mri.nii.gz +2. Query: "Skull stripping" +3. Agent detects: NIfTI, 3D volume, likely MRI +4. Results: Brain extraction tools supporting NIfTI +``` + +### TIFF Stack Workflow +``` +1. Upload: microscopy_stack.tif +2. Query: "Cell counting" +3. Agent detects: Multi-frame TIFF, 3D stack +4. Results: Microscopy analysis tools +``` + +## Next Steps + +- Learn about [Understanding Recommendations](recommendations.md) +- Explore [Running Demos](running-demos.md) +- Check [Advanced Features](advanced-features.md) diff --git a/docs/user-guide/recommendations.md b/docs/user-guide/recommendations.md index a477dfd082b92bcc76bef858d903d55236e67637..4070538e72cfc82eec4860c3e047652978a60640 100644 --- a/docs/user-guide/recommendations.md +++ b/docs/user-guide/recommendations.md @@ -1,329 +1,329 @@ -# Understanding Recommendations - -The AI Imaging Agent uses a sophisticated two-stage pipeline to provide ranked tool recommendations. This guide explains how recommendations are generated and how to interpret them. - -## How Recommendations Work - -### Two-Stage Pipeline - -```mermaid -graph TD - A[User Input: Image + Query] --> B[Stage 1: Retrieval] - B --> C[Candidate Tools] - C --> D[Stage 2: Agent Selection] - D --> E[Ranked Recommendations] -``` - -#### Stage 1: Retrieval (Fast Text Search) - -The retrieval stage quickly narrows down candidates: - -1. **Query Enhancement**: Your query is enriched with format tokens - ``` - Original: "segment lungs" - Enhanced: "segment lungs format:DICOM format:CT format:3D" - ``` - -2. **Embedding Search**: BGE-M3 model converts query to vector -3. **FAISS Vector Search**: Finds semantically similar tools -4. **CrossEncoder Reranking**: Re-scores candidates for better relevance -5. **Result**: Top-K candidates (default: 8) - -**No LLM calls** - this stage is fast and deterministic. - -#### Stage 2: Agent Selection (VLM-Powered) - -The agent analyzes candidates with full context: - -1. **Vision Analysis (only for VLM)**: GPT-4o/4o-mini (or your custom model) sees your image preview -2. **Context Integration**: Considers query + metadata + candidates -3. **Reasoning**: Explains why each tool matches -4. **Scoring**: Assigns accuracy scores (0-100%) -5. **Ranking**: Orders tools by relevance - -**Single VLM call** - comprehensive analysis with explanations. - -## Recommendation Format - -Each recommendation includes several components: - -### Header Information - -#### Rank Number -Position in the ranked list (1 = best match). - -``` -1๏ธโƒฃ TotalSegmentator -2๏ธโƒฃ MedSAM -3๏ธโƒฃ nnU-Net -``` - -#### Tool Name -The software or tool identifier, typically matching: -- GitHub repository name -- Published tool name -- Common community name - -#### Accuracy Score -Confidence level from 0-100%: - -- **90-100%**: Excellent match, highly confident -- **70-89%**: Good match, suitable for task -- **50-69%**: Moderate match, may need adaptation -- **Below 50%**: Weak match, alternative approach - -!!! note "Score Interpretation" - Scores reflect match quality for **your specific task and image**, not overall tool quality. - -### Body Content - -#### Description -Brief explanation of what the tool does: - -``` -TotalSegmentator: Automated multi-organ segmentation for CT scans supporting 104 anatomical structures. -``` - -#### Explanation -Why this tool matches your request: - -``` -Explanation: TotalSegmentator is specifically designed for whole-body CT segmentation including lung structures. It supports DICOM input and provides automated, accurate lung segmentation without manual intervention. -``` - -Key points in explanations: - -- **Task Alignment**: How well it matches your goal -- **Format Compatibility**: Support for your file format -- **Relevant Features**: Specific capabilities that help -- **Known Limitations**: Caveats or requirements - -#### Demo Link -Direct link to a runnable example: - -``` -๐Ÿš€ Demo: https://huggingface.co/spaces/example/totalsegmentator -``` - -Types of demos: - -- **HuggingFace Spaces**: Interactive Gradio/Streamlit apps -- **Colab Notebooks**: Jupyter notebooks you can run -- **Web Demos**: Hosted web interfaces -- **Documentation**: GitHub README with examples - -### Metadata Footer - -Technical details about the tool: - -#### Modality Support -Medical imaging modalities the tool works with: - -``` -Modalities: CT, MRI, X-ray -``` - -Common modalities: - -- **CT**: Computed Tomography -- **MRI**: Magnetic Resonance Imaging -- **XR**: X-ray radiography -- **US**: Ultrasound -- **PET**: Positron Emission Tomography -- **OCT**: Optical Coherence Tomography -- **Microscopy**: Various microscopy types - -#### Dimension Support -Image/volume dimensions supported: - -``` -Dimensions: 2D, 3D -``` - -- **2D**: Single slice images -- **3D**: Volumetric data -- **4D**: Time-series volumes - -#### Format Support -File formats the tool can process: - -``` -Formats: DICOM, NIfTI, PNG, JPEG -``` - -!!! tip "Format Importance" - Tools that support your **exact format** are prioritized in ranking. - -#### Tags -Categorization and keywords: - -``` -Tags: segmentation, medical-imaging, deep-learning, pytorch -``` - -Used for: -- Task categorization -- Technology stack -- Domain specificity -- Feature indicators - -## Scoring Factors - -The agent considers multiple factors when scoring: - -### Primary Factors (High Weight) - -1. **Task Match**: How well the tool's purpose aligns with your request -2. **Format Compatibility**: Support for your input format -3. **Image Content**: Visual analysis of what's in your image -4. **Dimension Match**: 2D tool for 2D images, 3D for volumes - -### Secondary Factors (Medium Weight) - -5. **Modality Specificity**: Tool designed for your imaging modality -6. **Feature Coverage**: Breadth of capabilities -7. **Stated Requirements**: Meets any specific requirements you mentioned -8. **Quality Indicators**: Stars, citations, community adoption - -### Tertiary Factors (Low Weight) - -9. **License**: Open-source vs. proprietary -10. **Recency**: Recently updated tools -11. **Documentation Quality**: Demo availability, examples -12. **Popularity**: Community usage and validation - -## Interpreting Results - -### High-Scoring Recommendations - -When you see scores above 85%: - -โœ… **Strong match** - Tool is designed for this task -โœ… **Format compatible** - Handles your file type -โœ… **Proven capability** - Demonstrated results in this domain - -**Action**: These are your best options. Try the top recommendation first. - -### Medium-Scoring Recommendations - -Scores 60-85%: - -โš ๏ธ **Good match** - Suitable but may need adaptation -โš ๏ธ **Possible format conversion** - Might require preprocessing -โš ๏ธ **Partial capability** - Covers some but not all requirements - -**Action**: Worth trying, especially if top choices don't work. Read explanations carefully. - -### Low-Scoring Recommendations - -Scores below 60%: - -โŒ **Weak match** - Limited alignment with task -โŒ **Format issues** - May not support your format -โŒ **Alternative approach** - Different methodology - -**Action**: Consider as fallback or for exploring alternative approaches. - -## Why Rankings Change - -Rankings depend on your specific context: - -### Same Tool, Different Queries - -"Segment lungs" vs "Detect tumors": -- Different tools excel at each task -- Rankings change based on task specificity - -### Same Task, Different Formats - -DICOM input vs PNG input: -- DICOM-compatible tools rank higher for DICOM -- General tools rank higher for standard images - -### Same Task, Different Images - -CT scan vs X-ray: -- Modality-specific tools get boosted -- Visual content influences selection - -## Common Patterns - -### All High Scores -Most recommendations >80%: - -- **Good news!** Multiple excellent options -- **Strategy**: Try top recommendation, then compare - -### Mixed Scores -Wide range (e.g., 90%, 65%, 45%): - -- **Top choice clear** - Focus on highest scorer -- **Strategy**: Try #1, fall back to #2 if needed - -### All Low Scores -All recommendations <60%: - -- **Limited options** - Task may be specialized -- **Strategy**: Try anyway, or rephrase query -- **Alternative**: Ask for suggestions - -## Acting on Recommendations - -### First Time with a Tool - -1. **Read the explanation** - Understand why it was recommended -2. **Check format compatibility** - Verify it supports your format -3. **Click demo link** - See it in action -4. **Try on your data** - Run if agent offers - -### Comparing Tools - -When choosing between similar scores: - -- **Check licenses** if redistribution matters -- **Compare formats** - prefer exact format match -- **Review tags** - match technology preferences -- **Demo availability** - easier to try - -### When Results Don't Match - -If recommendations seem wrong: - -1. **Provide more context**: "I need 3D volume support" -2. **Mention specific requirements**: "Must work with DICOM" -3. **Exclude irrelevant tools**: `[EXCLUDE:toolname]` -4. **Request alternatives**: "Can you search differently?" - -## Explanation Analysis - -Read explanations to understand: - -### Positive Indicators - -Look for phrases like: -- "Specifically designed for..." -- "Supports your exact format..." -- "Demonstrated accuracy on..." -- "Active development and maintained" - -### Caveats - -Watch for: -- "May require preprocessing..." -- "Limited to 2D images..." -- "Experimental feature..." -- "Requires specific environment..." - -### Requirements - -Note when explanations mention: -- "Needs GPU for inference" -- "Requires Python 3.8+" -- "DICOM headers must include..." -- "Minimum image resolution..." - -## Next Steps - -- Learn about [Running Demos](running-demos.md) -- Explore [Advanced Features](advanced-features.md) -- Understand the [Architecture Overview](../architecture/overview.md) +# Understanding Recommendations + +The AI Imaging Agent uses a sophisticated two-stage pipeline to provide ranked tool recommendations. This guide explains how recommendations are generated and how to interpret them. + +## How Recommendations Work + +### Two-Stage Pipeline + +```mermaid +graph TD + A[User Input: Image + Query] --> B[Stage 1: Retrieval] + B --> C[Candidate Tools] + C --> D[Stage 2: Agent Selection] + D --> E[Ranked Recommendations] +``` + +#### Stage 1: Retrieval (Fast Text Search) + +The retrieval stage quickly narrows down candidates: + +1. **Query Enhancement**: Your query is enriched with format tokens + ``` + Original: "segment lungs" + Enhanced: "segment lungs format:DICOM format:CT format:3D" + ``` + +2. **Embedding Search**: BGE-M3 model converts query to vector +3. **FAISS Vector Search**: Finds semantically similar tools +4. **CrossEncoder Reranking**: Re-scores candidates for better relevance +5. **Result**: Top-K candidates (default: 8) + +**No LLM calls** - this stage is fast and deterministic. + +#### Stage 2: Agent Selection (VLM-Powered) + +The agent analyzes candidates with full context: + +1. **Vision Analysis (only for VLM)**: GPT-4o/4o-mini (or your custom model) sees your image preview +2. **Context Integration**: Considers query + metadata + candidates +3. **Reasoning**: Explains why each tool matches +4. **Scoring**: Assigns accuracy scores (0-100%) +5. **Ranking**: Orders tools by relevance + +**Single VLM call** - comprehensive analysis with explanations. + +## Recommendation Format + +Each recommendation includes several components: + +### Header Information + +#### Rank Number +Position in the ranked list (1 = best match). + +``` +1๏ธโƒฃ TotalSegmentator +2๏ธโƒฃ MedSAM +3๏ธโƒฃ nnU-Net +``` + +#### Tool Name +The software or tool identifier, typically matching: +- GitHub repository name +- Published tool name +- Common community name + +#### Accuracy Score +Confidence level from 0-100%: + +- **90-100%**: Excellent match, highly confident +- **70-89%**: Good match, suitable for task +- **50-69%**: Moderate match, may need adaptation +- **Below 50%**: Weak match, alternative approach + +!!! note "Score Interpretation" + Scores reflect match quality for **your specific task and image**, not overall tool quality. + +### Body Content + +#### Description +Brief explanation of what the tool does: + +``` +TotalSegmentator: Automated multi-organ segmentation for CT scans supporting 104 anatomical structures. +``` + +#### Explanation +Why this tool matches your request: + +``` +Explanation: TotalSegmentator is specifically designed for whole-body CT segmentation including lung structures. It supports DICOM input and provides automated, accurate lung segmentation without manual intervention. +``` + +Key points in explanations: + +- **Task Alignment**: How well it matches your goal +- **Format Compatibility**: Support for your file format +- **Relevant Features**: Specific capabilities that help +- **Known Limitations**: Caveats or requirements + +#### Demo Link +Direct link to a runnable example: + +``` +๐Ÿš€ Demo: https://huggingface.co/spaces/example/totalsegmentator +``` + +Types of demos: + +- **HuggingFace Spaces**: Interactive Gradio/Streamlit apps +- **Colab Notebooks**: Jupyter notebooks you can run +- **Web Demos**: Hosted web interfaces +- **Documentation**: GitHub README with examples + +### Metadata Footer + +Technical details about the tool: + +#### Modality Support +Medical imaging modalities the tool works with: + +``` +Modalities: CT, MRI, X-ray +``` + +Common modalities: + +- **CT**: Computed Tomography +- **MRI**: Magnetic Resonance Imaging +- **XR**: X-ray radiography +- **US**: Ultrasound +- **PET**: Positron Emission Tomography +- **OCT**: Optical Coherence Tomography +- **Microscopy**: Various microscopy types + +#### Dimension Support +Image/volume dimensions supported: + +``` +Dimensions: 2D, 3D +``` + +- **2D**: Single slice images +- **3D**: Volumetric data +- **4D**: Time-series volumes + +#### Format Support +File formats the tool can process: + +``` +Formats: DICOM, NIfTI, PNG, JPEG +``` + +!!! tip "Format Importance" + Tools that support your **exact format** are prioritized in ranking. + +#### Tags +Categorization and keywords: + +``` +Tags: segmentation, medical-imaging, deep-learning, pytorch +``` + +Used for: +- Task categorization +- Technology stack +- Domain specificity +- Feature indicators + +## Scoring Factors + +The agent considers multiple factors when scoring: + +### Primary Factors (High Weight) + +1. **Task Match**: How well the tool's purpose aligns with your request +2. **Format Compatibility**: Support for your input format +3. **Image Content**: Visual analysis of what's in your image +4. **Dimension Match**: 2D tool for 2D images, 3D for volumes + +### Secondary Factors (Medium Weight) + +5. **Modality Specificity**: Tool designed for your imaging modality +6. **Feature Coverage**: Breadth of capabilities +7. **Stated Requirements**: Meets any specific requirements you mentioned +8. **Quality Indicators**: Stars, citations, community adoption + +### Tertiary Factors (Low Weight) + +9. **License**: Open-source vs. proprietary +10. **Recency**: Recently updated tools +11. **Documentation Quality**: Demo availability, examples +12. **Popularity**: Community usage and validation + +## Interpreting Results + +### High-Scoring Recommendations + +When you see scores above 85%: + +โœ… **Strong match** - Tool is designed for this task +โœ… **Format compatible** - Handles your file type +โœ… **Proven capability** - Demonstrated results in this domain + +**Action**: These are your best options. Try the top recommendation first. + +### Medium-Scoring Recommendations + +Scores 60-85%: + +โš ๏ธ **Good match** - Suitable but may need adaptation +โš ๏ธ **Possible format conversion** - Might require preprocessing +โš ๏ธ **Partial capability** - Covers some but not all requirements + +**Action**: Worth trying, especially if top choices don't work. Read explanations carefully. + +### Low-Scoring Recommendations + +Scores below 60%: + +โŒ **Weak match** - Limited alignment with task +โŒ **Format issues** - May not support your format +โŒ **Alternative approach** - Different methodology + +**Action**: Consider as fallback or for exploring alternative approaches. + +## Why Rankings Change + +Rankings depend on your specific context: + +### Same Tool, Different Queries + +"Segment lungs" vs "Detect tumors": +- Different tools excel at each task +- Rankings change based on task specificity + +### Same Task, Different Formats + +DICOM input vs PNG input: +- DICOM-compatible tools rank higher for DICOM +- General tools rank higher for standard images + +### Same Task, Different Images + +CT scan vs X-ray: +- Modality-specific tools get boosted +- Visual content influences selection + +## Common Patterns + +### All High Scores +Most recommendations >80%: + +- **Good news!** Multiple excellent options +- **Strategy**: Try top recommendation, then compare + +### Mixed Scores +Wide range (e.g., 90%, 65%, 45%): + +- **Top choice clear** - Focus on highest scorer +- **Strategy**: Try #1, fall back to #2 if needed + +### All Low Scores +All recommendations <60%: + +- **Limited options** - Task may be specialized +- **Strategy**: Try anyway, or rephrase query +- **Alternative**: Ask for suggestions + +## Acting on Recommendations + +### First Time with a Tool + +1. **Read the explanation** - Understand why it was recommended +2. **Check format compatibility** - Verify it supports your format +3. **Click demo link** - See it in action +4. **Try on your data** - Run if agent offers + +### Comparing Tools + +When choosing between similar scores: + +- **Check licenses** if redistribution matters +- **Compare formats** - prefer exact format match +- **Review tags** - match technology preferences +- **Demo availability** - easier to try + +### When Results Don't Match + +If recommendations seem wrong: + +1. **Provide more context**: "I need 3D volume support" +2. **Mention specific requirements**: "Must work with DICOM" +3. **Exclude irrelevant tools**: `[EXCLUDE:toolname]` +4. **Request alternatives**: "Can you search differently?" + +## Explanation Analysis + +Read explanations to understand: + +### Positive Indicators + +Look for phrases like: +- "Specifically designed for..." +- "Supports your exact format..." +- "Demonstrated accuracy on..." +- "Active development and maintained" + +### Caveats + +Watch for: +- "May require preprocessing..." +- "Limited to 2D images..." +- "Experimental feature..." +- "Requires specific environment..." + +### Requirements + +Note when explanations mention: +- "Needs GPU for inference" +- "Requires Python 3.8+" +- "DICOM headers must include..." +- "Minimum image resolution..." + +## Next Steps + +- Learn about [Running Demos](running-demos.md) +- Explore [Advanced Features](advanced-features.md) +- Understand the [Architecture Overview](../architecture/overview.md) diff --git a/docs/user-guide/running-demos.md b/docs/user-guide/running-demos.md index 52f12d879b3459c61ef66b022db31ee0d4cfc848..783454523d512bde7ea396ccb54d5ab96fea5bc2 100644 --- a/docs/user-guide/running-demos.md +++ b/docs/user-guide/running-demos.md @@ -1,372 +1,372 @@ -# Running Demos (This area is still under construction..) - -The AI Imaging Agent can execute tool demos directly on your uploaded images. This guide explains how demo execution works and how to use it effectively. - -## What Are Demos? - -Demos are **runnable examples** of imaging tools, typically hosted as: - -- **HuggingFace Spaces**: Interactive Gradio or Streamlit applications -- **Jupyter Notebooks**: Google Colab or similar notebook environments -- **Web Applications**: Hosted web interfaces -- **GitHub Examples**: Code repositories with example scripts - -## Demo Execution Flow - -### 1. Agent Offers to Run Demo - -After providing recommendations, the agent may offer: - -``` -Agent: Would you like me to run the demo with your image? -``` - -This appears when: - -- Tool has a compatible Gradio Space demo -- Your image format is compatible -- Demo's API is accessible - -### 2. You Confirm - -Respond with affirmative language: - -- "yes" -- "sure" -- "ok" -- "please" -- "go ahead" -- "run it" - -The agent detects these patterns and proceeds. - -### 3. Execution Happens - -The agent: - -1. Uploads your image to the demo space -2. Configures any required parameters -3. Triggers execution -4. Monitors progress -5. Retrieves results - -### 4. Results Display - -You receive: - -- **Success message** with output -- **Result images** or files -- **Execution trace** showing what happened - -## Demo Types - -### Gradio Space Demos - -**Best supported** - Direct API integration: - -``` -๐Ÿš€ Demo: https://huggingface.co/spaces/username/toolname -``` - -**Features**: - -- โœ… Automatic execution -- โœ… Progress monitoring -- โœ… Result retrieval -- โœ… Error handling - -**Example**: -``` -Running TotalSegmentator on your CT scan... -โœ“ Image uploaded -โœ“ Processing started -โœ“ Segmentation complete -โœ“ Results downloaded -``` - -### Notebook Demos - -**Partially supported** - Links provided for manual execution: - -``` -๐Ÿ““ Notebook: https://colab.research.google.com/... -``` - -**Process**: - -1. Click the notebook link -2. Open in Google Colab -3. Upload your image to the notebook -4. Run cells sequentially -5. Download results - -### Web Application Demos - -**Manual execution** - Opens in browser: - -``` -๐ŸŒ Web Demo: https://example.com/tool -``` - -**Process**: - -1. Click the demo link -2. Web app opens in new tab -3. Upload your image via the web UI -4. Configure settings -5. Run and download results - -### GitHub Repository Examples - -**Code-based** - Requires local setup: - -``` -๐Ÿ’ป Repository: https://github.com/user/repo -``` - -**Process**: - -1. Clone the repository -2. Install dependencies -3. Run example scripts -4. Adapt for your data - -## Execution Traces - -When demos run, you see detailed traces: - -```html -
-๐Ÿ”ง Tool Execution Trace - -Step 1: Uploading image to Gradio Space - โœ“ Connected to space: username/toolname - โœ“ Image uploaded: 2.3 MB - -Step 2: Configuring parameters - โœ“ Set task: lung-segmentation - โœ“ Set format: DICOM - -Step 3: Running inference - โณ Processing... (estimated 30s) - โœ“ Completed in 28s - -Step 4: Retrieving results - โœ“ Downloaded segmentation mask: 1.1 MB - โœ“ Downloaded visualization: 0.8 MB - -Status: โœ… Success -
-``` - -Click to expand and see full details. - -## Supported Gradio Spaces - -### Auto-Detected Parameters - -The agent automatically configures: - -#### Image Input -- Detects image input component(s) -- Uploads your file -- Converts format if needed - -#### Task Selection -Common task parameters: - -- **Task dropdown**: Matches your query to task option -- **Model selection**: Chooses appropriate model -- **Mode**: Inference, predict, analyze, etc. - -#### Format Options -- **Input format**: DICOM, NIfTI, PNG, etc. -- **Output format**: Segmentation mask, visualization, etc. -- **Data type**: 2D, 3D, specific modality - -### Manual Parameters - -Some demos require manual interaction: - -``` -Agent: This demo has additional parameters. Please visit the link to configure: -- Segmentation threshold: 0.5 -- Post-processing: enabled -``` - -## Demo Execution Best Practices - -!!! tip "Check Compatibility First" - Verify the tool supports your file format in the recommendation metadata. - -!!! tip "Use Standard Formats" - Demos work best with standard formats (PNG, JPEG for general; DICOM, NIfTI for medical). - -!!! tip "Be Patient" - Some demos take time, especially for: - - Large images or volumes - - Deep learning models - - 3D processing - - Typical times: 10 seconds to 2 minutes. - -!!! tip "Save Results Immediately" - Download result files promptly - they may not persist after closing the browser. - -!!! warning "Rate Limits" - Public Gradio Spaces may have rate limits or queue systems during high usage. - -## Troubleshooting Demo Execution - -### Demo Fails to Run - -**Error**: Connection timeout or failed upload - -**Solutions**: - -- Check internet connection -- Try again (server may be busy) -- Visit demo link manually -- Try alternative recommendation - -### Wrong Results - -**Error**: Output doesn't match expectations - -**Solutions**: - -- Check if correct parameters were used -- Verify image uploaded correctly -- Try adjusting task settings manually -- Compare with demo's example images - -### Incompatible Format - -**Error**: "Format not supported" - -**Solutions**: - -- Convert image to supported format -- Use tool that accepts your format -- Try alternative recommendation - -### Demo Link Broken - -**Error**: 404 or space not found - -**Solutions**: - -- Space may be temporarily down -- Check GitHub repo for alternative demo -- Try different tool recommendation -- Report broken link - -## Manual Demo Execution - -If automatic execution isn't available: - -### For Gradio Spaces - -1. Click the demo link -2. The space opens in your browser -3. Upload your image via the UI -4. Select appropriate options -5. Click "Submit" or "Run" -6. Download results - -### For Colab Notebooks - -1. Click the notebook link -2. Open in Google Colab -3. Run setup cells (install dependencies) -4. Upload your image when prompted: - ```python - from google.colab import files - uploaded = files.upload() - ``` -5. Run processing cells -6. Download results: - ```python - files.download('result.png') - ``` - -### For Local Execution - -1. Clone the repository: - ```bash - git clone https://github.com/user/repo - cd repo - ``` - -2. Install dependencies: - ```bash - pip install -r requirements.txt - ``` - -3. Run example script: - ```bash - python run_demo.py --input your_image.png --output result.png - ``` - -4. Check output directory for results - -## Understanding Results - -### Segmentation Results - -Typically includes: - -- **Segmentation mask**: Binary or multi-class mask -- **Overlay visualization**: Mask overlaid on original image -- **Statistics**: Volume, area, counts - -### Detection Results - -Usually provides: - -- **Bounding boxes**: Coordinates of detected objects -- **Annotated image**: Visual with boxes/labels -- **Confidence scores**: Detection confidence - -### Registration Results - -Common outputs: - -- **Transformed image**: Registered/aligned image -- **Transformation matrix**: Spatial transform parameters -- **Quality metrics**: Similarity scores - -### Classification Results - -Typical outputs: - -- **Class labels**: Predicted categories -- **Probabilities**: Confidence per class -- **Visualization**: Class activation maps - -## Demo Feedback - -Help improve the agent by reporting: - -### Successful Demos -When demos work well, this validates: - -- Tool compatibility -- Parameter auto-configuration -- Format handling - -### Issues -Report when: - -- Demo fails unexpectedly -- Results are incorrect -- Parameters were misconfigured -- Format conversion was wrong - -Feedback helps refine the agent's demo execution capabilities. - -## Next Steps - -- Explore [Advanced Features](advanced-features.md) -- Learn about the [Architecture](../architecture/overview.md) -- Check [CLI Commands](../reference/cli.md) +# Running Demos (This area is still under construction..) + +The AI Imaging Agent can execute tool demos directly on your uploaded images. This guide explains how demo execution works and how to use it effectively. + +## What Are Demos? + +Demos are **runnable examples** of imaging tools, typically hosted as: + +- **HuggingFace Spaces**: Interactive Gradio or Streamlit applications +- **Jupyter Notebooks**: Google Colab or similar notebook environments +- **Web Applications**: Hosted web interfaces +- **GitHub Examples**: Code repositories with example scripts + +## Demo Execution Flow + +### 1. Agent Offers to Run Demo + +After providing recommendations, the agent may offer: + +``` +Agent: Would you like me to run the demo with your image? +``` + +This appears when: + +- Tool has a compatible Gradio Space demo +- Your image format is compatible +- Demo's API is accessible + +### 2. You Confirm + +Respond with affirmative language: + +- "yes" +- "sure" +- "ok" +- "please" +- "go ahead" +- "run it" + +The agent detects these patterns and proceeds. + +### 3. Execution Happens + +The agent: + +1. Uploads your image to the demo space +2. Configures any required parameters +3. Triggers execution +4. Monitors progress +5. Retrieves results + +### 4. Results Display + +You receive: + +- **Success message** with output +- **Result images** or files +- **Execution trace** showing what happened + +## Demo Types + +### Gradio Space Demos + +**Best supported** - Direct API integration: + +``` +๐Ÿš€ Demo: https://huggingface.co/spaces/username/toolname +``` + +**Features**: + +- โœ… Automatic execution +- โœ… Progress monitoring +- โœ… Result retrieval +- โœ… Error handling + +**Example**: +``` +Running TotalSegmentator on your CT scan... +โœ“ Image uploaded +โœ“ Processing started +โœ“ Segmentation complete +โœ“ Results downloaded +``` + +### Notebook Demos + +**Partially supported** - Links provided for manual execution: + +``` +๐Ÿ““ Notebook: https://colab.research.google.com/... +``` + +**Process**: + +1. Click the notebook link +2. Open in Google Colab +3. Upload your image to the notebook +4. Run cells sequentially +5. Download results + +### Web Application Demos + +**Manual execution** - Opens in browser: + +``` +๐ŸŒ Web Demo: https://example.com/tool +``` + +**Process**: + +1. Click the demo link +2. Web app opens in new tab +3. Upload your image via the web UI +4. Configure settings +5. Run and download results + +### GitHub Repository Examples + +**Code-based** - Requires local setup: + +``` +๐Ÿ’ป Repository: https://github.com/user/repo +``` + +**Process**: + +1. Clone the repository +2. Install dependencies +3. Run example scripts +4. Adapt for your data + +## Execution Traces + +When demos run, you see detailed traces: + +```html +
+๐Ÿ”ง Tool Execution Trace + +Step 1: Uploading image to Gradio Space + โœ“ Connected to space: username/toolname + โœ“ Image uploaded: 2.3 MB + +Step 2: Configuring parameters + โœ“ Set task: lung-segmentation + โœ“ Set format: DICOM + +Step 3: Running inference + โณ Processing... (estimated 30s) + โœ“ Completed in 28s + +Step 4: Retrieving results + โœ“ Downloaded segmentation mask: 1.1 MB + โœ“ Downloaded visualization: 0.8 MB + +Status: โœ… Success +
+``` + +Click to expand and see full details. + +## Supported Gradio Spaces + +### Auto-Detected Parameters + +The agent automatically configures: + +#### Image Input +- Detects image input component(s) +- Uploads your file +- Converts format if needed + +#### Task Selection +Common task parameters: + +- **Task dropdown**: Matches your query to task option +- **Model selection**: Chooses appropriate model +- **Mode**: Inference, predict, analyze, etc. + +#### Format Options +- **Input format**: DICOM, NIfTI, PNG, etc. +- **Output format**: Segmentation mask, visualization, etc. +- **Data type**: 2D, 3D, specific modality + +### Manual Parameters + +Some demos require manual interaction: + +``` +Agent: This demo has additional parameters. Please visit the link to configure: +- Segmentation threshold: 0.5 +- Post-processing: enabled +``` + +## Demo Execution Best Practices + +!!! tip "Check Compatibility First" + Verify the tool supports your file format in the recommendation metadata. + +!!! tip "Use Standard Formats" + Demos work best with standard formats (PNG, JPEG for general; DICOM, NIfTI for medical). + +!!! tip "Be Patient" + Some demos take time, especially for: + - Large images or volumes + - Deep learning models + - 3D processing + + Typical times: 10 seconds to 2 minutes. + +!!! tip "Save Results Immediately" + Download result files promptly - they may not persist after closing the browser. + +!!! warning "Rate Limits" + Public Gradio Spaces may have rate limits or queue systems during high usage. + +## Troubleshooting Demo Execution + +### Demo Fails to Run + +**Error**: Connection timeout or failed upload + +**Solutions**: + +- Check internet connection +- Try again (server may be busy) +- Visit demo link manually +- Try alternative recommendation + +### Wrong Results + +**Error**: Output doesn't match expectations + +**Solutions**: + +- Check if correct parameters were used +- Verify image uploaded correctly +- Try adjusting task settings manually +- Compare with demo's example images + +### Incompatible Format + +**Error**: "Format not supported" + +**Solutions**: + +- Convert image to supported format +- Use tool that accepts your format +- Try alternative recommendation + +### Demo Link Broken + +**Error**: 404 or space not found + +**Solutions**: + +- Space may be temporarily down +- Check GitHub repo for alternative demo +- Try different tool recommendation +- Report broken link + +## Manual Demo Execution + +If automatic execution isn't available: + +### For Gradio Spaces + +1. Click the demo link +2. The space opens in your browser +3. Upload your image via the UI +4. Select appropriate options +5. Click "Submit" or "Run" +6. Download results + +### For Colab Notebooks + +1. Click the notebook link +2. Open in Google Colab +3. Run setup cells (install dependencies) +4. Upload your image when prompted: + ```python + from google.colab import files + uploaded = files.upload() + ``` +5. Run processing cells +6. Download results: + ```python + files.download('result.png') + ``` + +### For Local Execution + +1. Clone the repository: + ```bash + git clone https://github.com/user/repo + cd repo + ``` + +2. Install dependencies: + ```bash + pip install -r requirements.txt + ``` + +3. Run example script: + ```bash + python run_demo.py --input your_image.png --output result.png + ``` + +4. Check output directory for results + +## Understanding Results + +### Segmentation Results + +Typically includes: + +- **Segmentation mask**: Binary or multi-class mask +- **Overlay visualization**: Mask overlaid on original image +- **Statistics**: Volume, area, counts + +### Detection Results + +Usually provides: + +- **Bounding boxes**: Coordinates of detected objects +- **Annotated image**: Visual with boxes/labels +- **Confidence scores**: Detection confidence + +### Registration Results + +Common outputs: + +- **Transformed image**: Registered/aligned image +- **Transformation matrix**: Spatial transform parameters +- **Quality metrics**: Similarity scores + +### Classification Results + +Typical outputs: + +- **Class labels**: Predicted categories +- **Probabilities**: Confidence per class +- **Visualization**: Class activation maps + +## Demo Feedback + +Help improve the agent by reporting: + +### Successful Demos +When demos work well, this validates: + +- Tool compatibility +- Parameter auto-configuration +- Format handling + +### Issues +Report when: + +- Demo fails unexpectedly +- Results are incorrect +- Parameters were misconfigured +- Format conversion was wrong + +Feedback helps refine the agent's demo execution capabilities. + +## Next Steps + +- Explore [Advanced Features](advanced-features.md) +- Learn about the [Architecture](../architecture/overview.md) +- Check [CLI Commands](../reference/cli.md) diff --git a/mkdocs.yml b/mkdocs.yml index a5e0e8055ec4cfe23ee2adba4c984aafda0b6f7d..ecccf270c35c7cbc60ffdb0770616e91dc745ffa 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,108 +1,108 @@ -site_name: AI Imaging Agent -site_description: Intelligent RAG + AI agent for discovering imaging software -site_author: Imaging Plaza -site_url: https://imaging-plaza.github.io/ai-agent/ - -repo_name: imaging-plaza/ai-agent -repo_url: https://github.com/imaging-plaza/ai-agent -edit_uri: edit/main/docs/ - -theme: - name: material - palette: - # Palette toggle for light mode - - scheme: default - primary: teal - accent: green - toggle: - icon: material/brightness-7 - name: Switch to dark mode - # Palette toggle for dark mode - - scheme: slate - primary: teal - accent: green - toggle: - icon: material/brightness-4 - name: Switch to light mode - - font: - text: Roboto - code: Roboto Mono - - features: - - navigation.instant - - navigation.tracking - - navigation.tabs - - navigation.sections - - navigation.top - - navigation.footer - - search.suggest - - search.highlight - - content.code.copy - - content.code.annotate - - content.tabs.link - - icon: - repo: fontawesome/brands/github - -plugins: - - search - - tags - -markdown_extensions: - - admonition - - pymdownx.details - - pymdownx.superfences: - custom_fences: - - name: mermaid - class: mermaid - format: !!python/name:pymdownx.superfences.fence_code_format - - pymdownx.tabbed: - alternate_style: true - - pymdownx.highlight: - anchor_linenums: true - - pymdownx.inlinehilite - - pymdownx.snippets - - attr_list - - md_in_html - - pymdownx.emoji: - emoji_index: !!python/name:material.extensions.emoji.twemoji - emoji_generator: !!python/name:material.extensions.emoji.to_svg - - tables - - footnotes - - toc: - permalink: true - -nav: - - Home: index.md - - Getting Started: - - Installation: getting-started/installation.md - - Quick Start: getting-started/quickstart.md - - Configuration: getting-started/configuration.md - - User Guide: - - Using the Chat Interface: user-guide/chat-interface.md - - Supported File Formats: user-guide/file-formats.md - - Understanding Recommendations: user-guide/recommendations.md - - Running Demos: user-guide/running-demos.md - - Advanced Features: user-guide/advanced-features.md - - Architecture: - - Overview: architecture/overview.md - - Retrieval Pipeline: architecture/retrieval.md - - Agent & VLM Selection: architecture/agent.md - - Software Catalog: architecture/catalog.md - - Development: - - Project Guide: guide.md - - Project Structure: development/structure.md - - Contributing: development/contributing.md - - Testing: development/testing.md - - Reference: - - CLI Commands: reference/cli.md - - Environment Variables: reference/environment.md - - Changelog: reference/changelog.md - -extra: - social: - - icon: fontawesome/brands/github - link: https://github.com/imaging-plaza/ai-agent - -copyright: Copyright © 2024-2026 Imaging Plaza +site_name: AI Imaging Agent +site_description: Intelligent RAG + AI agent for discovering imaging software +site_author: Imaging Plaza +site_url: https://imaging-plaza.github.io/ai-agent/ + +repo_name: imaging-plaza/ai-agent +repo_url: https://github.com/imaging-plaza/ai-agent +edit_uri: edit/main/docs/ + +theme: + name: material + palette: + # Palette toggle for light mode + - scheme: default + primary: teal + accent: green + toggle: + icon: material/brightness-7 + name: Switch to dark mode + # Palette toggle for dark mode + - scheme: slate + primary: teal + accent: green + toggle: + icon: material/brightness-4 + name: Switch to light mode + + font: + text: Roboto + code: Roboto Mono + + features: + - navigation.instant + - navigation.tracking + - navigation.tabs + - navigation.sections + - navigation.top + - navigation.footer + - search.suggest + - search.highlight + - content.code.copy + - content.code.annotate + - content.tabs.link + + icon: + repo: fontawesome/brands/github + +plugins: + - search + - tags + +markdown_extensions: + - admonition + - pymdownx.details + - pymdownx.superfences: + custom_fences: + - name: mermaid + class: mermaid + format: !!python/name:pymdownx.superfences.fence_code_format + - pymdownx.tabbed: + alternate_style: true + - pymdownx.highlight: + anchor_linenums: true + - pymdownx.inlinehilite + - pymdownx.snippets + - attr_list + - md_in_html + - pymdownx.emoji: + emoji_index: !!python/name:material.extensions.emoji.twemoji + emoji_generator: !!python/name:material.extensions.emoji.to_svg + - tables + - footnotes + - toc: + permalink: true + +nav: + - Home: index.md + - Getting Started: + - Installation: getting-started/installation.md + - Quick Start: getting-started/quickstart.md + - Configuration: getting-started/configuration.md + - User Guide: + - Using the Chat Interface: user-guide/chat-interface.md + - Supported File Formats: user-guide/file-formats.md + - Understanding Recommendations: user-guide/recommendations.md + - Running Demos: user-guide/running-demos.md + - Advanced Features: user-guide/advanced-features.md + - Architecture: + - Overview: architecture/overview.md + - Retrieval Pipeline: architecture/retrieval.md + - Agent & VLM Selection: architecture/agent.md + - Software Catalog: architecture/catalog.md + - Development: + - Project Guide: guide.md + - Project Structure: development/structure.md + - Contributing: development/contributing.md + - Testing: development/testing.md + - Reference: + - CLI Commands: reference/cli.md + - Environment Variables: reference/environment.md + - Changelog: reference/changelog.md + +extra: + social: + - icon: fontawesome/brands/github + link: https://github.com/imaging-plaza/ai-agent + +copyright: Copyright © 2024-2026 Imaging Plaza diff --git a/pyproject.toml b/pyproject.toml index c64d8aafdc709c75972ac260545a4e7b33365419..f8f6a1391055a1889f51bbb25f1946842c9d144b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,10 @@ dependencies = [ "requests==2.32.4", "python-dotenv==1.1.1", "gradio==5.42.0", + "fastapi==0.118.0", + "uvicorn[standard]==0.36.0", + "sse-starlette==2.1.3", + "python-multipart==0.0.20", "json5==0.12.0", "pillow==11.3.0", "nibabel==5.3.2", diff --git a/src/ai_agent/agent/agent.py b/src/ai_agent/agent/agent.py index a7ef75db88d8fe1e74308bfbaf66dd7ebe636ce6..7431cc08a1307b3d46c620c2d44f24099d6be77a 100644 --- a/src/ai_agent/agent/agent.py +++ b/src/ai_agent/agent/agent.py @@ -1,611 +1,650 @@ -from __future__ import annotations - -import os -import logging -import time -import asyncio -import threading -from collections import OrderedDict -from datetime import datetime -from typing import List - -from pydantic_ai import Agent, RunContext -from pydantic_ai.usage import UsageLimits -from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIChatModel -from pydantic_ai.providers.openai import OpenAIProvider -from pydantic_ai.messages import BinaryContent - -from ai_agent.generator.prompts import get_agent_system_prompt -from ai_agent.generator.schema import ToolSelection, Conversation, ConversationStatus -from ai_agent.utils.config import get_config -from .models import AgentToolSelection, ToolRunLog, UsageStats -from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput -from ai_agent.agent.utils import coerce_github_url_or_none -from .tools.search_tool import tool_search_tools, SearchToolsInput -from .tools.search_alternative_tool import ( - tool_search_alternative, - SearchAlternativeInput, -) -from .tools.query_utils import sanitize_retrieval_query -from .utils import AgentState, limit_tool_calls, cap_prepare -from ai_agent.utils.image_meta import summarize_image_metadata, detect_ext_token - -log = logging.getLogger("agent.core") - -DEFAULT_NUM_CHOICES = int(os.getenv("NUM_CHOICES", "3")) - -# --------------------------------------------------------------------------- -# Dynamic agent instance cache -# Key: (model_name, base_url, api_key_env, num_choices) -# Avoids rebuilding Agent/OpenAIProvider/model objects on every request when -# the UI repeatedly uses the same custom endpoint + model combination. -# Bounded LRU (max AGENT_CACHE_MAX entries); protected by a lock so that -# concurrent requests cannot race while creating/inserting agents. -# --------------------------------------------------------------------------- -_AGENT_CACHE_MAX: int = int(os.getenv("AGENT_CACHE_MAX", "16")) -_AGENT_CACHE_LOCK: threading.Lock = threading.Lock() -_AGENT_CACHE: OrderedDict[tuple, "Agent"] = OrderedDict() - -# --------------------------------------------------------------------------- -# Model / provider setup -# --------------------------------------------------------------------------- -config = get_config() -agent_model_config = config.agent_model - -try: - api_key = agent_model_config.get_api_key() -except ValueError as e: - log.error(f"Failed to get API key for agent model: {e}") - raise - -log.info(f"Initializing agent model: {agent_model_config.name}") - -if agent_model_config.base_url: - log.info(f"Using custom OpenAI base URL: {agent_model_config.base_url}") - log.info("Using OpenAIChatModel (chat/completions API) for custom endpoint") - provider = OpenAIProvider( - base_url=agent_model_config.base_url, - api_key=api_key, - ) - openai_model = OpenAIChatModel( - model_name=agent_model_config.name, - provider=provider, - ) -else: - provider = OpenAIProvider(api_key=api_key) - openai_model = OpenAIResponsesModel( - model_name=agent_model_config.name, - provider=provider, - ) - -# --------------------------------------------------------------------------- -# Agent definition -# --------------------------------------------------------------------------- -agent = Agent( - model=openai_model, - system_prompt=get_agent_system_prompt(DEFAULT_NUM_CHOICES), - deps_type=AgentState, - output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), -) - -# --------------------------------------------------------------------------- -# Tool adapters for the agent -# --------------------------------------------------------------------------- - - -@agent.tool(retries=2, prepare=cap_prepare) -@limit_tool_calls("search_tools", cap=1) -async def search_tools( - ctx: RunContext[AgentState], - query: str, - excluded: List[str] | None = None, - top_k: int = 12, -) -> List[dict]: - """ - Agent-facing search tool. - - Delegates to tools.search_tool.tool_search_tools(), but automatically - injects: - - globally excluded tools (from ctx.deps.excluded_tools) - - image_paths and original_formats (from ctx.deps, set in run_agent) - so the language model never has to reason about file paths directly. - """ - # Merge explicit exclusions with global exclusions from AgentState - explicit_excluded = excluded or [] - global_excluded = getattr(ctx.deps, "excluded_tools", []) or [] - all_excluded = sorted(set(explicit_excluded + list(global_excluded))) - - original_formats = getattr(ctx.deps, "original_formats", []) or [] - image_paths = getattr(ctx.deps, "image_paths", []) or [] - - effective_top_k = ( - ctx.deps.override_top_k if ctx.deps.override_top_k is not None else top_k - ) - - started = time.perf_counter() - inp = SearchToolsInput( - query=sanitize_retrieval_query(query), - excluded=all_excluded, - top_k=effective_top_k, - original_formats=original_formats, - image_paths=image_paths, - ) - out = tool_search_tools(inp) - - ctx.deps.tool_calls.append( - { - "tool": "search_tools", - "query": query, - "count": len(out.candidates), - "duration_ms": round((time.perf_counter() - started) * 1000, 1), - "original_formats": original_formats, - "excluded": all_excluded, - "timestamp": datetime.now().isoformat(), - } - ) - - # Return plain dicts so the LLM sees a simple JSON-like structure. - return [c.model_dump(mode="python") for c in out.candidates] - - -@agent.tool(retries=2, prepare=cap_prepare) -@limit_tool_calls("search_alternative", cap=3) -async def search_alternative( - ctx: RunContext[AgentState], - alternative_query: str, - excluded: List[str] | None = None, - top_k: int = 12, -) -> List[dict]: - """ - Search with an alternative query formulation (includes automatic reranking). - """ - explicit_excluded = excluded or [] - global_excluded = getattr(ctx.deps, "excluded_tools", []) or [] - all_excluded = sorted(set(explicit_excluded + list(global_excluded))) - - original_formats = getattr(ctx.deps, "original_formats", []) or [] - image_paths = getattr(ctx.deps, "image_paths", []) or [] - - started = time.perf_counter() - inp = SearchAlternativeInput( - alternative_query=alternative_query, - excluded=all_excluded, - top_k=top_k, - original_formats=original_formats, - image_paths=image_paths, - ) - out = tool_search_alternative(inp) - - ctx.deps.tool_calls.append( - { - "tool": "search_alternative", - "alternative_query": alternative_query, - "query_used": out.query_used, - "count": len(out.candidates), - "duration_ms": round((time.perf_counter() - started) * 1000, 1), - "original_formats": original_formats, - "excluded": all_excluded, - "timestamp": datetime.now().isoformat(), - } - ) - - return [c.model_dump(mode="python") for c in out.candidates] - - -@agent.tool(retries=2, prepare=cap_prepare) -@limit_tool_calls("repo_info_batch", cap=4) -async def repo_info_batch( - ctx: RunContext[AgentState], - urls: List[str], -) -> List[dict]: - """Fetch repository summaries for multiple repositories in parallel.""" - started = time.perf_counter() - - if not urls: - return [] - - normalized: List[str] = [] - skipped: List[dict] = [] - seen: set[str] = set() - for raw in urls: - norm = coerce_github_url_or_none(raw) - if not norm: - skipped.append( - { - "url": raw, - "skipped": True, - "reason": "NON_GITHUB_URL", - } - ) - continue - if norm in seen: - continue - seen.add(norm) - normalized.append(norm) - - tasks = [tool_repo_summary(RepoSummaryInput(url=u)) for u in normalized] - outcomes = await asyncio.gather(*tasks, return_exceptions=True) - - results: List[dict] = [] - for url, outcome in zip(normalized, outcomes): - if isinstance(outcome, Exception): - results.append( - { - "url": url, - "source": "error", - "error": str(outcome), - } - ) - continue - payload = outcome.model_dump(mode="python") - payload["url"] = url - results.append(payload) - - if skipped: - results.extend(skipped) - - ctx.deps.tool_calls.append( - { - "tool": "repo_info_batch", - "requested": len(urls), - "normalized": len(normalized), - "returned": len(results), - "duration_ms": round((time.perf_counter() - started) * 1000, 1), - "timestamp": datetime.now().isoformat(), - } - ) - - return results - - -# --------------------------------------------------------------------------- -# High level entry point: run the agent on (text query + image) -# --------------------------------------------------------------------------- -def run_agent( - task: str, - image_paths: List[str], - excluded: List[str] | None = None, - conversation_history: List[str] | None = None, - *, - image_bytes: bytes | None = None, - model: str | None = None, - base_url: str | None = None, - api_key_env: str | None = None, - top_k: int | None = None, - num_choices: int | None = None, - image_metadata: str | None = None, -) -> AgentToolSelection: - """ - Execute the agent for a user task and at least one image path. - - - derive canonical original_formats (tiff / dicom / nifti / ...) - - build a compact image metadata summary (or use pre-computed one) - - pass both to the LLM as hidden context - - store image_paths/original_formats in deps so retrieval tools can use them - - optionally allow runtime model/base_url/top_k/num_choices overrides - - IMPORTANT: - The model only sees an actual image if `image_bytes` is provided. - `image_paths` are used for metadata + tool context only. - """ - run_started = time.perf_counter() - if not image_paths: - raise ValueError("run_agent requires at least one image path") - - tool_logs: List[ToolRunLog] = [] - - # ---- 1) Derive image-based metadata and format hints -------------------- - metadata_started = time.perf_counter() - meta_str = ( - image_metadata - if image_metadata is not None - else (summarize_image_metadata(image_paths) or "") - ) - fmt_str = detect_ext_token(image_paths) or "" - original_formats = [t.lower() for t in fmt_str.split()] if fmt_str else [] - metadata_duration_ms = round((time.perf_counter() - metadata_started) * 1000, 1) - - effective_top_k = top_k if top_k is not None else 12 - effective_num_choices = num_choices if num_choices is not None else 3 - - # ---- 2) Prepare dependency state passed to all tools -------------------- - deps = AgentState( - excluded_tools=excluded or [], - override_model=model, - override_base_url=base_url, - override_top_k=effective_top_k, - override_num_choices=effective_num_choices, - ) - - setattr(deps, "image_paths", list(image_paths)) - setattr(deps, "original_formats", original_formats) - - # ---- 3) Hidden metadata lines for the model ---------------------------- - hidden_meta = "" - if original_formats: - hidden_meta += "\n(Formats Hint: " + ",".join(original_formats) + ")" - if meta_str: - short_meta = " ".join(x.strip() for x in meta_str.splitlines() if x.strip()) - hidden_meta += ( - "\n(Image Metadata: " - + short_meta[:500] - + ("โ€ฆ" if len(short_meta) > 500 else "") - + ")" - ) - hidden_meta += f"\n(Search top_k: {effective_top_k})" - - extra_context = "\n\n**CRITICAL: Analyze the attached preview image showing the user's data.**\nUse visual observations (anatomy visible, image quality, dimensionality, contrast) combined with the metadata below to recommend tools. Reference what you see in your explanations." - - # ---- 4) Build the prompt (optionally including history) ---------------- - if conversation_history and len(conversation_history) > 0: - history_text = "\n".join(conversation_history) - prompt = ( - f"Previous conversation:\n{history_text}\n\n" - f"Current request: {task}{extra_context}{hidden_meta}" - ) - else: - prompt = task + extra_context + hidden_meta - - # ----------------------------------------------------------------------- - # Determine which agent instance to use - # ----------------------------------------------------------------------- - agent_instance = agent - effective_num_choices = num_choices if num_choices is not None else 3 - effective_model = model if model else agent_model_config.name - effective_top_k = top_k if top_k is not None else 12 - - # When model is provided from UI, base_url comes with it (can be None for OpenAI) - if model: - # Use api_key_env from config if provided, otherwise default to OPENAI_API_KEY - key_env_name = api_key_env if api_key_env else "OPENAI_API_KEY" - runtime_api_key = os.getenv(key_env_name) - if not runtime_api_key: - raise ValueError( - f"{key_env_name} not found in environment. Cannot use this model." - ) - effective_base_url = base_url # Can be None for OpenAI - log.info(f"โœ“ Using {key_env_name} for model {effective_model}") - log.debug(f"{key_env_name} is set: {bool(runtime_api_key)}") - else: - # No model override - use config defaults - effective_base_url = agent_model_config.base_url - runtime_api_key = api_key # Already loaded from config at startup - log.info(f"โœ“ Using API key from config for model {effective_model}") - - # Log runtime configuration - endpoint_display = effective_base_url if effective_base_url else "api.openai.com" - log.info( - f"๐Ÿค– Agent execution - Model: {effective_model}, endpoint: {endpoint_display}, " - f"top_k: {effective_top_k}, num_choices: {effective_num_choices}, excluded: {len(excluded or [])}" - ) - - needs_dynamic_agent = model is not None - - if needs_dynamic_agent: - cache_key = (effective_model, effective_base_url or "", api_key_env or "OPENAI_API_KEY", effective_num_choices) - with _AGENT_CACHE_LOCK: - agent_instance = _AGENT_CACHE.get(cache_key) - if agent_instance is not None: - _AGENT_CACHE.move_to_end(cache_key) - if agent_instance is None: - log.info( - f"๐Ÿ“ฆ Creating runtime agent with model={effective_model}, endpoint={effective_base_url or 'api.openai.com'}" - ) - - runtime_provider = OpenAIProvider( - base_url=effective_base_url, - api_key=runtime_api_key, - ) - - # Use OpenAIChatModel (chat/completions) for custom endpoints, OpenAIResponsesModel for default OpenAI - if effective_base_url: - log.info("Using OpenAIChatModel (chat/completions API) for custom endpoint") - runtime_model = OpenAIChatModel( - model_name=effective_model, provider=runtime_provider - ) - else: - runtime_model = OpenAIResponsesModel( - model_name=effective_model, provider=runtime_provider - ) - - agent_instance = Agent( - model=runtime_model, - system_prompt=get_agent_system_prompt(effective_num_choices), - deps_type=AgentState, - output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), - ) - - # Register tools on the dynamic agent - agent_instance.tool(search_tools, retries=2, prepare=cap_prepare) - agent_instance.tool(search_alternative, retries=2, prepare=cap_prepare) - agent_instance.tool(repo_info_batch, retries=2, prepare=cap_prepare) - - with _AGENT_CACHE_LOCK: - _AGENT_CACHE[cache_key] = agent_instance - _AGENT_CACHE.move_to_end(cache_key) - while len(_AGENT_CACHE) > _AGENT_CACHE_MAX: - _AGENT_CACHE.popitem(last=False) - else: - log.info( - f"โ™ป๏ธ Reusing cached dynamic agent (model: {effective_model}, num_choices: {effective_num_choices})" - ) - - elif ( - num_choices is not None and num_choices != DEFAULT_NUM_CHOICES - ): - cache_key = (effective_model, effective_base_url or "", api_key_env or "OPENAI_API_KEY", effective_num_choices) - with _AGENT_CACHE_LOCK: - agent_instance = _AGENT_CACHE.get(cache_key) - if agent_instance is not None: - _AGENT_CACHE.move_to_end(cache_key) - if agent_instance is None: - log.info( - f"๐Ÿ“ฆ Creating runtime agent with num_choices={effective_num_choices} (model: {effective_model})" - ) - agent_instance = Agent( - model=openai_model, - system_prompt=get_agent_system_prompt(effective_num_choices), - deps_type=AgentState, - output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), - ) - - # Register tools on the dynamic agent - agent_instance.tool(search_tools, retries=2, prepare=cap_prepare) - agent_instance.tool(search_alternative, retries=2, prepare=cap_prepare) - agent_instance.tool(repo_info_batch, retries=2, prepare=cap_prepare) - - with _AGENT_CACHE_LOCK: - _AGENT_CACHE[cache_key] = agent_instance - _AGENT_CACHE.move_to_end(cache_key) - while len(_AGENT_CACHE) > _AGENT_CACHE_MAX: - _AGENT_CACHE.popitem(last=False) - else: - log.info( - f"โ™ป๏ธ Reusing cached dynamic agent with num_choices={effective_num_choices} (model: {effective_model})" - ) - - else: - log.info( - f"โ™ป๏ธ Using global agent (model: {effective_model}, num_choices: {effective_num_choices})" - ) - - log.debug( - f"Prompt length: {len(prompt)} chars, has_image_paths: {bool(image_paths)}, has_image_bytes: {bool(image_bytes)}" - ) - - # ---- 5) Build multimodal prompt if image bytes provided ---------------- - if image_bytes: - log.info( - f"๐Ÿ–ผ๏ธ Sending image preview to model ({len(image_bytes)} bytes = {len(image_bytes)/1024:.1f}KB)" - ) - user_prompt = [ - prompt, - BinaryContent( - data=image_bytes, - media_type="image/png", - ), - ] - else: - log.warning( - "โš ๏ธ No image bytes provided - the model will not see the image preview" - ) - user_prompt = prompt - - # ---- 6) Run the agent -------------------------------------------------- - try: - llm_started = time.perf_counter() - run_result = agent_instance.run_sync( - user_prompt, - deps=deps, - output_type=ToolSelection, - usage_limits=UsageLimits(tool_calls_limit=20), - ) - llm_duration_ms = round((time.perf_counter() - llm_started) * 1000, 1) - result = run_result.output - - log.info( - f"โœ… Agent execution complete - choices returned: {len(result.choices)}" - ) - - # Log usage (helpful, but may not explicitly expose image-specific counters) - if run_result.usage: - usage = run_result.usage() - log.info( - f"๐Ÿ“Š Usage: total_tokens={usage.total_tokens}, " - f"input_tokens={usage.input_tokens}, output_tokens={usage.output_tokens}" - ) - - # Warn if using non-OpenAI endpoint with images - if image_bytes and effective_base_url: - log.warning( - "โš ๏ธ Using custom endpoint - confirm the selected model supports vision." - ) - - except Exception as e: - # Handle global tool quota limit (UsageLimitExceeded) and other errors gracefully - error_msg = str(e) - llm_duration_ms = round((time.perf_counter() - llm_started) * 1000, 1) - log.warning(f"โš ๏ธ Agent execution encountered an error: {error_msg}") - run_result = None # Ensure run_result is defined for usage stats extraction - - # Check if this is a usage limit error (global tool quota) - if ( - "UsageLimitExceeded" in str(type(e).__name__) - or "tool_calls_limit" in error_msg.lower() - ): - log.warning( - "Global tool call quota reached - continuing with partial results" - ) - - result = ToolSelection( - conversation=Conversation( - status=ConversationStatus.COMPLETE, - context="The agent reached the maximum number of tool calls allowed. Please try a more specific query or break down your request into smaller parts.", - question=None, - options=None, - ), - choices=[], - explanation="Tool call limit reached during execution. Try refining your query.", - reason=None, - ) - else: - raise - - # ---- 7) Convert raw tool call records into ToolRunLog objects ---------- - for tc in getattr(deps, "tool_calls", []): - tool_name = tc.get("tool") - timestamp = tc.get("timestamp") - error = tc.get("error") - inputs = { - k: v for k, v in tc.items() if k not in ("tool", "timestamp", "error") - } - tool_logs.append( - ToolRunLog( - tool=tool_name, - inputs=inputs, - timestamp=timestamp, - error=error, - ) - ) - - stage_counts: dict[str, int] = {} - stage_durations: dict[str, float] = {} - for tc in getattr(deps, "tool_calls", []): - name = tc.get("tool", "unknown") - stage_counts[name] = stage_counts.get(name, 0) + 1 - duration_ms = tc.get("duration_ms") - if isinstance(duration_ms, (int, float)): - stage_durations[name] = stage_durations.get(name, 0.0) + float(duration_ms) - - total_duration_ms = round((time.perf_counter() - run_started) * 1000, 1) - log.info( - "โฑ๏ธ Latency summary: total_ms=%s metadata_ms=%s llm_ms=%s tools=%s tool_ms=%s", - total_duration_ms, - metadata_duration_ms, - llm_duration_ms, - stage_counts, - {k: round(v, 1) for k, v in stage_durations.items()}, - ) - - # ---- 8) Extract usage statistics if available ------------------------- - usage_stats = None - if run_result and hasattr(run_result, "usage") and run_result.usage: - usage = run_result.usage() - usage_stats = UsageStats( - total_tokens=usage.total_tokens, - input_tokens=usage.input_tokens, - output_tokens=usage.output_tokens, - ) - - # ---- 9) Wrap into high-level AgentToolSelection ------------------------ - return AgentToolSelection( - conversation=result.conversation, - choices=result.choices, - explanation=result.explanation, - reason=result.reason, - tool_calls=tool_logs, - usage=usage_stats, - ) - - -__all__ = ["run_agent", "agent"] +from __future__ import annotations + +import os +import logging +import time +import asyncio +import threading +from collections import OrderedDict +from datetime import datetime +from typing import List + +from pydantic_ai import Agent, RunContext +from pydantic_ai.usage import UsageLimits +from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIChatModel +from pydantic_ai.providers.openai import OpenAIProvider +from pydantic_ai.messages import BinaryContent + +from ai_agent.generator.prompts import get_agent_system_prompt +from ai_agent.generator.schema import ToolSelection, Conversation, ConversationStatus +from ai_agent.utils.config import get_config +from .models import AgentToolSelection, ToolRunLog, UsageStats +from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput +from ai_agent.agent.utils import coerce_github_url_or_none +from .tools.search_tool import tool_search_tools, SearchToolsInput +from .tools.search_alternative_tool import ( + tool_search_alternative, + SearchAlternativeInput, +) +from .tools.sparql_tool import tool_sparql_query, SparqlQueryInput +from .tools.query_utils import sanitize_retrieval_query +from .utils import AgentState, limit_tool_calls, cap_prepare +from ai_agent.utils.image_meta import summarize_image_metadata, detect_ext_token + +log = logging.getLogger("agent.core") + +DEFAULT_NUM_CHOICES = int(os.getenv("NUM_CHOICES", "3")) + +# --------------------------------------------------------------------------- +# Dynamic agent instance cache +# Key: (model_name, base_url, api_key_env, num_choices) +# Avoids rebuilding Agent/OpenAIProvider/model objects on every request when +# the UI repeatedly uses the same custom endpoint + model combination. +# Bounded LRU (max AGENT_CACHE_MAX entries); protected by a lock so that +# concurrent requests cannot race while creating/inserting agents. +# --------------------------------------------------------------------------- +_AGENT_CACHE_MAX: int = int(os.getenv("AGENT_CACHE_MAX", "16")) +_AGENT_CACHE_LOCK: threading.Lock = threading.Lock() +_AGENT_CACHE: OrderedDict[tuple, "Agent"] = OrderedDict() + +# --------------------------------------------------------------------------- +# Model / provider setup +# --------------------------------------------------------------------------- +config = get_config() +agent_model_config = config.agent_model + +try: + api_key = agent_model_config.get_api_key() +except ValueError as e: + log.error(f"Failed to get API key for agent model: {e}") + raise + +log.info(f"Initializing agent model: {agent_model_config.name}") + +if agent_model_config.base_url: + log.info(f"Using custom OpenAI base URL: {agent_model_config.base_url}") + log.info("Using OpenAIChatModel (chat/completions API) for custom endpoint") + provider = OpenAIProvider( + base_url=agent_model_config.base_url, + api_key=api_key, + ) + openai_model = OpenAIChatModel( + model_name=agent_model_config.name, + provider=provider, + ) +else: + provider = OpenAIProvider(api_key=api_key) + openai_model = OpenAIResponsesModel( + model_name=agent_model_config.name, + provider=provider, + ) + +# --------------------------------------------------------------------------- +# Agent definition +# --------------------------------------------------------------------------- +agent = Agent( + model=openai_model, + system_prompt=get_agent_system_prompt(DEFAULT_NUM_CHOICES), + deps_type=AgentState, + output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), +) + +# --------------------------------------------------------------------------- +# Tool adapters for the agent +# --------------------------------------------------------------------------- + + +@agent.tool(retries=2, prepare=cap_prepare) +@limit_tool_calls("search_tools", cap=1) +async def search_tools( + ctx: RunContext[AgentState], + query: str, + excluded: List[str] | None = None, + top_k: int = 12, +) -> List[dict]: + """ + Agent-facing search tool. + + Delegates to tools.search_tool.tool_search_tools(), but automatically + injects: + - globally excluded tools (from ctx.deps.excluded_tools) + - image_paths and original_formats (from ctx.deps, set in run_agent) + so the language model never has to reason about file paths directly. + """ + # Merge explicit exclusions with global exclusions from AgentState + explicit_excluded = excluded or [] + global_excluded = getattr(ctx.deps, "excluded_tools", []) or [] + all_excluded = sorted(set(explicit_excluded + list(global_excluded))) + + original_formats = getattr(ctx.deps, "original_formats", []) or [] + image_paths = getattr(ctx.deps, "image_paths", []) or [] + + effective_top_k = ( + ctx.deps.override_top_k if ctx.deps.override_top_k is not None else top_k + ) + + started = time.perf_counter() + inp = SearchToolsInput( + query=sanitize_retrieval_query(query), + excluded=all_excluded, + top_k=effective_top_k, + original_formats=original_formats, + image_paths=image_paths, + ) + out = tool_search_tools(inp) + + ctx.deps.tool_calls.append( + { + "tool": "search_tools", + "query": query, + "count": len(out.candidates), + "duration_ms": round((time.perf_counter() - started) * 1000, 1), + "original_formats": original_formats, + "excluded": all_excluded, + "timestamp": datetime.now().isoformat(), + } + ) + + # Return plain dicts so the LLM sees a simple JSON-like structure. + return [c.model_dump(mode="python") for c in out.candidates] + + +@agent.tool(retries=2, prepare=cap_prepare) +@limit_tool_calls("search_alternative", cap=3) +async def search_alternative( + ctx: RunContext[AgentState], + alternative_query: str, + excluded: List[str] | None = None, + top_k: int = 12, +) -> List[dict]: + """ + Search with an alternative query formulation (includes automatic reranking). + """ + explicit_excluded = excluded or [] + global_excluded = getattr(ctx.deps, "excluded_tools", []) or [] + all_excluded = sorted(set(explicit_excluded + list(global_excluded))) + + original_formats = getattr(ctx.deps, "original_formats", []) or [] + image_paths = getattr(ctx.deps, "image_paths", []) or [] + + started = time.perf_counter() + inp = SearchAlternativeInput( + alternative_query=alternative_query, + excluded=all_excluded, + top_k=top_k, + original_formats=original_formats, + image_paths=image_paths, + ) + out = tool_search_alternative(inp) + + ctx.deps.tool_calls.append( + { + "tool": "search_alternative", + "alternative_query": alternative_query, + "query_used": out.query_used, + "count": len(out.candidates), + "duration_ms": round((time.perf_counter() - started) * 1000, 1), + "original_formats": original_formats, + "excluded": all_excluded, + "timestamp": datetime.now().isoformat(), + } + ) + + return [c.model_dump(mode="python") for c in out.candidates] + + +@agent.tool(retries=1, prepare=cap_prepare) +@limit_tool_calls("sparql_query", cap=3) +async def sparql_query( + ctx: RunContext[AgentState], + query: str, + limit: int = 50, +) -> dict: + """Run a read-only SPARQL SELECT / ASK against the GraphDB catalog. + + Use this for questions the semantic search can't answer cleanly โ€” + aggregates ("how many tools support DICOM?"), structural filters + ("tools that require a GPU AND are free"), distinct property values + ("which licences appear in the catalog?"), etc. + + UPDATE-style operations are rejected. Results capped at `limit` rows + (max 200). Returns columns + rows, plus a boolean for ASK queries. + """ + started = time.perf_counter() + inp = SparqlQueryInput(query=query, limit=max(1, min(200, int(limit)))) + out = tool_sparql_query(inp) + ctx.deps.tool_calls.append( + { + "tool": "sparql_query", + "query": query[:300], + "row_count": out.row_count, + "truncated": out.truncated, + "boolean": out.boolean, + "error": out.error, + "duration_ms": round((time.perf_counter() - started) * 1000, 1), + "timestamp": datetime.now().isoformat(), + } + ) + return out.model_dump(mode="python") + + +@agent.tool(retries=2, prepare=cap_prepare) +@limit_tool_calls("repo_info_batch", cap=4) +async def repo_info_batch( + ctx: RunContext[AgentState], + urls: List[str], +) -> List[dict]: + """Fetch repository summaries for multiple repositories in parallel.""" + started = time.perf_counter() + + if not urls: + return [] + + normalized: List[str] = [] + skipped: List[dict] = [] + seen: set[str] = set() + for raw in urls: + norm = coerce_github_url_or_none(raw) + if not norm: + skipped.append( + { + "url": raw, + "skipped": True, + "reason": "NON_GITHUB_URL", + } + ) + continue + if norm in seen: + continue + seen.add(norm) + normalized.append(norm) + + tasks = [tool_repo_summary(RepoSummaryInput(url=u)) for u in normalized] + outcomes = await asyncio.gather(*tasks, return_exceptions=True) + + results: List[dict] = [] + for url, outcome in zip(normalized, outcomes): + if isinstance(outcome, Exception): + results.append( + { + "url": url, + "source": "error", + "error": str(outcome), + } + ) + continue + payload = outcome.model_dump(mode="python") + payload["url"] = url + results.append(payload) + + if skipped: + results.extend(skipped) + + ctx.deps.tool_calls.append( + { + "tool": "repo_info_batch", + "requested": len(urls), + "normalized": len(normalized), + "returned": len(results), + "duration_ms": round((time.perf_counter() - started) * 1000, 1), + "timestamp": datetime.now().isoformat(), + } + ) + + return results + + +# --------------------------------------------------------------------------- +# High level entry point: run the agent on (text query + image) +# --------------------------------------------------------------------------- +def run_agent( + task: str, + image_paths: List[str], + excluded: List[str] | None = None, + conversation_history: List[str] | None = None, + *, + image_bytes: bytes | None = None, + model: str | None = None, + base_url: str | None = None, + api_key_env: str | None = None, + top_k: int | None = None, + num_choices: int | None = None, + image_metadata: str | None = None, +) -> AgentToolSelection: + """ + Execute the agent for a user task and at least one image path. + + - derive canonical original_formats (tiff / dicom / nifti / ...) + - build a compact image metadata summary (or use pre-computed one) + - pass both to the LLM as hidden context + - store image_paths/original_formats in deps so retrieval tools can use them + - optionally allow runtime model/base_url/top_k/num_choices overrides + + IMPORTANT: + The model only sees an actual image if `image_bytes` is provided. + `image_paths` are used for metadata + tool context only. + """ + run_started = time.perf_counter() + # image_paths may now be empty โ€” the agent runs in text-only mode and + # skips the VLM. Retrieval still works on the text query alone. + image_paths = list(image_paths or []) + + tool_logs: List[ToolRunLog] = [] + + # ---- 1) Derive image-based metadata and format hints -------------------- + metadata_started = time.perf_counter() + meta_str = ( + image_metadata + if image_metadata is not None + else (summarize_image_metadata(image_paths) or "") + ) + fmt_str = detect_ext_token(image_paths) or "" + original_formats = [t.lower() for t in fmt_str.split()] if fmt_str else [] + metadata_duration_ms = round((time.perf_counter() - metadata_started) * 1000, 1) + + effective_top_k = top_k if top_k is not None else 12 + effective_num_choices = num_choices if num_choices is not None else 3 + + # ---- 2) Prepare dependency state passed to all tools -------------------- + deps = AgentState( + excluded_tools=excluded or [], + override_model=model, + override_base_url=base_url, + override_top_k=effective_top_k, + override_num_choices=effective_num_choices, + ) + + setattr(deps, "image_paths", list(image_paths)) + setattr(deps, "original_formats", original_formats) + + # ---- 3) Hidden metadata lines for the model ---------------------------- + hidden_meta = "" + if original_formats: + hidden_meta += "\n(Formats Hint: " + ",".join(original_formats) + ")" + if meta_str: + short_meta = " ".join(x.strip() for x in meta_str.splitlines() if x.strip()) + hidden_meta += ( + "\n(Image Metadata: " + + short_meta[:500] + + ("โ€ฆ" if len(short_meta) > 500 else "") + + ")" + ) + hidden_meta += f"\n(Search top_k: {effective_top_k})" + + extra_context = "\n\n**CRITICAL: Analyze the attached preview image showing the user's data.**\nUse visual observations (anatomy visible, image quality, dimensionality, contrast) combined with the metadata below to recommend tools. Reference what you see in your explanations." + + # ---- 4) Build the prompt (optionally including history) ---------------- + if conversation_history and len(conversation_history) > 0: + history_text = "\n".join(conversation_history) + prompt = ( + f"Previous conversation:\n{history_text}\n\n" + f"Current request: {task}{extra_context}{hidden_meta}" + ) + else: + prompt = task + extra_context + hidden_meta + + # ----------------------------------------------------------------------- + # Determine which agent instance to use + # ----------------------------------------------------------------------- + agent_instance = agent + effective_num_choices = num_choices if num_choices is not None else 3 + effective_model = model if model else agent_model_config.name + effective_top_k = top_k if top_k is not None else 12 + + # When model is provided from UI, base_url comes with it (can be None for OpenAI) + if model: + # Use api_key_env from config if provided, otherwise default to OPENAI_API_KEY + key_env_name = api_key_env if api_key_env else "OPENAI_API_KEY" + runtime_api_key = os.getenv(key_env_name) + if not runtime_api_key: + raise ValueError( + f"{key_env_name} not found in environment. Cannot use this model." + ) + effective_base_url = base_url # Can be None for OpenAI + log.info(f"โœ“ Using {key_env_name} for model {effective_model}") + log.debug(f"{key_env_name} is set: {bool(runtime_api_key)}") + else: + # No model override - use config defaults + effective_base_url = agent_model_config.base_url + runtime_api_key = api_key # Already loaded from config at startup + log.info(f"โœ“ Using API key from config for model {effective_model}") + + # Log runtime configuration + endpoint_display = effective_base_url if effective_base_url else "api.openai.com" + log.info( + f"๐Ÿค– Agent execution - Model: {effective_model}, endpoint: {endpoint_display}, " + f"top_k: {effective_top_k}, num_choices: {effective_num_choices}, excluded: {len(excluded or [])}" + ) + + needs_dynamic_agent = model is not None + + if needs_dynamic_agent: + cache_key = (effective_model, effective_base_url or "", api_key_env or "OPENAI_API_KEY", effective_num_choices) + with _AGENT_CACHE_LOCK: + agent_instance = _AGENT_CACHE.get(cache_key) + if agent_instance is not None: + _AGENT_CACHE.move_to_end(cache_key) + if agent_instance is None: + log.info( + f"๐Ÿ“ฆ Creating runtime agent with model={effective_model}, endpoint={effective_base_url or 'api.openai.com'}" + ) + + runtime_provider = OpenAIProvider( + base_url=effective_base_url, + api_key=runtime_api_key, + ) + + # Use OpenAIChatModel (chat/completions) for custom endpoints, OpenAIResponsesModel for default OpenAI + if effective_base_url: + log.info("Using OpenAIChatModel (chat/completions API) for custom endpoint") + runtime_model = OpenAIChatModel( + model_name=effective_model, provider=runtime_provider + ) + else: + runtime_model = OpenAIResponsesModel( + model_name=effective_model, provider=runtime_provider + ) + + agent_instance = Agent( + model=runtime_model, + system_prompt=get_agent_system_prompt(effective_num_choices), + deps_type=AgentState, + output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), + ) + + # Register tools on the dynamic agent + agent_instance.tool(search_tools, retries=2, prepare=cap_prepare) + agent_instance.tool(search_alternative, retries=2, prepare=cap_prepare) + agent_instance.tool(repo_info_batch, retries=2, prepare=cap_prepare) + agent_instance.tool(sparql_query, retries=1, prepare=cap_prepare) + + with _AGENT_CACHE_LOCK: + _AGENT_CACHE[cache_key] = agent_instance + _AGENT_CACHE.move_to_end(cache_key) + while len(_AGENT_CACHE) > _AGENT_CACHE_MAX: + _AGENT_CACHE.popitem(last=False) + else: + log.info( + f"โ™ป๏ธ Reusing cached dynamic agent (model: {effective_model}, num_choices: {effective_num_choices})" + ) + + elif ( + num_choices is not None and num_choices != DEFAULT_NUM_CHOICES + ): + cache_key = (effective_model, effective_base_url or "", api_key_env or "OPENAI_API_KEY", effective_num_choices) + with _AGENT_CACHE_LOCK: + agent_instance = _AGENT_CACHE.get(cache_key) + if agent_instance is not None: + _AGENT_CACHE.move_to_end(cache_key) + if agent_instance is None: + log.info( + f"๐Ÿ“ฆ Creating runtime agent with num_choices={effective_num_choices} (model: {effective_model})" + ) + agent_instance = Agent( + model=openai_model, + system_prompt=get_agent_system_prompt(effective_num_choices), + deps_type=AgentState, + output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), + ) + + # Register tools on the dynamic agent + agent_instance.tool(search_tools, retries=2, prepare=cap_prepare) + agent_instance.tool(search_alternative, retries=2, prepare=cap_prepare) + agent_instance.tool(repo_info_batch, retries=2, prepare=cap_prepare) + agent_instance.tool(sparql_query, retries=1, prepare=cap_prepare) + + with _AGENT_CACHE_LOCK: + _AGENT_CACHE[cache_key] = agent_instance + _AGENT_CACHE.move_to_end(cache_key) + while len(_AGENT_CACHE) > _AGENT_CACHE_MAX: + _AGENT_CACHE.popitem(last=False) + else: + log.info( + f"โ™ป๏ธ Reusing cached dynamic agent with num_choices={effective_num_choices} (model: {effective_model})" + ) + + else: + log.info( + f"โ™ป๏ธ Using global agent (model: {effective_model}, num_choices: {effective_num_choices})" + ) + + log.debug( + f"Prompt length: {len(prompt)} chars, has_image_paths: {bool(image_paths)}, has_image_bytes: {bool(image_bytes)}" + ) + + # ---- 5) Build multimodal prompt if image bytes provided ---------------- + if image_bytes: + log.info( + f"๐Ÿ–ผ๏ธ Sending image preview to model ({len(image_bytes)} bytes = {len(image_bytes)/1024:.1f}KB)" + ) + user_prompt = [ + prompt, + BinaryContent( + data=image_bytes, + media_type="image/png", + ), + ] + else: + log.warning( + "โš ๏ธ No image bytes provided - the model will not see the image preview" + ) + user_prompt = prompt + + # ---- 6) Run the agent -------------------------------------------------- + try: + llm_started = time.perf_counter() + run_result = agent_instance.run_sync( + user_prompt, + deps=deps, + output_type=ToolSelection, + usage_limits=UsageLimits(tool_calls_limit=20), + ) + llm_duration_ms = round((time.perf_counter() - llm_started) * 1000, 1) + result = run_result.output + + log.info( + f"โœ… Agent execution complete - choices returned: {len(result.choices)}" + ) + + # Log usage (helpful, but may not explicitly expose image-specific counters) + if run_result.usage: + usage = run_result.usage() + log.info( + f"๐Ÿ“Š Usage: total_tokens={usage.total_tokens}, " + f"input_tokens={usage.input_tokens}, output_tokens={usage.output_tokens}" + ) + + # Warn if using non-OpenAI endpoint with images + if image_bytes and effective_base_url: + log.warning( + "โš ๏ธ Using custom endpoint - confirm the selected model supports vision." + ) + + except Exception as e: + # Handle global tool quota limit (UsageLimitExceeded) and other errors gracefully + error_msg = str(e) + llm_duration_ms = round((time.perf_counter() - llm_started) * 1000, 1) + log.warning(f"โš ๏ธ Agent execution encountered an error: {error_msg}") + run_result = None # Ensure run_result is defined for usage stats extraction + + # Check if this is a usage limit error (global tool quota) + if ( + "UsageLimitExceeded" in str(type(e).__name__) + or "tool_calls_limit" in error_msg.lower() + ): + log.warning( + "Global tool call quota reached - continuing with partial results" + ) + + result = ToolSelection( + conversation=Conversation( + status=ConversationStatus.COMPLETE, + context="The agent reached the maximum number of tool calls allowed. Please try a more specific query or break down your request into smaller parts.", + question=None, + options=None, + ), + choices=[], + explanation="Tool call limit reached during execution. Try refining your query.", + reason=None, + ) + else: + raise + + # ---- 7) Convert raw tool call records into ToolRunLog objects ---------- + for tc in getattr(deps, "tool_calls", []): + tool_name = tc.get("tool") + timestamp = tc.get("timestamp") + error = tc.get("error") + inputs = { + k: v for k, v in tc.items() if k not in ("tool", "timestamp", "error") + } + tool_logs.append( + ToolRunLog( + tool=tool_name, + inputs=inputs, + timestamp=timestamp, + error=error, + ) + ) + + stage_counts: dict[str, int] = {} + stage_durations: dict[str, float] = {} + for tc in getattr(deps, "tool_calls", []): + name = tc.get("tool", "unknown") + stage_counts[name] = stage_counts.get(name, 0) + 1 + duration_ms = tc.get("duration_ms") + if isinstance(duration_ms, (int, float)): + stage_durations[name] = stage_durations.get(name, 0.0) + float(duration_ms) + + total_duration_ms = round((time.perf_counter() - run_started) * 1000, 1) + log.info( + "โฑ๏ธ Latency summary: total_ms=%s metadata_ms=%s llm_ms=%s tools=%s tool_ms=%s", + total_duration_ms, + metadata_duration_ms, + llm_duration_ms, + stage_counts, + {k: round(v, 1) for k, v in stage_durations.items()}, + ) + + # ---- 8) Extract usage statistics if available ------------------------- + usage_stats = None + if run_result and hasattr(run_result, "usage") and run_result.usage: + usage = run_result.usage() + usage_stats = UsageStats( + total_tokens=usage.total_tokens, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + ) + + # ---- 9) Wrap into high-level AgentToolSelection ------------------------ + return AgentToolSelection( + conversation=result.conversation, + choices=result.choices, + explanation=result.explanation, + reason=result.reason, + tool_calls=tool_logs, + usage=usage_stats, + ) + + +__all__ = ["run_agent", "agent"] diff --git a/src/ai_agent/agent/tools/__init__.py b/src/ai_agent/agent/tools/__init__.py index 6b19ced4476fe5d5fb1f0c5823868219a5395a12..67be95e870634c5e5f58560c7eae9d67dbb44765 100644 --- a/src/ai_agent/agent/tools/__init__.py +++ b/src/ai_agent/agent/tools/__init__.py @@ -1,38 +1,38 @@ -"""Agent tools package.""" - -# Only export registry - tools will self-register when imported explicitly -from .mcp import ( - TOOL_REGISTRY, - get_tool, - register_tool, - list_tools, - ensure_mcp_tools_registered, -) - -# Import tools lazily to avoid loading heavy dependencies at package import -# Tools should be imported explicitly where needed, e.g.: -# from ai_agent.agent.tools.mcp.lungs_segmentation_tool import tool_lungs_segmentation - -__all__ = [ - "TOOL_REGISTRY", - "get_tool", - "register_tool", - "list_tools", - "ensure_tools_registered", -] - - -def ensure_tools_registered(): - """ - Import all tools to trigger their registration. - Call this once at app startup. - """ - from importlib import import_module - - import_module("ai_agent.agent.tools.search_tool") - import_module("ai_agent.agent.tools.search_alternative_tool") - import_module("ai_agent.agent.tools.repo_info_tool") - import_module("ai_agent.agent.tools.gradio_space_tool") - - # Import MCP tools - ensure_mcp_tools_registered() +"""Agent tools package.""" + +# Only export registry - tools will self-register when imported explicitly +from .mcp import ( + TOOL_REGISTRY, + get_tool, + register_tool, + list_tools, + ensure_mcp_tools_registered, +) + +# Import tools lazily to avoid loading heavy dependencies at package import +# Tools should be imported explicitly where needed, e.g.: +# from ai_agent.agent.tools.mcp.lungs_segmentation_tool import tool_lungs_segmentation + +__all__ = [ + "TOOL_REGISTRY", + "get_tool", + "register_tool", + "list_tools", + "ensure_tools_registered", +] + + +def ensure_tools_registered(): + """ + Import all tools to trigger their registration. + Call this once at app startup. + """ + from importlib import import_module + + import_module("ai_agent.agent.tools.search_tool") + import_module("ai_agent.agent.tools.search_alternative_tool") + import_module("ai_agent.agent.tools.repo_info_tool") + import_module("ai_agent.agent.tools.gradio_space_tool") + + # Import MCP tools + ensure_mcp_tools_registered() diff --git a/src/ai_agent/agent/tools/deepwiki_tool.py b/src/ai_agent/agent/tools/deepwiki_tool.py index 6e3bdeaa338d452d1fe4cdbbe25f17add8d481ea..3cc9692e74a8fb6c3579e94ed72d824e27022a70 100644 --- a/src/ai_agent/agent/tools/deepwiki_tool.py +++ b/src/ai_agent/agent/tools/deepwiki_tool.py @@ -1,103 +1,103 @@ -from __future__ import annotations - -import asyncio -import logging -import os -from typing import Optional - -from pydantic import BaseModel -from pydantic_ai.mcp import MCPServerStreamableHTTP - -from ai_agent.agent.tools.utils import _clip -from ai_agent.agent.utils import _coerce_owner_repo_ref - -log = logging.getLogger("agent.deepwiki") - -# DeepWiki MCP server endpoint (Streamable HTTP transport) -DEEPWIKI_HTTP_URL = "https://mcp.deepwiki.com/mcp" - -# Timeout for DeepWiki operations (seconds) -DEEPWIKI_TIMEOUT = int(os.getenv("DEEPWIKI_TIMEOUT", "20")) - - -class DeepWikiInput(BaseModel): - """Input for DeepWiki operations.""" - - url: str # GitHub repository URL or owner/repo format - - -class DeepWikiContentsOutput(BaseModel): - """Output from read_wiki_contents.""" - - success: bool - contents: Optional[str] = None - error: Optional[str] = None - truncated: bool = False - - -async def get_wiki_contents(input: DeepWikiInput) -> DeepWikiContentsOutput: - """ - Fetch repo docs from DeepWiki MCP (Streamable HTTP) and return a clipped string - to keep LLM token usage under control. - """ - owner, repo, _ = _coerce_owner_repo_ref(input.url) - repo = f"{owner}/{repo}" - - try: - server = MCPServerStreamableHTTP(DEEPWIKI_HTTP_URL) - - async with server: - result = await asyncio.wait_for( - server.direct_call_tool("read_wiki_contents", {"repoName": repo}), - timeout=DEEPWIKI_TIMEOUT, - ) - - # Handle different result types from MCP - text = None - if isinstance(result, str): - # Direct string result - text = result - elif hasattr(result, "content"): - # MCP ToolResult with content field - text_parts = [] - for item in result.content: - if hasattr(item, "text"): - text_parts.append(item.text) - elif isinstance(item, str): - text_parts.append(item) - text = "\n".join(text_parts) if text_parts else None - elif isinstance(result, list): - # List of strings or content items - text = "\n".join([str(p) for p in result if p]) or None - - if text and text.strip(): - clipped_text, truncated = _clip(text.strip()) - return DeepWikiContentsOutput( - success=True, - contents=clipped_text, - truncated=truncated, - ) - - return DeepWikiContentsOutput( - success=False, error="No content returned from DeepWiki" - ) - - except asyncio.TimeoutError: - log.warning(f"DeepWiki timed out after {DEEPWIKI_TIMEOUT}s for {repo}") - return DeepWikiContentsOutput( - success=False, - error=f"DeepWiki request timed out after {DEEPWIKI_TIMEOUT}s", - ) - except Exception as e: - log.error(f"Failed to get wiki contents for {repo}: {e}") - return DeepWikiContentsOutput( - success=False, - error=f"Failed to connect to DeepWiki: {str(e)}", - ) - - -__all__ = [ - "get_wiki_contents", - "DeepWikiInput", - "DeepWikiContentsOutput", -] +from __future__ import annotations + +import asyncio +import logging +import os +from typing import Optional + +from pydantic import BaseModel +from pydantic_ai.mcp import MCPServerStreamableHTTP + +from ai_agent.agent.tools.utils import _clip +from ai_agent.agent.utils import _coerce_owner_repo_ref + +log = logging.getLogger("agent.deepwiki") + +# DeepWiki MCP server endpoint (Streamable HTTP transport) +DEEPWIKI_HTTP_URL = "https://mcp.deepwiki.com/mcp" + +# Timeout for DeepWiki operations (seconds) +DEEPWIKI_TIMEOUT = int(os.getenv("DEEPWIKI_TIMEOUT", "20")) + + +class DeepWikiInput(BaseModel): + """Input for DeepWiki operations.""" + + url: str # GitHub repository URL or owner/repo format + + +class DeepWikiContentsOutput(BaseModel): + """Output from read_wiki_contents.""" + + success: bool + contents: Optional[str] = None + error: Optional[str] = None + truncated: bool = False + + +async def get_wiki_contents(input: DeepWikiInput) -> DeepWikiContentsOutput: + """ + Fetch repo docs from DeepWiki MCP (Streamable HTTP) and return a clipped string + to keep LLM token usage under control. + """ + owner, repo, _ = _coerce_owner_repo_ref(input.url) + repo = f"{owner}/{repo}" + + try: + server = MCPServerStreamableHTTP(DEEPWIKI_HTTP_URL) + + async with server: + result = await asyncio.wait_for( + server.direct_call_tool("read_wiki_contents", {"repoName": repo}), + timeout=DEEPWIKI_TIMEOUT, + ) + + # Handle different result types from MCP + text = None + if isinstance(result, str): + # Direct string result + text = result + elif hasattr(result, "content"): + # MCP ToolResult with content field + text_parts = [] + for item in result.content: + if hasattr(item, "text"): + text_parts.append(item.text) + elif isinstance(item, str): + text_parts.append(item) + text = "\n".join(text_parts) if text_parts else None + elif isinstance(result, list): + # List of strings or content items + text = "\n".join([str(p) for p in result if p]) or None + + if text and text.strip(): + clipped_text, truncated = _clip(text.strip()) + return DeepWikiContentsOutput( + success=True, + contents=clipped_text, + truncated=truncated, + ) + + return DeepWikiContentsOutput( + success=False, error="No content returned from DeepWiki" + ) + + except asyncio.TimeoutError: + log.warning(f"DeepWiki timed out after {DEEPWIKI_TIMEOUT}s for {repo}") + return DeepWikiContentsOutput( + success=False, + error=f"DeepWiki request timed out after {DEEPWIKI_TIMEOUT}s", + ) + except Exception as e: + log.error(f"Failed to get wiki contents for {repo}: {e}") + return DeepWikiContentsOutput( + success=False, + error=f"Failed to connect to DeepWiki: {str(e)}", + ) + + +__all__ = [ + "get_wiki_contents", + "DeepWikiInput", + "DeepWikiContentsOutput", +] diff --git a/src/ai_agent/agent/tools/mcp/__init__.py b/src/ai_agent/agent/tools/mcp/__init__.py index dd37cbbf01c24bd8cee1e605595851d3b34b4799..000f79b217dc2725fd9de62ce91ec9df39052f65 100644 --- a/src/ai_agent/agent/tools/mcp/__init__.py +++ b/src/ai_agent/agent/tools/mcp/__init__.py @@ -1,53 +1,53 @@ -""" -MCP (Model Context Protocol) tools package. - -This package contains registered imaging tools that require approval -and follow the tool registry pattern. -""" - -from .registry import ( - TOOL_REGISTRY, - CATALOG_NAME_TO_TOOL, - get_tool, - register_tool, - list_tools, - get_tool_display_name, - get_tool_icon, - extract_preview, - extract_downloads, - extract_metadata, - extract_output_field, - ToolConfig, -) - -from .base import BaseToolInput, BaseToolOutput, ImageToolInput - -__all__ = [ - # Registry - "TOOL_REGISTRY", - "CATALOG_NAME_TO_TOOL", - "get_tool", - "register_tool", - "list_tools", - "get_tool_display_name", - "get_tool_icon", - "extract_preview", - "extract_downloads", - "extract_metadata", - "extract_output_field", - "ToolConfig", - # Base models - "BaseToolInput", - "BaseToolOutput", - "ImageToolInput", -] - - -def ensure_mcp_tools_registered(): - """ - Import all MCP tools to trigger their registration. - Call this once at app startup. - """ - from importlib import import_module - - import_module("ai_agent.agent.tools.mcp.lungs_segmentation_tool") +""" +MCP (Model Context Protocol) tools package. + +This package contains registered imaging tools that require approval +and follow the tool registry pattern. +""" + +from .registry import ( + TOOL_REGISTRY, + CATALOG_NAME_TO_TOOL, + get_tool, + register_tool, + list_tools, + get_tool_display_name, + get_tool_icon, + extract_preview, + extract_downloads, + extract_metadata, + extract_output_field, + ToolConfig, +) + +from .base import BaseToolInput, BaseToolOutput, ImageToolInput + +__all__ = [ + # Registry + "TOOL_REGISTRY", + "CATALOG_NAME_TO_TOOL", + "get_tool", + "register_tool", + "list_tools", + "get_tool_display_name", + "get_tool_icon", + "extract_preview", + "extract_downloads", + "extract_metadata", + "extract_output_field", + "ToolConfig", + # Base models + "BaseToolInput", + "BaseToolOutput", + "ImageToolInput", +] + + +def ensure_mcp_tools_registered(): + """ + Import all MCP tools to trigger their registration. + Call this once at app startup. + """ + from importlib import import_module + + import_module("ai_agent.agent.tools.mcp.lungs_segmentation_tool") diff --git a/src/ai_agent/agent/tools/mcp/base.py b/src/ai_agent/agent/tools/mcp/base.py index 592cb20f45c256dbca027f3dee48a7a30841ad12..344045813a7988a67f43ec4b6f4ffe524c836387 100644 --- a/src/ai_agent/agent/tools/mcp/base.py +++ b/src/ai_agent/agent/tools/mcp/base.py @@ -1,88 +1,88 @@ -from __future__ import annotations - -from typing import Optional, Dict, Any -from pydantic import BaseModel, Field - - -class BaseToolInput(BaseModel): - """ - Base input model that tools can extend. - - Common patterns: - - image_path: Path to uploaded image/volume - - description: Optional context from agent - """ - - pass # Intentionally minimal - tools define their own inputs - - -class BaseToolOutput(BaseModel): - """ - Base output model that all tools should follow. - - This ensures consistent handling in the UI layer without - needing tool-specific code. - - Standard fields: - - success: bool - Whether execution succeeded - - error: Optional[str] - Error message if failed - - compute_time_seconds: float - Time taken by tool - - notes: Optional[str] - Additional info for user - - File outputs (at least one should be provided on success): - - result_preview: Optional[str] - PNG/GIF preview for inline display - - result_origin: Optional[str] - Original format file for download - - result_path: Optional[str] - Backward compat field - - Metadata: - - metadata_text: Optional[str] - Structured info about result - - metadata: Dict[str, Any] - Machine-readable metadata - - Tracking: - - endpoint_url: str - API endpoint used - - api_name: str - API method called - """ - - # Core status - success: bool = False - error: Optional[str] = None - compute_time_seconds: float = 0.0 - - # File outputs (tools should provide these for UI to display/download) - result_preview: Optional[str] = Field( - default=None, - description="Path to preview image (PNG/GIF) for inline display in chat", - ) - result_origin: Optional[str] = Field( - default=None, - description="Path to original format file (TIFF/NIfTI/DICOM) for download", - ) - result_path: Optional[str] = Field( - default=None, description="Backward compatibility: primary result path" - ) - - # Metadata - metadata_text: Optional[str] = Field( - default=None, description="Human-readable metadata about the result" - ) - metadata: Dict[str, Any] = Field( - default_factory=dict, description="Machine-readable metadata" - ) - notes: Optional[str] = Field( - default=None, description="Additional notes or context for user" - ) - - # Tracking/debugging - endpoint_url: str = "" - api_name: str = "" - - -class ImageToolInput(BaseToolInput): - """ - Common input pattern for image/volume processing tools. - """ - - image_path: str = Field(description="Path to the image/volume file") - description: Optional[str] = Field( - default=None, description="Optional context or notes from agent about the task" - ) +from __future__ import annotations + +from typing import Optional, Dict, Any +from pydantic import BaseModel, Field + + +class BaseToolInput(BaseModel): + """ + Base input model that tools can extend. + + Common patterns: + - image_path: Path to uploaded image/volume + - description: Optional context from agent + """ + + pass # Intentionally minimal - tools define their own inputs + + +class BaseToolOutput(BaseModel): + """ + Base output model that all tools should follow. + + This ensures consistent handling in the UI layer without + needing tool-specific code. + + Standard fields: + - success: bool - Whether execution succeeded + - error: Optional[str] - Error message if failed + - compute_time_seconds: float - Time taken by tool + - notes: Optional[str] - Additional info for user + + File outputs (at least one should be provided on success): + - result_preview: Optional[str] - PNG/GIF preview for inline display + - result_origin: Optional[str] - Original format file for download + - result_path: Optional[str] - Backward compat field + + Metadata: + - metadata_text: Optional[str] - Structured info about result + - metadata: Dict[str, Any] - Machine-readable metadata + + Tracking: + - endpoint_url: str - API endpoint used + - api_name: str - API method called + """ + + # Core status + success: bool = False + error: Optional[str] = None + compute_time_seconds: float = 0.0 + + # File outputs (tools should provide these for UI to display/download) + result_preview: Optional[str] = Field( + default=None, + description="Path to preview image (PNG/GIF) for inline display in chat", + ) + result_origin: Optional[str] = Field( + default=None, + description="Path to original format file (TIFF/NIfTI/DICOM) for download", + ) + result_path: Optional[str] = Field( + default=None, description="Backward compatibility: primary result path" + ) + + # Metadata + metadata_text: Optional[str] = Field( + default=None, description="Human-readable metadata about the result" + ) + metadata: Dict[str, Any] = Field( + default_factory=dict, description="Machine-readable metadata" + ) + notes: Optional[str] = Field( + default=None, description="Additional notes or context for user" + ) + + # Tracking/debugging + endpoint_url: str = "" + api_name: str = "" + + +class ImageToolInput(BaseToolInput): + """ + Common input pattern for image/volume processing tools. + """ + + image_path: str = Field(description="Path to the image/volume file") + description: Optional[str] = Field( + default=None, description="Optional context or notes from agent about the task" + ) diff --git a/src/ai_agent/agent/tools/mcp/lungs_segmentation_tool.py b/src/ai_agent/agent/tools/mcp/lungs_segmentation_tool.py index d95b55852c46e0559ff776f9f7f50c96bf8ead15..62350b3680281b7b9badc19374f952608d13b9c3 100644 --- a/src/ai_agent/agent/tools/mcp/lungs_segmentation_tool.py +++ b/src/ai_agent/agent/tools/mcp/lungs_segmentation_tool.py @@ -1,442 +1,442 @@ -from __future__ import annotations - -from typing import Optional, Any, Dict, Tuple -import os -import logging -import tempfile -from pathlib import Path -import time - -import requests -from gradio_client import Client, handle_file - -from ai_agent.utils.previews import _build_preview_for_vlm -from ai_agent.utils.temp_file_manager import register_temp_file -from ai_agent.agent.tools.mcp.registry import register_tool, ToolConfig -from ai_agent.agent.tools.mcp.base import BaseToolOutput, ImageToolInput - -log = logging.getLogger("agent.lungs_segmentation") - - -# --------------------------------------------------------------------- -# Models -# --------------------------------------------------------------------- -class LungsSegmentationInput(ImageToolInput): - """Input for 3D lungs segmentation tool.""" - - pass # Inherits image_path and description from ImageToolInput - - -class LungsSegmentationOutput(BaseToolOutput): - """Output from 3D lungs segmentation tool.""" - - # All standard fields inherited from BaseToolOutput: - # - success, error, compute_time_seconds, notes - # - result_preview, result_origin, result_path - # - metadata_text, endpoint_url, api_name - pass - - -# --------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------- -LUNGS_SEGMENTATION_ENDPOINT = "https://qchapp-3d-lungs-segmentation.hf.space/" -LUNGS_SEGMENTATION_API_NAME = "/segment" - -# Maximum file size for downloads (1GB for medical imaging) -MAX_DOWNLOAD_SIZE = 1024 * 1024 * 1024 # 1GB in bytes - - -# --------------------------------------------------------------------- -# Public tool -# --------------------------------------------------------------------- -def tool_lungs_segmentation(inp: LungsSegmentationInput) -> LungsSegmentationOutput: - """ - Run 3D lungs segmentation on a CT scan image via a Gradio Space. - - Materialization strategy (robust): - 1) If Space returns dict FileData (url/path/etc) -> download via URL. - 2) If Space returns URL string -> download. - 3) If Space returns local file -> use it. - 4) If Space returns server path (/tmp/...) -> try /gradio_api/file=... (may 403). - """ - start_time = time.time() - - if not os.path.exists(inp.image_path): - return LungsSegmentationOutput( - success=False, - error=f"Image file not found: {inp.image_path}", - endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, - api_name=LUNGS_SEGMENTATION_API_NAME, - ) - - hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") - - try: - log.info( - "Running lungs segmentation on %s (endpoint: %s)", - inp.image_path, - LUNGS_SEGMENTATION_ENDPOINT, - ) - - client = _make_gradio_client(LUNGS_SEGMENTATION_ENDPOINT, hf_token) - - # Call API - try: - result = client.predict( - file_obj=handle_file(inp.image_path), - api_name=LUNGS_SEGMENTATION_API_NAME, - ) - log.info("API returned type=%s value=%r", type(result), result) - except Exception as e: - return LungsSegmentationOutput( - success=False, - error=f"API call failed: {e}", - compute_time_seconds=time.time() - start_time, - endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, - api_name=LUNGS_SEGMENTATION_API_NAME, - ) - - # Materialize to local file - origin_path = _materialize_any(result, client=client, hf_token=hf_token) - - compute_time = time.time() - start_time - - if not origin_path or not os.path.exists(origin_path): - # This is the common case if the Space returns '/tmp/...' and Gradio blocks it (403). - return LungsSegmentationOutput( - success=False, - error="Could not materialize/download the result file.", - compute_time_seconds=compute_time, - endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, - api_name=LUNGS_SEGMENTATION_API_NAME, - notes=( - f"API returned: {result!r}. If this is a '/tmp/...' path and you see HTTP 403, " - "the Space must return a FileData/url (recommended) or whitelist the output directory " - "via allowed_paths / GRADIO_TEMP_DIR." - ), - ) - - # Build preview + metadata using your shared function - preview_path, meta_text = _safe_build_preview(origin_path) - - # Back-compat: prefer preview in result_path - result_path = preview_path or origin_path - - return LungsSegmentationOutput( - success=True, - result_path=result_path, - result_origin=origin_path, - result_preview=preview_path, - metadata_text=meta_text, - compute_time_seconds=compute_time, - endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, - api_name=LUNGS_SEGMENTATION_API_NAME, - notes=f"Successfully segmented lungs from {os.path.basename(inp.image_path)}", - ) - - except Exception as e: - log.exception("Lungs segmentation failed") - return LungsSegmentationOutput( - success=False, - error=str(e), - compute_time_seconds=time.time() - start_time, - endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, - api_name=LUNGS_SEGMENTATION_API_NAME, - ) - - -# --------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------- -def _make_gradio_client(endpoint: str, hf_token: Optional[str]) -> Client: - """ - Create a gradio_client.Client with best compatibility across versions. - """ - # Set extended timeout for both connection and operations (5 minutes for large files) - httpx_kwargs = {"timeout": 300.0} - - # Newer versions use token=, older versions used hf_token= - if hf_token: - try: - return Client(endpoint, hf_token=hf_token, httpx_kwargs=httpx_kwargs) - except TypeError: - # Fallback for very old versions without httpx_kwargs support - try: - return Client(endpoint, hf_token=hf_token) - except TypeError: - return Client(endpoint) - - try: - return Client(endpoint, httpx_kwargs=httpx_kwargs) - except TypeError: - # Fallback for very old versions - return Client(endpoint) - - -def _safe_build_preview(origin_path: str) -> Tuple[Optional[str], Optional[str]]: - """ - Wrapper around _build_preview_for_vlm so preview failures never break the tool. - """ - try: - preview_path, meta_text = _build_preview_for_vlm([origin_path]) - return preview_path, meta_text - except Exception as e: - log.debug("Preview build failed for %s: %r", origin_path, e) - return None, None - - -def _materialize_any( - obj: Any, client: Client, hf_token: Optional[str] = None, _depth: int = 0 -) -> Optional[str]: - """ - Convert common Gradio outputs into a local file path. - - Supported: - - local path string - - URL string - - dict (FileData-like) containing url/path/name/filepath - - list/tuple containing any of the above - - server path '/tmp/...' -> attempt Gradio file endpoint (may 403) - - Args: - obj: Object to materialize - client: Gradio client - hf_token: Optional HuggingFace token - _depth: Internal recursion depth counter (max 10) - """ - if obj is None or _depth > 10: - if _depth > 10: - log.warning("Recursion depth limit reached in _materialize_any, halting.") - return None - - # list/tuple: most Gradio outputs are single-element lists - if isinstance(obj, (list, tuple)) and obj: - return _materialize_any( - obj[0], client=client, hf_token=hf_token, _depth=_depth + 1 - ) - - # dict: FileData-like is best case (url provided) - if isinstance(obj, dict): - # Prefer URL if present - url = obj.get("url") - if isinstance(url, str) and url.startswith(("http://", "https://")): - log.info("Materialize: dict url=%s", url) - return _download_to_temp(url, hf_token=hf_token) - - # Fall back through common keys - for k in ("path", "filepath", "file", "name"): - v = obj.get(k) - if isinstance(v, str) and v: - return _materialize_any( - v, client=client, hf_token=hf_token, _depth=_depth + 1 - ) - - return None - - # string: local file, URL, or server path - if isinstance(obj, str): - s = obj.strip() - if not s: - return None - - # local file? - p = Path(s) - if p.exists() and p.is_file(): - log.info("Materialize: local file=%s", s) - return str(p) - - # URL? - if s.startswith(("http://", "https://")): - log.info("Materialize: url=%s", s) - return _download_to_temp(s, hf_token=hf_token) - - # server path? (e.g. /tmp/xxx_mask.tif) - if s.startswith("/"): - log.info("Materialize: server path=%s", s) - return _download_from_gradio_file_endpoint(client, s, hf_token=hf_token) - - return None - - -def _download_to_temp(url: str, hf_token: Optional[str] = None) -> Optional[str]: - """ - Download a URL to a temporary file (streaming) with size limit checks. - """ - headers: Dict[str, str] = {} - if hf_token: - headers["Authorization"] = f"Bearer {hf_token}" - - try: - with requests.get( - url, headers=headers, timeout=120, stream=True, allow_redirects=True - ) as r: - if r.status_code != 200: - log.error("Download failed: url=%s status=%s", url, r.status_code) - return None - - # Check Content-Length if available - content_length = r.headers.get("content-length") - if content_length and int(content_length) > MAX_DOWNLOAD_SIZE: - log.error( - "File too large: %s bytes (max %s)", - content_length, - MAX_DOWNLOAD_SIZE, - ) - return None - - ext = _guess_ext(url, r.headers.get("content-type", "")) - - with tempfile.NamedTemporaryFile( - delete=False, prefix="lungs_seg_", suffix=ext - ) as f: - downloaded_size = 0 - for chunk in r.iter_content(chunk_size=1024 * 1024): - if chunk: - downloaded_size += len(chunk) - if downloaded_size > MAX_DOWNLOAD_SIZE: - log.error( - "Download exceeded size limit: %s bytes", - downloaded_size, - ) - f.close() - os.remove(f.name) - return None - f.write(chunk) - log.info("Downloaded %s bytes: %s -> %s", downloaded_size, url, f.name) - return register_temp_file(f.name) - except Exception as e: - log.error("Failed to download %s: %r", url, e) - return None - - -def _download_from_gradio_file_endpoint( - client: Client, server_path: str, hf_token: Optional[str] = None -) -> Optional[str]: - """ - Last-resort fallback when API returns '/tmp/...' but no URL. - Often blocked with 403 unless Space allows that directory or writes into Gradio temp/cache. - Includes size limit checks. - """ - base = (getattr(client, "src", None) or LUNGS_SEGMENTATION_ENDPOINT).rstrip("/") - file_url = f"{base}/gradio_api/file={server_path}" - - headers: Dict[str, str] = {} - if hf_token: - headers["Authorization"] = f"Bearer {hf_token}" - - params: Dict[str, str] = {} - session_hash = getattr(client, "session_hash", None) - if session_hash: - params["session_hash"] = session_hash - - try: - r = requests.get( - file_url, headers=headers, params=params, timeout=60, stream=True - ) - if r.status_code == 403: - # Common: file exists but not allowed to be served - detail: Any - try: - detail = r.json() - except Exception: - detail = r.text[:200] - log.error("HTTP 403 from %s detail=%r", file_url, detail) - return None - - if r.status_code != 200: - log.error("HTTP %s from %s", r.status_code, file_url) - return None - - # Check Content-Length before downloading - content_length = r.headers.get("content-length") - if content_length and int(content_length) > MAX_DOWNLOAD_SIZE: - log.error( - "File too large: %s bytes (max %s)", content_length, MAX_DOWNLOAD_SIZE - ) - return None - - # Read content with size check - content = b"" - for chunk in r.iter_content(chunk_size=1024 * 1024): - if chunk: - content += chunk - if len(content) > MAX_DOWNLOAD_SIZE: - log.error("Download exceeded size limit: %s bytes", len(content)) - return None - - ct = r.headers.get("content-type", "") - if "html" in ct.lower() or content.startswith(b" %s", - len(content), - f.name, - ) - return register_temp_file(f.name) - - except Exception as e: - log.error("Failed gradio file endpoint download: %r", e) - return None - - -def _guess_ext(url: str, content_type: str) -> str: - """ - Guess file extension from URL path or Content-Type. - """ - from urllib.parse import urlparse - - path = urlparse(url).path.lower() - - if path.endswith(".nii.gz"): - return ".nii.gz" - - ext = os.path.splitext(path)[1] - if ext: - return ext - - ct = (content_type or "").lower() - if "tiff" in ct or "tif" in ct: - return ".tif" - if "png" in ct: - return ".png" - if "jpeg" in ct or "jpg" in ct: - return ".jpg" - if "gif" in ct: - return ".gif" - if "nifti" in ct or "nii" in ct: - return ".nii.gz" - return ".bin" - - -# --------------------------------------------------------------------- -# Tool Registration -# --------------------------------------------------------------------- -register_tool( - ToolConfig( - name="lungs_segmentation", - display_name="3D Lungs Segmentation", - icon="๐Ÿซ", - catalog_names=["lungs-segmentation"], # Catalog name from dataset/catalog.jsonl - input_model=LungsSegmentationInput, - output_model=LungsSegmentationOutput, - executor=tool_lungs_segmentation, - supports_images=True, - supports_files=True, - requires_approval=True, - preview_field="result_preview", - download_fields="result_origin", # Could also be ["result_origin", "other_file"] - metadata_field="metadata_text", - notes_field="notes", - success_field="success", - error_field="error", - compute_time_field="compute_time_seconds", - ) -) +from __future__ import annotations + +from typing import Optional, Any, Dict, Tuple +import os +import logging +import tempfile +from pathlib import Path +import time + +import requests +from gradio_client import Client, handle_file + +from ai_agent.utils.previews import _build_preview_for_vlm +from ai_agent.utils.temp_file_manager import register_temp_file +from ai_agent.agent.tools.mcp.registry import register_tool, ToolConfig +from ai_agent.agent.tools.mcp.base import BaseToolOutput, ImageToolInput + +log = logging.getLogger("agent.lungs_segmentation") + + +# --------------------------------------------------------------------- +# Models +# --------------------------------------------------------------------- +class LungsSegmentationInput(ImageToolInput): + """Input for 3D lungs segmentation tool.""" + + pass # Inherits image_path and description from ImageToolInput + + +class LungsSegmentationOutput(BaseToolOutput): + """Output from 3D lungs segmentation tool.""" + + # All standard fields inherited from BaseToolOutput: + # - success, error, compute_time_seconds, notes + # - result_preview, result_origin, result_path + # - metadata_text, endpoint_url, api_name + pass + + +# --------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------- +LUNGS_SEGMENTATION_ENDPOINT = "https://qchapp-3d-lungs-segmentation.hf.space/" +LUNGS_SEGMENTATION_API_NAME = "/segment" + +# Maximum file size for downloads (1GB for medical imaging) +MAX_DOWNLOAD_SIZE = 1024 * 1024 * 1024 # 1GB in bytes + + +# --------------------------------------------------------------------- +# Public tool +# --------------------------------------------------------------------- +def tool_lungs_segmentation(inp: LungsSegmentationInput) -> LungsSegmentationOutput: + """ + Run 3D lungs segmentation on a CT scan image via a Gradio Space. + + Materialization strategy (robust): + 1) If Space returns dict FileData (url/path/etc) -> download via URL. + 2) If Space returns URL string -> download. + 3) If Space returns local file -> use it. + 4) If Space returns server path (/tmp/...) -> try /gradio_api/file=... (may 403). + """ + start_time = time.time() + + if not os.path.exists(inp.image_path): + return LungsSegmentationOutput( + success=False, + error=f"Image file not found: {inp.image_path}", + endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, + api_name=LUNGS_SEGMENTATION_API_NAME, + ) + + hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") + + try: + log.info( + "Running lungs segmentation on %s (endpoint: %s)", + inp.image_path, + LUNGS_SEGMENTATION_ENDPOINT, + ) + + client = _make_gradio_client(LUNGS_SEGMENTATION_ENDPOINT, hf_token) + + # Call API + try: + result = client.predict( + file_obj=handle_file(inp.image_path), + api_name=LUNGS_SEGMENTATION_API_NAME, + ) + log.info("API returned type=%s value=%r", type(result), result) + except Exception as e: + return LungsSegmentationOutput( + success=False, + error=f"API call failed: {e}", + compute_time_seconds=time.time() - start_time, + endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, + api_name=LUNGS_SEGMENTATION_API_NAME, + ) + + # Materialize to local file + origin_path = _materialize_any(result, client=client, hf_token=hf_token) + + compute_time = time.time() - start_time + + if not origin_path or not os.path.exists(origin_path): + # This is the common case if the Space returns '/tmp/...' and Gradio blocks it (403). + return LungsSegmentationOutput( + success=False, + error="Could not materialize/download the result file.", + compute_time_seconds=compute_time, + endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, + api_name=LUNGS_SEGMENTATION_API_NAME, + notes=( + f"API returned: {result!r}. If this is a '/tmp/...' path and you see HTTP 403, " + "the Space must return a FileData/url (recommended) or whitelist the output directory " + "via allowed_paths / GRADIO_TEMP_DIR." + ), + ) + + # Build preview + metadata using your shared function + preview_path, meta_text = _safe_build_preview(origin_path) + + # Back-compat: prefer preview in result_path + result_path = preview_path or origin_path + + return LungsSegmentationOutput( + success=True, + result_path=result_path, + result_origin=origin_path, + result_preview=preview_path, + metadata_text=meta_text, + compute_time_seconds=compute_time, + endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, + api_name=LUNGS_SEGMENTATION_API_NAME, + notes=f"Successfully segmented lungs from {os.path.basename(inp.image_path)}", + ) + + except Exception as e: + log.exception("Lungs segmentation failed") + return LungsSegmentationOutput( + success=False, + error=str(e), + compute_time_seconds=time.time() - start_time, + endpoint_url=LUNGS_SEGMENTATION_ENDPOINT, + api_name=LUNGS_SEGMENTATION_API_NAME, + ) + + +# --------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------- +def _make_gradio_client(endpoint: str, hf_token: Optional[str]) -> Client: + """ + Create a gradio_client.Client with best compatibility across versions. + """ + # Set extended timeout for both connection and operations (5 minutes for large files) + httpx_kwargs = {"timeout": 300.0} + + # Newer versions use token=, older versions used hf_token= + if hf_token: + try: + return Client(endpoint, hf_token=hf_token, httpx_kwargs=httpx_kwargs) + except TypeError: + # Fallback for very old versions without httpx_kwargs support + try: + return Client(endpoint, hf_token=hf_token) + except TypeError: + return Client(endpoint) + + try: + return Client(endpoint, httpx_kwargs=httpx_kwargs) + except TypeError: + # Fallback for very old versions + return Client(endpoint) + + +def _safe_build_preview(origin_path: str) -> Tuple[Optional[str], Optional[str]]: + """ + Wrapper around _build_preview_for_vlm so preview failures never break the tool. + """ + try: + preview_path, meta_text = _build_preview_for_vlm([origin_path]) + return preview_path, meta_text + except Exception as e: + log.debug("Preview build failed for %s: %r", origin_path, e) + return None, None + + +def _materialize_any( + obj: Any, client: Client, hf_token: Optional[str] = None, _depth: int = 0 +) -> Optional[str]: + """ + Convert common Gradio outputs into a local file path. + + Supported: + - local path string + - URL string + - dict (FileData-like) containing url/path/name/filepath + - list/tuple containing any of the above + - server path '/tmp/...' -> attempt Gradio file endpoint (may 403) + + Args: + obj: Object to materialize + client: Gradio client + hf_token: Optional HuggingFace token + _depth: Internal recursion depth counter (max 10) + """ + if obj is None or _depth > 10: + if _depth > 10: + log.warning("Recursion depth limit reached in _materialize_any, halting.") + return None + + # list/tuple: most Gradio outputs are single-element lists + if isinstance(obj, (list, tuple)) and obj: + return _materialize_any( + obj[0], client=client, hf_token=hf_token, _depth=_depth + 1 + ) + + # dict: FileData-like is best case (url provided) + if isinstance(obj, dict): + # Prefer URL if present + url = obj.get("url") + if isinstance(url, str) and url.startswith(("http://", "https://")): + log.info("Materialize: dict url=%s", url) + return _download_to_temp(url, hf_token=hf_token) + + # Fall back through common keys + for k in ("path", "filepath", "file", "name"): + v = obj.get(k) + if isinstance(v, str) and v: + return _materialize_any( + v, client=client, hf_token=hf_token, _depth=_depth + 1 + ) + + return None + + # string: local file, URL, or server path + if isinstance(obj, str): + s = obj.strip() + if not s: + return None + + # local file? + p = Path(s) + if p.exists() and p.is_file(): + log.info("Materialize: local file=%s", s) + return str(p) + + # URL? + if s.startswith(("http://", "https://")): + log.info("Materialize: url=%s", s) + return _download_to_temp(s, hf_token=hf_token) + + # server path? (e.g. /tmp/xxx_mask.tif) + if s.startswith("/"): + log.info("Materialize: server path=%s", s) + return _download_from_gradio_file_endpoint(client, s, hf_token=hf_token) + + return None + + +def _download_to_temp(url: str, hf_token: Optional[str] = None) -> Optional[str]: + """ + Download a URL to a temporary file (streaming) with size limit checks. + """ + headers: Dict[str, str] = {} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" + + try: + with requests.get( + url, headers=headers, timeout=120, stream=True, allow_redirects=True + ) as r: + if r.status_code != 200: + log.error("Download failed: url=%s status=%s", url, r.status_code) + return None + + # Check Content-Length if available + content_length = r.headers.get("content-length") + if content_length and int(content_length) > MAX_DOWNLOAD_SIZE: + log.error( + "File too large: %s bytes (max %s)", + content_length, + MAX_DOWNLOAD_SIZE, + ) + return None + + ext = _guess_ext(url, r.headers.get("content-type", "")) + + with tempfile.NamedTemporaryFile( + delete=False, prefix="lungs_seg_", suffix=ext + ) as f: + downloaded_size = 0 + for chunk in r.iter_content(chunk_size=1024 * 1024): + if chunk: + downloaded_size += len(chunk) + if downloaded_size > MAX_DOWNLOAD_SIZE: + log.error( + "Download exceeded size limit: %s bytes", + downloaded_size, + ) + f.close() + os.remove(f.name) + return None + f.write(chunk) + log.info("Downloaded %s bytes: %s -> %s", downloaded_size, url, f.name) + return register_temp_file(f.name) + except Exception as e: + log.error("Failed to download %s: %r", url, e) + return None + + +def _download_from_gradio_file_endpoint( + client: Client, server_path: str, hf_token: Optional[str] = None +) -> Optional[str]: + """ + Last-resort fallback when API returns '/tmp/...' but no URL. + Often blocked with 403 unless Space allows that directory or writes into Gradio temp/cache. + Includes size limit checks. + """ + base = (getattr(client, "src", None) or LUNGS_SEGMENTATION_ENDPOINT).rstrip("/") + file_url = f"{base}/gradio_api/file={server_path}" + + headers: Dict[str, str] = {} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" + + params: Dict[str, str] = {} + session_hash = getattr(client, "session_hash", None) + if session_hash: + params["session_hash"] = session_hash + + try: + r = requests.get( + file_url, headers=headers, params=params, timeout=60, stream=True + ) + if r.status_code == 403: + # Common: file exists but not allowed to be served + detail: Any + try: + detail = r.json() + except Exception: + detail = r.text[:200] + log.error("HTTP 403 from %s detail=%r", file_url, detail) + return None + + if r.status_code != 200: + log.error("HTTP %s from %s", r.status_code, file_url) + return None + + # Check Content-Length before downloading + content_length = r.headers.get("content-length") + if content_length and int(content_length) > MAX_DOWNLOAD_SIZE: + log.error( + "File too large: %s bytes (max %s)", content_length, MAX_DOWNLOAD_SIZE + ) + return None + + # Read content with size check + content = b"" + for chunk in r.iter_content(chunk_size=1024 * 1024): + if chunk: + content += chunk + if len(content) > MAX_DOWNLOAD_SIZE: + log.error("Download exceeded size limit: %s bytes", len(content)) + return None + + ct = r.headers.get("content-type", "") + if "html" in ct.lower() or content.startswith(b" %s", + len(content), + f.name, + ) + return register_temp_file(f.name) + + except Exception as e: + log.error("Failed gradio file endpoint download: %r", e) + return None + + +def _guess_ext(url: str, content_type: str) -> str: + """ + Guess file extension from URL path or Content-Type. + """ + from urllib.parse import urlparse + + path = urlparse(url).path.lower() + + if path.endswith(".nii.gz"): + return ".nii.gz" + + ext = os.path.splitext(path)[1] + if ext: + return ext + + ct = (content_type or "").lower() + if "tiff" in ct or "tif" in ct: + return ".tif" + if "png" in ct: + return ".png" + if "jpeg" in ct or "jpg" in ct: + return ".jpg" + if "gif" in ct: + return ".gif" + if "nifti" in ct or "nii" in ct: + return ".nii.gz" + return ".bin" + + +# --------------------------------------------------------------------- +# Tool Registration +# --------------------------------------------------------------------- +register_tool( + ToolConfig( + name="lungs_segmentation", + display_name="3D Lungs Segmentation", + icon="๐Ÿซ", + catalog_names=["lungs-segmentation"], # Catalog name from dataset/catalog.jsonl + input_model=LungsSegmentationInput, + output_model=LungsSegmentationOutput, + executor=tool_lungs_segmentation, + supports_images=True, + supports_files=True, + requires_approval=True, + preview_field="result_preview", + download_fields="result_origin", # Could also be ["result_origin", "other_file"] + metadata_field="metadata_text", + notes_field="notes", + success_field="success", + error_field="error", + compute_time_field="compute_time_seconds", + ) +) diff --git a/src/ai_agent/agent/tools/mcp/registry.py b/src/ai_agent/agent/tools/mcp/registry.py index 73a78ae78a0db7a3406b2041a7f4a7483dd60884..09201219703a728cbd96251ef229c31aba19ee85 100644 --- a/src/ai_agent/agent/tools/mcp/registry.py +++ b/src/ai_agent/agent/tools/mcp/registry.py @@ -1,203 +1,203 @@ -from __future__ import annotations - -from typing import Dict, Type, Callable, Optional, List, Any -from pydantic import BaseModel -from dataclasses import dataclass - - -@dataclass -class ToolConfig: - """ - Declarative configuration for a tool. - - Tools register themselves with this config, and the UI uses it - to generically handle execution, display, and file management. - """ - - # Core identification - name: str # Internal name (e.g., "lungs_segmentation") - display_name: str # User-facing name (e.g., "3D Lungs Segmentation") - icon: str # Emoji for UI display - - # Type information - input_model: Type[BaseModel] # Pydantic model for inputs - output_model: Type[BaseModel] # Pydantic model for outputs - executor: Callable # Function that takes input_model and returns output_model - - # Capability flags - catalog_names: Optional[List[str]] = ( - None # Catalog names for this tool (e.g., ["lungs-segmentation"]) - ) - supports_images: bool = True - supports_files: bool = True - requires_approval: bool = True # Whether to show approval button - - # Output field mappings (how to extract results from output_model) - # These map generic concepts to tool-specific field names - preview_field: str = "result_preview" # Field containing preview image path - download_fields: List[str] | str = ( - "result_origin" # Field(s) for downloadable files - ) - metadata_field: Optional[str] = "metadata_text" # Optional metadata text - notes_field: str = "notes" # Field containing execution notes - - # Success detection - success_field: str = "success" # Field indicating success/failure - error_field: str = "error" # Field containing error message - compute_time_field: str = "compute_time_seconds" # Field with timing info - - -# Global tool registry -TOOL_REGISTRY: Dict[str, ToolConfig] = {} - -# Reverse mapping from catalog names to tool names -CATALOG_NAME_TO_TOOL: Dict[str, str] = {} - - -def register_tool(config: ToolConfig) -> None: - """ - Register a tool with the global registry. - - Args: - config: Tool configuration - - Raises: - ValueError: If tool name already registered or catalog name collision - """ - if config.name in TOOL_REGISTRY: - raise ValueError(f"Tool '{config.name}' is already registered") - - # Check for catalog name collisions before registering - if config.catalog_names: - for catalog_name in config.catalog_names: - if ( - catalog_name in CATALOG_NAME_TO_TOOL - and CATALOG_NAME_TO_TOOL[catalog_name] != config.name - ): - raise ValueError( - f"Catalog name '{catalog_name}' already registered to " - f"'{CATALOG_NAME_TO_TOOL[catalog_name]}'" - ) - - TOOL_REGISTRY[config.name] = config - - # Register catalog name mappings - if config.catalog_names: - for catalog_name in config.catalog_names: - CATALOG_NAME_TO_TOOL[catalog_name] = config.name - - -def get_tool(name: str) -> Optional[ToolConfig]: - """ - Get tool configuration by name. - - Args: - name: Tool name (registry name or catalog name) - - Returns: - ToolConfig if found, None otherwise - - Note: - This function checks both the tool registry name and catalog names. - """ - # First try direct registry lookup - config = TOOL_REGISTRY.get(name) - if config: - return config - - # Try catalog name mapping - tool_name = CATALOG_NAME_TO_TOOL.get(name) - if tool_name: - return TOOL_REGISTRY.get(tool_name) - - return None - - -def list_tools() -> List[str]: - """Get list of all registered tool names.""" - return list(TOOL_REGISTRY.keys()) - - -def get_tool_display_name(name: str) -> str: - """ - Get display name for a tool, with fallback to name. - - Args: - name: Tool name - - Returns: - Display name or formatted version of name - """ - tool = get_tool(name) - if tool: - return tool.display_name - # Fallback: format name nicely - return name.replace("_", " ").title() - - -def get_tool_icon(name: str) -> str: - """ - Get icon for a tool, with fallback. - - Args: - name: Tool name - - Returns: - Icon emoji or default - """ - tool = get_tool(name) - if tool: - return tool.icon - return "๐Ÿ”ง" # Default tool icon - - -def extract_output_field(output: BaseModel, field_name: str) -> Any: - """ - Safely extract a field from tool output. - - Args: - output: Tool output object - field_name: Field name to extract - - Returns: - Field value or None if not found - """ - return getattr(output, field_name, None) - - -def extract_preview(output: BaseModel, tool_name: str) -> Optional[str]: - """Extract preview image path from tool output.""" - tool = get_tool(tool_name) - if not tool: - return None - return extract_output_field(output, tool.preview_field) - - -def extract_downloads(output: BaseModel, tool_name: str) -> List[str]: - """Extract downloadable file paths from tool output.""" - tool = get_tool(tool_name) - if not tool: - return [] - - download_fields = tool.download_fields - if isinstance(download_fields, str): - download_fields = [download_fields] - - downloads = [] - for field in download_fields: - value = extract_output_field(output, field) - if value: - if isinstance(value, list): - downloads.extend([v for v in value if v]) - elif isinstance(value, str): - downloads.append(value) - - return [d for d in downloads if d] # Filter None/empty - - -def extract_metadata(output: BaseModel, tool_name: str) -> Optional[str]: - """Extract metadata text from tool output.""" - tool = get_tool(tool_name) - if not tool or not tool.metadata_field: - return None - return extract_output_field(output, tool.metadata_field) +from __future__ import annotations + +from typing import Dict, Type, Callable, Optional, List, Any +from pydantic import BaseModel +from dataclasses import dataclass + + +@dataclass +class ToolConfig: + """ + Declarative configuration for a tool. + + Tools register themselves with this config, and the UI uses it + to generically handle execution, display, and file management. + """ + + # Core identification + name: str # Internal name (e.g., "lungs_segmentation") + display_name: str # User-facing name (e.g., "3D Lungs Segmentation") + icon: str # Emoji for UI display + + # Type information + input_model: Type[BaseModel] # Pydantic model for inputs + output_model: Type[BaseModel] # Pydantic model for outputs + executor: Callable # Function that takes input_model and returns output_model + + # Capability flags + catalog_names: Optional[List[str]] = ( + None # Catalog names for this tool (e.g., ["lungs-segmentation"]) + ) + supports_images: bool = True + supports_files: bool = True + requires_approval: bool = True # Whether to show approval button + + # Output field mappings (how to extract results from output_model) + # These map generic concepts to tool-specific field names + preview_field: str = "result_preview" # Field containing preview image path + download_fields: List[str] | str = ( + "result_origin" # Field(s) for downloadable files + ) + metadata_field: Optional[str] = "metadata_text" # Optional metadata text + notes_field: str = "notes" # Field containing execution notes + + # Success detection + success_field: str = "success" # Field indicating success/failure + error_field: str = "error" # Field containing error message + compute_time_field: str = "compute_time_seconds" # Field with timing info + + +# Global tool registry +TOOL_REGISTRY: Dict[str, ToolConfig] = {} + +# Reverse mapping from catalog names to tool names +CATALOG_NAME_TO_TOOL: Dict[str, str] = {} + + +def register_tool(config: ToolConfig) -> None: + """ + Register a tool with the global registry. + + Args: + config: Tool configuration + + Raises: + ValueError: If tool name already registered or catalog name collision + """ + if config.name in TOOL_REGISTRY: + raise ValueError(f"Tool '{config.name}' is already registered") + + # Check for catalog name collisions before registering + if config.catalog_names: + for catalog_name in config.catalog_names: + if ( + catalog_name in CATALOG_NAME_TO_TOOL + and CATALOG_NAME_TO_TOOL[catalog_name] != config.name + ): + raise ValueError( + f"Catalog name '{catalog_name}' already registered to " + f"'{CATALOG_NAME_TO_TOOL[catalog_name]}'" + ) + + TOOL_REGISTRY[config.name] = config + + # Register catalog name mappings + if config.catalog_names: + for catalog_name in config.catalog_names: + CATALOG_NAME_TO_TOOL[catalog_name] = config.name + + +def get_tool(name: str) -> Optional[ToolConfig]: + """ + Get tool configuration by name. + + Args: + name: Tool name (registry name or catalog name) + + Returns: + ToolConfig if found, None otherwise + + Note: + This function checks both the tool registry name and catalog names. + """ + # First try direct registry lookup + config = TOOL_REGISTRY.get(name) + if config: + return config + + # Try catalog name mapping + tool_name = CATALOG_NAME_TO_TOOL.get(name) + if tool_name: + return TOOL_REGISTRY.get(tool_name) + + return None + + +def list_tools() -> List[str]: + """Get list of all registered tool names.""" + return list(TOOL_REGISTRY.keys()) + + +def get_tool_display_name(name: str) -> str: + """ + Get display name for a tool, with fallback to name. + + Args: + name: Tool name + + Returns: + Display name or formatted version of name + """ + tool = get_tool(name) + if tool: + return tool.display_name + # Fallback: format name nicely + return name.replace("_", " ").title() + + +def get_tool_icon(name: str) -> str: + """ + Get icon for a tool, with fallback. + + Args: + name: Tool name + + Returns: + Icon emoji or default + """ + tool = get_tool(name) + if tool: + return tool.icon + return "๐Ÿ”ง" # Default tool icon + + +def extract_output_field(output: BaseModel, field_name: str) -> Any: + """ + Safely extract a field from tool output. + + Args: + output: Tool output object + field_name: Field name to extract + + Returns: + Field value or None if not found + """ + return getattr(output, field_name, None) + + +def extract_preview(output: BaseModel, tool_name: str) -> Optional[str]: + """Extract preview image path from tool output.""" + tool = get_tool(tool_name) + if not tool: + return None + return extract_output_field(output, tool.preview_field) + + +def extract_downloads(output: BaseModel, tool_name: str) -> List[str]: + """Extract downloadable file paths from tool output.""" + tool = get_tool(tool_name) + if not tool: + return [] + + download_fields = tool.download_fields + if isinstance(download_fields, str): + download_fields = [download_fields] + + downloads = [] + for field in download_fields: + value = extract_output_field(output, field) + if value: + if isinstance(value, list): + downloads.extend([v for v in value if v]) + elif isinstance(value, str): + downloads.append(value) + + return [d for d in downloads if d] # Filter None/empty + + +def extract_metadata(output: BaseModel, tool_name: str) -> Optional[str]: + """Extract metadata text from tool output.""" + tool = get_tool(tool_name) + if not tool or not tool.metadata_field: + return None + return extract_output_field(output, tool.metadata_field) diff --git a/src/ai_agent/agent/tools/query_utils.py b/src/ai_agent/agent/tools/query_utils.py index f9fff5b6d939c0c8758ed430e77a1683f7c652ca..fdc19263507815dd8acb678138207ed60f26e5a6 100644 --- a/src/ai_agent/agent/tools/query_utils.py +++ b/src/ai_agent/agent/tools/query_utils.py @@ -1,128 +1,128 @@ -from __future__ import annotations - -import re -from typing import List - -FORMAT_TOKEN_MAP = { - "tif": "TIFF", - "tiff": "TIFF", - "nii": "NIfTI", - "nii.gz": "NIfTI", - "dcm": "DICOM", - "dicom": "DICOM", - "nrrd": "NRRD", - "png": "PNG", - "jpg": "JPEG", - "jpeg": "JPEG", -} - -_REPO_DRIFT_TERMS = { - "github", - "repository", - "repo", - "official", - "readme", - "docs", - "documentation", - "source", - "sourcecode", -} - -_LOW_SIGNAL_TERMS = _REPO_DRIFT_TERMS | { - "tool", - "tools", - "project", - "framework", -} - - -def _tokenize_query(query: str) -> List[str]: - return [t for t in re.findall(r"[a-z0-9_+-]+", (query or "").lower()) if t] - - -def normalize_formats(formats: List[str]) -> List[str]: - seen = set() - out: List[str] = [] - for ext in formats: - norm = (ext or "").strip().lower() - if not norm or norm in seen: - continue - seen.add(norm) - out.append(norm) - return out - - -def append_format_tokens(query: str, formats: List[str]) -> str: - fmt_tokens: List[str] = [] - for ext in normalize_formats(formats): - canon = FORMAT_TOKEN_MAP.get(ext, ext.upper()) - if canon not in fmt_tokens: - fmt_tokens.append(canon) - - if not fmt_tokens: - return query.strip() - return (query.strip() + " " + " ".join(f"format:{t}" for t in fmt_tokens)).strip() - - -def strip_legacy_original_formats_line(query: str) -> tuple[str, List[str]]: - """Parse and remove legacy OriginalFormats: line from query text.""" - original_formats: List[str] = [] - clean_lines = [] - - for line in (query or "").splitlines(): - if line.lower().startswith("originalformats:"): - parts = line.split(":", 1)[1].strip().split() - original_formats.extend(parts) - continue - clean_lines.append(line) - - base_query = " ".join(ln.strip() for ln in clean_lines if ln.strip()) - return base_query, normalize_formats(original_formats) - - -def sanitize_retrieval_query( - query: str, - known_tool_names: List[str] | None = None, - fallback_query: str | None = None, -) -> str: - """ - Sanitize LLM-generated retrieval queries by removing repository drift terms. - - If the query collapses into a tool-name-only or low-signal query, fallback - to the previous task-centric query when provided. - """ - raw = (query or "").strip() - if not raw: - return (fallback_query or "").strip() - - # Remove URLs and punctuation-heavy fragments. - s = re.sub(r"https?://\S+", " ", raw, flags=re.IGNORECASE) - s = re.sub(r"www\.\S+", " ", s, flags=re.IGNORECASE) - - tokens = _tokenize_query(s) - if not tokens: - return (fallback_query or raw).strip() - - filtered = [t for t in tokens if t not in _REPO_DRIFT_TERMS] - if not filtered: - return (fallback_query or raw).strip() - - # Detect tool-name-only drift (e.g., "dhsegment official github repository"). - if known_tool_names: - token_set = set(filtered) - for nm in known_tool_names: - nm_tokens = set(_tokenize_query(nm)) - if not nm_tokens: - continue - if token_set.issubset(nm_tokens) and len(token_set) <= 3: - return (fallback_query or " ".join(filtered)).strip() - - low_signal = all(t in _LOW_SIGNAL_TERMS for t in filtered) - if low_signal: - return (fallback_query or " ".join(filtered)).strip() - - # If it became too short and a previous task query exists, prefer that. - if len(filtered) <= 2 and fallback_query: - return fallback_query.strip() - - return " ".join(filtered).strip() +from __future__ import annotations + +import re +from typing import List + +FORMAT_TOKEN_MAP = { + "tif": "TIFF", + "tiff": "TIFF", + "nii": "NIfTI", + "nii.gz": "NIfTI", + "dcm": "DICOM", + "dicom": "DICOM", + "nrrd": "NRRD", + "png": "PNG", + "jpg": "JPEG", + "jpeg": "JPEG", +} + +_REPO_DRIFT_TERMS = { + "github", + "repository", + "repo", + "official", + "readme", + "docs", + "documentation", + "source", + "sourcecode", +} + +_LOW_SIGNAL_TERMS = _REPO_DRIFT_TERMS | { + "tool", + "tools", + "project", + "framework", +} + + +def _tokenize_query(query: str) -> List[str]: + return [t for t in re.findall(r"[a-z0-9_+-]+", (query or "").lower()) if t] + + +def normalize_formats(formats: List[str]) -> List[str]: + seen = set() + out: List[str] = [] + for ext in formats: + norm = (ext or "").strip().lower() + if not norm or norm in seen: + continue + seen.add(norm) + out.append(norm) + return out + + +def append_format_tokens(query: str, formats: List[str]) -> str: + fmt_tokens: List[str] = [] + for ext in normalize_formats(formats): + canon = FORMAT_TOKEN_MAP.get(ext, ext.upper()) + if canon not in fmt_tokens: + fmt_tokens.append(canon) + + if not fmt_tokens: + return query.strip() + return (query.strip() + " " + " ".join(f"format:{t}" for t in fmt_tokens)).strip() + + +def strip_legacy_original_formats_line(query: str) -> tuple[str, List[str]]: + """Parse and remove legacy OriginalFormats: line from query text.""" + original_formats: List[str] = [] + clean_lines = [] + + for line in (query or "").splitlines(): + if line.lower().startswith("originalformats:"): + parts = line.split(":", 1)[1].strip().split() + original_formats.extend(parts) + continue + clean_lines.append(line) + + base_query = " ".join(ln.strip() for ln in clean_lines if ln.strip()) + return base_query, normalize_formats(original_formats) + + +def sanitize_retrieval_query( + query: str, + known_tool_names: List[str] | None = None, + fallback_query: str | None = None, +) -> str: + """ + Sanitize LLM-generated retrieval queries by removing repository drift terms. + + If the query collapses into a tool-name-only or low-signal query, fallback + to the previous task-centric query when provided. + """ + raw = (query or "").strip() + if not raw: + return (fallback_query or "").strip() + + # Remove URLs and punctuation-heavy fragments. + s = re.sub(r"https?://\S+", " ", raw, flags=re.IGNORECASE) + s = re.sub(r"www\.\S+", " ", s, flags=re.IGNORECASE) + + tokens = _tokenize_query(s) + if not tokens: + return (fallback_query or raw).strip() + + filtered = [t for t in tokens if t not in _REPO_DRIFT_TERMS] + if not filtered: + return (fallback_query or raw).strip() + + # Detect tool-name-only drift (e.g., "dhsegment official github repository"). + if known_tool_names: + token_set = set(filtered) + for nm in known_tool_names: + nm_tokens = set(_tokenize_query(nm)) + if not nm_tokens: + continue + if token_set.issubset(nm_tokens) and len(token_set) <= 3: + return (fallback_query or " ".join(filtered)).strip() + + low_signal = all(t in _LOW_SIGNAL_TERMS for t in filtered) + if low_signal: + return (fallback_query or " ".join(filtered)).strip() + + # If it became too short and a previous task query exists, prefer that. + if len(filtered) <= 2 and fallback_query: + return fallback_query.strip() + + return " ".join(filtered).strip() diff --git a/src/ai_agent/agent/tools/repo_info_tool.py b/src/ai_agent/agent/tools/repo_info_tool.py index 82b4d588064f1be213635a6df87d464f95e5d34a..038b2a17cb515bbdef98dc541ae0b002e94c3892 100644 --- a/src/ai_agent/agent/tools/repo_info_tool.py +++ b/src/ai_agent/agent/tools/repo_info_tool.py @@ -2,7 +2,6 @@ from __future__ import annotations import asyncio import os -import time from typing import Optional import repocards @@ -10,6 +9,7 @@ from pydantic import BaseModel from .deepwiki_tool import get_wiki_contents, DeepWikiInput from .utils import _clip, get_catalog_docs, _is_github_url +from ai_agent.utils.cache_db import get_cache_db import logging @@ -18,7 +18,9 @@ log = logging.getLogger("agent.repo_info") REPO_INFO_CACHE_TTL_SECONDS = int(os.getenv("REPO_INFO_CACHE_TTL_SECONDS", "3600")) REPO_INFO_CACHE_MAX_ENTRIES = int(os.getenv("REPO_INFO_CACHE_MAX_ENTRIES", "256")) -_REPO_INFO_CACHE: dict[str, tuple[float, "RepoSummaryOutput"]] = {} +_REPO_INFO_NS = "repo_info" + +# In-flight request deduplication must remain in-memory (holds asyncio.Future objects) _REPO_INFO_INFLIGHT: dict[str, asyncio.Future["RepoSummaryOutput"]] = {} _REPO_INFO_LOCK = asyncio.Lock() @@ -41,25 +43,9 @@ def _normalize_cache_key(url: str) -> str: return url.strip().lower() -def _evict_expired_entries(now: float) -> None: - expired = [k for k, (exp, _) in _REPO_INFO_CACHE.items() if exp <= now] - for key in expired: - _REPO_INFO_CACHE.pop(key, None) - - -def _enforce_cache_capacity() -> None: - if len(_REPO_INFO_CACHE) <= REPO_INFO_CACHE_MAX_ENTRIES: - return - # Evict oldest expiration first. Good enough for bounded in-memory cache. - keys_by_expiry = sorted(_REPO_INFO_CACHE.items(), key=lambda item: item[1][0]) - over = len(_REPO_INFO_CACHE) - REPO_INFO_CACHE_MAX_ENTRIES - for key, _ in keys_by_expiry[:over]: - _REPO_INFO_CACHE.pop(key, None) - - def _clear_repo_summary_cache_for_tests() -> None: """Test helper to avoid state leakage across test cases.""" - _REPO_INFO_CACHE.clear() + get_cache_db().clear(_REPO_INFO_NS) _REPO_INFO_INFLIGHT.clear() @@ -94,17 +80,31 @@ async def tool_repo_summary(input: RepoSummaryInput) -> RepoSummaryOutput: log.warning(f"Failed to lookup repo URL from catalog: {e}") cache_key = _normalize_cache_key(effective_url) - now = time.monotonic() + db = get_cache_db() + # --- Cache read outside the lock so we don't block other coroutines --- + try: + raw = await asyncio.to_thread(db.get, _REPO_INFO_NS, cache_key) + except Exception: + log.warning("Repo info cache read failed; treating as cache miss.", exc_info=True) + raw = None + if raw is not None: + log.info(f"Repo info cache hit for {effective_url}") + return RepoSummaryOutput.model_validate_json(raw) + + # --- Lock only for inflight bookkeeping (pure in-memory, no I/O) --- await _REPO_INFO_LOCK.acquire() try: - _evict_expired_entries(now) - - cached_entry = _REPO_INFO_CACHE.get(cache_key) - if cached_entry: - _, cached = cached_entry - log.info(f"Repo info cache hit for {effective_url}") - return cached.model_copy(deep=True) + # Re-check the cache inside the lock in case another coroutine wrote it + # between our lockless read above and acquiring the lock now. + try: + raw = await asyncio.to_thread(db.get, _REPO_INFO_NS, cache_key) + except Exception: + log.warning("Repo info cache read (after lock) failed; treating as cache miss.", exc_info=True) + raw = None + if raw is not None: + log.info(f"Repo info cache hit (after lock) for {effective_url}") + return RepoSummaryOutput.model_validate_json(raw) inflight = _REPO_INFO_INFLIGHT.get(cache_key) if inflight is None: @@ -133,19 +133,30 @@ async def tool_repo_summary(input: RepoSummaryInput) -> RepoSummaryOutput: _REPO_INFO_LOCK.release() raise + # --- Resolve the in-flight future immediately so waiters are unblocked + # before the (potentially slow) DB write happens. --- await _REPO_INFO_LOCK.acquire() try: - if result.source != "error" and REPO_INFO_CACHE_TTL_SECONDS > 0: - expires_at = time.monotonic() + REPO_INFO_CACHE_TTL_SECONDS - _REPO_INFO_CACHE[cache_key] = (expires_at, result) - _enforce_cache_capacity() - if not inflight.done(): inflight.set_result(result) _REPO_INFO_INFLIGHT.pop(cache_key, None) finally: _REPO_INFO_LOCK.release() + # --- DB write after waiters are already released --- + if result.source != "error" and REPO_INFO_CACHE_TTL_SECONDS > 0: + try: + await asyncio.to_thread( + db.set, + _REPO_INFO_NS, + cache_key, + result.model_dump_json(), + ttl_seconds=REPO_INFO_CACHE_TTL_SECONDS, + max_entries=REPO_INFO_CACHE_MAX_ENTRIES, + ) + except Exception: + log.warning("Repo info cache write failed; result will not be cached.", exc_info=True) + return result.model_copy(deep=True) diff --git a/src/ai_agent/agent/tools/search_alternative_tool.py b/src/ai_agent/agent/tools/search_alternative_tool.py index 618b77c89f45722ee52ca97c658e0c57b0385365..ecceaded241078a30a9f419f2423d3cdead0c5bb 100644 --- a/src/ai_agent/agent/tools/search_alternative_tool.py +++ b/src/ai_agent/agent/tools/search_alternative_tool.py @@ -1,73 +1,73 @@ -from __future__ import annotations - -from typing import List -from pydantic import BaseModel, Field - -from ai_agent.generator.schema import CandidateDoc -from .utils import get_known_names, get_pipeline -from .query_utils import append_format_tokens, normalize_formats, sanitize_retrieval_query - - -class SearchAlternativeInput(BaseModel): - """ - Input for searching with an alternative query formulation. - - Use this when initial search results are insufficient and you want to - try a different phrasing or broader/narrower terms. - """ - - alternative_query: str = Field( - description="Alternative query phrasing to try (can be similar terms, broader/narrower, etc.)" - ) - excluded: List[str] = Field(default_factory=list) - top_k: int = 12 - original_formats: List[str] = Field(default_factory=list) - image_paths: List[str] = Field(default_factory=list) - - -class SearchAlternativeOutput(BaseModel): - candidates: List[CandidateDoc] - query_used: str - - -def tool_search_alternative(inp: SearchAlternativeInput) -> SearchAlternativeOutput: - """ - Search with an alternative query formulation, with automatic reranking. - - This tool allows the agent to explicitly try a different search approach - when initial results are not satisfactory. - """ - pipe = get_pipeline() - - # Use the alternative query directly - query = sanitize_retrieval_query( - inp.alternative_query.strip(), known_tool_names=get_known_names() - ) - - # Normalize formats - original_formats: List[str] = normalize_formats(inp.original_formats) - query = append_format_tokens(query, original_formats) - - # Call retrieve() which includes automatic reranking - hits = pipe.retrieve( - query, - image_paths=inp.image_paths or None, - exclusions=inp.excluded, - top_k=inp.top_k, - ) - - # Convert hits to CandidateDoc objects - candidates: List[CandidateDoc] = [] - for h in hits: - d = h.get("doc") - if not d: - continue - try: - candidates.append(CandidateDoc.model_validate(d.model_dump(mode="python"))) - except Exception: - continue - - return SearchAlternativeOutput( - candidates=candidates, - query_used=query, - ) +from __future__ import annotations + +from typing import List +from pydantic import BaseModel, Field + +from ai_agent.generator.schema import CandidateDoc +from .utils import get_known_names, get_pipeline +from .query_utils import append_format_tokens, normalize_formats, sanitize_retrieval_query + + +class SearchAlternativeInput(BaseModel): + """ + Input for searching with an alternative query formulation. + + Use this when initial search results are insufficient and you want to + try a different phrasing or broader/narrower terms. + """ + + alternative_query: str = Field( + description="Alternative query phrasing to try (can be similar terms, broader/narrower, etc.)" + ) + excluded: List[str] = Field(default_factory=list) + top_k: int = 12 + original_formats: List[str] = Field(default_factory=list) + image_paths: List[str] = Field(default_factory=list) + + +class SearchAlternativeOutput(BaseModel): + candidates: List[CandidateDoc] + query_used: str + + +def tool_search_alternative(inp: SearchAlternativeInput) -> SearchAlternativeOutput: + """ + Search with an alternative query formulation, with automatic reranking. + + This tool allows the agent to explicitly try a different search approach + when initial results are not satisfactory. + """ + pipe = get_pipeline() + + # Use the alternative query directly + query = sanitize_retrieval_query( + inp.alternative_query.strip(), known_tool_names=get_known_names() + ) + + # Normalize formats + original_formats: List[str] = normalize_formats(inp.original_formats) + query = append_format_tokens(query, original_formats) + + # Call retrieve() which includes automatic reranking + hits = pipe.retrieve( + query, + image_paths=inp.image_paths or None, + exclusions=inp.excluded, + top_k=inp.top_k, + ) + + # Convert hits to CandidateDoc objects + candidates: List[CandidateDoc] = [] + for h in hits: + d = h.get("doc") + if not d: + continue + try: + candidates.append(CandidateDoc.model_validate(d.model_dump(mode="python"))) + except Exception: + continue + + return SearchAlternativeOutput( + candidates=candidates, + query_used=query, + ) diff --git a/src/ai_agent/agent/tools/sparql_tool.py b/src/ai_agent/agent/tools/sparql_tool.py new file mode 100644 index 0000000000000000000000000000000000000000..f869ffadfcb0aa54406333cea6b5edab446fa09d --- /dev/null +++ b/src/ai_agent/agent/tools/sparql_tool.py @@ -0,0 +1,178 @@ +"""SPARQL query tool โ€” lets the agent ask the catalog's GraphDB anything the +RAG can't answer cleanly (counts, distincts, structural filters, etc.). + +Safety: + ยท read-only by syntax check โ€” UPDATE / INSERT / DELETE / DROP / CLEAR / + CREATE / LOAD / COPY / MOVE / ADD are rejected before the query goes + anywhere near the endpoint + ยท result row count is capped (we inject / override LIMIT) + ยท network and query time are bounded (SPARQLWrapper timeout) + ยท auth comes from the same GRAPHDB_USER / GRAPHDB_PASSWORD env the + catalog sync uses +""" + +from __future__ import annotations + +import logging +import os +import re +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field +from SPARQLWrapper import JSON, SPARQLWrapper + +log = logging.getLogger("agent.tools.sparql") + +# Default ceiling for any SELECT โ€” overrideable via input. +DEFAULT_LIMIT = 50 +HARD_MAX_LIMIT = 200 + +# Update keywords that must NOT appear at the top level of the query. We +# match on word boundaries (case-insensitive) outside of `<...>` IRIs and +# quoted strings โ€” a real parser would be safer but for an agent-authored +# query this catches the common cases. +_UPDATE_KEYWORDS = re.compile( + r"\b(" + r"INSERT|DELETE|UPDATE|DROP|CLEAR|CREATE|LOAD|COPY|MOVE|ADD|MODIFY" + r")\b", + re.IGNORECASE, +) + +_LIMIT_RE = re.compile(r"\blimit\s+(\d+)\b", re.IGNORECASE) + + +class SparqlQueryInput(BaseModel): + query: str = Field(..., description="A SPARQL SELECT / ASK query.") + limit: int = Field( + default=DEFAULT_LIMIT, + ge=1, + le=HARD_MAX_LIMIT, + description="Maximum row count (we'll add / cap LIMIT in the query).", + ) + + +class SparqlQueryOutput(BaseModel): + columns: List[str] = Field(default_factory=list) + rows: List[Dict[str, Any]] = Field(default_factory=list) + row_count: int = 0 + truncated: bool = False + boolean: Optional[bool] = None # only for ASK queries + query_executed: str = "" + error: Optional[str] = None + + +def _strip_strings_and_iris(q: str) -> str: + """Remove `"..."`, `'...'`, and `<...>` so keyword scanning isn't fooled + by literal text inside the query.""" + q = re.sub(r'"(?:[^"\\]|\\.)*"', '""', q) + q = re.sub(r"'(?:[^'\\]|\\.)*'", "''", q) + q = re.sub(r"<[^>]*>", "<>", q) + return q + + +def _looks_readonly(query: str) -> bool: + scrubbed = _strip_strings_and_iris(query) + return _UPDATE_KEYWORDS.search(scrubbed) is None + + +def _enforce_limit(query: str, limit: int) -> str: + """Inject LIMIT, or cap an existing one to `limit`. SELECT only โ€” ASK, + CONSTRUCT, DESCRIBE pass through unchanged (LIMIT is moot for ASK; the + other two are blocked upstream).""" + if not re.search(r"\bselect\b", query, re.IGNORECASE): + return query + m = _LIMIT_RE.search(query) + if m: + requested = int(m.group(1)) + if requested <= limit: + return query + return _LIMIT_RE.sub(f"LIMIT {limit}", query) + trimmed = query.rstrip(" \t\n;") + return trimmed + "\nLIMIT " + str(limit) + + +def tool_sparql_query(inp: SparqlQueryInput) -> SparqlQueryOutput: + """Execute a SPARQL SELECT or ASK against GRAPHDB_URL and return a flat + table of bindings.""" + endpoint = (os.getenv("GRAPHDB_URL") or "").strip() + if not endpoint: + return SparqlQueryOutput( + error="GRAPHDB_URL is not configured on the server.", + query_executed=inp.query, + ) + + if not _looks_readonly(inp.query): + return SparqlQueryOutput( + error="rejected_update_query โ€” only read-only SPARQL is allowed", + query_executed=inp.query, + ) + + # We currently only normalise SELECT/ASK. CONSTRUCT/DESCRIBE return RDF + # graphs which would need a different output schema โ€” bail with a clear + # message rather than silently misformatting. + qkind = re.search( + r"\b(select|ask|construct|describe)\b", inp.query, re.IGNORECASE + ) + if not qkind: + return SparqlQueryOutput( + error="missing_query_form โ€” start the query with SELECT or ASK", + query_executed=inp.query, + ) + form = qkind.group(1).lower() + if form in ("construct", "describe"): + return SparqlQueryOutput( + error=f"unsupported_form โ€” {form.upper()} not supported by this tool yet; rephrase as SELECT", + query_executed=inp.query, + ) + + final_query = _enforce_limit(inp.query, inp.limit) if form == "select" else inp.query + + sparql = SPARQLWrapper(endpoint) + sparql.setQuery(final_query) + sparql.setReturnFormat(JSON) + user = os.getenv("GRAPHDB_USER") + pwd = os.getenv("GRAPHDB_PASSWORD") + if user and pwd: + sparql.setCredentials(user=user, passwd=pwd) + sparql.setTimeout(20) + + try: + result = sparql.query().convert() + except Exception as exc: + log.warning("sparql_query failed: %s", exc) + return SparqlQueryOutput( + error=f"sparql_error: {exc}", query_executed=final_query + ) + + # ASK โ†’ boolean + if "boolean" in result: + return SparqlQueryOutput( + boolean=bool(result["boolean"]), + row_count=0, + query_executed=final_query, + ) + + cols: List[str] = list(result.get("head", {}).get("vars", [])) + bindings = result.get("results", {}).get("bindings", []) or [] + rows: List[Dict[str, Any]] = [] + for b in bindings: + row: Dict[str, Any] = {} + for c in cols: + cell = b.get(c) + if cell is None: + row[c] = None + else: + row[c] = cell.get("value") + rows.append(row) + + truncated = len(rows) >= inp.limit + return SparqlQueryOutput( + columns=cols, + rows=rows, + row_count=len(rows), + truncated=truncated, + query_executed=final_query, + ) + + +__all__ = ["SparqlQueryInput", "SparqlQueryOutput", "tool_sparql_query"] diff --git a/src/ai_agent/api/deps.py b/src/ai_agent/api/deps.py new file mode 100644 index 0000000000000000000000000000000000000000..b81147ad1d09f495015435ab0063a14ab16dcb6b --- /dev/null +++ b/src/ai_agent/api/deps.py @@ -0,0 +1,103 @@ +"""Shared FastAPI dependencies. + + - ``require_auth``: simple cookie-based gate using ``APP_PASSWORD`` env var. + If ``APP_PASSWORD`` is empty/unset, auth is disabled entirely (dev mode). + - ``get_pipeline``: returns the singleton ``RAGImagingPipeline`` (lazy-init, + shared with the Gradio path via ``ai_agent.core.pipeline_registry``). + - ``get_doc_index``: name -> SoftwareDoc, derived from the pipeline. +""" + +from __future__ import annotations + +import hashlib +import hmac +import logging +import os +from typing import Dict, Optional + +from fastapi import Cookie, Depends, HTTPException, status + +from ai_agent.api.pipeline import RAGImagingPipeline +from ai_agent.core.pipeline_registry import get_pipeline as _shared_pipeline +from ai_agent.retriever.software_doc import SoftwareDoc + +log = logging.getLogger("api.deps") + +AUTH_COOKIE_NAME = "ai_agent_auth" +_AUTH_VERSION = "v1" + + +def _expected_cookie() -> Optional[str]: + """The cookie value clients must present, or None if auth is disabled.""" + pw = os.getenv("APP_PASSWORD") or "" + if not pw: + return None + return hmac.new( + pw.encode("utf-8"), + _AUTH_VERSION.encode("utf-8"), + hashlib.sha256, + ).hexdigest() + + +def make_cookie_value(password: str) -> str: + """Compute the cookie value for a successful login.""" + return hmac.new( + password.encode("utf-8"), + _AUTH_VERSION.encode("utf-8"), + hashlib.sha256, + ).hexdigest() + + +def verify_password(password: str) -> bool: + expected = os.getenv("APP_PASSWORD") or "" + if not expected: + # Auth disabled โ€” accept any password (login is then effectively a noop). + return True + return hmac.compare_digest(password, expected) + + +def auth_disabled() -> bool: + return not (os.getenv("APP_PASSWORD") or "") + + +async def require_auth( + ai_agent_auth: Optional[str] = Cookie(default=None, alias=AUTH_COOKIE_NAME), +) -> None: + """FastAPI dependency: 401 if cookie missing or wrong (when auth on).""" + expected = _expected_cookie() + if expected is None: + return # auth disabled + if not ai_agent_auth or not hmac.compare_digest(ai_agent_auth, expected): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="not_authenticated" + ) + + +# --------------------------------------------------------------------------- +# Pipeline / doc index providers (re-export so routers don't import internals). +# --------------------------------------------------------------------------- +def get_pipeline() -> RAGImagingPipeline: + index_dir = os.getenv("RAG_INDEX_DIR", "artifacts/rag_index") + return _shared_pipeline(index_dir=index_dir) + + +def get_doc_index( + pipe: RAGImagingPipeline = Depends(get_pipeline), +) -> Dict[str, SoftwareDoc]: + try: + docs = list(pipe.index.docs.values()) + return {d.name: d for d in docs if getattr(d, "name", None)} + except Exception: + log.exception("Failed to build doc_index from pipeline") + return {} + + +__all__ = [ + "AUTH_COOKIE_NAME", + "auth_disabled", + "get_doc_index", + "get_pipeline", + "make_cookie_value", + "require_auth", + "verify_password", +] diff --git a/src/ai_agent/api/routers/__init__.py b/src/ai_agent/api/routers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/ai_agent/api/routers/auth.py b/src/ai_agent/api/routers/auth.py new file mode 100644 index 0000000000000000000000000000000000000000..394bf43fdf358d5645d44c10e8901978eb9c62b7 --- /dev/null +++ b/src/ai_agent/api/routers/auth.py @@ -0,0 +1,61 @@ +"""Password login routes. + +The frontend POSTs ``{password}`` to ``/api/auth/login``; if the value +matches ``APP_PASSWORD`` (env), an HMAC token is set as an httpOnly cookie +and all subsequent ``/api/*`` requests authenticate via that cookie. + +When ``APP_PASSWORD`` is unset, login still works (any password accepted) +and the cookie is not strictly required โ€” useful for local dev or when the +deployment sits behind a Cloudflare tunnel without per-user auth. +""" + +from __future__ import annotations + +import os + +from fastapi import APIRouter, HTTPException, Response, status + +from ai_agent.api.deps import ( + AUTH_COOKIE_NAME, + auth_disabled, + make_cookie_value, + verify_password, +) +from ai_agent.api.schemas import LoginRequest, LoginResponse + +router = APIRouter(prefix="/api/auth", tags=["auth"]) + + +@router.get("/status") +def status_route() -> dict: + """Tells the frontend whether auth is enforced at all. + + Used by the login page to skip itself when ``APP_PASSWORD`` is empty. + """ + return {"required": not auth_disabled()} + + +@router.post("/login", response_model=LoginResponse) +def login(body: LoginRequest, response: Response) -> LoginResponse: + if not verify_password(body.password): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid_password" + ) + expected_pw = os.getenv("APP_PASSWORD") or "" + if expected_pw: + response.set_cookie( + key=AUTH_COOKIE_NAME, + value=make_cookie_value(expected_pw), + httponly=True, + samesite="lax", + secure=False, # tunnel terminates TLS; cookie travels over plain HTTP internally + max_age=60 * 60 * 24 * 7, + path="/", + ) + return LoginResponse(ok=True) + + +@router.post("/logout", response_model=LoginResponse) +def logout(response: Response) -> LoginResponse: + response.delete_cookie(AUTH_COOKIE_NAME, path="/") + return LoginResponse(ok=True) diff --git a/src/ai_agent/api/routers/catalog.py b/src/ai_agent/api/routers/catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..160447988eeb03fc6e332e084fa68e44b0644e86 --- /dev/null +++ b/src/ai_agent/api/routers/catalog.py @@ -0,0 +1,34 @@ +"""Return the FAISS catalog as a list of compact CatalogEntry rows.""" + +from __future__ import annotations + +from typing import Dict, List + +from fastapi import APIRouter, Depends + +from ai_agent.api.deps import get_doc_index, require_auth +from ai_agent.api.schemas import CatalogEntry +from ai_agent.retriever.software_doc import SoftwareDoc + +router = APIRouter( + prefix="/api/catalog", tags=["catalog"], dependencies=[Depends(require_auth)] +) + + +@router.get("", response_model=List[CatalogEntry]) +def list_catalog( + doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index), +) -> List[CatalogEntry]: + out: List[CatalogEntry] = [] + for name, doc in sorted(doc_index.items(), key=lambda kv: kv[0].lower()): + out.append( + CatalogEntry( + name=name, + description=getattr(doc, "description", "") or "", + modality=list(getattr(doc, "modality", []) or []), + license=getattr(doc, "license", None), + dims=list(getattr(doc, "dims", []) or []), + keywords=list(getattr(doc, "keywords", []) or [])[:20], + ) + ) + return out diff --git a/src/ai_agent/api/routers/chat.py b/src/ai_agent/api/routers/chat.py new file mode 100644 index 0000000000000000000000000000000000000000..1d541a14fd93c7f46f4ae719b9b9c374d8ec6b19 --- /dev/null +++ b/src/ai_agent/api/routers/chat.py @@ -0,0 +1,249 @@ +"""Chat router with SSE streaming. + +Event protocol (named events; payloads are JSON): + + ``session`` ``{"session_id": "..."}`` once, first + ``user_message`` ``{"text": "...", "asset_ids": [...]}`` echo + ``status`` ``{"phase": "thinking" | "agent_running" | ...}`` occasional + ``text`` ``{"content": "..."}`` main reply + ``recommendation`` ``{"item": {rank, name, accuracy, why, ...}}`` per card + ``tool_trace`` ``{"trace": {...}}`` per call + ``pending_action`` ``{"action": {type, tool_name, ...}}`` if pause + ``clarification`` ``{"question": "...", "options": [...]}`` if needed + ``images`` ``{"paths": [...]}`` tool-exec + ``files`` ``{"items": [{path, label}, ...]}`` tool-exec + ``usage`` ``{"total": n, "input": n, "output": n}`` once + ``error`` ``{"message": "...", "code": "..."}`` fatal + ``done`` ``{}`` terminal + +The underlying agent (``services.chat.process_turn``) is still synchronous; +this router translates the final ``ChatTurnResult`` into a sequence of +events. As soon as pydantic-ai's ``run_stream`` is wired in, the same +endpoint will emit ``text`` deltas live. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any, AsyncIterator, Dict, List + +from fastapi import APIRouter, Depends, HTTPException, status +from sse_starlette.sse import EventSourceResponse + +from ai_agent.api.deps import get_doc_index, require_auth +from ai_agent.api.schemas import ChatStartBody +from ai_agent.retriever.software_doc import SoftwareDoc +from ai_agent.services.chat import ( + ChatRequest, + ChatTurnResult, + approve_pending, + decline_pending, + process_turn, +) +from ai_agent.services.sessions import Session, get_session_store + +log = logging.getLogger("api.routers.chat") + +router = APIRouter(prefix="/api/chat", tags=["chat"], dependencies=[Depends(require_auth)]) + + +def _event(name: str, data: Any) -> Dict[str, Any]: + return {"event": name, "data": json.dumps(data, default=str)} + + +# Status messages emitted on a heartbeat while the (synchronous) agent runs. +# The list is intentionally aligned with the typical agent timeline: +# pre-flight -> retrieval -> rerank -> VLM call -> shaping the response. +_STATUS_BEATS: List[str] = [ + "analyzing your requestโ€ฆ", + "searching software catalogโ€ฆ", + "ranking candidates against your imageโ€ฆ", + "consulting the agent โ€” this can take a few secondsโ€ฆ", + "fetching repo metadataโ€ฆ", + "preparing recommendationsโ€ฆ", + "still working โ€” large models are slower on first callโ€ฆ", +] +_STATUS_HEARTBEAT_SECONDS = 1.6 + + +def _stream_result( + session: Session, result: ChatTurnResult +) -> AsyncIterator[Dict[str, Any]]: + """Translate a finished ChatTurnResult into a sequence of SSE events.""" + + async def gen() -> AsyncIterator[Dict[str, Any]]: + # Always lead with the session id so the client can persist it. + yield _event("session", {"session_id": session.session_id}) + + if result.text: + yield _event("text", {"content": result.text}) + + for rec in result.recommendations: + yield _event( + "recommendation", + { + "rank": rec.rank, + "name": rec.name, + "accuracy": rec.accuracy, + "why": rec.why, + "doc": rec.doc, + "demo_url": rec.demo_url, + }, + ) + + for trace in result.tool_traces: + yield _event("tool_trace", {"trace": trace}) + + if result.clarification: + yield _event( + "clarification", + { + "question": result.clarification.question, + "context": result.clarification.context, + "options": result.clarification.options, + }, + ) + + if result.pending_action: + pa = result.pending_action + yield _event( + "pending_action", + { + "type": pa.type, + "tool_name": pa.tool_name, + "display_name": pa.display_name, + "icon": pa.icon, + "image_name": pa.image_name, + "demo_url": pa.demo_url, + "prompt": pa.prompt, + }, + ) + + if result.images: + yield _event("images", {"paths": result.images}) + if result.files: + yield _event( + "files", + {"items": [{"path": p, "label": label} for p, label in result.files]}, + ) + + if result.usage: + yield _event("usage", result.usage) + + if result.status == "error" and result.error: + yield _event("error", {"message": result.text, "code": result.error}) + + yield _event("done", {"status": result.status}) + + return gen() + + +def _resolve_or_error(session_id: str) -> Session: + session = get_session_store().get(session_id) + if session is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="session_not_found" + ) + return session + + +def _with_heartbeat( + session: Session, + coro_factory, +) -> AsyncIterator[Dict[str, Any]]: + """Run a blocking callable in a thread; emit ``status`` heartbeats while + it executes; then stream the real result. + + ``coro_factory`` is a zero-arg lambda that returns the ``ChatTurnResult`` + when called from a worker thread. + """ + + async def gen() -> AsyncIterator[Dict[str, Any]]: + yield _event("session", {"session_id": session.session_id}) + yield _event("status", {"phase": "started", "message": _STATUS_BEATS[0]}) + + task = asyncio.create_task(asyncio.to_thread(coro_factory)) + i = 1 + while True: + try: + result = await asyncio.wait_for( + asyncio.shield(task), timeout=_STATUS_HEARTBEAT_SECONDS + ) + break + except asyncio.TimeoutError: + msg = _STATUS_BEATS[min(i, len(_STATUS_BEATS) - 1)] + yield _event( + "status", + {"phase": f"beat_{i}", "message": msg}, + ) + i += 1 + + # Drain the (now-finished) result through the existing translator, + # but skip its leading ``session`` event since we already emitted one. + translator = _stream_result(session, result) + first = True + async for ev in translator: + if first and ev.get("event") == "session": + first = False + continue + first = False + yield ev + + return gen() + + +@router.post("") +async def start_chat( + body: ChatStartBody, + doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index), +): + store = get_session_store() + session = store.get_or_create(body.session_id) + + # Restore prior conversation transcript into a fresh server-side session. + if body.seed_history and not session.conversation_history: + session.conversation_history.extend(body.seed_history) + + request = ChatRequest( + message=body.message, + asset_ids=list(body.asset_ids), + model=body.model, + top_k=body.top_k, + num_choices=body.num_choices, + ) + return EventSourceResponse( + _with_heartbeat(session, lambda: process_turn(session, request, doc_index)) + ) + + +@router.post("/{session_id}/approve") +async def approve( + session_id: str, + doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index), +): + session = _resolve_or_error(session_id) + return EventSourceResponse( + _with_heartbeat(session, lambda: approve_pending(session)) + ) + + +@router.post("/{session_id}/decline") +async def decline(session_id: str): + session = _resolve_or_error(session_id) + result = decline_pending(session) + return EventSourceResponse(_stream_result(session, result)) + + +@router.post("/{session_id}/confirm-demo") +async def confirm_demo( + session_id: str, + doc_index: Dict[str, SoftwareDoc] = Depends(get_doc_index), +): + """Resume a turn that asked the user to confirm running a demo URL.""" + session = _resolve_or_error(session_id) + request = ChatRequest(message="yes") + return EventSourceResponse( + _with_heartbeat(session, lambda: process_turn(session, request, doc_index)) + ) diff --git a/src/ai_agent/api/routers/files.py b/src/ai_agent/api/routers/files.py new file mode 100644 index 0000000000000000000000000000000000000000..f8956a0d468affcd3960be47ec1408cc69089044 --- /dev/null +++ b/src/ai_agent/api/routers/files.py @@ -0,0 +1,306 @@ +"""Multipart file upload + per-asset preview serving. + +POST /api/files multipart, returns AssetResponse[] +GET /api/files/preview/{asset_id} serves the cached preview PNG + +Uploaded files are saved into a per-session temp directory so the existing +preview/metadata helpers (which take filesystem paths) keep working. +""" + +from __future__ import annotations + +import logging +import os +import tempfile +import uuid +from pathlib import Path +from typing import List, Optional + +from fastapi import ( + APIRouter, + Depends, + File, + Form, + HTTPException, + UploadFile, + status, +) +from fastapi.responses import FileResponse + +from fastapi import Query +from fastapi.responses import Response + +from ai_agent.api.deps import require_auth +from ai_agent.api.schemas import AssetResponse, FilesUploadResponse +from ai_agent.services.files import ingest_files +from ai_agent.services.sessions import Asset, get_session_store +from ai_agent.services.views import ( + extract_mip, + extract_slice, + get_info, + get_volume_bytes, +) + + +_MIME_BY_EXT: Dict[str, str] = { + "png": "image/png", + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "tif": "image/tiff", + "tiff": "image/tiff", + "webp": "image/webp", + "bmp": "image/bmp", + "gif": "image/gif", + "nii": "application/octet-stream", + "nii.gz": "application/octet-stream", + "dcm": "application/dicom", + "mp4": "video/mp4", + "mov": "video/quicktime", + "webm": "video/webm", + "mkv": "video/x-matroska", + "mp3": "audio/mpeg", + "wav": "audio/wav", + "ogg": "audio/ogg", + "flac": "audio/flac", + "m4a": "audio/mp4", + "pdf": "application/pdf", +} + +log = logging.getLogger("api.routers.files") + +router = APIRouter(prefix="/api/files", tags=["files"], dependencies=[Depends(require_auth)]) + +_UPLOAD_ROOT = Path( + os.getenv("UPLOAD_ROOT") + or os.path.join(tempfile.gettempdir(), "ai_agent_uploads") +) +_UPLOAD_ROOT.mkdir(parents=True, exist_ok=True) + + +def _session_dir(session_id: str) -> Path: + path = _UPLOAD_ROOT / session_id + path.mkdir(parents=True, exist_ok=True) + return path + + +def _safe_filename(name: str) -> str: + # Basic sanitization โ€” keep extension, replace path separators. + base = os.path.basename(name) or "upload" + base = base.replace("\x00", "") + return base[:200] + + +def _to_asset_response(asset: Asset) -> AssetResponse: + preview_url = ( + f"/api/files/preview/{asset.asset_id}" if asset.preview_path else None + ) + return AssetResponse( + asset_id=asset.asset_id, + display_name=asset.display_name, + original_format=asset.original_format, + preview_url=preview_url, + metadata_text=asset.metadata_text, + created_at=asset.created_at, + ) + + +@router.post("", response_model=FilesUploadResponse) +async def upload_files( + files: List[UploadFile] = File(...), + session_id: Optional[str] = Form(None), +) -> FilesUploadResponse: + store = get_session_store() + session = store.get_or_create(session_id) + target_dir = _session_dir(session.session_id) + + saved_paths: List[str] = [] + for f in files: + safe = _safe_filename(f.filename or f"upload-{uuid.uuid4().hex}") + # Avoid collisions in the session dir by suffixing a short uuid + stem, ext = os.path.splitext(safe) + target = target_dir / f"{stem}-{uuid.uuid4().hex[:8]}{ext}" + with target.open("wb") as out: + while chunk := await f.read(1024 * 1024): + out.write(chunk) + saved_paths.append(str(target)) + + result = ingest_files(session, saved_paths) + if result.validation_errors and not result.assets: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"errors": result.validation_errors}, + ) + + return FilesUploadResponse( + session_id=session.session_id, + assets=[_to_asset_response(a) for a in result.assets], + ) + + +@router.get("/sessions/{session_id}", response_model=List[AssetResponse]) +def list_session_assets(session_id: str) -> List[AssetResponse]: + """List every asset currently registered on the given session. + + Powers the "gallery" picker in the frontend so users can re-attach files + they've already uploaded in the same conversation without re-picking from + disk. + """ + store = get_session_store() + session = store.get(session_id) + if session is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="session_not_found" + ) + items: List[AssetResponse] = [ + _to_asset_response(a) for a in session.assets.values() + ] + # Most-recent uploads first. + items.sort( + key=lambda a: a.created_at if a.created_at is not None else 0.0, + reverse=True, + ) + return items + + +def _find_asset(asset_id: str) -> Optional[Asset]: + store = get_session_store() + for sid in list(store._sessions.keys()): # noqa: SLF001 โ€” intentional + session = store.get(sid) + if not session: + continue + asset = session.assets.get(asset_id) + if asset is not None: + return asset + return None + + +@router.get("/asset/{asset_id}/raw") +def asset_raw(asset_id: str): + """Serve the original uploaded file with its native MIME type. + + Used by the frontend to render