The AI Imaging Agent is organized into modular components with clear separation of concerns.
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
Main package containing all application code.
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 definitionsmodels.py: Agent output and tool usage schemasutils.py:AgentStateplus call caps/prepare hookstools/: Tool implementations (search, alternatives, repo info, mcp tools)
Dependencies: api/, utils/
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/
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/
VLM selection schemas and types.
generator/
├── __init__.py
└── schema.py # Pydantic models for responses
Models:
ToolRecommendation: Individual tool recommendationAgentResponse: Complete response with statusConversationStatus: Enum for conversation statesToolReason: Enum for recommendation reasons
Dependencies: None (pure schemas)
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:
text_embedder.py: Embed queryvector_index.py: FAISS searchreranker.py: CrossEncoder reranking- Output: Top-K candidates
Dependencies: None (pure retrieval)
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 interfacehandlers.py:respond()function - core interaction logicformatters.py: Format recommendations as markdown/cardscomponents.py: Reusable Gradio components
Dependencies: agent/, api/
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: Loadconfig.yamlwith Pydantic validationfile_validator.py: Size limits, format checksimage_meta.py: Extract DICOM/NIfTI/TIFF metadatapreviews.py: Convert medical images to PNGtags.py: Parse exclusion tags and strip control tags from queries
Dependencies: None (pure utilities)
Command-line interface entry point.
def main():
# Parse arguments
# Route to chat or syncCommands:
ai_agent chat: Launch UIai_agent sync: Sync catalog
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)
Python package metadata and dependencies.
[project]
name = "ai_agent"
version = "0.1.1"
dependencies = [...]
[project.scripts]
ai_agent = "ai_agent.cli:main"Model configuration.
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"
...Documentation configuration.
site_name: AI Imaging Agent
theme:
name: material
nav: [...]Environment variables (not committed).
OPENAI_API_KEY=sk-xxxx
SOFTWARE_CATALOG=dataset/catalog.jsonlSoftware catalog in JSON Lines format.
Each line is a complete JSON object following schema.org SoftwareSourceCode.
Pre-built FAISS index and metadata.
artifacts/rag_index/
├── index.faiss # FAISS binary index
└── meta.json # Tool IDs, config, timestamps
Clear separation prevents circular dependencies:
ui/ → agent/ → api/ → retriever/
→ generator/
→ utils/
Rules:
utils/: No dependencies on other modulesretriever/: Pure retrieval, no generationgenerator/: Pure schemas, no retrievalapi/: Orchestrates retriever + generatoragent/: Uses api for tool callsui/: Top-level, depends on agent + api
All imports use absolute paths from ai_agent:
from ai_agent.retriever.vector_index import VectorIndex
from ai_agent.utils.config import load_config
from ai_agent.agent.utils import AgentStateNever use relative imports like from ..utils import ...
Add tool adapters to agent/agent.py and implement logic in agent/tools/:
@agent.tool
async def new_tool(ctx: RunContext[AgentState], param: str) -> str:
"""Tool description."""
# Implementation
return resultAdd to utils/image_meta.py:
def extract_custom_format(file_path: str) -> dict:
"""Extract metadata from custom format."""
# Implementation
return metadataReplace in retriever/text_embedder.py:
class TextEmbedder:
def __init__(self, model_name="new-embedding-model"):
self.model = SentenceTransformer(model_name)- Learn about Contributing
- Explore Testing
- Return to Architecture Overview