---
title: AI Agent
emoji: ๐ค
colorFrom: blue
colorTo: green
sdk: docker
app_port: 7860
pinned: false
---
# AI Imaging Agent (Imaging Plaza)
[](https://www.python.org/downloads/)
[](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.