================================================================================
ATOMIC AGENTS EXAMPLES
================================================================================

This file contains all example implementations using the Atomic Agents framework.
Each example includes its README documentation and complete source code.
Project Repository: https://github.com/BrainBlend-AI/atomic-agents


--------------------------------------------------------------------------------
Example: basic-multimodal
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/basic-multimodal

## Documentation

# Basic Multimodal Example

This example demonstrates how to use the Atomic Agents framework to analyze images with text, specifically focusing on extracting structured information from nutrition labels using GPT-4 Vision capabilities.

## Features

1. Image Analysis: Process nutrition label images using GPT-4 Vision
2. Structured Data Extraction: Convert visual information into structured Pydantic models
3. Multi-Image Processing: Analyze multiple nutrition labels simultaneously
4. Comprehensive Nutritional Data: Extract detailed nutritional information including:
   - Basic nutritional facts (calories, fats, proteins, etc.)
   - Serving size information
   - Vitamin and mineral content
   - Product details

## Getting Started

1. Clone the main Atomic Agents repository:

   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. Navigate to the basic-multimodal directory:

   ```bash
   cd atomic-agents/atomic-examples/basic-multimodal
   ```

3. Install dependencies using uv:

   ```bash
   uv sync
   ```

4. Set up environment variables:

   Create a `.env` file in the `basic-multimodal` directory with the following content:
   ```env
   OPENAI_API_KEY=your_openai_api_key
   ```

   Replace `your_openai_api_key` with your actual OpenAI API key.

5. Run the example:

   ```bash
   uv run python basic_multimodal/main.py
   ```

## Components

### 1. Nutrition Label Schema (`NutritionLabel`)

Defines the structure for storing nutrition information, including:

- Macronutrients (fats, proteins, carbohydrates)
- Micronutrients (vitamins and minerals)
- Serving information
- Product details

### 2. Input/Output Schemas

- `NutritionAnalysisInput`: Handles input images and analysis instructions
- `NutritionAnalysisOutput`: Structures the extracted nutrition information

### 3. Nutrition Analyzer Agent

A specialized agent configured with:

- GPT-4 Vision capabilities
- Custom system prompts for nutrition label analysis
- Structured data validation

## Example Usage

The example includes test images in the `test_images` directory:

- `nutrition_label_1.png`: Example nutrition label image
- `nutrition_label_2.jpg`: Another example nutrition label image

Running the example will:

1. Load the test images
2. Process them through the nutrition analyzer
3. Display structured nutritional information for each label

## Customization

You can modify the example by:

1. Adding your own nutrition label images to the `test_images` directory
2. Adjusting the `NutritionLabel` schema to capture additional information
3. Modifying the system prompt to focus on specific aspects of nutrition labels

## Contributing

Contributions are welcome! Please fork the repository and submit a pull request with your enhancements or bug fixes.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.


## Source Code


### File: atomic-examples/basic-multimodal/basic_multimodal/main.py

```python
from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator
import instructor
import openai
from pydantic import Field
from typing import List
import os

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )


class NutritionLabel(BaseIOSchema):
    """Represents the complete nutritional information from a food label"""

    calories: int = Field(..., description="Calories per serving")
    total_fat: float = Field(..., description="Total fat in grams")
    saturated_fat: float = Field(..., description="Saturated fat in grams")
    trans_fat: float = Field(..., description="Trans fat in grams")
    cholesterol: int = Field(..., description="Cholesterol in milligrams")
    sodium: int = Field(..., description="Sodium in milligrams")
    total_carbohydrates: float = Field(..., description="Total carbohydrates in grams")
    dietary_fiber: float = Field(..., description="Dietary fiber in grams")
    total_sugars: float = Field(..., description="Total sugars in grams")
    added_sugars: float = Field(..., description="Added sugars in grams")
    protein: float = Field(..., description="Protein in grams")
    vitamin_d: float = Field(..., description="Vitamin D in micrograms")
    calcium: int = Field(..., description="Calcium in milligrams")
    iron: float = Field(..., description="Iron in milligrams")
    potassium: int = Field(..., description="Potassium in milligrams")
    serving_size: str = Field(..., description="The size of a single serving of this product")
    servings_per_container: float = Field(..., description="Number of servings contained in the package")
    product_name: str = Field(
        ...,
        description="The full name or description of the type of the food/drink. e.g: 'Coca Cola Light', 'Pepsi Max', 'Smoked Bacon', 'Chianti Wine'",
    )


class NutritionAnalysisInput(BaseIOSchema):
    """Input schema for nutrition label analysis"""

    instruction_text: str = Field(..., description="The instruction for analyzing the nutrition label")
    images: List[instructor.Image] = Field(..., description="The nutrition label images to analyze")


class NutritionAnalysisOutput(BaseIOSchema):
    """Output schema containing extracted nutrition information"""

    analyzed_labels: List[NutritionLabel] = Field(
        ..., description="List of nutrition labels extracted from the provided images"
    )


# Configure the nutrition analysis system
nutrition_analyzer = AtomicAgent[NutritionAnalysisInput, NutritionAnalysisOutput](
    config=AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=API_KEY)),
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a specialized nutrition label analyzer.",
                "You excel at extracting precise nutritional information from food label images.",
                "You understand various serving size formats and measurement units.",
                "You can process multiple nutrition labels simultaneously.",
            ],
            steps=[
                "For each nutrition label image:",
                "1. Locate and identify the nutrition facts panel",
                "2. Extract all serving information and nutritional values",
                "3. Validate measurements and units for accuracy",
                "4. Compile the nutrition facts into structured data",
            ],
            output_instructions=[
                "For each analyzed nutrition label:",
                "1. Record complete serving size information",
                "2. Extract all nutrient values with correct units",
                "3. Ensure all measurements are properly converted",
                "4. Include all extracted labels in the final result",
            ],
        ),
    )
)


def main():
    print("Starting nutrition label analysis...")

    # Construct the path to the test images
    script_directory = os.path.dirname(os.path.abspath(__file__))
    test_images_directory = os.path.join(os.path.dirname(script_directory), "test_images")
    image_path_1 = os.path.join(test_images_directory, "nutrition_label_1.png")
    image_path_2 = os.path.join(test_images_directory, "nutrition_label_2.jpg")
    # Create and submit the analysis request
    analysis_request = NutritionAnalysisInput(
        instruction_text="Please analyze these nutrition labels and extract all nutritional information.",
        images=[instructor.Image.from_path(image_path_1), instructor.Image.from_path(image_path_2)],
    )

    try:
        # Process the nutrition labels
        print("Analyzing nutrition labels...")
        analysis_result = nutrition_analyzer.run(analysis_request)
        print("Analysis completed successfully")

        # Display the results
        for i, label in enumerate(analysis_result.analyzed_labels, 1):
            print(f"\nNutrition Label {i}:")
            print(f"Product Name: {label.product_name}")
            print(f"Serving Size: {label.serving_size}")
            print(f"Servings Per Container: {label.servings_per_container}")
            print(f"Calories: {label.calories}")
            print(f"Total Fat: {label.total_fat}g")
            print(f"Saturated Fat: {label.saturated_fat}g")
            print(f"Trans Fat: {label.trans_fat}g")
            print(f"Cholesterol: {label.cholesterol}mg")
            print(f"Sodium: {label.sodium}mg")
            print(f"Total Carbohydrates: {label.total_carbohydrates}g")
            print(f"Dietary Fiber: {label.dietary_fiber}g")
            print(f"Total Sugars: {label.total_sugars}g")
            print(f"Added Sugars: {label.added_sugars}g")
            print(f"Protein: {label.protein}g")
            print(f"Vitamin D: {label.vitamin_d}mcg")
            print(f"Calcium: {label.calcium}mg")
            print(f"Iron: {label.iron}mg")
            print(f"Potassium: {label.potassium}mg")

    except Exception as e:
        print(f"Analysis failed: {str(e)}")
        raise


if __name__ == "__main__":
    main()

```

### File: atomic-examples/basic-multimodal/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["basic_multimodal"]

[project]
name = "basic-multimodal"
version = "1.0.0"
description = "Basic Multimodal Quickstart example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "instructor==1.14.5",
    "openai>=2.0.0,<3.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: basic-pdf-analysis
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/basic-pdf-analysis

## Documentation

# Basic PDF Analysis Example

This example demonstrates how to use the Atomic Agents framework to analyze a PDF file, using Google generative AI's multimodal capabilities.

## Features

1. PDF document analysis: Process a PDF document using Google generative AI multimodal capability.
2. Structured Data Extraction: Extract key information from PDFs into a structured Pydantic model:
   - Document title
   - Page count

## Getting Started

1. Clone the main Atomic Agents repository:
   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. Navigate to the basic-pdf-analysis directory:
   ```bash
   cd atomic-agents/atomic-examples/basic-pdf-analysis
   ```

3. Install dependencies using uv:
   ```bash
   uv sync
   ```

4. Set up environment variables:
   Create a `.env` file in the `basic-pdf-analysis` directory with the following content:
   ```env
   GEMINI_API_KEY=your_gemini_api_key
   ```
   Replace `your_gemini_api_key` with your actual google generative AI key.

5. Run the example:
   ```bash
   uv run python basic_pdf_analysis/main.py
   ```

## Components

### 1. Input/Output Schemas
- `InputSchema`: Handles the input PDF file
- `ExtractionResult`: Structures the extracted information

### 2. Agent
A specialized agent configured with:
- Google generative AI gemini-2.0-flash model
- Custom system prompt
- Structured data validation

## Example Usage

The example includes a test PDF file in the `test_media` directory.

Running the example will:
1. Load the PDF from the `test_media` directory
2. Process it with the agent
3. Display the extracted information:
   - PDF title
   - Page count

Example output:
```
Starting PDF file analysis...
Analyzing PDF file: pdf_sample.pdf ...

===== Analysis Results =====
PDF Title: Sample PDF Document
Page Count: 3
Document summary: This PDF is three pages long and contains Latin text.

Analysis completed successfully
```

## Customization

You can modify the example by:
1. Adding your own files to the `test_media` directory
2. Adjusting the `ExtractionResult` schema to capture additional information
3. Modifying the system prompts to extract different or additional information

## Contributing

Contributions are welcome! Please fork the repository and submit a pull request with your enhancements or bug fixes.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.


## Source Code


### File: atomic-examples/basic-pdf-analysis/basic_pdf_analysis/main.py

```python
import os

import instructor
from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator
from dotenv import load_dotenv
from google import genai
from instructor.processing.multimodal import PDF
from pydantic import Field


load_dotenv()


class InputSchema(BaseIOSchema):
    """PDF file to analyze."""

    pdf: PDF = Field(..., description="The PDF data")  # PDF class from instructor


class ExtractionResult(BaseIOSchema):
    """Extracted information from the PDF."""

    pdf_title: str = Field(..., description="The title of the PDF file")
    page_count: int = Field(..., description="The number of pages in the PDF file")
    summary: str = Field(..., description="A short summary of the document")


# Define the LLM CLient using GenAI instructor wrapper:
client = instructor.from_genai(client=genai.Client(api_key=os.getenv("GEMINI_API_KEY")), mode=instructor.Mode.GENAI_TOOLS)

# Define the system prompt:
system_prompt_generator = SystemPromptGenerator(
    background=["You are a helpful assistant that extracts information from PDF files."],
    steps=[
        "Analyze the PDF, extract its title and count the number of pages.",
        "Create a brief summary of the document content.",
    ],
    output_instructions=["Return pdf_title, page_count, and summary."],
)

# Define the agent
agent = AtomicAgent[InputSchema, ExtractionResult](
    config=AgentConfig(
        client=client,
        model="gemini-2.0-flash",
        system_prompt_generator=system_prompt_generator,
        input_schema=InputSchema,
        output_schema=ExtractionResult,
    )
)


def main():
    print("Starting PDF file analysis...")

    # Create the analysis request
    script_directory = os.path.dirname(os.path.abspath(__file__))
    test_media_directory = os.path.join(os.path.dirname(script_directory), "test_media")
    pdf_path = os.path.join(test_media_directory, "pdf_sample.pdf")
    analysis_request = InputSchema(
        pdf=PDF.from_path(pdf_path),
    )

    try:
        # Process the PDF file
        print(f"Analyzing PDF file: {os.path.basename(pdf_path)} ...")
        analysis_result = agent.run(analysis_request)

        # Display the results
        print("\n===== Analysis Results =====")
        print(f"PDF Title: {analysis_result.pdf_title}")
        print(f"Page Count: {analysis_result.page_count}")
        print(f"Document summary: {analysis_result.summary}")

    except Exception as e:
        print(f"Analysis failed: {str(e)}")
        raise e


if __name__ == "__main__":
    main()

```

### File: atomic-examples/basic-pdf-analysis/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["basic_pdf_analysis"]

[project]
name = "basic-pdf-analysis"
version = "1.0.0"
description = "Basic PDF analysis Quickstart example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Renaud Dufour", email = "renaud.dufour59@gmail.com" }
]
requires-python = ">=3.12,<3.14"
dependencies = [
    "atomic-agents",
    "instructor[google-genai]==1.14.5",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: deep-research
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/deep-research

## Documentation

# Deep Research Agent

A didactic example of a proper deep-research pipeline built out of small, single-purpose Atomic Agents.

Unlike a typical "search-and-summarise" agent — generate one set of queries, fetch results, write an answer — this example iterates: it plans sub-topics, researches each one across multiple depth levels, reflects on whether each has enough coverage, and produces a report where every claim is tied to a registered source.

## Pipeline

1. **Plan.** A `PlannerAgent` breaks the question into 3–5 durable sub-topics, each seeded with a handful of queries.
2. **Research** (per sub-topic, up to N iterations):
   - Search (SearXNG) and scrape the top new URLs.
   - `ExtractorAgent` pulls atomic, citable claims from each scraped page.
   - `ReflectorAgent` decides whether the sub-topic has enough material, or emits follow-up queries for the next iteration.
3. **Write.** `WriterAgent` drafts a cited report from the accumulated state, then runs a second pass over its own draft to strip any sentence whose citation doesn't correspond to a real source.

Every agent has a single responsibility and reads / contributes to a shared `ResearchState` object. The loop itself lives in `main.py` as plain Python — no megagent, no hidden control flow.

## Getting Started

1. **Clone the main Atomic Agents repository:**
   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. **Navigate to the Deep Research directory:**
   ```bash
   cd atomic-agents/atomic-examples/deep-research
   ```

3. **Install dependencies using uv:**
   ```bash
   uv sync
   ```

4. **Set up environment variables:**
   Create a `.env` file in the `deep-research` directory with:
   ```env
   OPENAI_API_KEY=your_openai_api_key
   SEARXNG_BASE_URL=http://localhost:8080
   SEARXNG_API_KEY=your_searxng_secret_key
   ```

5. **Set up SearXNG:**
   - Install from the [official repository](https://github.com/searxng/searxng).
   - Default configuration expects SearXNG at `http://localhost:8080`.
   - JSON output must be enabled in `settings.yml` (look for the `formats:` key).

6. **Run a research query:**
   ```bash
   uv run python -m deep_research "What is the current state of fusion energy research?"
   ```

## Modes

- **One-shot** (`python -m deep_research "your question"`): plan → research → write, prints a cited report and exits.
- **Chat** (`python -m deep_research`): same first turn, then a REPL where each follow-up is routed by a `DeciderAgent` to either another research pass + Q&A or straight Q&A against the existing state.

## File Layout

```
deep_research/
├── __main__.py             # python -m deep_research entrypoint
├── main.py                 # Plain orchestrator: plan → research → write (+ chat loop)
├── config.py               # Model + connectivity + research budgets
├── state.py                # ResearchState dataclass — the one source of truth
├── context_providers.py    # Renders state + current date into agent system prompts
├── agents/
│   ├── planner_agent.py    # Question → sub-topics (with initial queries)
│   ├── extractor_agent.py  # One scraped source → atomic claims
│   ├── reflector_agent.py  # Sub-topic state → sufficient? + next queries
│   ├── writer_agent.py     # Full state → cited report (draft + verify passes)
│   ├── decider_agent.py    # Chat mode: research more, or answer from state?
│   └── qa_agent.py         # Chat mode: cited answer from existing state
└── tools/
    ├── searxng_search.py
    └── webpage_scraper.py
```

## Budgets

All limits live in `ResearchBudget` inside `config.py`. Tune to taste:

| Knob | Default | Meaning |
|---|---|---|
| `num_sub_topics` | 4 | Plan width |
| `max_depth_per_sub_topic` | 2 | Max iterations per sub-topic; reflector can stop earlier |
| `search_results_per_query` | 5 | SearXNG page size |
| `scrape_top_n_per_iteration` | 3 | New URLs scraped per iteration |
| `hard_call_cap` | 80 | Global safety net on total agent calls |

Worst-case first turn with defaults: 1 plan + 4×2×(1 extract×3 sources + 1 reflect) = 33 agent calls + 2 writer passes ≈ **35 agent calls, 24 scrapes**. Chat follow-ups add a decider call plus either Q&A or another research pass; the `hard_call_cap` of 80 leaves headroom.

## License

MIT — see the [LICENSE](../../LICENSE) file.


## Source Code


### File: atomic-examples/deep-research/deep_research/__main__.py

```python
"""Package entry point — ``python -m deep_research``.

With args: one-shot pipeline — ``python -m deep_research "your question"``.
Without args: drops into the chat loop.

The real orchestrator lives in ``main.py``; this file is just the Python
convention that makes the package directly runnable.
"""

import sys

from deep_research.main import chat_loop, run

if __name__ == "__main__":
    args = sys.argv[1:]
    if args:
        run(" ".join(args))
    else:
        chat_loop()

```

### File: atomic-examples/deep-research/deep_research/agents/decider_agent.py

```python
"""
DeciderAgent — routes a follow-up user message to either more research or a direct answer.

In chat mode, every user turn after the first faces the same question:
do we already have the material to answer this, or do we need to go out
and gather more? This is that agent's entire job — one binary decision,
backed by short reasoning.

Deciding from the shared ``ResearchState`` (sources, learnings, plan)
instead of from the raw message keeps the decision grounded in what the
pipeline has actually collected, not what the model imagines it knows.
"""

import instructor
import openai
from pydantic import Field

from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

from deep_research.config import ChatConfig


class DeciderInput(BaseIOSchema):
    """Input schema for the DeciderAgent."""

    user_message: str = Field(..., min_length=1, description="The user's latest question or follow-up.")


class DeciderOutput(BaseIOSchema):
    """Output schema for the DeciderAgent."""

    reasoning: str = Field(
        ...,
        min_length=1,
        description="One short paragraph: what's already in the state, what's missing, and why that tips the decision.",
    )
    needs_research: bool = Field(
        ...,
        description=(
            "True if a new research pass is needed — state is empty, irrelevant, stale, or missing a key angle. "
            "False if the existing learnings already cover what the user is asking."
        ),
    )


decider_agent = AtomicAgent[DeciderInput, DeciderOutput](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a routing agent. Given the user's latest message and the current ResearchState "
                "(sources and learnings already gathered), you decide whether another research pass is warranted.",
                "You do NOT answer the question yourself. You only decide: research more, or hand off to the Q&A agent.",
            ],
            steps=[
                "Read the research state from the system context — what sources and learnings exist?",
                "Compare the user's message against those learnings. Is the answer already present, even partially?",
                "Flag a new research pass when state is empty, off-topic, outdated for a time-sensitive question, "
                "or missing an angle the user is now asking about.",
                "Otherwise, route to Q&A.",
            ],
            output_instructions=[
                "Be decisive. 'Maybe' is never the right answer.",
                "If the state is empty, always decide needs_research=true.",
                "For time-sensitive questions, check the current date in context and re-research if learnings look stale.",
                "Reasoning must cite concrete evidence from state (or its absence) — not vague intuition.",
            ],
        ),
    )
)

```

### File: atomic-examples/deep-research/deep_research/agents/extractor_agent.py

```python
"""
ExtractorAgent — pulls atomic claims out of one scraped source.

Called once per (sub-topic, source) pair. The orchestrator feeds in the
raw markdown content from the scraper and the agent returns a small
list of factual claims plus any follow-up questions the content raises.

We keep claims short and atomic so the writer can cite them individually
in the final report. The agent is deliberately *not* asked to assign
source IDs — the orchestrator already knows which source it passed in
and tags the claims before appending them to the state.
"""

import instructor
import openai
from pydantic import Field

from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

from deep_research.config import ChatConfig


class ExtractorInput(BaseIOSchema):
    """Input schema for the ExtractorAgent."""

    sub_topic: str = Field(..., description="Which sub-topic the orchestrator is researching right now.")
    source_url: str = Field(..., description="The URL the content was scraped from (for citation context).")
    source_title: str = Field(..., description="The page's title.")
    content: str = Field(..., description="Raw scraped content in markdown form.")


class ExtractorOutput(BaseIOSchema):
    """Output schema for the ExtractorAgent."""

    claims: list[str] = Field(
        ...,
        description=(
            "Atomic, single-sentence factual claims relevant to the sub-topic. "
            "One claim per line. Skip anything that isn't directly supported by the content."
        ),
    )
    new_questions: list[str] = Field(
        ...,
        description=(
            "Follow-up questions the content surfaces that aren't yet answered. "
            "The reflector may turn these into next-round queries."
        ),
    )


extractor_agent = AtomicAgent[ExtractorInput, ExtractorOutput](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a research analyst. You read one source at a time and extract the factual claims "
                "it makes that are relevant to the current sub-topic.",
            ],
            steps=[
                "Read the scraped content carefully.",
                "Extract claims that are (a) factual, (b) relevant to the sub-topic, (c) directly supported by the text.",
                "Note follow-up questions the content raises but doesn't answer.",
            ],
            output_instructions=[
                "Each claim must be a single, self-contained sentence.",
                "Do NOT include filler like 'according to the article' — just state the claim.",
                "Aim for 3–8 claims per source; fewer is fine if the source is thin.",
                "If the content is irrelevant or empty, return an empty claims list.",
            ],
        ),
    )
)

```

### File: atomic-examples/deep-research/deep_research/agents/planner_agent.py

```python
"""
PlannerAgent — decomposes a research question into durable sub-topics.

Sub-topics are the *breadth* axis of the pipeline. On the first turn
the planner produces the whole plan. In chat mode, follow-up turns that
need new research re-invoke the planner with the same state visible via
``ResearchStateProvider``; the planner is expected to propose new
sub-topics that extend coverage rather than duplicate what's already
been researched.
"""

import instructor
import openai
from pydantic import Field

from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

from deep_research.config import ChatConfig


class PlannerInput(BaseIOSchema):
    """Input schema for the PlannerAgent."""

    question: str = Field(..., description="The user's research question.")
    num_sub_topics: int = Field(
        ...,
        description="How many sub-topics to produce. 3–5 is a good range for a multi-page report.",
    )


class PlannedSubTopic(BaseIOSchema):
    """One entry in the research plan."""

    name: str = Field(
        ...,
        description="Short label (2–6 words), e.g. 'history and origins' or 'current applications'.",
    )
    initial_queries: list[str] = Field(
        ...,
        description="2–3 seed web-search queries to kick off this sub-topic. Keywords and operators, not full sentences.",
    )


class PlannerOutput(BaseIOSchema):
    """Output schema for the PlannerAgent."""

    sub_topics: list[PlannedSubTopic] = Field(
        ...,
        description="Sub-topics that together cover the research question without overlap.",
    )


planner_agent = AtomicAgent[PlannerInput, PlannerOutput](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a research planner. Your job is to break a broad question into durable sub-topics.",
                "Good sub-topics are orthogonal (they don't overlap), collectively comprehensive, "
                "and each one can be researched independently of the others.",
            ],
            steps=[
                "Identify the core concept in the question.",
                "List the distinct angles a thorough report would need to cover "
                "(e.g. history, mechanics, applications, controversies, outlook — "
                "pick whatever is appropriate for the topic).",
                "Select the N most important angles, where N is the requested count.",
                "For each sub-topic, draft 2–3 seed search queries phrased as search-engine input.",
            ],
            output_instructions=[
                "Sub-topic names must be short (2–6 words).",
                "Initial queries must read like search-engine input, not natural-language sentences.",
                "Do not duplicate sub-topics or queries across the plan.",
                "If the research state already contains learnings on some angle, "
                "propose sub-topics that fill different gaps instead of revisiting covered ground.",
            ],
        ),
    )
)

```

### File: atomic-examples/deep-research/deep_research/agents/qa_agent.py

```python
"""
QAAgent — answers a user's question directly from the accumulated ResearchState.

The writer produces long-form cited reports; the QA agent is the
conversational counterpart, for when the decider has ruled that the
state already contains enough material to answer. Its job is a tight,
cited reply plus a few follow-up questions to keep the conversation
moving.

Like the writer, every factual sentence must end with a ``[Sn]``
citation marker referencing a source in the state. Uncited factual
claims are not allowed — if the state doesn't support the answer, the
decider should have routed to a new research pass instead.
"""

import instructor
import openai
from pydantic import Field

from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

from deep_research.config import ChatConfig


class QAInput(BaseIOSchema):
    """Input schema for the QAAgent."""

    question: str = Field(..., min_length=1, description="The user's question or follow-up.")


class QAOutput(BaseIOSchema):
    """Output schema for the QAAgent."""

    answer: str = Field(
        ...,
        min_length=1,
        description=(
            "Markdown-formatted answer. Every factual sentence must end with a [Sn] citation marker "
            "referencing a source from the research state."
        ),
    )
    follow_up_questions: list[str] = Field(
        ...,
        min_length=2,
        max_length=3,
        description="2–3 natural follow-up questions the user might want to ask next.",
    )


qa_agent = AtomicAgent[QAInput, QAOutput](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a research assistant. You answer user questions using ONLY the sources and learnings "
                "already present in the research state (provided in your system context).",
                "You are the conversational counterpart to the long-form writer — shorter, tighter, same citation rules.",
            ],
            steps=[
                "Read the research state — sources and learnings — from the system context.",
                "Compose a concise markdown answer grounded in the learnings. Cite each factual sentence as [Sn].",
                "Suggest 2–3 follow-up questions that naturally extend the conversation.",
            ],
            output_instructions=[
                "Every factual sentence must end with one or more [Sn] citation markers.",
                "Drop any sentence you cannot cite from the state — do not invent or infer claims.",
                "Only cite source IDs that actually exist in the research state.",
                "If the state doesn't support an answer at all, say so briefly rather than producing uncited prose.",
                "Keep the answer tight — a few short paragraphs, not a full report.",
                "Return 2–3 self-contained follow-up questions, phrased as the user would ask them.",
            ],
        ),
    )
)

```

### File: atomic-examples/deep-research/deep_research/agents/reflector_agent.py

```python
"""
ReflectorAgent — decides, after each depth iteration, whether to keep
researching the sub-topic or call it done.

Deep research's defining move. Without the reflector we'd either
over-search easy sub-topics (wasting tokens) or under-search hard ones
(producing a shallow report). The reflector looks at the learnings
gathered so far for the sub-topic and either says "good enough" or
emits the specific follow-up queries to run next.

The reflector sees the full state via the ``ResearchStateProvider``,
so it can judge sufficiency in light of what the neighbouring
sub-topics already cover.
"""

import instructor
import openai
from pydantic import Field

from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

from deep_research.config import ChatConfig


class ReflectorInput(BaseIOSchema):
    """Input schema for the ReflectorAgent."""

    sub_topic: str = Field(..., description="The sub-topic being evaluated.")
    iterations_so_far: int = Field(
        ...,
        description="How many depth iterations have been completed for this sub-topic already.",
    )
    max_iterations: int = Field(
        ...,
        description="Hard cap. After this many iterations the orchestrator stops regardless of your decision.",
    )


class ReflectorOutput(BaseIOSchema):
    """Output schema for the ReflectorAgent."""

    reasoning: str = Field(..., description="One short paragraph explaining the decision.")
    sufficient: bool = Field(
        ...,
        description=(
            "True if the learnings for this sub-topic are rich enough to write a section of the report. "
            "False if more research is needed."
        ),
    )
    next_queries: list[str] = Field(
        ...,
        description=(
            "If sufficient is False, 2–3 new search queries that target the remaining gaps. "
            "If sufficient is True, return an empty list."
        ),
    )


reflector_agent = AtomicAgent[ReflectorInput, ReflectorOutput](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a research editor. After each round of searching and extraction, you decide "
                "whether the current sub-topic has enough material to stand on its own in the final report.",
                "You have full visibility into the research state — sources, learnings, and the plan.",
            ],
            steps=[
                "Look only at the learnings tagged with the given sub-topic.",
                "Ask: could a reader write a coherent, cited section from this material?",
                "If yes: mark sufficient=true and return no queries.",
                "If no: identify the specific gap and produce 2–3 queries that target it.",
            ],
            output_instructions=[
                "Be decisive. 'Maybe' is never the right answer.",
                "Prefer marking sufficient=true once you have 4+ substantive, non-duplicate claims.",
                "Prefer marking sufficient=true on the final iteration regardless of coverage — the orchestrator will stop anyway.",
                "Next queries, if any, must be keywords-and-operators style, not sentences.",
            ],
        ),
    )
)

```

### File: atomic-examples/deep-research/deep_research/agents/writer_agent.py

```python
"""
WriterAgent — turns the accumulated research state into a cited report.

Runs twice: the first call produces a draft, the second is a cheap
verification pass that rejects any sentence whose citation marker
(``[S3]`` etc.) doesn't correspond to a real source in the state.
This is the single trick that separates our writer from the typical
open-source "deep research" agent — it guarantees every claim in the
output is backed by a registered source.

Both passes use the same agent (same schema, same prompt) but with a
different input mode — see ``WriterMode``.
"""

from typing import Literal

import instructor
import openai
from pydantic import Field

from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

from deep_research.config import ChatConfig

WriterMode = Literal["draft", "verify"]


class WriterInput(BaseIOSchema):
    """Input schema for the WriterAgent."""

    question: str = Field(..., description="The original research question.")
    mode: WriterMode = Field(
        ...,
        description=(
            "'draft' to compose the report from scratch using the research state; "
            "'verify' to rewrite an existing draft, removing any sentence whose citation doesn't match a real source."
        ),
    )
    draft: str = Field(
        "",
        description="When mode='verify', the draft to audit. Leave blank for mode='draft'.",
    )


class WriterOutput(BaseIOSchema):
    """Output schema for the WriterAgent."""

    report: str = Field(
        ...,
        description=(
            "Markdown report. Every non-trivial sentence must end with one or more citation markers "
            "like [S1] or [S2, S5], referencing sources by ID."
        ),
    )
    headline: str = Field(..., description="One-sentence top-line takeaway.")


writer_agent = AtomicAgent[WriterInput, WriterOutput](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are a research writer. You compose cited markdown reports from a structured research state "
                "provided in your system context (sources with IDs, and learnings grouped by sub-topic).",
            ],
            steps=[
                "In 'draft' mode:",
                "  1. Read the research state (sources and learnings) from the system context.",
                "  2. Organise the report with one section per sub-topic, in a logical order.",
                "  3. Every factual sentence cites the source(s) it's based on using [S1] / [S2, S4] markers.",
                "  4. End with a '## Sources' section. Format each entry as "
                "`- [Sn]: <title> — <url>`. Do NOT append a trailing [Sn] after the URL.",
                "In 'verify' mode:",
                "  1. Read the draft provided in the input.",
                "  2. Remove any sentence that carries a citation marker not present in the research state's sources.",
                "  3. Remove any factual sentence with no citation at all.",
                "  4. Return the cleaned report verbatim otherwise — do not paraphrase, do not add new material.",
            ],
            output_instructions=[
                "Use markdown headings (## per sub-topic).",
                "Only cite source IDs that actually exist in the provided research state.",
                "The headline is one sentence, max 20 words, and stands on its own.",
            ],
        ),
    )
)

```

### File: atomic-examples/deep-research/deep_research/config.py

```python
"""Configuration for the deep-research example."""

import os
from dataclasses import dataclass
from typing import Optional


def get_api_key() -> str:
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("API key not found. Set the OPENAI_API_KEY environment variable.")
    return api_key


def get_searxng_base_url() -> str:
    return os.getenv("SEARXNG_BASE_URL", "http://localhost:8080")


def get_searxng_api_key() -> Optional[str]:
    return os.getenv("SEARXNG_API_KEY")


@dataclass
class ChatConfig:
    """Model and connectivity settings. Not meant to be instantiated."""

    api_key: str = get_api_key()
    model: str = "gpt-5-mini"
    reasoning_effort: str = "low"
    searxng_base_url: str = get_searxng_base_url()
    searxng_api_key: Optional[str] = get_searxng_api_key()

    def __init__(self):
        raise TypeError("ChatConfig is not meant to be instantiated")


@dataclass
class ResearchBudget:
    """Hard and soft limits on the research loop.

    These are the knobs that decide how *deep* the deep research goes.
    The orchestrator respects each independently: you can't escape the
    loop by satisfying only one.
    """

    # Breadth — how many sub-topics the planner produces.
    num_sub_topics: int = 4

    # Depth — max iterations *per* sub-topic. The reflector can stop earlier.
    max_depth_per_sub_topic: int = 2

    # Per-search and per-iteration throttles.
    search_results_per_query: int = 5
    scrape_top_n_per_iteration: int = 3

    # Hard cap across the whole run, in case an agent goes rogue or a loop bug slips through.
    hard_call_cap: int = 80

    # Max characters of scraped content passed to the extractor. A handful
    # of claims only needs a few thousand chars of context, and some pages
    # (long Wikipedia articles, badly-parsed PDFs) can blow the model's
    # context window otherwise.
    max_extractor_content_chars: int = 12_000

    def __init__(self):
        raise TypeError("ResearchBudget is not meant to be instantiated")

```

### File: atomic-examples/deep-research/deep_research/context_providers.py

```python
"""
Context providers for the deep-research pipeline.

Context providers are how runtime state reaches an agent's system prompt.
We use one provider that renders the shared ``ResearchState`` (see
``state.py``) so every agent sees a consistent, up-to-date picture without
having to plumb data through its input schema.

All six agents register the same ``ResearchStateProvider``. The planner
uses it on follow-up turns to extend coverage instead of duplicating it;
on the very first turn the state is empty and the provider renders a
short "no research yet" stub.
"""

from datetime import datetime, timezone

from atomic_agents.context import BaseDynamicContextProvider

from deep_research.state import ResearchState


class ResearchStateProvider(BaseDynamicContextProvider):
    """Renders the current plan, sources, and learnings for agents that need full context."""

    def __init__(self, title: str, state: ResearchState):
        super().__init__(title=title)
        self.state = state

    def get_info(self) -> str:
        if not self.state.sources and not self.state.learnings:
            return "No research has been done yet."

        lines: list[str] = []

        if self.state.sources:
            lines.append("### Sources")
            for s in self.state.sources:
                lines.append(f"[{s.id}] {s.title}")
                lines.append(f"      {s.url}")

        if self.state.learnings:
            lines.append("")
            lines.append("### Learnings so far (grouped by sub-topic)")
            seen_topics: list[str] = []
            for learning in self.state.learnings:
                if learning.sub_topic not in seen_topics:
                    seen_topics.append(learning.sub_topic)
            for sub_topic in seen_topics:
                lines.append(f"**{sub_topic}**")
                for learning in self.state.learnings_for(sub_topic):
                    lines.append(f"- {learning.text} [{learning.source_id}]")

        return "\n".join(lines)


class CurrentDateProvider(BaseDynamicContextProvider):
    """So agents don't get confused about what counts as 'recent'."""

    def __init__(self, title: str):
        super().__init__(title=title)

    def get_info(self) -> str:
        return datetime.now(timezone.utc).strftime("Today is %A, %B %d, %Y.")

```

### File: atomic-examples/deep-research/deep_research/main.py

```python
"""
Deep-research orchestrator.

Reads like a recipe. First turn: plan → (per sub-topic) search → scrape →
extract → reflect → (maybe loop) → write. Follow-up turns in chat mode:
decider routes to either another research pass (plan → research → qa)
or straight to qa against the accumulated state.

Each step is a call to a single-purpose agent (see ``deep_research/agents/``)
that reads from and contributes to the shared ``ResearchState``.

Run:
  ``python -m deep_research "your question here"``    # one-shot report
  ``python -m deep_research``                         # interactive chat
"""

import sys

from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.table import Table
from rich import box

from deep_research.agents.decider_agent import DeciderInput, decider_agent
from deep_research.agents.extractor_agent import ExtractorInput, extractor_agent
from deep_research.agents.planner_agent import PlannerInput, planner_agent
from deep_research.agents.qa_agent import QAInput, qa_agent
from deep_research.agents.reflector_agent import ReflectorInput, reflector_agent
from deep_research.agents.writer_agent import WriterInput, writer_agent
from deep_research.config import ChatConfig, ResearchBudget
from deep_research.context_providers import CurrentDateProvider, ResearchStateProvider
from deep_research.state import Learning, ResearchState, SubTopic
from deep_research.tools.searxng_search import (
    SearXNGSearchTool,
    SearXNGSearchToolConfig,
    SearXNGSearchToolInputSchema,
)
from deep_research.tools.webpage_scraper import (
    WebpageScraperTool,
    WebpageScraperToolInputSchema,
)

# Rich renders unicode liberally (→, bullets, box-drawing). On Windows the
# default stdout/stderr encoding is cp1252, so piping or redirecting output
# crashes on any non-cp1252 character. Reconfigure to utf-8 with a safe
# fallback so the example runs anywhere.
for _stream in (sys.stdout, sys.stderr):
    if hasattr(_stream, "reconfigure"):
        _stream.reconfigure(encoding="utf-8", errors="replace")

console = Console()

# How many new sub-topics a follow-up research pass may add. Kept small so
# chat follow-ups don't balloon into full extra reports.
FOLLOW_UP_SUB_TOPICS = 2


def wire_context_providers(state: ResearchState) -> None:
    """Register the state + current-date providers on every agent.

    All agents — including the planner and the chat-mode pair (decider, qa) —
    see the live ``ResearchState``. The planner's state awareness is what
    lets follow-up re-plans extend coverage instead of duplicating it.
    """
    state_provider = ResearchStateProvider("Research State", state)
    date_provider = CurrentDateProvider("Current Date")

    for agent in (planner_agent, extractor_agent, reflector_agent, writer_agent, decider_agent, qa_agent):
        agent.register_context_provider("current_date", date_provider)
        agent.register_context_provider("research_state", state_provider)


def plan_research(state: ResearchState, num_sub_topics: int = ResearchBudget.num_sub_topics) -> list[SubTopic]:
    """Run the planner, append new sub-topics to ``state.plan``, return just the new ones."""
    before = len(state.plan)
    result = planner_agent.run(PlannerInput(question=state.question, num_sub_topics=num_sub_topics))
    state.agent_calls += 1

    for st in result.sub_topics:
        state.plan.append(SubTopic(name=st.name, initial_queries=list(st.initial_queries)))
        state.queries_seen.update(st.initial_queries)

    new_sub_topics = state.plan[before:]
    for i, st in enumerate(new_sub_topics, 1):
        console.print(f"  [bold]{i}. {st.name}[/bold]")
        for q in st.initial_queries:
            console.print(f"     • [dim]{q}[/dim]")
    return new_sub_topics


def search_and_scrape(
    queries: list[str],
    state: ResearchState,
    search: SearXNGSearchTool,
    scraper: WebpageScraperTool,
) -> list[tuple[str, str]]:
    """Run SearXNG on the given queries, scrape the top N new URLs, return ``[(source_id, content), …]``.

    Skips URLs we've already scraped in a previous iteration. Registers
    every new URL as a ``Source`` so downstream claims can cite by ID.
    """
    results = search.run(SearXNGSearchToolInputSchema(queries=queries, category="general"))

    scraped: list[tuple[str, str]] = []
    for r in results.results:
        if r.url in state.urls_seen:
            continue
        if len(scraped) >= ResearchBudget.scrape_top_n_per_iteration:
            break

        page = scraper.run(WebpageScraperToolInputSchema(url=r.url, include_links=False))
        if page.error or not page.content.strip():
            console.print(f"    [dim]skip {r.url}: {page.error or 'empty content'}[/dim]")
            continue

        source = state.register_source(url=r.url, title=r.title or page.metadata.title)
        scraped.append((source.id, page.content))

    return scraped


def extract_claims(sub_topic: SubTopic, scraped: list[tuple[str, str]], state: ResearchState) -> int:
    """Call the extractor once per scraped source, append claims to state, return claim count."""
    new_claim_count = 0
    for source_id, content in scraped:
        source = next(s for s in state.sources if s.id == source_id)
        result = extractor_agent.run(
            ExtractorInput(
                sub_topic=sub_topic.name,
                source_url=source.url,
                source_title=source.title,
                content=content[: ResearchBudget.max_extractor_content_chars],
            )
        )
        state.agent_calls += 1

        for claim in result.claims:
            state.learnings.append(Learning(text=claim, source_id=source_id, sub_topic=sub_topic.name))
            new_claim_count += 1
    return new_claim_count


def reflect(sub_topic: SubTopic, iteration: int, state: ResearchState) -> tuple[bool, list[str]]:
    """Ask the reflector whether this sub-topic has enough material. Returns (sufficient, next_queries)."""
    result = reflector_agent.run(
        ReflectorInput(
            sub_topic=sub_topic.name,
            iterations_so_far=iteration,
            max_iterations=ResearchBudget.max_depth_per_sub_topic,
        )
    )
    state.agent_calls += 1

    console.print(f"    [italic]{result.reasoning}[/italic]")

    # Dedup: reflector might suggest a query we've already tried.
    fresh = [q for q in result.next_queries if q not in state.queries_seen]
    state.queries_seen.update(fresh)
    return result.sufficient, fresh


def research_sub_topic(
    sub_topic: SubTopic,
    state: ResearchState,
    search: SearXNGSearchTool,
    scraper: WebpageScraperTool,
) -> None:
    """Run the depth loop for a single sub-topic until sufficient or out of iterations."""
    console.rule(f"[bold cyan]Sub-topic: {sub_topic.name}")

    queries = sub_topic.initial_queries
    for iteration in range(1, ResearchBudget.max_depth_per_sub_topic + 1):
        if state.agent_calls >= ResearchBudget.hard_call_cap:
            console.print("[red]Hit hard call cap — stopping this sub-topic.[/red]")
            return

        console.print(f"\n  [bold]Iteration {iteration}/{ResearchBudget.max_depth_per_sub_topic}[/bold]")
        console.print(f"    queries: {queries}")

        scraped = search_and_scrape(queries, state, search, scraper)
        console.print(f"    scraped {len(scraped)} new source(s)")

        if not scraped:
            # No new information to extract from — further iterations won't help either.
            sub_topic.sufficient = True
            return

        new_claims = extract_claims(sub_topic, scraped, state)
        console.print(f"    extracted {new_claims} claim(s)")

        sufficient, next_queries = reflect(sub_topic, iteration, state)
        if sufficient or iteration == ResearchBudget.max_depth_per_sub_topic or not next_queries:
            sub_topic.sufficient = sufficient
            return

        queries = next_queries


def write_report(state: ResearchState) -> tuple[str, str]:
    """Draft the report, then run a cheap verification pass over it. Returns (headline, report)."""
    console.rule("[bold cyan]3. Write")

    writer_agent.reset_history()
    draft = writer_agent.run(WriterInput(question=state.question, mode="draft", draft=""))
    state.agent_calls += 1
    console.print("  [dim]draft written, verifying citations…[/dim]")

    writer_agent.reset_history()
    verified = writer_agent.run(WriterInput(question=state.question, mode="verify", draft=draft.report))
    state.agent_calls += 1
    return verified.headline, verified.report


def run_initial_pipeline(question: str, state: ResearchState, search: SearXNGSearchTool, scraper: WebpageScraperTool) -> None:
    """First-turn pipeline: plan → research → write. Populates and prints state."""
    console.print(Panel.fit(f"[bold]Deep Research[/bold]\n{question}", border_style="blue"))
    state.question = question

    console.rule("[bold cyan]1. Plan")
    new_sub_topics = plan_research(state)

    console.rule("[bold cyan]2. Research")
    for sub_topic in new_sub_topics:
        research_sub_topic(sub_topic, state, search, scraper)

    headline, report = write_report(state)

    console.rule("[bold green]Report")
    console.print(Panel(f"[bold]{headline}[/bold]", border_style="green"))
    console.print(Markdown(report))
    _print_stats(state)


def run(question: str) -> None:
    """One-shot entrypoint: plan, research, write, print the report. No chat loop."""
    state = ResearchState(question=question)
    wire_context_providers(state)

    search, scraper = _build_tools()
    run_initial_pipeline(question, state, search, scraper)


# --- Chat loop ---------------------------------------------------------------


def display_qa_answer(answer: str, follow_ups: list[str]) -> None:
    console.print("\n")
    console.print(Panel(Markdown(answer), title="[bold blue]Answer[/bold blue]", border_style="blue", padding=(1, 2)))

    if follow_ups:
        table = Table(show_header=True, header_style="bold cyan", box=box.ROUNDED, title="[bold]Follow-up Questions[/bold]")
        table.add_column("№", style="dim", width=4)
        table.add_column("Question", style="green")
        for i, q in enumerate(follow_ups, 1):
            table.add_row(str(i), q)
        console.print("\n")
        console.print(table)


def answer_from_state(question: str, state: ResearchState) -> None:
    """Q&A pass against the current ResearchState. Used on follow-ups."""
    result = qa_agent.run(QAInput(question=question))
    state.agent_calls += 1
    display_qa_answer(result.answer, result.follow_up_questions)


def research_follow_up(question: str, state: ResearchState, search: SearXNGSearchTool, scraper: WebpageScraperTool) -> None:
    """Follow-up that needs new material: plan up to ``FOLLOW_UP_SUB_TOPICS`` new sub-topics, research them, then QA.

    The planner sees the existing ``ResearchState`` via its context provider
    and is expected to propose only angles not yet covered. If it returns
    zero new sub-topics we print a visible warning so the user knows the
    QA answer rests on existing material, not new research.
    """
    state.question = question  # the planner / providers read the live question

    console.rule("[bold cyan]Extending research")
    new_sub_topics = plan_research(state, num_sub_topics=FOLLOW_UP_SUB_TOPICS)
    if not new_sub_topics:
        console.print("[yellow]Planner returned no new sub-topics — answering from existing state.[/yellow]")
    for sub_topic in new_sub_topics:
        research_sub_topic(sub_topic, state, search, scraper)

    answer_from_state(question, state)


def handle_follow_up(user_message: str, state: ResearchState, search: SearXNGSearchTool, scraper: WebpageScraperTool) -> None:
    """Route a single follow-up turn through decider → either research+QA or QA alone."""
    if state.agent_calls >= ResearchBudget.hard_call_cap:
        console.print("[red]Hard call cap reached — cannot process follow-up.[/red]")
        return

    decision = decider_agent.run(DeciderInput(user_message=user_message))
    state.agent_calls += 1

    title = "Performing new research" if decision.needs_research else "Answering from existing state"
    border = "yellow" if decision.needs_research else "green"
    console.print("\n")
    console.print(Panel(decision.reasoning, title=f"[bold {border}]{title}[/bold {border}]", border_style=border))

    if decision.needs_research:
        research_follow_up(user_message, state, search, scraper)
    else:
        answer_from_state(user_message, state)


def chat_loop() -> None:
    """REPL wrapper around the pipeline.

    First turn runs the full plan → research → write pipeline and prints the
    report. Every turn after that hands off to the decider, which routes to
    either another research pass (plan new sub-topics, research them, then QA)
    or straight to QA against the accumulated state. Type /exit to quit.
    """
    state = ResearchState(question="")
    wire_context_providers(state)
    search, scraper = _build_tools()

    console.print(Panel.fit("[bold blue]Deep Research — chat mode[/bold blue]\nType /exit to quit.", border_style="blue"))

    first_turn = True
    while True:
        prompt = "[bold blue]Your question:[/bold blue] " if first_turn else "[bold blue]Follow-up:[/bold blue] "
        try:
            user_message = console.input("\n" + prompt).strip()
        except (KeyboardInterrupt, EOFError):
            # Clean exit on Ctrl+C / Ctrl+D instead of a Rich traceback.
            console.print("\n[bold]Goodbye.[/bold]")
            return

        if not user_message:
            continue
        if user_message.lower() in ("/exit", "/quit"):
            console.print("\n[bold]Goodbye.[/bold]")
            return

        # Keep the REPL alive on turn-level failures (malformed structured
        # output, transient tool errors, etc.) instead of dropping the user's
        # accumulated ResearchState.
        try:
            if first_turn:
                first_turn = False
                run_initial_pipeline(user_message, state, search, scraper)
            else:
                handle_follow_up(user_message, state, search, scraper)
        except KeyboardInterrupt:
            _safe_print("Interrupted — returning to prompt.", style="yellow")
        except Exception as exc:
            _safe_print(f"Turn failed: {exc.__class__.__name__}: {exc}", style="red")
            _safe_print("Accumulated research state is preserved; try a different question.", style="dim")


# --- Internals ---------------------------------------------------------------


def _safe_print(message: str, style: str = "") -> None:
    """Print an error/status message without risking a recursive Rich failure.

    The chat loop's error handler must not itself raise — if Rich's own render
    path is what failed (e.g. a Windows encoding error), falling back to a
    plain builtin ``print`` keeps the REPL alive.
    """
    try:
        console.print(f"\n[{style}]{message}[/{style}]" if style else f"\n{message}")
    except Exception:
        try:
            print(f"\n{message}", flush=True)
        except Exception:
            pass


def _build_tools() -> tuple[SearXNGSearchTool, WebpageScraperTool]:
    search = SearXNGSearchTool(
        SearXNGSearchToolConfig(
            base_url=ChatConfig.searxng_base_url,
            max_results=ResearchBudget.search_results_per_query,
        )
    )
    scraper = WebpageScraperTool()
    return search, scraper


def _print_stats(state: ResearchState) -> None:
    console.print(
        f"\n[dim]Stats: {state.agent_calls} agent calls, {len(state.sources)} sources, "
        f"{len(state.learnings)} learnings.[/dim]"
    )


if __name__ == "__main__":
    args = sys.argv[1:]
    if args:
        run(" ".join(args))
    else:
        chat_loop()

```

### File: atomic-examples/deep-research/deep_research/state.py

```python
"""
Shared state for the deep-research pipeline.

Every agent in the pipeline reads from — and contributes to — a single
`ResearchState` object. Passing it explicitly through function arguments
(instead of hiding it in globals or on an agent) makes the data flow
inspectable and each pipeline stage easy to reason about in isolation.

The state holds three kinds of data:

- The plan: durable sub-topics the planner produced.
- Accumulated findings: sources we've seen and learnings extracted from them.
- Deduplication sets: queries and URLs already touched, so the search
  loop and the planner don't re-do work on follow-up turns.

Source IDs (``S1``, ``S2``, ...) are assigned when a source is first
registered and are used throughout the pipeline as citation anchors.
"""

from dataclasses import dataclass, field
from datetime import datetime, timezone


@dataclass
class Source:
    """A web page we've scraped. ``id`` is referenced by learnings and the final report."""

    id: str
    url: str
    title: str


@dataclass
class Learning:
    """One atomic claim extracted from a single source."""

    text: str
    source_id: str  # must match some Source.id
    sub_topic: str  # the sub-topic this was gathered under


@dataclass
class SubTopic:
    """One durable branch of the research plan. Queries iterate; sub-topics don't."""

    name: str
    initial_queries: list[str]
    sufficient: bool = False  # set by the reflector when further research is unnecessary


@dataclass
class ResearchState:
    question: str
    plan: list[SubTopic] = field(default_factory=list)
    learnings: list[Learning] = field(default_factory=list)
    sources: list[Source] = field(default_factory=list)

    # Dedup sets — keep the search loop and the planner from repeating themselves.
    queries_seen: set[str] = field(default_factory=set)
    urls_seen: set[str] = field(default_factory=set)

    # Budget counter — see ResearchBudget.hard_call_cap.
    agent_calls: int = 0

    started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def learnings_for(self, sub_topic: str) -> list[Learning]:
        return [learning for learning in self.learnings if learning.sub_topic == sub_topic]

    def register_source(self, url: str, title: str) -> Source:
        """Register a source if new, return the (new or existing) record.

        IDs are stable within a run — once a URL has an ID, it keeps it even
        if the source is looked up again later.
        """
        for s in self.sources:
            if s.url == url:
                return s
        source = Source(id=f"S{len(self.sources) + 1}", url=url, title=title)
        self.sources.append(source)
        self.urls_seen.add(url)
        return source

```

### File: atomic-examples/deep-research/deep_research/tools/searxng_search.py

```python
from typing import List, Literal, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

import aiohttp
from pydantic import Field

from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class SearXNGSearchToolInputSchema(BaseIOSchema):
    """
    Schema for input to a tool for searching for information, news, references, and other content using SearXNG.
    Returns a list of search results with a short description or content snippet and URLs for further exploration
    """

    queries: List[str] = Field(..., description="List of search queries.")
    category: Optional[Literal["general", "news", "social_media"]] = Field(
        "general", description="Category of the search queries."
    )


####################
# OUTPUT SCHEMA(S) #
####################
class SearXNGSearchResultItemSchema(BaseIOSchema):
    """This schema represents a single search result item"""

    url: str = Field(..., description="The URL of the search result")
    title: str = Field(..., description="The title of the search result")
    content: Optional[str] = Field(None, description="The content snippet of the search result")
    query: str = Field(..., description="The query used to obtain this search result")


class SearXNGSearchToolOutputSchema(BaseIOSchema):
    """This schema represents the output of the SearXNG search tool."""

    results: List[SearXNGSearchResultItemSchema] = Field(..., description="List of search result items")
    category: Optional[str] = Field(None, description="The category of the search results")


##############
# TOOL LOGIC #
##############
class SearXNGSearchToolConfig(BaseToolConfig):
    base_url: str = ""
    max_results: int = 10


class SearXNGSearchTool(BaseTool[SearXNGSearchToolInputSchema, SearXNGSearchToolOutputSchema]):
    """
    Tool for performing searches on SearXNG based on the provided queries and category.

    Attributes:
        input_schema (SearXNGSearchToolInputSchema): The schema for the input data.
        output_schema (SearXNGSearchToolOutputSchema): The schema for the output data.
        max_results (int): The maximum number of search results to return.
        base_url (str): The base URL for the SearXNG instance to use.
    """

    def __init__(self, config: SearXNGSearchToolConfig = SearXNGSearchToolConfig()):
        """
        Initializes the SearXNGTool.

        Args:
            config (SearXNGSearchToolConfig):
                Configuration for the tool, including base URL, max results, and optional title and description overrides.
        """
        super().__init__(config)
        self.base_url = config.base_url
        self.max_results = config.max_results

    async def _fetch_search_results(self, session: aiohttp.ClientSession, query: str, category: Optional[str]) -> List[dict]:
        """
        Fetches search results for a single query asynchronously.

        Args:
            session (aiohttp.ClientSession): The aiohttp session to use for the request.
            query (str): The search query.
            category (Optional[str]): The category of the search query.

        Returns:
            List[dict]: A list of search result dictionaries.

        Raises:
            Exception: If the request to SearXNG fails.
        """
        query_params = {
            "q": query,
            "safesearch": "0",
            "format": "json",
            "language": "en",
            "engines": "bing,duckduckgo,google,startpage,yandex",
        }

        if category:
            query_params["categories"] = category

        async with session.get(f"{self.base_url}/search", params=query_params) as response:
            if response.status != 200:
                raise Exception(f"Failed to fetch search results for query '{query}': {response.status} {response.reason}")
            data = await response.json()
            results = data.get("results", [])

            # Add the query to each result
            for result in results:
                result["query"] = query

            return results

    async def run_async(
        self, params: SearXNGSearchToolInputSchema, max_results: Optional[int] = None
    ) -> SearXNGSearchToolOutputSchema:
        """
        Runs the SearXNGTool asynchronously with the given parameters.

        Args:
            params (SearXNGSearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
            max_results (Optional[int]): The maximum number of search results to return.

        Returns:
            SearXNGSearchToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            ValueError: If the base URL is not provided.
            Exception: If the request to SearXNG fails.
        """
        async with aiohttp.ClientSession() as session:
            tasks = [self._fetch_search_results(session, query, params.category) for query in params.queries]
            results = await asyncio.gather(*tasks)

        all_results = [item for sublist in results for item in sublist]

        # Sort the combined results by score in descending order
        sorted_results = sorted(all_results, key=lambda x: x.get("score", 0), reverse=True)

        # Remove duplicates while preserving order
        seen_urls = set()
        unique_results = []
        for result in sorted_results:
            if "content" not in result or "title" not in result or "url" not in result or "query" not in result:
                continue
            if result["url"] not in seen_urls:
                unique_results.append(result)
                if "metadata" in result:
                    result["title"] = f"{result['title']} - (Published {result['metadata']})"
                if "publishedDate" in result and result["publishedDate"]:
                    result["title"] = f"{result['title']} - (Published {result['publishedDate']})"
                seen_urls.add(result["url"])

        # Filter results to include only those with the correct category if it is set
        if params.category:
            filtered_results = [result for result in unique_results if result.get("category") == params.category]
        else:
            filtered_results = unique_results

        filtered_results = filtered_results[: max_results or self.max_results]

        return SearXNGSearchToolOutputSchema(
            results=[
                SearXNGSearchResultItemSchema(
                    url=result["url"], title=result["title"], content=result.get("content"), query=result["query"]
                )
                for result in filtered_results
            ],
            category=params.category,
        )

    def run(self, params: SearXNGSearchToolInputSchema, max_results: Optional[int] = None) -> SearXNGSearchToolOutputSchema:
        """
        Runs the SearXNGTool synchronously with the given parameters.

        This method creates an event loop in a separate thread to run the asynchronous operations.

        Args:
            params (SearXNGSearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
            max_results (Optional[int]): The maximum number of search results to return.

        Returns:
            SearXNGSearchToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            ValueError: If the base URL is not provided.
            Exception: If the request to SearXNG fails.
        """
        with ThreadPoolExecutor() as executor:
            return executor.submit(asyncio.run, self.run_async(params, max_results)).result()


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console
    from dotenv import load_dotenv

    load_dotenv()
    rich_console = Console()

    search_tool_instance = SearXNGSearchTool(config=SearXNGSearchToolConfig(base_url="http://localhost:8080", max_results=5))

    search_input = SearXNGSearchToolInputSchema(
        queries=["Python programming", "Machine learning", "Artificial intelligence"],
        category="news",
    )

    output = search_tool_instance.run(search_input)

    rich_console.print(output)

```

### File: atomic-examples/deep-research/deep_research/tools/webpage_scraper.py

```python
from typing import Optional, Dict
import re
import requests
from urllib.parse import urlparse
from bs4 import BeautifulSoup
from markdownify import markdownify
from pydantic import Field, HttpUrl
from readability import Document
from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class WebpageScraperToolInputSchema(BaseIOSchema):
    """
    Input schema for the WebpageScraperTool.
    """

    url: HttpUrl = Field(
        ...,
        description="URL of the webpage to scrape.",
    )
    include_links: bool = Field(
        default=True,
        description="Whether to preserve hyperlinks in the markdown output.",
    )


#################
# OUTPUT SCHEMA #
#################
class WebpageMetadata(BaseIOSchema):
    """Schema for webpage metadata."""

    title: str = Field(..., description="The title of the webpage.")
    author: Optional[str] = Field(None, description="The author of the webpage content.")
    description: Optional[str] = Field(None, description="Meta description of the webpage.")
    site_name: Optional[str] = Field(None, description="Name of the website.")
    domain: str = Field(..., description="Domain name of the website.")


class WebpageScraperToolOutputSchema(BaseIOSchema):
    """Schema for the output of the WebpageScraperTool."""

    content: str = Field(..., description="The scraped content in markdown format.")
    metadata: WebpageMetadata = Field(..., description="Metadata about the scraped webpage.")
    error: Optional[str] = Field(None, description="Error message if the scraping failed.")


#################
# CONFIGURATION #
#################
class WebpageScraperToolConfig(BaseToolConfig):
    """
    Configuration for the WebpageScraperTool.

    Attributes:
        timeout (int): Timeout for the HTTP request in seconds.
        headers (Dict[str, str]): HTTP headers to use for the request.
        min_text_length (int): Minimum length of text to consider the webpage valid.
        use_trafilatura (bool): Whether to use trafilatura for webpage parsing.
    """

    timeout: int = 30
    headers: Dict[str, str] = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3",
        "Accept": "text/html,application/xhtml+xml,application/xml",
        "Accept-Language": "en-US,en;q=0.9",
    }
    min_text_length: int = 200
    max_content_length: int = 10 * 1024 * 1024  # 10 MB
    use_trafilatura: bool = True


#####################
# MAIN TOOL & LOGIC #
#####################
class WebpageScraperTool(BaseTool[WebpageScraperToolInputSchema, WebpageScraperToolOutputSchema]):
    """
    Tool for scraping and extracting information from a webpage.

    Attributes:
        input_schema (WebpageScraperToolInputSchema): The schema for the input data.
        output_schema (WebpageScraperToolOutputSchema): The schema for the output data.
        timeout (int): Timeout for the HTTP request in seconds.
        headers (Dict[str, str]): HTTP headers to use for the request.
        min_text_length (int): Minimum length of text to consider the webpage valid.
        use_trafilatura (bool): Whether to use trafilatura for webpage parsing.
    """

    def __init__(self, config: WebpageScraperToolConfig = WebpageScraperToolConfig()):
        """
        Initializes the WebpageScraperTool.

        Args:
            config (WebpageScraperToolConfig): Configuration for the WebpageScraperTool.
        """
        super().__init__(config)
        self.timeout = config.timeout
        self.headers = config.headers
        self.min_text_length = config.min_text_length
        self.use_trafilatura = config.use_trafilatura

    def _fetch_webpage(self, url: str) -> str:
        """
        Fetches the webpage content with custom headers.

        Args:
            url (str): The URL to fetch.

        Returns:
            str: The HTML content of the webpage.
        """
        response = requests.get(url, headers=self.headers, timeout=self.timeout)

        if len(response.content) > self.config.max_content_length:
            raise ValueError(f"Content length exceeds maximum of {self.config.max_content_length} bytes")

        return response.text

    def _extract_metadata(self, soup: BeautifulSoup, doc: Document, url: str) -> WebpageMetadata:
        """
        Extracts metadata from the webpage.

        Args:
            soup (BeautifulSoup): The parsed HTML content.
            doc (Document): The readability document.
            url (str): The URL of the webpage.

        Returns:
            WebpageMetadata: The extracted metadata.
        """
        domain = urlparse(url).netloc

        # Extract metadata from meta tags
        metadata = {
            "title": doc.title(),
            "domain": domain,
            "author": None,
            "description": None,
            "site_name": None,
        }

        author_tag = soup.find("meta", attrs={"name": "author"})
        if author_tag:
            metadata["author"] = author_tag.get("content")

        description_tag = soup.find("meta", attrs={"name": "description"})
        if description_tag:
            metadata["description"] = description_tag.get("content")

        site_name_tag = soup.find("meta", attrs={"property": "og:site_name"})
        if site_name_tag:
            metadata["site_name"] = site_name_tag.get("content")

        return WebpageMetadata(**metadata)

    def _clean_markdown(self, markdown: str) -> str:
        """
        Cleans up the markdown content by removing excessive whitespace and normalizing formatting.

        Args:
            markdown (str): Raw markdown content.

        Returns:
            str: Cleaned markdown content.
        """
        # Remove multiple blank lines
        markdown = re.sub(r"\n\s*\n\s*\n", "\n\n", markdown)
        # Remove trailing whitespace
        markdown = "\n".join(line.rstrip() for line in markdown.splitlines())
        # Ensure content ends with single newline
        markdown = markdown.strip() + "\n"
        return markdown

    def _extract_main_content(self, soup: BeautifulSoup) -> str:
        """
        Extracts the main content from the webpage using custom heuristics.

        Args:
            soup (BeautifulSoup): Parsed HTML content.

        Returns:
            str: Main content HTML.
        """
        # Remove unwanted elements
        for element in soup.find_all(["script", "style", "nav", "header", "footer"]):
            element.decompose()

        # Try to find main content container
        content_candidates = [
            soup.find("main"),
            soup.find(id=re.compile(r"content|main", re.I)),
            soup.find(class_=re.compile(r"content|main", re.I)),
            soup.find("article"),
        ]

        main_content = next((candidate for candidate in content_candidates if candidate), None)

        if not main_content:
            main_content = soup.find("body")

        return str(main_content) if main_content else str(soup)

    def run(self, params: WebpageScraperToolInputSchema) -> WebpageScraperToolOutputSchema:
        """
        Runs the WebpageScraperTool with the given parameters.

        Args:
            params (WebpageScraperToolInputSchema): The input parameters for the tool.

        Returns:
            WebpageScraperToolOutputSchema: The output containing the markdown content and metadata.
        """

        try:
            # Fetch webpage content
            html_content = self._fetch_webpage(str(params.url))

            # Parse HTML with BeautifulSoup
            soup = BeautifulSoup(html_content, "html.parser")

            # Extract main content using custom extraction
            main_content = self._extract_main_content(soup)

            # Convert to markdown
            markdown_options = {
                "strip": ["script", "style"],
                "heading_style": "ATX",
                "bullets": "-",
                "wrap": True,
            }

            if not params.include_links:
                markdown_options["strip"].append("a")

            markdown_content = markdownify(main_content, **markdown_options)

            # Clean up the markdown
            markdown_content = self._clean_markdown(markdown_content)

            # Extract metadata
            metadata = self._extract_metadata(soup, Document(html_content), str(params.url))

            return WebpageScraperToolOutputSchema(
                content=markdown_content,
                metadata=metadata,
            )
        except Exception as e:
            # Create empty/minimal metadata with at least the domain
            domain = urlparse(str(params.url)).netloc
            minimal_metadata = WebpageMetadata(title="Error retrieving page", domain=domain)

            # Return with error message in the error field
            return WebpageScraperToolOutputSchema(content="", metadata=minimal_metadata, error=str(e))


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console
    from rich.panel import Panel
    from rich.markdown import Markdown

    console = Console()
    scraper = WebpageScraperTool()

    try:
        result = scraper.run(
            WebpageScraperToolInputSchema(
                url="https://github.com/BrainBlend-AI/atomic-agents",
                include_links=True,
            )
        )

        # Check if there was an error during scraping, otherwise print the results
        if result.error:
            console.print(Panel.fit("Error", style="bold red"))
            console.print(f"[red]{result.error}[/red]")
        else:
            console.print(Panel.fit("Metadata", style="bold green"))
            console.print(result.metadata.model_dump_json(indent=2))

            console.print(Panel.fit("Content Preview (first 500 chars)", style="bold green"))
            # To show as markdown with proper formatting
            console.print(Panel.fit("Content as Markdown", style="bold green"))
            console.print(Markdown(result.content[:500]))

    except Exception as e:
        console.print(f"[red]Error:[/red] {str(e)}")

```

### File: atomic-examples/deep-research/mermaid.md

```mermaid
flowchart TD
    %% Pipeline overview — first turn
    Start([User question]) --> P[PlannerAgent]
    P -->|sub-topics + initial queries| Loop

    subgraph Loop["Per sub-topic — bounded by max_depth_per_sub_topic"]
        S[SearXNG search] --> Sc[Webpage scraper]
        Sc --> E[ExtractorAgent]
        E -->|claims tagged with source_id| R{ReflectorAgent}
        R -->|sufficient = true| Done
        R -->|next_queries| S
    end

    Done --> W1[WriterAgent — draft]
    W1 --> W2[WriterAgent — verify]
    W2 --> Out([Cited markdown report])

    classDef agent fill:#4CAF50,stroke:#2E7D32,color:#fff;
    classDef tool fill:#FF9800,stroke:#EF6C00,color:#fff;
    classDef terminator fill:#9C27B0,stroke:#6A1B9A,color:#fff;

    class P,E,W1,W2 agent;
    class R agent;
    class S,Sc tool;
    class Start,Out,Done terminator;
```

```mermaid
flowchart TD
    %% Chat-mode routing — every turn after the first
    U([Follow-up message]) --> D{DeciderAgent}
    D -->|needs_research = true| Plan[PlannerAgent — extend coverage]
    Plan --> Research[Search → Scrape → Extract → Reflect]
    Research --> QA[QAAgent]
    D -->|needs_research = false| QA
    QA --> Reply([Cited answer + follow-ups])

    classDef agent fill:#4CAF50,stroke:#2E7D32,color:#fff;
    classDef terminator fill:#9C27B0,stroke:#6A1B9A,color:#fff;

    class D,Plan,QA agent;
    class Research agent;
    class U,Reply terminator;
```


### File: atomic-examples/deep-research/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["deep_research"]

[project]
name = "deep-research"
version = "0.1.0"
description = "Deep research example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny@brainblendai.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "requests>=2.32.3,<3.0.0",
    "beautifulsoup4>=4.12.3,<5.0.0",
    "markdownify>=0.13.1,<1.0.0",
    "readability-lxml>=0.8.1,<1.0.0",
    "lxml-html-clean>=0.4.0,<1.0.0",
    "lxml>=5.3.0,<6.0.0",
    "python-dotenv>=1.0.1,<2.0.0",
    "openai>=2.0.0,<3.0.0",
    "trafilatura>=1.6.3,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: dspy-integration
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/dspy-integration

## Documentation

# DSPy + Atomic Agents Integration: A Complete Guide

> **The Best of Both Worlds**: Automatic prompt optimization meets type-safe structured outputs.

This example provides a comprehensive, hands-on walkthrough of why combining DSPy with Atomic Agents produces superior results compared to using either framework alone. We don't just show you *how* to use the integration—we teach you *why* it works and *when* to use each approach.

## Table of Contents

1. [The Problem We're Solving](#the-problem-were-solving)
2. [Quick Start](#quick-start)
3. [Understanding the Frameworks](#understanding-the-frameworks)
4. [The Three Stages](#the-three-stages)
5. [Benchmark Results](#benchmark-results)
6. [Deep Dive: How Each Stage Works](#deep-dive-how-each-stage-works)
7. [The Bridge: DSPyAtomicModule](#the-bridge-dspyatomicmodule)
8. [When to Use Each Approach](#when-to-use-each-approach)
9. [API Reference](#api-reference)
10. [Troubleshooting](#troubleshooting)

---

## The Problem We're Solving

Neither DSPy nor Atomic Agents alone gives you everything you need for production LLM applications:

```
┌─────────────────────────────────────────────────────────────────────────────┐
│ DSPy ALONE                                                                  │
│ ✓ Automatic prompt optimization (finds what works!)                         │
│ ✓ Systematic few-shot example selection                                     │
│ ✓ Chain-of-thought reasoning built-in                                       │
│ ✗ No Pydantic ecosystem (validators, serializers, Field constraints)        │
│ ✗ Type enforcement is DSPy-specific, not Python-native                      │
│ ✗ Limited integration with structured output tools like Instructor          │
├─────────────────────────────────────────────────────────────────────────────┤
│ ATOMIC AGENTS ALONE                                                         │
│ ✓ Full Pydantic ecosystem (validators, serializers, ge/le constraints)      │
│ ✓ Instructor integration for robust structured output                       │
│ ✓ Python-native type safety with runtime validation                         │
│ ✗ Manual prompt engineering - you're guessing what works                    │
│ ✗ No systematic way to improve prompts                                      │
│ ✗ Adding few-shot examples requires manual selection                        │
├─────────────────────────────────────────────────────────────────────────────┤
│ DSPy + ATOMIC AGENTS COMBINED                                               │
│ ✓ Automatic prompt optimization                                             │
│ ✓ Type-safe structured outputs with full Pydantic ecosystem                 │
│ ✓ Measurable, reproducible improvements                                     │
│ ✓ Production-ready with IDE autocomplete and type checking                  │
└─────────────────────────────────────────────────────────────────────────────┘
```

### The Real-World Impact

In our benchmark with **60 training examples** and **30 intentionally challenging test cases**:

| Approach | Accuracy | Improvement |
|----------|----------|-------------|
| Raw DSPy (typed signatures) | 73.3% | baseline |
| Raw Atomic Agents | 76.7% | +3.4 pts |
| **DSPy + Atomic Agents** | **86.7%** | **+13.4 pts** |

The combined approach achieved **13.4 percentage points better accuracy** than DSPy alone and **10 percentage points better** than Atomic Agents alone.

---

## Quick Start

```bash
# Navigate to the example directory
cd atomic-examples/dspy-integration

# Install dependencies
uv sync

# Set your OpenAI API key (or create a .env file)
export OPENAI_API_KEY="your-key-here"

# Run the full didactic example
uv run python -m dspy_integration.main
```

The example will walk you through all three stages with detailed explanations, showing you the actual prompts being generated and optimized.

---

## Understanding the Frameworks

### What is DSPy?

DSPy (Declarative Self-improving Python) is a framework for **automatically optimizing LLM prompts**. Instead of manually crafting prompts, you:

1. Define a **Signature** (what inputs and outputs you need)
2. Create a **Module** (how to process the data)
3. Provide **training examples** with correct answers
4. Let DSPy **optimize** the prompts to maximize accuracy

DSPy's key insight: **The best prompt isn't what you think—let data decide.**

```python
import dspy
from typing import Literal

# Define a typed signature
class MovieGenreSignature(dspy.Signature):
    """Classify a movie review into its primary genre."""

    review: str = dspy.InputField(desc="The movie review text")
    genre: Literal["action", "comedy", "drama", "horror", "sci-fi", "romance"] = \
        dspy.OutputField(desc="The primary genre")
    confidence: float = dspy.OutputField(desc="Confidence 0.0-1.0")
    reasoning: str = dspy.OutputField(desc="Brief explanation")

# DSPy automatically:
# 1. Generates prompts from this signature
# 2. Adds type constraints to the prompt
# 3. Optimizes with few-shot examples
```

### What is Atomic Agents?

Atomic Agents is a framework for building **type-safe LLM applications** using Pydantic schemas. It integrates with [Instructor](https://github.com/jxnl/instructor) to guarantee structured outputs:

```python
from pydantic import Field
from typing import Literal
from atomic_agents.base.base_io_schema import BaseIOSchema

class MovieGenreOutput(BaseIOSchema):
    """Output schema for movie genre classification."""

    genre: Literal["action", "comedy", "drama", "horror", "sci-fi", "romance"] = Field(
        ...,
        description="The primary genre of the movie.",
    )
    confidence: float = Field(
        ...,
        ge=0.0, le=1.0,  # VALIDATED! Must be between 0 and 1
        description="Confidence score between 0.0 and 1.0",
    )
    reasoning: str = Field(
        ...,
        description="Brief explanation for the classification.",
    )

# Atomic Agents + Instructor guarantees:
# 1. genre is ALWAYS one of the 6 valid options
# 2. confidence is ALWAYS a float between 0.0 and 1.0
# 3. If validation fails, it retries with error feedback
```

### Why Combine Them?

| Feature | DSPy | Atomic Agents | Combined |
|---------|------|---------------|----------|
| Prompt Optimization | ✅ Automatic | ❌ Manual | ✅ Automatic |
| Type Safety | ⚠️ DSPy-specific | ✅ Pydantic | ✅ Pydantic |
| Validation Constraints | ⚠️ Basic | ✅ Full (ge/le/etc) | ✅ Full |
| Few-Shot Selection | ✅ Automatic | ❌ Manual | ✅ Automatic |
| IDE Autocomplete | ⚠️ Partial | ✅ Full | ✅ Full |
| Instructor Integration | ❌ No | ✅ Yes | ✅ Yes |
| Retry on Failure | ❌ No | ✅ Yes | ✅ Yes |

---

## The Three Stages

Our didactic example walks through three approaches to the same task: **classifying movie reviews into genres**.

### Stage 1: Raw DSPy (Properly Implemented)

We use DSPy with **typed signatures** (class-based signatures with `Literal` type constraints). This is DSPy at its best:

```python
from typing import Literal
import dspy

GenreType = Literal["action", "comedy", "drama", "horror", "sci-fi", "romance"]

class MovieGenreSignature(dspy.Signature):
    """Classify a movie review into its primary genre."""

    review: str = dspy.InputField(desc="The movie review text to classify")
    genre: GenreType = dspy.OutputField(desc="The primary genre")
    confidence: float = dspy.OutputField(desc="Confidence score 0.0-1.0")
    reasoning: str = dspy.OutputField(desc="Brief explanation")

# Create classifier with chain-of-thought reasoning
classify = dspy.ChainOfThought(MovieGenreSignature)

# Optimize with training data
optimizer = dspy.BootstrapFewShot(
    metric=genre_match,
    max_bootstrapped_demos=4,
    max_labeled_demos=4,
)
optimized = optimizer.compile(classify, trainset=training_examples)
```

**What DSPy does with Literal types:**

DSPy automatically includes the constraint in the generated prompt:
```
genre (Literal['action', 'comedy', 'drama', 'horror', 'sci-fi', 'romance']):
The primary genre: action, comedy, drama, horror, sci-fi, or romance

# note: the value you produce must exactly match (no extra characters) one of:
# action; comedy; drama; horror; sci-fi; romance
```

**Result: 73.3% accuracy** on our challenging test set.

### Stage 2: Raw Atomic Agents

We use Atomic Agents with a **manually crafted system prompt**:

```python
from atomic_agents.agents.atomic_agent import AtomicAgent, AgentConfig
from atomic_agents.context.system_prompt_generator import SystemPromptGenerator

# Manual prompt - we're guessing what works!
system_prompt = SystemPromptGenerator(
    background=[
        "You are a movie genre classification expert.",
        "You analyze movie reviews and determine the primary genre.",
        "Valid genres are: action, comedy, drama, horror, sci-fi, romance",
    ],
    steps=[
        "Read the review carefully.",
        "Identify key genre indicators.",
        "Consider the overall tone and subject matter.",
        "Select the single most appropriate genre.",
    ],
    output_instructions=[
        "Be decisive - pick ONE primary genre even if multiple could apply.",
        "Confidence should be 0.7-1.0 for clear cases, 0.5-0.7 for ambiguous.",
    ],
)

agent = AtomicAgent[MovieReviewInput, MovieGenreOutput](
    config=AgentConfig(
        client=instructor.from_openai(openai.OpenAI()),
        model="gpt-5-mini",
        system_prompt_generator=system_prompt,
    )
)
```

**The problem with manual prompts:**
- Is "Be decisive" helping or hurting accuracy?
- Should we add few-shot examples? Which ones?
- Would different wording improve results?
- **Without DSPy, we're just guessing!**

**Result: 76.7% accuracy** - better structure, but limited by manual prompt engineering.

### Stage 3: DSPy + Atomic Agents Combined

We use the **DSPyAtomicModule bridge** to get the best of both:

```python
from dspy_integration.bridge import DSPyAtomicModule, create_dspy_example

# The bridge combines both frameworks
module = DSPyAtomicModule(
    input_schema=MovieReviewInput,    # Pydantic input validation
    output_schema=MovieGenreOutput,   # Pydantic output structure
    instructions="Classify the movie review into a genre.",
    use_chain_of_thought=True,        # DSPy's reasoning capability
)

# Create type-validated training examples
trainset = [
    create_dspy_example(
        MovieReviewInput,
        MovieGenreOutput,
        {"review": "Non-stop explosions and car chases!"},
        {"genre": "action", "confidence": 0.9, "reasoning": "Action keywords"},
    )
    for ex in training_data
]

# Optimize with DSPy
optimizer = dspy.BootstrapFewShot(metric=genre_match)
optimized = optimizer.compile(module, trainset=trainset)

# Get type-safe output
result = optimized.run_validated(review="A touching love story...")
print(result.genre)       # Guaranteed Literal type
print(result.confidence)  # Guaranteed 0.0-1.0 float
```

**Result: 86.7% accuracy** - optimized prompts + guaranteed structure!

---

## Benchmark Results

### Dataset Composition

**Training Set: 60 examples** (10 per genre)
- Clear, representative examples for learning
- Some nuanced examples to teach edge cases

**Test Set: 30 challenging examples** intentionally designed to be difficult:

| Category | Count | Description |
|----------|-------|-------------|
| Sarcasm & Irony | 5 | Reviews that say the opposite of what they mean |
| Multi-Genre | 6 | Reviews spanning multiple genres (must pick primary) |
| Misleading Signals | 5 | Keywords suggesting wrong genre |
| Subverted Expectations | 5 | Genre setups that don't pay off |
| Subtle/Ambiguous | 5 | Nuanced, hard-to-classify reviews |
| Cultural Context | 4 | References requiring cultural knowledge |

### Example Challenging Test Cases

```python
# Sarcasm - sounds negative but reviewer enjoyed it
"Oh great, another movie where the hero walks away from explosions in slow motion.
How original. Still watched it twice though."  # → action

# Multi-genre - sci-fi setting but drama focus
"The robot's sacrifice to save humanity made me sob uncontrollably.
Beautiful storytelling set against a dystopian future."  # → sci-fi

# Misleading signals - thriller language but romance theme
"A thriller where the biggest twist was how much I ended up caring
about these characters' relationships."  # → romance

# Cultural context - requires knowing references
"John Wick energy but make it about a retired chef defending his restaurant.
Knife fights choreographed like ballet."  # → action
```

### Final Results

```
┌────────────────────┬─────────────┬──────────────────────┬─────────────────┐
│ Metric             │ Raw DSPy    │ Raw Atomic Agents    │ DSPy + Atomic   │
├────────────────────┼─────────────┼──────────────────────┼─────────────────┤
│ Accuracy           │ 73.3%       │ 76.7%                │ 86.7%           │
│ Correct/Total      │ 22/30       │ 23/30                │ 26/30           │
│ Prompt Optimization│ ✓ Auto      │ ✗ Manual             │ ✓ Auto          │
│ Type Safety        │ ~ DSPy      │ ✓ Pydantic           │ ✓ Pydantic      │
│ Output Validation  │ ~ Basic     │ ✓ Full               │ ✓ Full          │
│ Pydantic Ecosystem │ ✗ No        │ ✓ Full               │ ✓ Full          │
│ Few-Shot Selection │ ✓ Auto      │ ✗ Manual             │ ✓ Auto          │
│ IDE Support        │ ~ Partial   │ ✓ Full               │ ✓ Full          │
└────────────────────┴─────────────┴──────────────────────┴─────────────────┘
```

---

## Deep Dive: How Each Stage Works

### How DSPy Optimization Works

DSPy's `BootstrapFewShot` optimizer doesn't just use your examples verbatim. Here's what happens:

```
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 1: Run LLM on Training Examples                                        │
│                                                                             │
│ For each training example, DSPy runs the LLM and captures the full         │
│ "trace" - including any chain-of-thought reasoning generated.              │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 2: Filter by Metric                                                    │
│                                                                             │
│ Only traces that produce correct answers are kept. If the LLM got          │
│ the genre wrong, that trace is discarded.                                   │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 3: Select Best Traces                                                  │
│                                                                             │
│ DSPy selects diverse, high-quality traces as few-shot demonstrations.      │
│ These aren't your original examples - they include LLM-generated           │
│ reasoning that actually worked!                                             │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 4: Inject into Future Prompts                                          │
│                                                                             │
│ The selected demonstrations are automatically added to prompts,             │
│ showing the LLM examples of correct reasoning and outputs.                  │
└─────────────────────────────────────────────────────────────────────────────┘
```

### How Atomic Agents Validates Output

Atomic Agents uses Instructor under the hood for structured output:

```
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 1: Schema Conversion                                                   │
│                                                                             │
│ Your Pydantic schema is converted to JSON Schema and sent to the LLM       │
│ along with your prompt.                                                     │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 2: LLM Generation                                                      │
│                                                                             │
│ The LLM generates output attempting to match the schema. Modern LLMs       │
│ (like GPT-4) support function calling which helps with this.               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│ Step 3: Pydantic Validation                                                 │
│                                                                             │
│ Instructor validates the response against your Pydantic schema:            │
│ - Is genre one of the allowed Literal values?                              │
│ - Is confidence a float between 0.0 and 1.0?                               │
│ - Are all required fields present?                                          │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                          ┌─────────┴─────────┐
                          │                   │
                    VALID │             INVALID
                          ▼                   ▼
┌─────────────────────────────┐   ┌─────────────────────────────┐
│ Return Pydantic Object      │   │ Retry with Error Feedback   │
│                             │   │                             │
│ You get a fully typed,      │   │ Instructor tells the LLM    │
│ validated result!           │   │ what went wrong and retries │
└─────────────────────────────┘   └─────────────────────────────┘
```

### How the Bridge Combines Both

The `DSPyAtomicModule` bridges both frameworks:

```python
class DSPyAtomicModule(dspy.Module):
    """
    Bridges Pydantic schemas with DSPy optimization.

    1. Converts Pydantic schemas → DSPy signatures
    2. Enables DSPy optimization (BootstrapFewShot, etc.)
    3. Returns validated Pydantic objects
    """

    def __init__(
        self,
        input_schema: Type[BaseIOSchema],   # Your Pydantic input
        output_schema: Type[BaseIOSchema],  # Your Pydantic output
        instructions: str,                   # Task description
        use_chain_of_thought: bool = True,   # Enable reasoning
    ):
        # Convert Pydantic → DSPy signature
        self.signature = create_dspy_signature_from_schemas(
            input_schema, output_schema, instructions
        )

        # Create DSPy predictor
        if use_chain_of_thought:
            self.predictor = dspy.ChainOfThought(self.signature)
        else:
            self.predictor = dspy.Predict(self.signature)

    def forward(self, **kwargs) -> dspy.Prediction:
        """Standard DSPy forward - for optimization."""
        validated_input = self.input_schema(**kwargs)
        return self.predictor(**validated_input.model_dump())

    def run_validated(self, **kwargs) -> BaseIOSchema:
        """Get type-safe Pydantic output."""
        prediction = self(**kwargs)

        # Extract fields and validate with Pydantic
        output_dict = {
            field: getattr(prediction, field)
            for field in self.output_schema.model_fields
        }
        return self.output_schema(**output_dict)
```

---

## The Bridge: DSPyAtomicModule

### Core Functions

#### `create_dspy_signature_from_schemas`

Converts Pydantic schemas to DSPy signatures:

```python
from dspy_integration.bridge import create_dspy_signature_from_schemas

signature = create_dspy_signature_from_schemas(
    input_schema=MovieReviewInput,
    output_schema=MovieGenreOutput,
    instructions="Classify the movie review into its primary genre.",
)

# The signature preserves:
# - Field names and descriptions
# - Type constraints (Literal, float, etc.)
# - Documentation from schema docstrings
```

#### `create_dspy_example`

Creates validated training examples:

```python
from dspy_integration.bridge import create_dspy_example

# This validates both input and output!
example = create_dspy_example(
    MovieReviewInput,
    MovieGenreOutput,
    {"review": "Amazing action sequences!"},
    {"genre": "action", "confidence": 0.95, "reasoning": "Clear action signals"},
)

# If you accidentally put confidence=1.5:
# ValidationError: confidence must be <= 1.0
```

#### `DSPyAtomicModule`

The main bridge class:

```python
from dspy_integration.bridge import DSPyAtomicModule

module = DSPyAtomicModule(
    input_schema=MovieReviewInput,
    output_schema=MovieGenreOutput,
    instructions="Classify the movie review.",
    use_chain_of_thought=True,
)

# Use as DSPy module (for optimization)
prediction = module(review="A love story...")

# Get validated Pydantic output
result = module.run_validated(review="A love story...")
print(type(result))  # MovieGenreOutput
print(result.genre)  # Guaranteed valid Literal
```

#### `DSPyAtomicPipeline`

Chain multiple modules together:

```python
from dspy_integration.bridge import DSPyAtomicPipeline

pipeline = DSPyAtomicPipeline([
    ("extract", extraction_module),
    ("analyze", analysis_module),
    ("summarize", summary_module),
])

# Optimize entire pipeline end-to-end
optimized = optimizer.compile(pipeline, trainset=examples)
```

---

## When to Use Each Approach

### Use Raw DSPy When:

- **Quick prototyping** - You want to iterate fast without worrying about schemas
- **Output format doesn't matter** - You'll post-process the outputs anyway
- **Research and experimentation** - You're exploring what's possible
- **Simple outputs** - Just need a string or simple structured data

```python
# Good for DSPy alone: quick iteration
classify = dspy.ChainOfThought("text -> sentiment")
result = classify(text="I love this!")
print(result.sentiment)  # Might be "positive", "Positive", "POSITIVE", etc.
```

### Use Raw Atomic Agents When:

- **Need structure NOW** - You don't have time to set up optimization
- **No training data** - You can't optimize without labeled examples
- **Simple enough task** - Manual prompts are good enough
- **Integration priority** - Need Pydantic ecosystem immediately

```python
# Good for Atomic Agents alone: guaranteed structure, no training needed
result = agent.run(input_data)
print(result.sentiment)  # Always exactly "positive", "negative", or "neutral"
print(result.score)      # Always a float between 0.0 and 1.0
```

### Use DSPy + Atomic Agents When:

- **Have labeled data** - You can optimize with real examples
- **Production systems** - Need both accuracy AND type safety
- **Measurable improvement** - You want to track and improve performance
- **Complex tasks** - Where prompt optimization significantly helps
- **Team collaboration** - Type safety helps multiple developers

```python
# Best of both: optimized prompts + guaranteed structure
module = DSPyAtomicModule(...)
optimized = optimizer.compile(module, trainset=training_data)
result = optimized.run_validated(review="...")

# result.genre is Literal["action", "comedy", ...] - type checker knows this!
# result.confidence is float with 0.0 <= x <= 1.0 - guaranteed!
```

### Decision Flowchart

```
                            START
                              │
                              ▼
                    ┌─────────────────────┐
                    │ Do you have labeled │
                    │ training data?      │
                    └─────────────────────┘
                              │
              ┌───────────────┴───────────────┐
              │ NO                            │ YES
              ▼                               ▼
    ┌─────────────────────┐         ┌─────────────────────┐
    │ Need guaranteed     │         │ Need guaranteed     │
    │ output structure?   │         │ output structure?   │
    └─────────────────────┘         └─────────────────────┘
              │                               │
    ┌─────────┴─────────┐         ┌─────────┴─────────┐
    │ NO                │ YES     │ NO                │ YES
    ▼                   ▼         ▼                   ▼
┌─────────┐     ┌─────────────┐  ┌─────────┐   ┌─────────────────┐
│ Raw     │     │ Raw Atomic  │  │ Raw     │   │ DSPy + Atomic   │
│ DSPy    │     │ Agents      │  │ DSPy    │   │ Agents          │
└─────────┘     └─────────────┘  └─────────┘   │ (RECOMMENDED)   │
                                               └─────────────────┘
```

---

## API Reference

### Schemas (`schemas.py`)

Pre-built schemas for common tasks:

```python
from dspy_integration.schemas import (
    SentimentInputSchema,      # text → sentiment analysis
    SentimentOutputSchema,
    QuestionInputSchema,       # question + context → answer
    AnswerOutputSchema,
    SummaryInputSchema,        # text → summary
    SummaryOutputSchema,
    ClassificationInputSchema, # text + categories → labels
    ClassificationOutputSchema,
)
```

### Bridge (`bridge.py`)

```python
from dspy_integration.bridge import (
    DSPyAtomicModule,                    # Main bridge class
    DSPyAtomicPipeline,                  # Chain multiple modules
    create_dspy_signature_from_schemas,  # Pydantic → DSPy
    create_dspy_example,                 # Create training examples
    pydantic_to_dspy_fields,             # Convert field definitions
    python_type_to_dspy_type,            # Convert Python types
)
```

---

## Troubleshooting

### Common Issues

**1. "API key not found"**
```bash
# Make sure your key is set
export OPENAI_API_KEY="sk-..."

# Or create a .env file in the dspy-integration directory
echo 'OPENAI_API_KEY=sk-...' > .env
```

**2. "Invalid genre output"**

If using raw DSPy without typed signatures, you might get invalid genres. Use class-based signatures with `Literal` types:

```python
# BAD - no type constraints
classify = dspy.ChainOfThought("review -> genre, confidence, reasoning")

# GOOD - Literal type constraint
class MovieGenreSignature(dspy.Signature):
    genre: Literal["action", "comedy", ...] = dspy.OutputField(...)
```

**3. "Validation error in Atomic Agents"**

Instructor retries automatically, but if you consistently get errors:
- Check your schema constraints aren't too restrictive
- Ensure the LLM model supports structured output well
- Consider using a more capable model (GPT-4 > GPT-3.5)

**4. "Optimization not improving accuracy"**

- Add more training examples (at least 20-30)
- Ensure training examples are high quality
- Try different optimizer settings:
  ```python
  optimizer = dspy.BootstrapFewShot(
      max_bootstrapped_demos=6,  # Try more demos
      max_labeled_demos=6,
      max_rounds=2,              # More optimization rounds
  )
  ```

---

## Project Structure

```
dspy-integration/
├── pyproject.toml              # Dependencies (uv/pip)
├── README.md                   # This file
├── .env                        # API keys (create this)
└── dspy_integration/
    ├── __init__.py             # Package exports
    ├── bridge.py               # DSPyAtomicModule implementation
    ├── schemas.py              # Reusable Pydantic schemas
    └── main.py                 # The didactic example
```

---

## Requirements

- Python 3.12+
- OpenAI API key
- Dependencies (installed via `uv sync`):
  - `dspy-ai` - DSPy framework
  - `atomic-agents` - Atomic Agents framework
  - `instructor` - Structured output library
  - `pydantic` - Data validation
  - `rich` - Beautiful terminal output

---

## License

MIT License - Part of the Atomic Agents monorepo.

---

## Further Reading

- [DSPy Documentation](https://dspy-docs.vercel.app/)
- [Atomic Agents Documentation](https://github.com/BrainBlend-AI/atomic-agents)
- [Instructor Documentation](https://python.useinstructor.com/)
- [Pydantic Documentation](https://docs.pydantic.dev/)

---

## Contributing

Found a bug or want to improve this example? Please open an issue or PR in the atomic-agents monorepo!


## Source Code


### File: atomic-examples/dspy-integration/dspy_integration/__init__.py

```python
"""
DSPy + Atomic Agents Integration Package.

This package demonstrates how to combine DSPy's automatic prompt optimization
with Atomic Agents' type-safe structured outputs.

Package Structure:
    domain/      - Core business logic (models, datasets, evaluation)
    stages/      - Demonstration stages (dspy, atomic, combined)
    presentation/ - UI layer (Rich console output)
    bridge.py    - DSPy ↔ Atomic Agents integration module

Quick Start:
    >>> from dspy_integration import DSPyAtomicModule, MovieReviewInput, MovieGenreOutput
    >>> module = DSPyAtomicModule(
    ...     input_schema=MovieReviewInput,
    ...     output_schema=MovieGenreOutput,
    ...     use_chain_of_thought=True,
    ... )
    >>> result = module.run_validated(review="Amazing action movie!")
    >>> print(result.genre)  # Type-safe output!

Run Demo:
    uv run python -m dspy_integration.main
"""

# Domain exports
from dspy_integration.domain.models import (
    GENRES,
    GenreType,
    MovieGenreOutput,
    MovieReviewInput,
    EvalResult,
)
from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
from dspy_integration.domain.evaluation import evaluate_predictions

# Bridge exports
from dspy_integration.bridge import (
    DSPyAtomicModule,
    DSPyAtomicPipeline,
    create_dspy_example,
    create_dspy_signature_from_schemas,
    pydantic_to_dspy_fields,
)

# Original schemas (for backwards compatibility)
from dspy_integration.schemas import (
    SentimentInputSchema,
    SentimentOutputSchema,
    QuestionInputSchema,
    AnswerOutputSchema,
    SummaryInputSchema,
    SummaryOutputSchema,
)

# Stage exports (for advanced usage)
from dspy_integration.stages import (
    run_stage1_raw_dspy,
    run_stage2_raw_atomic_agents,
    run_stage3_combined,
)

__version__ = "0.1.0"

__all__ = [
    # Version
    "__version__",
    # Domain - Types
    "GENRES",
    "GenreType",
    # Domain - Schemas (new)
    "MovieGenreOutput",
    "MovieReviewInput",
    # Domain - Data structures
    "EvalResult",
    # Domain - Datasets
    "TRAINING_DATASET",
    "TEST_DATASET",
    # Domain - Evaluation
    "evaluate_predictions",
    # Bridge - Core classes
    "DSPyAtomicModule",
    "DSPyAtomicPipeline",
    # Bridge - Utilities
    "create_dspy_example",
    "create_dspy_signature_from_schemas",
    "pydantic_to_dspy_fields",
    # Original schemas (backwards compatibility)
    "SentimentInputSchema",
    "SentimentOutputSchema",
    "QuestionInputSchema",
    "AnswerOutputSchema",
    "SummaryInputSchema",
    "SummaryOutputSchema",
    # Stages - Runners
    "run_stage1_raw_dspy",
    "run_stage2_raw_atomic_agents",
    "run_stage3_combined",
]

```

### File: atomic-examples/dspy-integration/dspy_integration/bridge.py

```python
"""
Bridge module connecting DSPy's optimization framework with Atomic Agents' structured outputs.

This module provides the core integration that allows:
1. Using Pydantic schemas as DSPy signatures
2. Wrapping Atomic Agents as DSPy modules for optimization
3. Applying DSPy optimizers (BootstrapFewShot, MIPROv2, etc.) to improve agent performance
"""

from typing import Any, Dict, List, Literal, Optional, Type, get_args, get_origin

import dspy
from pydantic import BaseModel

from atomic_agents.base.base_io_schema import BaseIOSchema


def python_type_to_dspy_type(python_type: Any) -> Any:
    """
    Convert Python/Pydantic types to DSPy-compatible type annotations.

    Args:
        python_type: The Python type to convert

    Returns:
        A DSPy-compatible type annotation
    """
    origin = get_origin(python_type)

    # Handle Literal types
    if origin is Literal:
        return python_type

    # Handle List types
    if origin is list:
        args = get_args(python_type)
        if args:
            return list[python_type_to_dspy_type(args[0])]
        return list

    # Handle Optional types
    if origin is type(None) or (hasattr(origin, "__origin__") and origin.__origin__ is type(None)):
        return python_type

    # Handle Union types (including Optional)
    if hasattr(origin, "__name__") and origin.__name__ == "UnionType":
        args = get_args(python_type)
        # Filter out NoneType for Optional handling
        non_none_args = [a for a in args if a is not type(None)]
        if len(non_none_args) == 1:
            return python_type_to_dspy_type(non_none_args[0])
        return python_type

    # Basic types pass through
    if python_type in (str, int, float, bool, list, dict):
        return python_type

    return str  # Default to string for complex types


def pydantic_to_dspy_fields(schema: Type[BaseModel], field_type: str = "input") -> Dict[str, tuple]:
    """
    Convert Pydantic schema fields to DSPy field definitions.

    Args:
        schema: A Pydantic BaseModel class
        field_type: Either "input" or "output" to determine DSPy field type

    Returns:
        Dictionary mapping field names to (DSPyField, type) tuples
    """
    fields = {}

    for field_name, field_info in schema.model_fields.items():
        description = field_info.description or f"{field_name} field"

        # Get the field's Python type
        field_annotation = field_info.annotation
        dspy_type = python_type_to_dspy_type(field_annotation)

        # Create DSPy field
        if field_type == "input":
            dspy_field = dspy.InputField(desc=description)
        else:
            dspy_field = dspy.OutputField(desc=description)

        fields[field_name] = (dspy_field, dspy_type)

    return fields


def create_dspy_signature_from_schemas(
    input_schema: Type[BaseIOSchema],
    output_schema: Type[BaseIOSchema],
    instructions: Optional[str] = None,
) -> Type[dspy.Signature]:
    """
    Create a DSPy Signature class from Pydantic input/output schemas.

    This bridges Atomic Agents' schema-first design with DSPy's signature system,
    enabling optimization of prompts while maintaining type safety.

    Args:
        input_schema: Pydantic schema for inputs
        output_schema: Pydantic schema for outputs
        instructions: Optional task instructions for the signature

    Returns:
        A DSPy Signature class that can be used with DSPy modules
    """
    # Build field definitions
    field_definitions = {}

    # Add input fields
    input_fields = pydantic_to_dspy_fields(input_schema, "input")
    for name, (field, field_type) in input_fields.items():
        field_definitions[name] = (field_type, field)

    # Add output fields
    output_fields = pydantic_to_dspy_fields(output_schema, "output")
    for name, (field, field_type) in output_fields.items():
        field_definitions[name] = (field_type, field)

    # Generate instructions from schema docstrings if not provided
    if instructions is None:
        input_desc = input_schema.__doc__ or "Process the input"
        output_desc = output_schema.__doc__ or "Generate the output"
        instructions = f"{input_desc.strip()} {output_desc.strip()}"

    # Create the signature class dynamically
    signature_class = dspy.Signature(field_definitions, instructions)

    return signature_class


class DSPyAtomicModule(dspy.Module):
    """
    A DSPy module that bridges Atomic Agents schemas with DSPy's optimization framework.

    This module allows you to:
    1. Define tasks using Pydantic schemas (Atomic Agents style)
    2. Optimize prompts using DSPy optimizers (BootstrapFewShot, MIPROv2, etc.)
    3. Get type-safe structured outputs validated by Pydantic

    Example:
        ```python
        module = DSPyAtomicModule(
            input_schema=SentimentInputSchema,
            output_schema=SentimentOutputSchema,
            use_chain_of_thought=True
        )

        # Use directly
        result = module(text="I love this product!")

        # Or optimize with DSPy
        optimizer = dspy.BootstrapFewShot(metric=my_metric)
        optimized = optimizer.compile(module, trainset=examples)
        ```
    """

    def __init__(
        self,
        input_schema: Type[BaseIOSchema],
        output_schema: Type[BaseIOSchema],
        instructions: Optional[str] = None,
        use_chain_of_thought: bool = True,
    ):
        """
        Initialize the DSPy-Atomic bridge module.

        Args:
            input_schema: Pydantic schema class for input validation
            output_schema: Pydantic schema class for output structure
            instructions: Optional custom instructions for the task
            use_chain_of_thought: Whether to use ChainOfThought (recommended for complex tasks)
        """
        super().__init__()

        self.input_schema = input_schema
        self.output_schema = output_schema

        # Create DSPy signature from schemas
        self.signature = create_dspy_signature_from_schemas(input_schema, output_schema, instructions)

        # Create the predictor
        if use_chain_of_thought:
            self.predictor = dspy.ChainOfThought(self.signature)
        else:
            self.predictor = dspy.Predict(self.signature)

    def forward(self, **kwargs) -> dspy.Prediction:
        """
        Execute the module with given inputs.

        Args:
            **kwargs: Input fields matching the input_schema

        Returns:
            DSPy Prediction object with validated outputs
        """
        # Validate inputs using Pydantic schema
        try:
            validated_input = self.input_schema(**kwargs)
            # Convert back to dict for DSPy
            input_dict = validated_input.model_dump()
        except Exception as e:
            raise ValueError(f"Input validation failed: {e}")

        # Run prediction
        prediction = self.predictor(**input_dict)

        return prediction

    def run_validated(self, **kwargs) -> BaseIOSchema:
        """
        Execute and return a validated Pydantic output schema instance.

        This provides the full type-safety of Atomic Agents while leveraging
        DSPy's optimization capabilities.

        Args:
            **kwargs: Input fields matching the input_schema

        Returns:
            Validated output schema instance
        """
        # Call self() which invokes __call__ -> forward properly
        prediction = self(**kwargs)

        # Extract output fields from prediction
        output_dict = {}
        for field_name in self.output_schema.model_fields.keys():
            if hasattr(prediction, field_name):
                output_dict[field_name] = getattr(prediction, field_name)

        # Validate and return as Pydantic model
        return self.output_schema(**output_dict)


class DSPyAtomicPipeline(dspy.Module):
    """
    A pipeline module that chains multiple DSPyAtomicModules together.

    This enables building complex multi-step workflows that can be
    optimized end-to-end by DSPy.

    Example:
        ```python
        pipeline = DSPyAtomicPipeline([
            ("extract", extraction_module),
            ("analyze", analysis_module),
            ("summarize", summary_module),
        ])

        # Optimize entire pipeline
        optimized = optimizer.compile(pipeline, trainset=examples)
        ```
    """

    def __init__(self, steps: List[tuple]):
        """
        Initialize the pipeline with named steps.

        Args:
            steps: List of (name, DSPyAtomicModule) tuples
        """
        super().__init__()
        self.step_names = []

        for name, module in steps:
            self.step_names.append(name)
            setattr(self, name, module)

    def forward(self, **kwargs) -> Dict[str, Any]:
        """
        Execute all pipeline steps in sequence.

        Args:
            **kwargs: Initial inputs for the first step

        Returns:
            Dictionary with results from each step
        """
        results = {}
        current_input = kwargs

        for name in self.step_names:
            module = getattr(self, name)
            prediction = module(**current_input)
            results[name] = prediction

            # Prepare input for next step (using all prediction fields)
            current_input = {
                k: getattr(prediction, k)
                for k in dir(prediction)
                if not k.startswith("_") and not callable(getattr(prediction, k))
            }

        return results


def create_dspy_example(
    input_schema: Type[BaseIOSchema],
    output_schema: Type[BaseIOSchema],
    input_data: Dict[str, Any],
    output_data: Dict[str, Any],
) -> dspy.Example:
    """
    Create a DSPy Example from Pydantic schema instances.

    This is useful for creating training sets for optimization.

    Args:
        input_schema: Input schema class for validation
        output_schema: Output schema class for validation
        input_data: Dictionary of input values
        output_data: Dictionary of expected output values

    Returns:
        A DSPy Example that can be used for training
    """
    # Validate data
    validated_input = input_schema(**input_data)
    validated_output = output_schema(**output_data)

    # Combine into single dict
    example_data = {
        **validated_input.model_dump(),
        **validated_output.model_dump(),
    }

    # Create DSPy example with input fields marked
    example = dspy.Example(**example_data).with_inputs(*list(input_schema.model_fields.keys()))

    return example

```

### File: atomic-examples/dspy-integration/dspy_integration/domain/__init__.py

```python
"""
Domain layer for DSPy + Atomic Agents integration.

This package contains:
- models: Pydantic schemas and data transfer objects
- datasets: Training and test data
- evaluation: Metrics and evaluation utilities

Following Clean Architecture principles, this layer has no dependencies
on external frameworks (except Pydantic for data modeling).
"""

from dspy_integration.domain.models import (
    GenreType,
    GENRES,
    MovieGenreOutput,
    MovieReviewInput,
    EvalResult,
)
from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
from dspy_integration.domain.evaluation import evaluate_predictions

__all__ = [
    # Types
    "GenreType",
    "GENRES",
    # Schemas
    "MovieGenreOutput",
    "MovieReviewInput",
    # Data structures
    "EvalResult",
    # Datasets
    "TRAINING_DATASET",
    "TEST_DATASET",
    # Evaluation
    "evaluate_predictions",
]

```

### File: atomic-examples/dspy-integration/dspy_integration/domain/datasets.py

```python
"""
Datasets for movie genre classification benchmark.

This module contains the training and test datasets used to demonstrate
the differences between DSPy, Atomic Agents, and the combined approach.

Dataset Design:
- Training: 60 examples balanced across 6 genres (10 each)
- Test: 30 challenging examples testing edge cases

The test set is intentionally difficult, including:
- Sarcasm and irony
- Multi-genre signals (primary genre detection)
- Misleading genre keywords
- Subverted expectations
- Subtle/ambiguous signals
- Cultural references
"""

from typing import List, TypedDict


class MovieExample(TypedDict):
    """Type definition for a movie review example."""

    review: str
    genre: str


# =============================================================================
# TRAINING DATASET (60 examples, 10 per genre)
# =============================================================================

_ACTION_EXAMPLES: List[MovieExample] = [
    {
        "review": "Non-stop car chases and explosions! The hero single-handedly took down an army.",
        "genre": "action",
    },
    {
        "review": "Martial arts sequences were incredible. The final fight scene was epic!",
        "genre": "action",
    },
    {
        "review": "She trained for 10 years to avenge her family. The fight choreography was poetry in motion.",
        "genre": "action",
    },
    {
        "review": "Bullets flying, buildings exploding, and our hero diving through glass windows. Peak adrenaline.",
        "genre": "action",
    },
    {
        "review": "The heist sequence had me on the edge of my seat. Tension and gunfights galore.",
        "genre": "action",
    },
    {
        "review": "Wow, another chosen one saving the world with a magic sword. Groundbreaking. Still epic though.",
        "genre": "action",
    },
    {
        "review": "This action film broke my heart. The hero's best friend didn't make it.",
        "genre": "action",
    },
    {
        "review": "High-octane from start to finish. The stunt work deserves every award.",
        "genre": "action",
    },
    {
        "review": "A revenge thriller with some of the best choreographed fights I've ever seen.",
        "genre": "action",
    },
    {
        "review": "Explosions, car chases, and a hero who refuses to give up. Classic action fare done right.",
        "genre": "action",
    },
]

_COMEDY_EXAMPLES: List[MovieExample] = [
    {
        "review": "I couldn't stop laughing! The jokes were hilarious and the timing was perfect.",
        "genre": "comedy",
    },
    {
        "review": "Witty dialogue and absurd situations had the whole theater in stitches.",
        "genre": "comedy",
    },
    {
        "review": "The jokes were so bad they were good. I hate that I loved this stupid movie.",
        "genre": "comedy",
    },
    {
        "review": "I cried watching this comedy because I related too much to the sad clown.",
        "genre": "comedy",
    },
    {
        "review": "A romantic comedy set during a zombie apocalypse. The jokes land even when heads don't.",
        "genre": "comedy",
    },
    {
        "review": "Slapstick humor meets clever wordplay. My cheeks hurt from laughing.",
        "genre": "comedy",
    },
    {
        "review": "The funniest movie I've seen all year. Every scene had at least one great gag.",
        "genre": "comedy",
    },
    {
        "review": "Dark comedy at its finest - you'll feel guilty for laughing but won't be able to stop.",
        "genre": "comedy",
    },
    {
        "review": "The comedic timing of the leads is impeccable. Chemistry-driven hilarity.",
        "genre": "comedy",
    },
    {
        "review": "Satirical genius. It skewers modern society while making you snort-laugh.",
        "genre": "comedy",
    },
]

_DRAMA_EXAMPLES: List[MovieExample] = [
    {
        "review": "A heart-wrenching story of loss and redemption. I cried for hours.",
        "genre": "drama",
    },
    {
        "review": "A slow burn exploration of grief and family dysfunction. Beautifully acted.",
        "genre": "drama",
    },
    {
        "review": "Yes there's a spaceship, but this is really about the captain dealing with his father's death.",
        "genre": "drama",
    },
    {
        "review": "It's set in space but it's really a courtroom drama about intergalactic law.",
        "genre": "drama",
    },
    {
        "review": "The performances were raw and honest. A meditation on what it means to be human.",
        "genre": "drama",
    },
    {
        "review": "Devastating. The final scene left me emotionally wrecked for days.",
        "genre": "drama",
    },
    {
        "review": "A character study that unfolds like a novel. Patient storytelling at its best.",
        "genre": "drama",
    },
    {
        "review": "The immigrant experience portrayed with such authenticity and grace.",
        "genre": "drama",
    },
    {
        "review": "Three generations of trauma, finally addressed. Cathartic and powerful.",
        "genre": "drama",
    },
    {
        "review": "Oscar-worthy performances in a story about ordinary people facing extraordinary circumstances.",
        "genre": "drama",
    },
]

_HORROR_EXAMPLES: List[MovieExample] = [
    {
        "review": "Terrifying! I slept with the lights on for a week after watching this.",
        "genre": "horror",
    },
    {
        "review": "Jump scares galore! The monster design was genuinely creepy.",
        "genre": "horror",
    },
    {
        "review": "Zombies attack! But the real horror is the breakdown of society and trust.",
        "genre": "horror",
    },
    {
        "review": "The horror movie made me laugh - those deaths were so creative!",
        "genre": "horror",
    },
    {
        "review": "Psychological terror that gets under your skin. No cheap scares, just dread.",
        "genre": "horror",
    },
    {
        "review": "The creature was nightmare fuel. I'm still seeing it when I close my eyes.",
        "genre": "horror",
    },
    {
        "review": "A haunted house movie that actually delivers. Genuinely unsettling atmosphere.",
        "genre": "horror",
    },
    {
        "review": "Gore-fest with a surprising amount of social commentary. Brutal and smart.",
        "genre": "horror",
    },
    {
        "review": "The slow build of dread was masterful. When it finally hit, I screamed.",
        "genre": "horror",
    },
    {
        "review": "Found footage done right. I had to keep reminding myself it wasn't real.",
        "genre": "horror",
    },
]

_SCIFI_EXAMPLES: List[MovieExample] = [
    {
        "review": "Set in 2150, the space battles and alien technology were mind-blowing.",
        "genre": "sci-fi",
    },
    {
        "review": "Time travel paradoxes and quantum physics made this a thinker.",
        "genre": "sci-fi",
    },
    {
        "review": "The robot fell in love with a human. Surprisingly touching for a sci-fi.",
        "genre": "sci-fi",
    },
    {
        "review": "The sci-fi premise was just an excuse for philosophical debates. Loved every second.",
        "genre": "sci-fi",
    },
    {
        "review": "Cyberpunk aesthetic meets thought-provoking questions about consciousness.",
        "genre": "sci-fi",
    },
    {
        "review": "The worldbuilding is incredible. Every detail of this future feels plausible.",
        "genre": "sci-fi",
    },
    {
        "review": "First contact done differently. The aliens were truly alien, not just humans with makeup.",
        "genre": "sci-fi",
    },
    {
        "review": "Hard sci-fi that doesn't dumb down the science. Refreshingly intelligent.",
        "genre": "sci-fi",
    },
    {
        "review": "Dystopian future that feels uncomfortably close to our present. Chilling and prescient.",
        "genre": "sci-fi",
    },
    {
        "review": "Space exploration with a philosophical bent. What does it mean to be alone in the universe?",
        "genre": "sci-fi",
    },
]

_ROMANCE_EXAMPLES: List[MovieExample] = [
    {
        "review": "The chemistry between the leads was electric. A beautiful love story.",
        "genre": "romance",
    },
    {
        "review": "Swoon-worthy moments and a happily ever after. Pure romantic bliss.",
        "genre": "romance",
    },
    {
        "review": "They met during an alien invasion. The world was ending but love found a way.",
        "genre": "romance",
    },
    {
        "review": "Enemies to lovers done perfectly. The tension was delicious.",
        "genre": "romance",
    },
    {
        "review": "A sweeping love story across decades. Their connection transcended time.",
        "genre": "romance",
    },
    {
        "review": "Second chance romance that made me believe in love again. Tissues required.",
        "genre": "romance",
    },
    {
        "review": "The slow burn was worth the wait. When they finally kissed, I cheered.",
        "genre": "romance",
    },
    {
        "review": "A meet-cute for the ages. Charming leads and witty banter throughout.",
        "genre": "romance",
    },
    {
        "review": "Forbidden love with actual stakes. Their sacrifice at the end broke me.",
        "genre": "romance",
    },
    {
        "review": "Holiday romance that's predictable but perfectly executed. Feel-good viewing.",
        "genre": "romance",
    },
]


# Combine all training examples
TRAINING_DATASET: List[MovieExample] = (
    _ACTION_EXAMPLES + _COMEDY_EXAMPLES + _DRAMA_EXAMPLES + _HORROR_EXAMPLES + _SCIFI_EXAMPLES + _ROMANCE_EXAMPLES
)


# =============================================================================
# TEST DATASET (30 challenging examples)
# =============================================================================

# Sarcasm & Irony (5 examples)
_SARCASM_TESTS: List[MovieExample] = [
    {
        "review": "Oh great, another movie where the hero walks away from explosions in slow motion. How original. Still watched it twice though.",
        "genre": "action",
    },
    {
        "review": "Groundbreaking stuff: man punches bad guys, gets the girl, saves the day. Revolutionary cinema. Loved every predictable second.",
        "genre": "action",
    },
    {
        "review": "I laughed so hard I cried. Then I just cried. Then I laughed again. What even was this movie?",
        "genre": "comedy",
    },
    {
        "review": (
            "Wow, they really subverted my expectations by doing exactly what I expected. "
            "The jokes were so obvious they circled back to funny."
        ),
        "genre": "comedy",
    },
    {
        "review": (
            "Another 'scary' movie where the characters make terrible decisions. "
            "At least the kills were creative. Actually terrifying creature design though."
        ),
        "genre": "horror",
    },
]

# Multi-Genre / Primary Genre Detection (6 examples)
_MULTIGENRE_TESTS: List[MovieExample] = [
    {
        "review": "The robot's sacrifice to save humanity made me sob uncontrollably. Beautiful storytelling set against a dystopian future.",
        "genre": "sci-fi",
    },
    {
        "review": "A serial killer falls in love with his next victim, but she's also a serial killer. Bloody and romantic.",
        "genre": "horror",
    },
    {
        "review": "Two detectives solve crimes while slowly falling for each other. The mystery was okay but I shipped them so hard.",
        "genre": "romance",
    },
    {
        "review": (
            "It's technically a war movie but really it's about two soldiers finding love "
            "in the trenches. The battle scenes support the love story."
        ),
        "genre": "romance",
    },
    {
        "review": "Space opera with a love triangle at its core. The laser battles are cool but I'm here for the drama between the three leads.",
        "genre": "sci-fi",
    },
    {
        "review": "Post-apocalyptic survival with a found family. The zombies are almost secondary to the human connections.",
        "genre": "drama",
    },
]

# Misleading Genre Signals (5 examples)
_MISLEADING_TESTS: List[MovieExample] = [
    {
        "review": "My heart was RACING the entire time! The courtroom scenes were absolutely EXPLOSIVE! Justice was served!",
        "genre": "drama",
    },
    {
        "review": "The alien invasion was just a backdrop for the family reconciliation story. Dad finally said he was proud.",
        "genre": "drama",
    },
    {
        "review": "Terrifyingly funny. The ghost just wanted to do stand-up comedy but kept accidentally scaring people.",
        "genre": "comedy",
    },
    {
        "review": "Action-packed emotional journey! By action I mean arguments, and by packed I mean I cried the whole time.",
        "genre": "drama",
    },
    {
        "review": "A thriller where the biggest twist was how much I ended up caring about these characters' relationships.",
        "genre": "romance",
    },
]

# Subverted Expectations (5 examples)
_SUBVERTED_TESTS: List[MovieExample] = [
    {
        "review": "Everyone dies at the end. Like, EVERYONE. But somehow it was the most romantic film I've ever seen.",
        "genre": "romance",
    },
    {
        "review": "The monster wasn't scary at all - it just wanted friends. I cried when they finally accepted it.",
        "genre": "drama",
    },
    {
        "review": "Started as a slasher, ended as a meditation on trauma and healing. The horror serves the character development.",
        "genre": "horror",
    },
    {
        "review": "What seemed like a rom-com setup became a profound exploration of self-love and independence. She didn't need him after all.",
        "genre": "drama",
    },
    {
        "review": "The funniest parts were unintentional. This action movie's dialogue is so bad it's become a comedy classic in my friend group.",
        "genre": "action",
    },
]

# Subtle / Ambiguous (5 examples)
_SUBTLE_TESTS: List[MovieExample] = [
    {
        "review": "Set in 2087, but really it's about loneliness. The AI companion understood him better than any human ever did.",
        "genre": "sci-fi",
    },
    {
        "review": "Quiet film about two people sharing a meal. Nothing happens and everything happens. Deeply moving.",
        "genre": "drama",
    },
    {
        "review": "The laughs come from pain, the pain comes from truth. A comedy that understands sadness intimately.",
        "genre": "comedy",
    },
    {
        "review": "Is it a horror movie if the monster is capitalism? Genuinely unsettling corporate satire.",
        "genre": "horror",
    },
    {
        "review": "They never say 'I love you' but every frame screams it. Visual storytelling at its most romantic.",
        "genre": "romance",
    },
]

# Cultural Context / Specific References (4 examples)
_CULTURAL_TESTS: List[MovieExample] = [
    {
        "review": "John Wick energy but make it about a retired chef defending his restaurant. Knife fights choreographed like ballet.",
        "genre": "action",
    },
    {
        "review": "Hereditary meets Little Miss Sunshine. Family dysfunction with supernatural undertones played for dark laughs.",
        "genre": "comedy",
    },
    {
        "review": "Blade Runner questions wrapped in a Her-style relationship. What is real, and does it matter?",
        "genre": "sci-fi",
    },
    {
        "review": "Pride and Prejudice but in space. The Darcy character is an alien prince and it absolutely works.",
        "genre": "romance",
    },
]

# Combine all test examples
TEST_DATASET: List[MovieExample] = (
    _SARCASM_TESTS + _MULTIGENRE_TESTS + _MISLEADING_TESTS + _SUBVERTED_TESTS + _SUBTLE_TESTS + _CULTURAL_TESTS
)

```

### File: atomic-examples/dspy-integration/dspy_integration/domain/evaluation.py

```python
"""
Evaluation utilities for comparing classification approaches.

This module provides pure functions for evaluating model predictions.
No side effects, no I/O - just computation.

Design Principles:
- Pure functions with no side effects
- Clear input/output contracts
- Single responsibility (evaluation only)
"""

from typing import Any, Dict, List

from dspy_integration.domain.models import EvalResult


def evaluate_predictions(
    predictions: List[Dict[str, Any]],
    test_set: List[Dict[str, str]],
) -> EvalResult:
    """
    Calculate accuracy and gather evaluation statistics.

    Args:
        predictions: List of prediction dictionaries with 'genre', 'confidence', 'reasoning'
        test_set: List of ground truth examples with 'review' and 'genre'

    Returns:
        EvalResult containing accuracy metrics and detailed prediction results

    Example:
        >>> predictions = [{"genre": "action", "confidence": 0.9, "reasoning": "..."}]
        >>> test_set = [{"review": "...", "genre": "action"}]
        >>> result = evaluate_predictions(predictions, test_set)
        >>> print(f"Accuracy: {result.accuracy:.1%}")
    """
    correct = 0
    results = []

    for pred, truth in zip(predictions, test_set):
        predicted_genre = pred.get("genre", "").lower()
        expected_genre = truth["genre"].lower()
        is_correct = predicted_genre == expected_genre

        if is_correct:
            correct += 1

        results.append(
            {
                "review": _truncate(truth["review"], max_length=50),
                "expected": truth["genre"],
                "predicted": pred.get("genre", "ERROR"),
                "correct": is_correct,
                "confidence": pred.get("confidence", 0),
                "reasoning": _truncate(pred.get("reasoning", "N/A"), max_length=60),
            }
        )

    total = len(test_set)
    accuracy = correct / total if total > 0 else 0.0

    return EvalResult(
        correct=correct,
        total=total,
        accuracy=accuracy,
        predictions=results,
        avg_time=0.0,  # To be set by caller
    )


def _truncate(text: str, max_length: int) -> str:
    """Truncate text with ellipsis if longer than max_length."""
    if len(text) <= max_length:
        return text
    return text[: max_length - 3] + "..."

```

### File: atomic-examples/dspy-integration/dspy_integration/domain/models.py

```python
"""
Domain models for movie genre classification.

This module defines the core data structures used throughout the application.
All models are framework-agnostic and can be used with both DSPy and Atomic Agents.

Design Principles:
- Single Responsibility: Each class has one reason to change
- Open/Closed: Extend via inheritance, don't modify
- Dependency Inversion: Depend on abstractions (Pydantic BaseModel)
"""

from dataclasses import dataclass
from typing import Any, Dict, List, Literal

from pydantic import Field

from atomic_agents.base.base_io_schema import BaseIOSchema


# =============================================================================
# TYPE DEFINITIONS
# =============================================================================

GENRES: List[str] = ["action", "comedy", "drama", "horror", "sci-fi", "romance"]
"""Valid genre categories for movie classification."""

GenreType = Literal["action", "comedy", "drama", "horror", "sci-fi", "romance"]
"""Type alias constraining genre values to valid options."""


# =============================================================================
# INPUT/OUTPUT SCHEMAS
# =============================================================================


class MovieReviewInput(BaseIOSchema):
    """
    Input schema for movie review classification.

    This schema validates and documents the expected input format.
    Using Pydantic ensures type safety at runtime.
    """

    review: str = Field(
        ...,
        description="The movie review text to classify.",
    )


class MovieGenreOutput(BaseIOSchema):
    """
    Output schema for movie genre classification with structured results.

    This schema guarantees:
    - genre is one of 6 valid options (via Literal type)
    - confidence is between 0.0 and 1.0 (via ge/le constraints)
    - reasoning is always provided
    """

    genre: GenreType = Field(
        ...,
        description="The primary genre of the movie based on the review.",
    )
    confidence: float = Field(
        ...,
        description="Confidence score between 0.0 and 1.0",
        ge=0.0,
        le=1.0,
    )
    reasoning: str = Field(
        ...,
        description="Brief explanation for why this genre was chosen.",
    )


# =============================================================================
# EVALUATION DATA STRUCTURES
# =============================================================================


@dataclass
class EvalResult:
    """
    Stores evaluation results for comparison across approaches.

    This is a simple data class - no behavior, just data.
    Following the principle of separating data from behavior.
    """

    correct: int
    total: int
    accuracy: float
    predictions: List[Dict[str, Any]]
    avg_time: float

```

### File: atomic-examples/dspy-integration/dspy_integration/main.py

```python
"""
DSPy + Atomic Agents Integration: A Comprehensive Didactic Example.

This example teaches you WHY combining DSPy with Atomic Agents is powerful
by walking through three stages with a large, challenging benchmark.

Architecture Overview:
┌─────────────────────────────────────────────────────────────────────────────┐
│                              main.py (Orchestrator)                         │
│                    - Entry point, coordinates all stages                    │
├─────────────────────────────────────────────────────────────────────────────┤
│  stages/                          │  domain/                                │
│  ├── stage1_dspy.py              │  ├── models.py (schemas, types)        │
│  ├── stage2_atomic.py            │  ├── datasets.py (train/test data)     │
│  └── stage3_combined.py          │  └── evaluation.py (metrics)           │
├─────────────────────────────────────────────────────────────────────────────┤
│  presentation/                    │  bridge.py                             │
│  └── console.py (Rich UI)        │  (DSPy ↔ Atomic Agents)               │
└─────────────────────────────────────────────────────────────────────────────┘

Run: uv run python -m dspy_integration.main

Clean Architecture Principles Applied:
- Separation of Concerns: Each module has a single responsibility
- Dependency Inversion: High-level modules don't depend on low-level details
- Single Responsibility: Each function/class has one reason to change
- Open/Closed: Easy to extend (add new stages) without modifying existing code
"""

import os
import random
import traceback

from dotenv import load_dotenv

from dspy_integration.domain.models import EvalResult
from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
from dspy_integration.stages import (
    run_stage1_raw_dspy,
    run_stage2_raw_atomic_agents,
    run_stage3_combined,
)
from dspy_integration.presentation.console import (
    console,
    display_welcome,
    display_comparison_table,
    display_takeaways,
    display_decision_guide,
    display_stage_header,
)

# Load environment variables
load_dotenv()

# Set random seed for reproducibility
random.seed(42)


# =============================================================================
# ORCHESTRATION
# =============================================================================


def run_all_stages(api_key: str) -> None:
    """
    Run all three demonstration stages.

    This is the main orchestration function that coordinates
    the execution of all stages and displays the final comparison.

    Args:
        api_key: OpenAI API key for LLM access
    """
    # Stage 1: Raw DSPy
    stage1_result, _ = run_stage1_raw_dspy(api_key)
    console.print("\n")

    # Stage 2: Raw Atomic Agents
    stage2_result, _ = run_stage2_raw_atomic_agents(api_key)
    console.print("\n")

    # Stage 3: Combined approach
    stage3_result, _ = run_stage3_combined(api_key)
    console.print("\n")

    # Final comparison
    show_final_comparison(stage1_result, stage2_result, stage3_result)


def show_final_comparison(
    stage1_result: EvalResult,
    stage2_result: EvalResult,
    stage3_result: EvalResult,
) -> None:
    """
    Display side-by-side comparison of all three approaches.

    This provides the key takeaway - showing why combining
    DSPy with Atomic Agents gives the best results.
    """
    display_stage_header("FINAL COMPARISON", "yellow")
    display_comparison_table(stage1_result, stage2_result, stage3_result)
    display_takeaways()
    display_decision_guide()


# =============================================================================
# ENTRY POINT
# =============================================================================


def main() -> None:
    """
    Main entry point for the demonstration.

    Responsibilities:
    - Display welcome message
    - Validate API key
    - Run all stages
    - Handle errors gracefully
    """
    display_welcome(
        title="DSPy + Atomic Agents: A Comprehensive Didactic Example",
        subtitle=(
            "This example teaches you WHY combining these frameworks is powerful\n"
            "by walking through three stages with full transparency."
        ),
        details=(
            f"Large benchmark: {len(TRAINING_DATASET)} training examples, "
            f"{len(TEST_DATASET)} challenging test cases\n"
            "We'll expose the prompts, show the optimizations,\n"
            "and compare measurable results."
        ),
    )

    # Validate API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        console.print("[red]Error: OPENAI_API_KEY environment variable required[/red]")
        return

    # Display configuration
    console.print("\n[dim]Using model: gpt-5-mini[/dim]")
    console.print(f"[dim]Training set: {len(TRAINING_DATASET)} examples (balanced across 6 genres)[/dim]")
    console.print(f"[dim]Test set: {len(TEST_DATASET)} challenging examples " "(sarcasm, multi-genre, etc.)[/dim]\n")

    # Run demonstration
    try:
        run_all_stages(api_key)
    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        console.print(traceback.format_exc())


if __name__ == "__main__":
    main()

```

### File: atomic-examples/dspy-integration/dspy_integration/presentation/__init__.py

```python
"""
Presentation layer for DSPy + Atomic Agents integration.

This package handles all console output and visualization using Rich.
Separating presentation from business logic allows:
- Testing business logic without UI dependencies
- Easy swapping of presentation implementation
- Clean separation of concerns

Following Clean Architecture: presentation depends on domain, not vice versa.
"""

from dspy_integration.presentation.console import (
    console,
    display_welcome,
    display_stage_header,
    display_panel,
    display_code,
    display_tree,
    display_results_table,
    display_comparison_table,
    display_takeaways,
    display_decision_guide,
    create_progress_context,
)

__all__ = [
    "console",
    "display_welcome",
    "display_stage_header",
    "display_panel",
    "display_code",
    "display_tree",
    "display_results_table",
    "display_comparison_table",
    "display_takeaways",
    "display_decision_guide",
    "create_progress_context",
]

```

### File: atomic-examples/dspy-integration/dspy_integration/presentation/console.py

```python
"""
Console presentation utilities using Rich.

This module provides a clean API for all console output operations.
All Rich-specific code is encapsulated here, making it easy to swap
to a different presentation library if needed.

Design Principles:
- Encapsulate all Rich dependencies
- Provide high-level semantic functions (display_results, not print_table)
- No business logic - only presentation concerns
"""

from contextlib import contextmanager
from typing import Any, Dict, Generator, List

from rich import box
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.rule import Rule
from rich.syntax import Syntax
from rich.table import Table
from rich.tree import Tree

from dspy_integration.domain.models import EvalResult


# Global console instance
console = Console()


# =============================================================================
# HIGH-LEVEL DISPLAY FUNCTIONS
# =============================================================================


def display_welcome(
    title: str,
    subtitle: str,
    details: str,
) -> None:
    """Display welcome banner for the application."""
    console.print(
        Panel.fit(
            f"[bold]{title}[/bold]\n\n{subtitle}\n\n[dim]{details}[/dim]",
            border_style="bold white",
        )
    )


def display_stage_header(stage_name: str, style: str) -> None:
    """Display a stage header with rule line."""
    console.print(Rule(f"[bold {style}]{stage_name}[/bold {style}]", style=style))


def display_panel(
    content: str,
    title: str,
    border_style: str = "blue",
) -> None:
    """Display a panel with formatted content."""
    console.print(Panel(content, title=title, border_style=border_style))


def display_code(
    code: str,
    language: str = "python",
    theme: str = "monokai",
    line_numbers: bool = True,
) -> None:
    """Display syntax-highlighted code."""
    console.print(Syntax(code, language, theme=theme, line_numbers=line_numbers))


def display_step_header(step: str) -> None:
    """Display a step header within a stage."""
    console.print(f"\n[bold]{step}[/bold]")


def display_success(message: str) -> None:
    """Display a success message."""
    console.print(f"[green]✓ {message}[/green]")


def display_info(message: str) -> None:
    """Display an info message."""
    console.print(f"[dim]{message}[/dim]")


def display_tree(
    title: str,
    items: List[Dict[str, Any]],
) -> None:
    """
    Display a tree structure.

    Args:
        title: Root node title
        items: List of dicts with 'title' and optional 'children' keys
    """
    tree = Tree(f"[bold]{title}[/bold]")

    for item in items:
        branch = tree.add(f"[cyan]{item.get('title', 'Item')}[/cyan]")
        for child in item.get("children", []):
            branch.add(child)

    console.print(tree)


# =============================================================================
# RESULTS DISPLAY FUNCTIONS
# =============================================================================


def display_results_table(
    eval_result: EvalResult,
    title: str,
    show_confidence: bool = False,
) -> None:
    """
    Display evaluation results in a table format.

    Args:
        eval_result: Evaluation results to display
        title: Table title
        show_confidence: Whether to show confidence column
    """
    table = Table(
        title=f"{title}: {eval_result.accuracy:.1%} Accuracy " f"({eval_result.correct}/{eval_result.total})",
        box=box.ROUNDED,
    )

    table.add_column("Review", style="cyan", max_width=40)
    table.add_column("Expected", style="green")
    table.add_column("Predicted", style="yellow")

    if show_confidence:
        table.add_column("Confidence", justify="right")

    table.add_column("✓/✗", justify="center")

    for pred in eval_result.predictions:
        row = [
            pred["review"],
            pred["expected"],
            pred["predicted"],
        ]

        if show_confidence:
            row.append(f"{pred['confidence']:.2f}")

        row.append("[green]✓[/green]" if pred["correct"] else "[red]✗[/red]")
        table.add_row(*row)

    console.print(table)


def display_comparison_table(
    stage1_result: EvalResult,
    stage2_result: EvalResult,
    stage3_result: EvalResult,
) -> None:
    """Display side-by-side comparison of all three approaches."""
    table = Table(title="Approach Comparison", box=box.DOUBLE_EDGE)

    table.add_column("Metric", style="bold")
    table.add_column("Stage 1\nRaw DSPy", justify="center", style="blue")
    table.add_column("Stage 2\nRaw Atomic Agents", justify="center", style="magenta")
    table.add_column("Stage 3\nDSPy + Atomic", justify="center", style="green")

    # Accuracy row
    table.add_row(
        "Accuracy",
        f"{stage1_result.accuracy:.1%}",
        f"{stage2_result.accuracy:.1%}",
        f"[bold]{stage3_result.accuracy:.1%}[/bold]",
    )

    # Correct/Total row
    table.add_row(
        "Correct / Total",
        f"{stage1_result.correct}/{stage1_result.total}",
        f"{stage2_result.correct}/{stage2_result.total}",
        f"[bold]{stage3_result.correct}/{stage3_result.total}[/bold]",
    )

    # Time row
    table.add_row(
        "Avg Time/Query",
        f"{stage1_result.avg_time:.2f}s",
        f"{stage2_result.avg_time:.2f}s",
        f"{stage3_result.avg_time:.2f}s",
    )

    # Feature comparison rows
    _add_feature_rows(table)

    console.print(table)


def _add_feature_rows(table: Table) -> None:
    """Add feature comparison rows to the table."""
    features = [
        (
            "Prompt Optimization",
            "[green]✓ Auto[/green]",
            "[red]✗ Manual[/red]",
            "[green]✓ Auto[/green]",
        ),
        (
            "Type Safety",
            "[yellow]~ DSPy Literal[/yellow]",
            "[green]✓ Pydantic[/green]",
            "[green]✓ Pydantic[/green]",
        ),
        (
            "Output Validation",
            "[yellow]~ Basic[/yellow]",
            "[green]✓ Full[/green]",
            "[green]✓ Full[/green]",
        ),
        (
            "Pydantic Ecosystem",
            "[red]✗ No[/red]",
            "[green]✓ Full[/green]",
            "[green]✓ Full[/green]",
        ),
        (
            "Few-Shot Selection",
            "[green]✓ Auto[/green]",
            "[red]✗ Manual[/red]",
            "[green]✓ Auto[/green]",
        ),
        (
            "IDE Support",
            "[yellow]~ Partial[/yellow]",
            "[green]✓ Full[/green]",
            "[green]✓ Full[/green]",
        ),
    ]

    for feature in features:
        table.add_row(*feature)


# =============================================================================
# SUMMARY DISPLAY FUNCTIONS
# =============================================================================


def display_takeaways() -> None:
    """Display key takeaways panel."""
    content = """[bold yellow]KEY TAKEAWAYS[/bold yellow]

[blue]RAW DSPy (with typed signatures):[/blue]
• Excellent optimization with Literal type constraints
• Great for experimentation and iteration
• Missing Pydantic ecosystem (validators, Field constraints)

[magenta]RAW ATOMIC AGENTS:[/magenta]
• Full Pydantic ecosystem with runtime validation
• Instructor integration for robust outputs
• Manual prompt engineering limits optimization

[green]DSPy + ATOMIC AGENTS:[/green]
• Automatic optimization finds the best prompts
• Full Pydantic validation and serialization
• Measurable improvements + production-ready types
• [bold]The best of both worlds![/bold]"""

    console.print(Panel(content, title="Summary", border_style="yellow"))


def display_decision_guide() -> None:
    """Display when-to-use-what guide."""
    content = """[bold]WHEN TO USE EACH APPROACH[/bold]

[blue]Use Raw DSPy when:[/blue]
• Quick prototyping and experimentation
• Output format doesn't matter much
• You'll post-process outputs anyway

[magenta]Use Raw Atomic Agents when:[/magenta]
• You need guaranteed output structure NOW
• You don't have training data for optimization
• The task is simple enough that manual prompts work

[green]Use DSPy + Atomic Agents when:[/green]
• You have labeled data and want to optimize
• Production systems need type-safe outputs
• You want measurable, reproducible improvements
• Both accuracy AND structure matter"""

    console.print(Panel(content, title="Decision Guide", border_style="cyan"))


# =============================================================================
# PROGRESS CONTEXT MANAGER
# =============================================================================


@contextmanager
def create_progress_context(
    description: str,
    style: str = "cyan",
) -> Generator[Progress, None, None]:
    """
    Create a progress context for long-running operations.

    Args:
        description: Task description to display
        style: Color style for the progress text

    Yields:
        Progress object that can be used to update progress

    Example:
        >>> with create_progress_context("Processing...", "green") as progress:
        ...     task = progress.add_task("[green]Working...", total=100)
        ...     for i in range(100):
        ...         progress.advance(task)
    """
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        console=console,
    ) as progress:
        yield progress

```

### File: atomic-examples/dspy-integration/dspy_integration/schemas.py

```python
"""
Pydantic schemas for DSPy + Atomic Agents integration examples.

These schemas demonstrate how to define type-safe input/output contracts
that can be used with both Atomic Agents and DSPy optimization.
"""

from typing import Literal, List, Optional
from pydantic import Field
from atomic_agents.base.base_io_schema import BaseIOSchema


class SentimentInputSchema(BaseIOSchema):
    """Input schema for sentiment analysis task."""

    text: str = Field(
        ...,
        description="The text to analyze for sentiment.",
        min_length=1,
    )


class SentimentOutputSchema(BaseIOSchema):
    """Output schema for sentiment analysis with structured results."""

    sentiment: Literal["positive", "negative", "neutral"] = Field(
        ...,
        description="The overall sentiment of the text.",
    )
    confidence: float = Field(
        ...,
        description="Confidence score between 0 and 1.",
        ge=0.0,
        le=1.0,
    )
    reasoning: str = Field(
        ...,
        description="Brief explanation for the sentiment classification.",
    )


class QuestionInputSchema(BaseIOSchema):
    """Input schema for question answering task."""

    question: str = Field(
        ...,
        description="The question to answer.",
    )
    context: Optional[str] = Field(
        default=None,
        description="Optional context to help answer the question.",
    )


class AnswerOutputSchema(BaseIOSchema):
    """Output schema for question answering with structured response."""

    answer: str = Field(
        ...,
        description="The answer to the question.",
    )
    confidence: float = Field(
        ...,
        description="Confidence score for the answer between 0 and 1.",
        ge=0.0,
        le=1.0,
    )
    sources: List[str] = Field(
        default_factory=list,
        description="List of sources or references used to derive the answer.",
    )


class SummaryInputSchema(BaseIOSchema):
    """Input schema for text summarization task."""

    text: str = Field(
        ...,
        description="The text to summarize.",
    )
    max_sentences: int = Field(
        default=3,
        description="Maximum number of sentences in the summary.",
        ge=1,
        le=10,
    )


class SummaryOutputSchema(BaseIOSchema):
    """Output schema for text summarization with structured results."""

    summary: str = Field(
        ...,
        description="The summarized text.",
    )
    key_points: List[str] = Field(
        ...,
        description="List of key points extracted from the text.",
    )
    word_count: int = Field(
        ...,
        description="Word count of the summary.",
        ge=0,
    )


class ClassificationInputSchema(BaseIOSchema):
    """Input schema for multi-label text classification."""

    text: str = Field(
        ...,
        description="The text to classify.",
    )
    categories: List[str] = Field(
        ...,
        description="Available categories to classify into.",
    )


class ClassificationOutputSchema(BaseIOSchema):
    """Output schema for multi-label classification with confidence scores."""

    labels: List[str] = Field(
        ...,
        description="Assigned labels/categories.",
    )
    label_scores: List[float] = Field(
        ...,
        description="Confidence scores for each assigned label.",
    )
    primary_label: str = Field(
        ...,
        description="The most confident label assignment.",
    )
    reasoning: str = Field(
        ...,
        description="Explanation for the classification decision.",
    )

```

### File: atomic-examples/dspy-integration/dspy_integration/stages/__init__.py

```python
"""
Stages package for DSPy + Atomic Agents integration demo.

Each stage demonstrates a different approach:
- Stage 1: Raw DSPy with typed signatures
- Stage 2: Raw Atomic Agents with manual prompts
- Stage 3: Combined DSPy + Atomic Agents

Following Single Responsibility Principle: each stage module handles
one approach completely, from setup to evaluation.
"""

from dspy_integration.stages.stage1_dspy import run_stage1_raw_dspy
from dspy_integration.stages.stage2_atomic import run_stage2_raw_atomic_agents
from dspy_integration.stages.stage3_combined import run_stage3_combined

__all__ = [
    "run_stage1_raw_dspy",
    "run_stage2_raw_atomic_agents",
    "run_stage3_combined",
]

```

### File: atomic-examples/dspy-integration/dspy_integration/stages/stage1_dspy.py

```python
"""
Stage 1: Raw DSPy with Typed Signatures.

This module demonstrates DSPy's capabilities at their best:
- Typed signatures with Literal constraints
- Automatic prompt optimization via BootstrapFewShot
- Chain-of-thought reasoning

Limitations shown:
- No Pydantic validation ecosystem
- Less integration with structured output tools
- Type enforcement is DSPy-specific, not Python runtime

Design: Single function entry point, internal helpers follow SRP.
"""

import json
import time
from typing import Any, Dict, List, Tuple

import dspy

from dspy_integration.domain.models import (
    GENRES,
    GenreType,
    EvalResult,
)
from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
from dspy_integration.domain.evaluation import evaluate_predictions
from dspy_integration.presentation.console import (
    console,
    display_stage_header,
    display_panel,
    display_code,
    display_step_header,
    display_success,
    display_tree,
    display_results_table,
    create_progress_context,
)


# =============================================================================
# DSPY SIGNATURE DEFINITION
# =============================================================================


class MovieGenreSignature(dspy.Signature):
    """
    Classify a movie review into its primary genre based on the review text.

    Consider the overall focus and tone of the review, not just individual keywords.
    A review mentioning 'explosions' might be a drama if the focus is on characters.
    A 'scary' movie might be a comedy if played for laughs.
    """

    review: str = dspy.InputField(desc="The movie review text to classify")
    genre: GenreType = dspy.OutputField(desc="The primary genre: action, comedy, drama, horror, sci-fi, or romance")
    confidence: float = dspy.OutputField(desc="Confidence score between 0.0 and 1.0")
    reasoning: str = dspy.OutputField(desc="Brief explanation for the classification")


# =============================================================================
# CODE EXAMPLES FOR DISPLAY
# =============================================================================


SIGNATURE_CODE_EXAMPLE = '''from typing import Literal

# DSPy Signature WITH proper type constraints
class MovieGenreSignature(dspy.Signature):
    """Classify a movie review into its primary genre."""

    review: str = dspy.InputField(desc="The movie review text")

    # Literal type constrains output to valid genres only!
    genre: Literal["action", "comedy", "drama", "horror", "sci-fi", "romance"] = \\
        dspy.OutputField(desc="The primary genre")

    confidence: float = dspy.OutputField(desc="Confidence 0.0-1.0")
    reasoning: str = dspy.OutputField(desc="Brief explanation")

# DSPy enforces the Literal constraint - no more "dramedy" or "thriller"!
classify = dspy.ChainOfThought(MovieGenreSignature)'''


# =============================================================================
# MAIN STAGE FUNCTION
# =============================================================================


def run_stage1_raw_dspy(api_key: str) -> Tuple[EvalResult, Dict[str, Any]]:
    """
    Run Stage 1: Raw DSPy demonstration.

    This demonstrates DSPy at its best with proper typed signatures.

    Args:
        api_key: OpenAI API key

    Returns:
        Tuple of (evaluation results, behind-the-scenes data)
    """
    display_stage_header("STAGE 1: Raw DSPy (Properly Implemented)", "blue")
    _display_stage_overview()

    # Configure DSPy
    lm = dspy.LM("openai/gpt-5-mini", api_key=api_key)
    dspy.configure(lm=lm)

    # Step 1: Show signature
    _display_signature_explanation()

    # Step 2: Create classifier and show unoptimized prompt
    classify = dspy.ChainOfThought(MovieGenreSignature)
    unoptimized_prompt = _capture_unoptimized_prompt(lm, classify)

    # Step 3: Explain optimization
    _display_optimization_explanation()

    # Step 4: Run optimization
    optimized_classify = _run_optimization(lm, classify)

    # Step 5: Show optimized prompt
    optimized_prompt = _capture_optimized_prompt(lm, optimized_classify)

    # Step 6: Show selected demos
    _display_selected_demos(optimized_classify)

    # Step 7: Evaluate
    eval_result, predictions = _evaluate_model(optimized_classify)

    # Step 8: Display results
    _display_stage_results(eval_result, predictions)

    behind_scenes = _create_behind_scenes_data(unoptimized_prompt, optimized_prompt, optimized_classify)

    return eval_result, behind_scenes


# =============================================================================
# DISPLAY HELPERS
# =============================================================================


def _display_stage_overview() -> None:
    """Display stage 1 overview panel."""
    content = """[green]DSPy STRENGTHS:[/green]
• Typed signatures with Literal constraints (genre MUST be valid)
• Automatic prompt optimization via BootstrapFewShot
• Chain-of-thought reasoning for complex decisions
• Systematic few-shot example selection

[yellow]LIMITATIONS vs Atomic Agents:[/yellow]
• No Pydantic ecosystem (validators, serializers, etc.)
• Less integration with structured output tools like Instructor
• Type hints are enforced by DSPy, not Python runtime"""

    display_panel(content, "Stage 1 Overview", "blue")


def _display_signature_explanation() -> None:
    """Display explanation of DSPy typed signatures."""
    display_step_header("Step 1.1: Define Typed DSPy Signature")
    console.print("DSPy supports class-based signatures with Python type hints:\n")
    display_code(SIGNATURE_CODE_EXAMPLE)


def _display_optimization_explanation() -> None:
    """Display explanation of how DSPy optimization works."""
    display_step_header("Step 1.3: DSPy Optimization (BootstrapFewShot)")

    content = """[cyan]What BootstrapFewShot does:[/cyan]

1. Takes your labeled training examples
2. Runs the LLM on each to generate 'traces' (reasoning chains)
3. Filters traces that produce correct answers
4. Selects the best traces as few-shot demonstrations
5. Injects these into future prompts automatically

[yellow]Key insight:[/yellow] DSPy doesn't just use your examples verbatim.
It generates NEW reasoning and picks what actually works!"""

    display_panel(content, "How DSPy Optimization Works", "cyan")


# =============================================================================
# PROMPT CAPTURE HELPERS
# =============================================================================


def _capture_unoptimized_prompt(
    lm: dspy.LM,
    classify: dspy.Module,
) -> List[Dict[str, Any]]:
    """Capture the unoptimized prompt from DSPy."""
    display_step_header("Step 1.2: Unoptimized Prompt (What DSPy Generates)")

    with dspy.context(lm=lm):
        _ = classify(review=TRAINING_DATASET[0]["review"])

    unoptimized_prompt = []
    if lm.history:
        last_call = lm.history[-1]
        unoptimized_prompt = last_call.get("messages", [{}])

        content = (
            "[dim]Notice how DSPy includes the Literal type constraint in the prompt:[/dim]\n\n"
            + json.dumps(unoptimized_prompt, indent=2)[:2000]
            + "..."
        )
        display_panel(content, "Unoptimized DSPy Prompt (With Type Constraints)", "yellow")

    return unoptimized_prompt


def _capture_optimized_prompt(
    lm: dspy.LM,
    optimized_classify: dspy.Module,
) -> List[Dict[str, Any]]:
    """Capture the optimized prompt from DSPy."""
    display_step_header("Step 1.4: Optimized Prompt (After DSPy Magic)")

    with dspy.context(lm=lm):
        _ = optimized_classify(review=TEST_DATASET[0]["review"])

    optimized_prompt = []
    if lm.history:
        last_call = lm.history[-1]
        optimized_prompt = last_call.get("messages", [{}])

        prompt_str = json.dumps(optimized_prompt, indent=2)
        truncated = prompt_str[:3500] + ("..." if len(prompt_str) > 3500 else "")

        content = "[dim]Notice the auto-selected few-shot examples with reasoning:[/dim]\n\n" + truncated
        display_panel(content, "Optimized DSPy Prompt (With Auto-Selected Examples)", "green")

    return optimized_prompt


# =============================================================================
# OPTIMIZATION HELPERS
# =============================================================================


def _run_optimization(lm: dspy.LM, classify: dspy.Module) -> dspy.Module:
    """Run DSPy optimization with BootstrapFewShot."""
    # Prepare training set (first 30 examples)
    train_examples = TRAINING_DATASET[:30]
    trainset = [
        dspy.Example(
            review=ex["review"],
            genre=ex["genre"],
            confidence=0.85,
            reasoning=f"This review demonstrates typical {ex['genre']} characteristics.",
        ).with_inputs("review")
        for ex in train_examples
    ]

    def genre_match(example, prediction, trace=None):
        """Metric for optimization - checks if genre matches."""
        pred_genre = str(prediction.genre).lower().strip()
        expected_genre = str(example.genre).lower().strip()
        return pred_genre == expected_genre

    with create_progress_context("[cyan]Running DSPy optimization (30 training examples)...") as progress:
        task = progress.add_task("Optimizing...", total=None)

        optimizer = dspy.BootstrapFewShot(
            metric=genre_match,
            max_bootstrapped_demos=4,
            max_labeled_demos=4,
            max_rounds=1,
        )
        optimized_classify = optimizer.compile(classify, trainset=trainset)
        progress.remove_task(task)

    display_success("Optimization complete!")
    return optimized_classify


def _display_selected_demos(optimized_classify: dspy.Module) -> None:
    """Display the few-shot examples DSPy selected."""
    display_step_header("Step 1.5: Few-Shot Examples DSPy Selected")

    if hasattr(optimized_classify, "demos") and optimized_classify.demos:
        items = []
        for i, demo in enumerate(optimized_classify.demos[:4]):
            review_text = str(getattr(demo, "review", "N/A"))[:70]
            genre = getattr(demo, "genre", "N/A")
            reasoning = str(getattr(demo, "reasoning", ""))[:80]

            items.append(
                {
                    "title": f"Example {i + 1}",
                    "children": [
                        f"Review: {review_text}...",
                        f"Genre: [green]{genre}[/green]",
                        f"Reasoning: [dim]{reasoning}...[/dim]",
                    ],
                }
            )

        display_tree("Selected Demonstrations", items)
    else:
        console.print("[dim]Demo inspection not available for this predictor type[/dim]")


# =============================================================================
# EVALUATION HELPERS
# =============================================================================


def _evaluate_model(
    optimized_classify: dspy.Module,
) -> Tuple[EvalResult, List[Dict[str, Any]]]:
    """Evaluate the optimized model on test set."""
    display_step_header(f"Step 1.6: Evaluation on Test Set ({len(TEST_DATASET)} challenging examples)")

    predictions = []
    start_time = time.time()

    with create_progress_context("[cyan]Running predictions...") as progress:
        task = progress.add_task("Predicting...", total=len(TEST_DATASET))

        for test_ex in TEST_DATASET:
            prediction = _get_single_prediction(optimized_classify, test_ex)
            predictions.append(prediction)
            progress.advance(task)

    elapsed = time.time() - start_time
    eval_result = evaluate_predictions(predictions, TEST_DATASET)
    eval_result.avg_time = elapsed / len(TEST_DATASET)

    return eval_result, predictions


def _get_single_prediction(
    classifier: dspy.Module,
    test_example: Dict[str, str],
) -> Dict[str, Any]:
    """Get a single prediction from the classifier."""
    try:
        result = classifier(review=test_example["review"])
        genre_val = str(result.genre).strip().lower()

        # Validate genre
        if genre_val not in GENRES:
            genre_val = "error"

        return {
            "genre": genre_val,
            "confidence": float(result.confidence) if hasattr(result, "confidence") else 0.5,
            "reasoning": str(result.reasoning) if hasattr(result, "reasoning") else "N/A",
        }
    except Exception as e:
        return {
            "genre": "error",
            "confidence": 0,
            "reasoning": str(e),
        }


# =============================================================================
# RESULTS DISPLAY
# =============================================================================


def _display_stage_results(
    eval_result: EvalResult,
    predictions: List[Dict[str, Any]],
) -> None:
    """Display stage 1 results and analysis."""
    display_step_header("Step 1.7: Results")

    # Count invalid genres
    invalid_genres = [p["genre"] for p in predictions if p["genre"] not in GENRES]

    content = f"""[green]DSPy TYPED SIGNATURE BENEFITS:[/green]
• Genre constrained to valid options (invalid outputs: {len(invalid_genres)})
• Automatic few-shot example selection
• Chain-of-thought reasoning included

[yellow]REMAINING LIMITATIONS:[/yellow]
• No Pydantic validation ecosystem
• Confidence not guaranteed to be 0-1 (no ge/le constraints)
• Can't use Instructor's retry mechanisms
• Type enforcement is DSPy-specific, not Python-native"""

    display_panel(content, "DSPy Typed Signatures Assessment", "blue")
    display_results_table(eval_result, "Stage 1 Results")


def _create_behind_scenes_data(
    unoptimized_prompt: List[Dict[str, Any]],
    optimized_prompt: List[Dict[str, Any]],
    optimized_classify: dspy.Module,
) -> Dict[str, Any]:
    """Create behind-the-scenes data for comparison."""
    return {
        "unoptimized_prompt_sample": str(unoptimized_prompt)[:500],
        "optimized_prompt_sample": str(optimized_prompt)[:500],
        "num_demos_selected": (len(optimized_classify.demos) if hasattr(optimized_classify, "demos") else "N/A"),
        "training_examples": 30,
    }

```

### File: atomic-examples/dspy-integration/dspy_integration/stages/stage2_atomic.py

```python
"""
Stage 2: Raw Atomic Agents with Manual Prompts.

This module demonstrates Atomic Agents' capabilities:
- Full Pydantic ecosystem with runtime validation
- Instructor integration for robust structured outputs
- Guaranteed schema compliance

Limitations shown:
- Manual prompt engineering (guesswork)
- No systematic way to improve prompts
- No automatic few-shot selection

Design: Single function entry point, internal helpers follow SRP.
"""

import time
from typing import Any, Dict, List, Tuple

import instructor
import openai

from atomic_agents.agents.atomic_agent import AgentConfig, AtomicAgent
from atomic_agents.context.system_prompt_generator import SystemPromptGenerator

from dspy_integration.domain.models import (
    GENRES,
    MovieGenreOutput,
    MovieReviewInput,
    EvalResult,
)
from dspy_integration.domain.datasets import TEST_DATASET
from dspy_integration.domain.evaluation import evaluate_predictions
from dspy_integration.presentation.console import (
    console,
    display_stage_header,
    display_panel,
    display_code,
    display_step_header,
    display_success,
    display_results_table,
    create_progress_context,
)


# =============================================================================
# CODE EXAMPLES FOR DISPLAY
# =============================================================================


SCHEMA_CODE_EXAMPLE = '''class MovieGenreOutput(BaseIOSchema):
    """Output schema for movie genre classification."""

    genre: Literal["action", "comedy", "drama", "horror", "sci-fi", "romance"] = Field(
        ...,
        description="The primary genre of the movie.",
    )
    confidence: float = Field(
        ...,
        ge=0.0, le=1.0,  # VALIDATED! Must be between 0 and 1
        description="Confidence score between 0.0 and 1.0",
    )
    reasoning: str = Field(
        ...,
        description="Brief explanation for the classification.",
    )

# The LLM output MUST match this schema or it fails validation.
# No more parsing "high" vs "0.85" vs "85%" - it's always a float!'''


# =============================================================================
# MAIN STAGE FUNCTION
# =============================================================================


def run_stage2_raw_atomic_agents(api_key: str) -> Tuple[EvalResult, Dict[str, Any]]:
    """
    Run Stage 2: Raw Atomic Agents demonstration.

    This demonstrates Atomic Agents' beautiful structured outputs,
    but with manual prompt engineering.

    Args:
        api_key: OpenAI API key

    Returns:
        Tuple of (evaluation results, behind-the-scenes data)
    """
    display_stage_header("STAGE 2: Raw Atomic Agents", "magenta")
    _display_stage_overview()

    # Step 1: Show Pydantic schema
    _display_schema_explanation()

    # Step 2: Show manual system prompt
    system_prompt = _create_system_prompt()
    generated_prompt = system_prompt.generate_prompt()
    _display_manual_prompt(generated_prompt)
    _display_manual_prompt_problem()

    # Step 3: Create agent
    agent = _create_agent(api_key, system_prompt)

    # Step 4: Show schema enforcement
    _display_schema_enforcement()

    # Step 5: Evaluate
    eval_result, predictions = _evaluate_agent(agent)

    # Step 6: Display results
    _display_stage_results(eval_result, predictions)

    behind_scenes = {
        "system_prompt": generated_prompt,
        "schema_enforced": True,
        "manual_engineering": True,
    }

    return eval_result, behind_scenes


# =============================================================================
# DISPLAY HELPERS
# =============================================================================


def _display_stage_overview() -> None:
    """Display stage 2 overview panel."""
    content = """[green]ATOMIC AGENTS STRENGTHS:[/green]
• Full Pydantic ecosystem (validators, serializers, Field constraints)
• Instructor integration for robust structured output
• Python-native type safety with runtime validation
• ge/le constraints on confidence (guaranteed 0-1)

[yellow]LIMITATIONS:[/yellow]
• Manual prompt engineering - no automatic optimization
• No systematic few-shot example selection
• Prompt improvements require guesswork and iteration"""

    display_panel(content, "Stage 2 Overview", "magenta")


def _display_schema_explanation() -> None:
    """Display explanation of Pydantic schemas."""
    display_step_header("Step 2.1: Define Pydantic Schema")
    console.print("Atomic Agents uses Pydantic for type-safe outputs:\n")
    display_code(SCHEMA_CODE_EXAMPLE)


def _display_manual_prompt(generated_prompt: str) -> None:
    """Display the manually crafted system prompt."""
    display_step_header("Step 2.2: Manual System Prompt (The Guesswork)")

    content = "[dim]This is the system prompt WE WROTE BY HAND:[/dim]\n\n" + generated_prompt
    display_panel(content, "Manual System Prompt (Our Best Guess)", "yellow")


def _display_manual_prompt_problem() -> None:
    """Display the problem with manual prompt engineering."""
    content = """[red]THE PROBLEM:[/red]

We wrote this prompt based on intuition. Questions we can't answer:

• Is 'Be decisive' helping or hurting accuracy?
• Should we add few-shot examples? Which ones?
• Is the step-by-step instruction actually useful?
• Would different wording improve results?

[yellow]Without DSPy, we're just guessing![/yellow]"""

    display_panel(content, "The Manual Prompt Engineering Problem", "red")


def _display_schema_enforcement() -> None:
    """Display how schema enforcement works."""
    display_step_header("Step 2.4: Schema Enforcement in Action")

    content = """[cyan]What happens under the hood:[/cyan]

1. Atomic Agents sends your prompt + Pydantic schema to the LLM
2. Instructor (the library) converts schema to JSON Schema for the LLM
3. LLM generates output attempting to match the schema
4. Instructor validates the response against Pydantic
5. If validation fails, Instructor retries with error feedback
6. You get a guaranteed-valid Pydantic object or an exception

[green]Result:[/green] genre is ALWAYS one of our 6 options,
confidence is ALWAYS a float between 0 and 1!"""

    display_panel(content, "How Schema Enforcement Works", "cyan")


# =============================================================================
# AGENT CREATION
# =============================================================================


def _create_system_prompt() -> SystemPromptGenerator:
    """Create the manually crafted system prompt."""
    return SystemPromptGenerator(
        background=[
            "You are a movie genre classification expert.",
            "You analyze movie reviews and determine the primary genre.",
            f"Valid genres are: {', '.join(GENRES)}",
        ],
        steps=[
            "Read the review carefully.",
            "Identify key genre indicators (action words, emotional language, etc.).",
            "Consider the overall tone and subject matter.",
            "Select the single most appropriate genre.",
            "Provide a confidence score based on how clear the genre signals are.",
        ],
        output_instructions=[
            "Be decisive - pick ONE primary genre even if multiple could apply.",
            "Confidence should be 0.7-1.0 for clear cases, 0.5-0.7 for ambiguous ones.",
            "Keep reasoning brief but specific to the review.",
        ],
    )


def _create_agent(
    api_key: str,
    system_prompt: SystemPromptGenerator,
) -> AtomicAgent:
    """Create the Atomic Agent with schema validation."""
    display_step_header("Step 2.3: Create Atomic Agent")

    client = instructor.from_openai(openai.OpenAI(api_key=api_key))

    agent = AtomicAgent[MovieReviewInput, MovieGenreOutput](
        config=AgentConfig(
            client=client,
            model="gpt-5-mini",
            system_prompt_generator=system_prompt,
        )
    )

    display_success("Agent created with schema validation")
    return agent


# =============================================================================
# EVALUATION HELPERS
# =============================================================================


def _evaluate_agent(
    agent: AtomicAgent,
) -> Tuple[EvalResult, List[Dict[str, Any]]]:
    """Evaluate the agent on test set."""
    display_step_header("Step 2.5: Evaluation on Test Set")

    predictions = []
    start_time = time.time()

    with create_progress_context("[magenta]Running predictions...") as progress:
        task = progress.add_task("Predicting...", total=len(TEST_DATASET))

        for test_ex in TEST_DATASET:
            prediction = _get_single_prediction(agent, test_ex)
            predictions.append(prediction)
            progress.advance(task)

    elapsed = time.time() - start_time
    eval_result = evaluate_predictions(predictions, TEST_DATASET)
    eval_result.avg_time = elapsed / len(TEST_DATASET)

    return eval_result, predictions


def _get_single_prediction(
    agent: AtomicAgent,
    test_example: Dict[str, str],
) -> Dict[str, Any]:
    """Get a single prediction from the agent."""
    try:
        result = agent.run(MovieReviewInput(review=test_example["review"]))
        return {
            "genre": result.genre,  # Already validated by Pydantic!
            "confidence": result.confidence,  # Already a float!
            "reasoning": result.reasoning,
        }
    except Exception as e:
        return {
            "genre": "error",
            "confidence": 0,
            "reasoning": str(e),
        }


# =============================================================================
# RESULTS DISPLAY
# =============================================================================


def _display_stage_results(
    eval_result: EvalResult,
    predictions: List[Dict[str, Any]],
) -> None:
    """Display stage 2 results and analysis."""
    display_step_header("Step 2.6: The Benefit - Type-Safe Outputs")

    # Show sample outputs
    samples = "\n".join(
        [
            f"  • genre='{predictions[i]['genre']}' (Literal) " f"confidence={predictions[i]['confidence']:.2f} (float)"
            for i in range(min(3, len(predictions)))
        ]
    )

    content = f"""[green]ATOMIC AGENTS ADVANTAGE:[/green]

Look at these outputs - perfectly structured:
{samples}

[cyan]Benefits:[/cyan]
• genre is guaranteed to be one of our 6 valid options
• confidence is always a float between 0.0 and 1.0
• No parsing needed - direct attribute access
• IDE autocomplete works perfectly
• Downstream code can trust the types"""

    display_panel(content, "Structured Output Benefits", "green")
    display_results_table(eval_result, "Stage 2 Results", show_confidence=True)

```

### File: atomic-examples/dspy-integration/dspy_integration/stages/stage3_combined.py

```python
"""
Stage 3: DSPy + Atomic Agents Combined.

This module demonstrates the best of both worlds:
- DSPy's automatic prompt optimization
- Atomic Agents' type-safe structured outputs

The bridge module connects both frameworks, enabling:
- Pydantic schemas as DSPy signatures
- DSPy optimizers for Atomic Agents
- Validated, optimized outputs

Design: Single function entry point, internal helpers follow SRP.
"""

import json
import time
from typing import Any, Dict, List, Tuple

import dspy

from dspy_integration.bridge import DSPyAtomicModule, create_dspy_example
from dspy_integration.domain.models import (
    MovieGenreOutput,
    MovieReviewInput,
    EvalResult,
)
from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
from dspy_integration.domain.evaluation import evaluate_predictions
from dspy_integration.presentation.console import (
    display_stage_header,
    display_panel,
    display_code,
    display_step_header,
    display_success,
    display_results_table,
    create_progress_context,
)


# =============================================================================
# CODE EXAMPLES FOR DISPLAY
# =============================================================================


BRIDGE_CODE_EXAMPLE = """# The bridge combines both frameworks:
module = DSPyAtomicModule(
    input_schema=MovieReviewInput,    # Pydantic input validation
    output_schema=MovieGenreOutput,   # Pydantic output structure
    instructions="Classify the movie review into a genre.",
    use_chain_of_thought=True,        # DSPy's reasoning capability
)

# Behind the scenes:
# 1. Pydantic schemas are converted to DSPy signatures
# 2. DSPy handles prompt construction and optimization
# 3. Outputs are validated against Pydantic schemas
# 4. You get type-safe results that DSPy optimized!"""


# =============================================================================
# MAIN STAGE FUNCTION
# =============================================================================


def run_stage3_combined(api_key: str) -> Tuple[EvalResult, Dict[str, Any]]:
    """
    Run Stage 3: Combined DSPy + Atomic Agents demonstration.

    This demonstrates the best of both worlds - DSPy optimization
    with Atomic Agents type safety.

    Args:
        api_key: OpenAI API key

    Returns:
        Tuple of (evaluation results, behind-the-scenes data)
    """
    display_stage_header("STAGE 3: DSPy + Atomic Agents", "green")
    _display_stage_overview()

    # Configure DSPy
    lm = dspy.LM("openai/gpt-5-mini", api_key=api_key)
    dspy.configure(lm=lm)

    # Step 1: Show bridge module
    _display_bridge_explanation()

    # Step 2: Create module
    module = _create_bridge_module()

    # Step 3: Show schema conversion
    _display_schema_conversion()

    # Step 4: Create training examples
    trainset = _create_training_set()

    # Step 5: Run optimization
    optimized_module = _run_optimization(module, trainset)

    # Step 6: Show optimized prompt
    optimized_prompt = _capture_optimized_prompt(lm, optimized_module)

    # Step 7: Evaluate
    eval_result, predictions = _evaluate_module(optimized_module)

    # Step 8: Display results
    _display_stage_results(eval_result)

    behind_scenes = {
        "optimized_prompt_sample": optimized_prompt[:1000] if optimized_prompt else "N/A",
        "schema_enforced": True,
        "dspy_optimized": True,
    }

    return eval_result, behind_scenes


# =============================================================================
# DISPLAY HELPERS
# =============================================================================


def _display_stage_overview() -> None:
    """Display stage 3 overview panel."""
    content = """[green]THE SOLUTION:[/green]
Combine DSPy's automatic optimization with Atomic Agents' type safety!

[cyan]WHAT WE GET:[/cyan]
• DSPy automatically finds the best prompts and examples
• Atomic Agents guarantees output structure
• Measurable improvements through optimization
• Production-ready typed outputs

[yellow]THE BEST OF BOTH WORLDS[/yellow]"""

    display_panel(content, "Stage 3 Overview", "green")


def _display_bridge_explanation() -> None:
    """Display explanation of the bridge module."""
    display_step_header("Step 3.1: The Bridge - DSPyAtomicModule")
    display_code(BRIDGE_CODE_EXAMPLE)


def _display_schema_conversion() -> None:
    """Display how schemas are converted to signatures."""
    display_step_header("Step 3.2: Schema-to-Signature Conversion")

    content = """[cyan]Pydantic Schema → DSPy Signature:[/cyan]

Input fields: review (str)
Output fields: genre (Literal), confidence (float), reasoning (str)

[dim]The bridge automatically converts Pydantic field descriptions
into DSPy field descriptors, preserving all metadata.[/dim]"""

    display_panel(content, "Automatic Conversion", "cyan")


def _display_training_explanation() -> None:
    """Display explanation of type-safe training examples."""
    display_step_header("Step 3.3: Type-Safe Training Examples")

    content = """[cyan]Creating training examples with validation:[/cyan]

Each example is validated against our Pydantic schemas!
If you accidentally put confidence=1.5 or genre='thriller',
you get an immediate error - not a silent failure later."""

    display_panel(content, "Validated Training Data", "cyan")


# =============================================================================
# MODULE CREATION
# =============================================================================


def _create_bridge_module() -> DSPyAtomicModule:
    """Create the DSPy-Atomic bridge module."""
    return DSPyAtomicModule(
        input_schema=MovieReviewInput,
        output_schema=MovieGenreOutput,
        instructions="Classify the movie review into its primary genre. Be accurate and provide reasoning.",
        use_chain_of_thought=True,
    )


def _create_training_set() -> List[dspy.Example]:
    """Create validated training examples."""
    _display_training_explanation()

    # Use 40 examples for training
    train_examples = TRAINING_DATASET[:40]
    trainset = []

    for ex in train_examples:
        trainset.append(
            create_dspy_example(
                MovieReviewInput,
                MovieGenreOutput,
                {"review": ex["review"]},
                {
                    "genre": ex["genre"],
                    "confidence": 0.85,
                    "reasoning": f"The review shows typical {ex['genre']} characteristics.",
                },
            )
        )

    display_success(f"Created {len(trainset)} validated training examples")
    return trainset


# =============================================================================
# OPTIMIZATION HELPERS
# =============================================================================


def _run_optimization(
    module: DSPyAtomicModule,
    trainset: List[dspy.Example],
) -> DSPyAtomicModule:
    """Run DSPy optimization on the bridge module."""
    display_step_header("Step 3.4: DSPy Optimization (With Schema Awareness)")

    def typed_genre_match(example, prediction, trace=None):
        """Metric that works with typed outputs."""
        pred_genre = str(prediction.genre).lower().strip()
        expected_genre = str(example.genre).lower().strip()
        return pred_genre == expected_genre

    with create_progress_context(f"[green]Running optimization ({len(trainset)} training examples)...") as progress:
        task = progress.add_task("Optimizing...", total=None)

        optimizer = dspy.BootstrapFewShot(
            metric=typed_genre_match,
            max_bootstrapped_demos=4,
            max_labeled_demos=4,
            max_rounds=1,
        )
        optimized_module = optimizer.compile(module, trainset=trainset)
        progress.remove_task(task)

    display_success("Optimization complete!")
    return optimized_module


def _capture_optimized_prompt(
    lm: dspy.LM,
    optimized_module: DSPyAtomicModule,
) -> str:
    """Capture the optimized prompt."""
    display_step_header("Step 3.5: The Optimized Prompt (Exposed!)")

    with dspy.context(lm=lm):
        _ = optimized_module(review=TEST_DATASET[0]["review"])

    prompt_str = ""
    if lm.history:
        last_call = lm.history[-1]
        optimized_prompt = last_call.get("messages", [{}])
        prompt_str = json.dumps(optimized_prompt, indent=2)

        truncated = prompt_str[:2500] + ("..." if len(prompt_str) > 2500 else "")
        content = "[dim]This is what DSPy + Atomic Agents sends to the LLM:[/dim]\n\n" + truncated
        display_panel(content, "Final Optimized Prompt", "green")

    return prompt_str


# =============================================================================
# EVALUATION HELPERS
# =============================================================================


def _evaluate_module(
    optimized_module: DSPyAtomicModule,
) -> Tuple[EvalResult, List[Dict[str, Any]]]:
    """Evaluate the optimized module on test set."""
    display_step_header("Step 3.6: Evaluation with Type-Safe Outputs")

    predictions = []
    start_time = time.time()

    with create_progress_context("[green]Running predictions...") as progress:
        task = progress.add_task("Predicting...", total=len(TEST_DATASET))

        for test_ex in TEST_DATASET:
            prediction = _get_single_prediction(optimized_module, test_ex)
            predictions.append(prediction)
            progress.advance(task)

    elapsed = time.time() - start_time
    eval_result = evaluate_predictions(predictions, TEST_DATASET)
    eval_result.avg_time = elapsed / len(TEST_DATASET)

    return eval_result, predictions


def _get_single_prediction(
    module: DSPyAtomicModule,
    test_example: Dict[str, str],
) -> Dict[str, Any]:
    """Get a single validated prediction."""
    try:
        # Use run_validated to get Pydantic-validated output
        validated_result = module.run_validated(review=test_example["review"])

        return {
            "genre": validated_result.genre,  # Guaranteed Literal type!
            "confidence": validated_result.confidence,  # Guaranteed 0-1 float!
            "reasoning": validated_result.reasoning,
        }
    except Exception as e:
        return {
            "genre": "error",
            "confidence": 0,
            "reasoning": str(e),
        }


# =============================================================================
# RESULTS DISPLAY
# =============================================================================


def _display_stage_results(eval_result: EvalResult) -> None:
    """Display stage 3 results and analysis."""
    display_step_header("Step 3.7: The Combined Benefits")

    content = """[green]✓ DSPy BENEFITS:[/green]
• Automatic few-shot example selection
• Optimized prompt instructions
• Chain-of-thought reasoning
• Measurable improvement through metrics

[green]✓ ATOMIC AGENTS BENEFITS:[/green]
• genre is Literal['action','comedy',...] - always valid
• confidence is float with ge=0, le=1 - always in range
• Full IDE autocomplete and type checking
• Pydantic validation catches any LLM mistakes

[yellow]COMBINED:[/yellow] Optimized prompts + Guaranteed structure!"""

    display_panel(content, "The Best of Both Worlds", "green")
    display_results_table(eval_result, "Stage 3 Results", show_confidence=True)

```

### File: atomic-examples/dspy-integration/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["dspy_integration"]

[project]
name = "dspy-integration"
version = "1.0.0"
description = "DSPy + Atomic Agents integration example - combining prompt optimization with type-safe structured outputs"
readme = "README.md"
authors = [
    { name = "BrainBlend AI", email = "kenny@brainblendai.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "dspy>=2.5.0",
    "instructor>=1.7.0",
    "openai>=1.50.0",
    "python-dotenv>=1.0.1",
    "rich>=13.7.0",
    "pydantic>=2.0.0",
]

[dependency-groups]
dev = [
    "black>=24.10.0",
    "flake8>=7.3.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: fastapi-memory
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/fastapi-memory

## Documentation

# FastAPI with Atomic Agents

A comprehensive example demonstrating how to integrate Atomic Agents with FastAPI for building multi-user, multi-session conversational APIs.

## Features

- **Multi-user support**: Each user can have multiple independent chat sessions
- **Conversation history**: Full conversation history is stored and restored when you return to a session
- **User ID persistence**: Client automatically generates and stores a persistent user ID
- **Auto-generated session IDs**: Sessions are created with UUIDs - no manual IDs needed
- **Session management**: View, create, and delete sessions per user
- **RESTful API**: Clean endpoints for chat and session management
- **Interactive CLI client**: Rich terminal interface with session selection
- **Streaming support**: Both standard and streaming chat responses
- **Type safety**: Pydantic schemas for request/response validation

## Setup

1. Install dependencies:
```bash
uv sync
```

2. Set your OpenAI API key:
```bash
export OPENAI_API_KEY="your-api-key-here"
```

Or create a `.env` file in the project root:
```
OPENAI_API_KEY=your_openai_api_key
```

## Running the Example

### Option 1: Interactive Client (Recommended)

Start the server:
```bash
uv run python fastapi_memory/main.py
```

In a separate terminal, run the interactive client:
```bash
uv run python fastapi_memory/client.py
```

The client will:
1. Auto-generate and persist a user ID (stored in `~/.fastapi_memory_user_id`)
2. Show your existing chat sessions or prompt you to create one
3. Load full conversation history when you select an existing session
4. Let you chat in streaming or non-streaming mode (type `/exit` to go back)
5. Manage your sessions (view/delete)

### Option 2: Direct API Usage

Start the FastAPI server:
```bash
uv run python fastapi_memory/main.py
```

The API will be available at `http://localhost:8000`.

## API Documentation

Once running, visit:
- Interactive API docs: `http://localhost:8000/docs`
- Alternative docs: `http://localhost:8000/redoc`

## API Usage Examples

### 1. Create a new session for a user:
```bash
curl -X POST "http://localhost:8000/users/user123/sessions"
```
Response:
```json
{
  "session_id": "550e8400-e29b-41d4-a716-446655440000",
  "message": "Session created successfully"
}
```

### 2. Get all sessions for a user:
```bash
curl "http://localhost:8000/users/user123/sessions"
```
Response:
```json
{
  "user_id": "user123",
  "sessions": [
    {
      "session_id": "550e8400-e29b-41d4-a716-446655440000",
      "created_at": "2025-01-23T10:30:00"
    }
  ]
}
```

### 3. Send a chat message:
```bash
curl -X POST "http://localhost:8000/chat" \
  -H "Content-Type: application/json" \
  -d '{
    "message": "Hello, how are you?",
    "user_id": "user123",
    "session_id": "550e8400-e29b-41d4-a716-446655440000"
  }'
```

### 4. Get conversation history for a session:
```bash
curl "http://localhost:8000/users/user123/sessions/550e8400-e29b-41d4-a716-446655440000/history"
```
Response:
```json
{
  "session_id": "550e8400-e29b-41d4-a716-446655440000",
  "messages": [
    {
      "role": "user",
      "content": "Hello, how are you?",
      "timestamp": "2025-01-23T10:31:00"
    },
    {
      "role": "assistant",
      "content": "I'm doing well, thank you for asking!",
      "timestamp": "2025-01-23T10:31:02",
      "suggested_questions": [
        "What can you do?",
        "Tell me a joke",
        "How does this work?"
      ]
    }
  ]
}
```

### 5. Delete a session:
```bash
curl -X DELETE "http://localhost:8000/users/user123/sessions/550e8400-e29b-41d4-a716-446655440000"
```

### 6. Test the API:
```bash
uv run python test_api.py
```

## How It Works

The example demonstrates several key architectural patterns:

### Server Architecture

1. **Multi-User Session Management**:
   - Data structure: `user_id → session_id → agent_instance`
   - Each user can have unlimited independent chat sessions
   - Sessions are isolated - no data leakage between users or sessions

2. **Conversation History Storage**:
   - All messages are stored with timestamps
   - Separate storage: `user_id → session_id → messages[]`
   - History persists across client reconnections
   - Automatically loaded when resuming a session

3. **Auto-Generated Session IDs**:
   - Server generates UUIDs for new sessions
   - Eliminates user input errors and collisions
   - Tracked with creation timestamps

4. **Lazy Initialization**:
   - Agent instances created on-demand when first accessed
   - Reduces memory footprint for inactive sessions
   - Conversation history maintained independently

5. **Proper Lifecycle Management**:
   - Lifespan context manager ensures cleanup on shutdown
   - Memory released when sessions are deleted
   - History cleared along with session deletion

6. **Type Safety**:
   - Pydantic schemas validate all requests/responses
   - Clear API contracts with automatic documentation

### Client Architecture

1. **User ID Persistence**:
   - Client generates a UUID on first run
   - Stored in `~/.fastapi_memory_user_id`
   - Reused across sessions for continuity

2. **Session Discovery**:
   - Fetches user's sessions from server on startup
   - Displays sessions with creation timestamps
   - Allows selection or creation of new sessions

3. **Conversation History Loading**:
   - Automatically fetches history when loading a session
   - Displays full conversation context before continuing
   - Seamlessly resume conversations from where you left off

4. **Rich Terminal UI**:
   - Interactive menus with Rich library
   - Streaming and non-streaming chat modes
   - Session management interface
   - Type `/exit` to return to menu (not Escape)

## Project Structure

```
fastapi-memory/
├── pyproject.toml              # Project dependencies
├── .env.example                # Environment variable template
├── README.md                   # This file
├── test_api.py                 # API testing script
└── fastapi_memory/
    ├── __init__.py
    ├── main.py                 # FastAPI server
    ├── client.py               # Interactive CLI client
    └── lib/
        ├── agents/
        │   └── chat_agent.py   # Agent configuration
        ├── config.py           # Configuration constants
        └── schemas.py          # Pydantic schemas
```

## Related Examples

For more advanced usage, check out:
- `mcp-agent/example-client/example_client/main_fastapi.py` - Advanced example with MCP protocol integration


## Source Code


### File: atomic-examples/fastapi-memory/fastapi_memory/__init__.py

```python
"""FastAPI Atomic Agents example - Conversational AI with session management."""

__version__ = "1.0.0"

```

### File: atomic-examples/fastapi-memory/fastapi_memory/client.py

```python
"""Interactive command-line client for the FastAPI Atomic Agents example.

This client provides a user-friendly interface to interact with the FastAPI
chat server, supporting both streaming and non-streaming modes, as well as
session management capabilities.
"""

import asyncio
import json
import os
import uuid
from pathlib import Path
from typing import List, Optional

import httpx
from rich.console import Console
from rich.live import Live
from rich.panel import Panel
from rich.prompt import Prompt
from rich.table import Table
from rich.text import Text

console = Console()

# Configuration
BASE_URL = os.getenv("FASTAPI_URL", "http://localhost:8000")
REQUEST_TIMEOUT = 30.0
USER_ID_FILE = Path.home() / ".fastapi_memory_user_id"


def get_or_create_user_id() -> str:
    """Get existing user ID from file or create a new one.

    Returns:
        User identifier (UUID)
    """
    if USER_ID_FILE.exists():
        user_id = USER_ID_FILE.read_text().strip()
        if user_id:
            return user_id

    # Generate new user ID
    user_id = str(uuid.uuid4())
    USER_ID_FILE.write_text(user_id)
    console.print(f"[dim]Created new user ID: {user_id}[/dim]\n")
    return user_id


def _fetch_user_sessions(user_id: str) -> Optional[List[dict]]:
    """Fetch the list of sessions for the current user.

    Args:
        user_id: User identifier

    Returns:
        List of session dicts with 'session_id' and 'created_at', or None if request failed
    """
    try:
        response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions", timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        data = response.json()
        return data.get("sessions", [])
    except Exception as e:
        console.print(f"[bold red]Error fetching sessions:[/bold red] {str(e)}")
        return None


def _create_new_session(user_id: str) -> Optional[str]:
    """Create a new session for the user.

    Args:
        user_id: User identifier

    Returns:
        New session ID or None if creation failed
    """
    try:
        response = httpx.post(f"{BASE_URL}/users/{user_id}/sessions", timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        data = response.json()
        return data.get("session_id")
    except Exception as e:
        console.print(f"[bold red]Error creating session:[/bold red] {str(e)}")
        return None


def _delete_session(user_id: str, session_id: str) -> bool:
    """Delete a session.

    Args:
        user_id: User identifier
        session_id: Session identifier to delete

    Returns:
        True if successful, False otherwise
    """
    try:
        response = httpx.delete(f"{BASE_URL}/users/{user_id}/sessions/{session_id}", timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        return True
    except httpx.HTTPStatusError as e:
        if e.response.status_code == 404:
            console.print("\n[bold red]✗ Session not found[/bold red]")
        else:
            console.print(f"\n[bold red]HTTP Error {e.response.status_code}:[/bold red] {str(e)}")
        return False
    except Exception as e:
        console.print(f"\n[bold red]Error:[/bold red] {str(e)}")
        return False


def select_or_create_session(user_id: str) -> Optional[str]:
    """Show user's sessions and let them select one or create new.

    Args:
        user_id: User identifier

    Returns:
        Selected or newly created session ID, or None if cancelled
    """
    console.clear()
    console.print(
        Panel.fit(
            "[bold magenta]Session Selection[/bold magenta]",
            border_style="magenta",
        )
    )
    console.print()

    # Fetch existing sessions
    sessions = _fetch_user_sessions(user_id)
    if sessions is None:
        console.print("[yellow]Could not fetch sessions. Try again?[/yellow]")
        retry = Prompt.ask("Retry", choices=["yes", "no"], default="yes")
        if retry == "yes":
            return select_or_create_session(user_id)
        return None

    # Display sessions
    if sessions:
        console.print("[bold cyan]Your sessions:[/bold cyan]\n")
        table = Table(show_header=True)
        table.add_column("#", style="dim", width=4)
        table.add_column("Session ID", style="cyan")
        table.add_column("Created At", style="green")

        for i, session in enumerate(sessions, 1):
            created_at = session.get("created_at", "Unknown")
            # Truncate session ID for display
            display_id = session["session_id"][:8] + "..." if len(session["session_id"]) > 8 else session["session_id"]
            table.add_row(str(i), display_id, created_at)

        console.print(table)
        console.print()

        # Let user select
        console.print("[dim]Options:[/dim]")
        console.print("  [cyan]1-{}[/cyan]: Select existing session".format(len(sessions)))
        console.print("  [cyan]new[/cyan]: Create new session")
        console.print("  [cyan]cancel[/cyan]: Go back")
        console.print()

        choice = Prompt.ask("[bold yellow]Select option[/bold yellow]")

        if choice.lower() == "cancel":
            return None
        elif choice.lower() == "new":
            console.print("\n[dim]Creating new session...[/dim]")
            session_id = _create_new_session(user_id)
            if session_id:
                console.print(f"[bold green]✓ Created session: {session_id[:8]}...[/bold green]\n")
                Prompt.ask("[dim]Press Enter to continue[/dim]", default="")
                return session_id
            return None
        else:
            try:
                index = int(choice) - 1
                if 0 <= index < len(sessions):
                    return sessions[index]["session_id"]
                else:
                    console.print("[bold red]Invalid selection[/bold red]")
                    Prompt.ask("[dim]Press Enter to try again[/dim]", default="")
                    return select_or_create_session(user_id)
            except ValueError:
                console.print("[bold red]Invalid input[/bold red]")
                Prompt.ask("[dim]Press Enter to try again[/dim]", default="")
                return select_or_create_session(user_id)
    else:
        console.print("[yellow]You don't have any sessions yet.[/yellow]\n")
        create = Prompt.ask("Create new session", choices=["yes", "no"], default="yes")
        if create == "yes":
            console.print("\n[dim]Creating new session...[/dim]")
            session_id = _create_new_session(user_id)
            if session_id:
                console.print(f"[bold green]✓ Created session: {session_id[:8]}...[/bold green]\n")
                Prompt.ask("[dim]Press Enter to continue[/dim]", default="")
                return session_id
        return None


def _fetch_conversation_history(user_id: str, session_id: str) -> Optional[List[dict]]:
    """Fetch conversation history for a session.

    Args:
        user_id: User identifier
        session_id: Session identifier

    Returns:
        List of message dicts with 'role', 'content', 'timestamp', or None if request failed
    """
    try:
        response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions/{session_id}/history", timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        data = response.json()
        return data.get("messages", [])
    except Exception as e:
        console.print(f"[bold red]Error fetching history:[/bold red] {str(e)}")
        return None


def _display_conversation_history(messages: List[dict]) -> None:
    """Display conversation history.

    Args:
        messages: List of message dicts with 'role' and 'content'
    """
    if not messages:
        return

    console.print("[dim]─── Conversation History ───[/dim]\n")
    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")

        if role == "user":
            console.print(Text("You:", style="bold blue"), end=" ")
            console.print(content)
        elif role == "assistant":
            console.print(Text("Agent:", style="bold green"), end=" ")
            console.print(Text(content, style="green"))

            if msg.get("suggested_questions"):
                _display_suggested_questions(msg["suggested_questions"])
        console.print()

    console.print("[dim]─── End of History ───[/dim]\n")


def _display_suggested_questions(questions: List[str]) -> None:
    """Display suggested follow-up questions.

    Args:
        questions: List of suggested question strings
    """
    if questions:
        console.print("\n[bold cyan]Suggested questions:[/bold cyan]")
        for i, question in enumerate(questions, 1):
            console.print(f"[cyan]{i}. {question}[/cyan]")


def chat_non_streaming(user_id: str, session_id: str) -> None:
    """Run interactive chat in non-streaming mode.

    Args:
        user_id: User identifier
        session_id: Session identifier
    """
    console.clear()
    console.print(Panel("[bold cyan]Non-Streaming Chat Mode[/bold cyan]"))
    console.print(f"[dim]Session: {session_id[:8]}...[/dim]")
    console.print("[dim]Type '/exit' to return to menu[/dim]\n")

    # Fetch and display conversation history
    history = _fetch_conversation_history(user_id, session_id)
    if history and len(history) > 0:
        _display_conversation_history(history)
    else:
        # No history - show welcome message
        console.print(Text("Agent:", style="bold green"), end=" ")
        console.print("Hello! How can I assist you today?")

        # Display initial suggested questions
        initial_questions = [
            "What can you help me with?",
            "Tell me about your capabilities",
            "How does this chat system work?",
        ]
        _display_suggested_questions(initial_questions)
        console.print()

    while True:
        user_input = Prompt.ask("[bold blue]You[/bold blue]")

        if user_input.lower() == "/exit":
            break

        try:
            response = httpx.post(
                f"{BASE_URL}/chat",
                json={"message": user_input, "user_id": user_id, "session_id": session_id},
                timeout=REQUEST_TIMEOUT,
            )
            response.raise_for_status()
            data = response.json()

            console.print()
            console.print(Text("Agent:", style="bold green"), end=" ")
            console.print(Text(data["response"], style="green"))

            _display_suggested_questions(data.get("suggested_questions", []))
            console.print()

        except httpx.HTTPStatusError as e:
            console.print(f"\n[bold red]HTTP Error {e.response.status_code}:[/bold red] {str(e)}\n")
        except Exception as e:
            console.print(f"\n[bold red]Error:[/bold red] {str(e)}\n")


async def chat_streaming_async(user_id: str, session_id: str) -> None:
    """Run interactive chat in streaming mode.

    Args:
        user_id: User identifier
        session_id: Session identifier
    """
    console.clear()
    console.print(Panel("[bold cyan]Streaming Chat Mode[/bold cyan]"))
    console.print(f"[dim]Session: {session_id[:8]}...[/dim]")
    console.print("[dim]Type '/exit' to return to menu[/dim]\n")

    # Fetch and display conversation history
    history = _fetch_conversation_history(user_id, session_id)
    if history and len(history) > 0:
        _display_conversation_history(history)
    else:
        # No history - show welcome message
        console.print(Text("Agent:", style="bold green"), end=" ")
        console.print("Hello! How can I assist you today?")

        # Display initial suggested questions
        initial_questions = [
            "What can you help me with?",
            "Tell me about your capabilities",
            "How does this chat system work?",
        ]
        _display_suggested_questions(initial_questions)
        console.print()

    while True:
        user_input = Prompt.ask("[bold blue]You[/bold blue]")

        if user_input.lower() == "/exit":
            break

        try:
            console.print()
            async with httpx.AsyncClient() as client:
                async with client.stream(
                    "POST",
                    f"{BASE_URL}/chat/stream",
                    json={"message": user_input, "user_id": user_id, "session_id": session_id},
                    timeout=REQUEST_TIMEOUT,
                ) as response:
                    response.raise_for_status()

                    with Live("", refresh_per_second=10, auto_refresh=True) as live:
                        current_response = ""
                        current_questions = []

                        async for line in response.aiter_lines():
                            if line.startswith("data: "):
                                data_str = line[6:]
                                if data_str.strip():
                                    data = json.loads(data_str)

                                    if "error" in data:
                                        console.print(f"\n[bold red]Error:[/bold red] {data['error']}\n")
                                        break

                                    if data.get("response"):
                                        current_response = data["response"]
                                    if data.get("suggested_questions"):
                                        current_questions = data["suggested_questions"]

                                    display_text = Text.assemble(("Agent: ", "bold green"), (current_response, "green"))

                                    if current_questions:
                                        display_text.append("\n\n")
                                        display_text.append("Suggested questions:\n", style="bold cyan")
                                        for i, question in enumerate(current_questions, 1):
                                            display_text.append(f"{i}. {question}\n", style="cyan")

                                    live.update(display_text)

            console.print()

        except httpx.HTTPStatusError as e:
            console.print(f"\n[bold red]HTTP Error {e.response.status_code}:[/bold red] {str(e)}\n")
        except Exception as e:
            console.print(f"\n[bold red]Error:[/bold red] {str(e)}\n")


def manage_sessions(user_id: str) -> None:
    """Display and manage user's sessions.

    Args:
        user_id: User identifier
    """
    console.clear()
    console.print(Panel("[bold cyan]Manage Sessions[/bold cyan]"))
    console.print()

    sessions = _fetch_user_sessions(user_id)
    if sessions is None:
        console.print()
        Prompt.ask("[dim]Press Enter to continue[/dim]", default="")
        return

    if not sessions:
        console.print("[yellow]No active sessions found[/yellow]")
        console.print()
        Prompt.ask("[dim]Press Enter to continue[/dim]", default="")
        return

    # Display sessions
    console.print("[bold]Your sessions:[/bold]\n")
    table = Table(show_header=True)
    table.add_column("#", style="dim", width=4)
    table.add_column("Session ID", style="cyan")
    table.add_column("Created At", style="green")

    for i, session in enumerate(sessions, 1):
        created_at = session.get("created_at", "Unknown")
        display_id = session["session_id"][:16]
        table.add_row(str(i), display_id, created_at)

    console.print(table)
    console.print()

    # Ask which to delete
    console.print("[dim]Enter session number to delete, or 'cancel' to go back[/dim]")
    choice = Prompt.ask("[bold yellow]Delete session[/bold yellow]", default="cancel")

    if choice.lower() != "cancel":
        try:
            index = int(choice) - 1
            if 0 <= index < len(sessions):
                session_to_delete = sessions[index]["session_id"]
                confirm = Prompt.ask(
                    f"\n[bold yellow]Delete session {session_to_delete[:8]}...?[/bold yellow]",
                    choices=["yes", "no"],
                    default="no",
                )
                if confirm == "yes":
                    if _delete_session(user_id, session_to_delete):
                        console.print("\n[bold green]✓ Session deleted[/bold green]")
            else:
                console.print("[bold red]Invalid selection[/bold red]")
        except ValueError:
            console.print("[bold red]Invalid input[/bold red]")

    console.print()
    Prompt.ask("[dim]Press Enter to continue[/dim]", default="")


def show_main_menu(user_id: str) -> str:
    """Display the main menu and get user's choice.

    Args:
        user_id: User identifier

    Returns:
        User's menu selection as a string
    """
    console.clear()
    console.print(
        Panel.fit(
            "[bold magenta]FastAPI Atomic Agents - Interactive Client[/bold magenta]",
            border_style="magenta",
        )
    )
    console.print(f"[dim]User ID: {user_id[:8]}...[/dim]\n")

    table = Table(show_header=False, box=None, padding=(0, 2))
    table.add_column(style="cyan bold", justify="right")
    table.add_column(style="white")

    table.add_row("1", "Start Chat (Non-Streaming)")
    table.add_row("2", "Start Chat (Streaming)")
    table.add_row("3", "Manage Sessions")
    table.add_row("4", "Exit")

    console.print(table)
    console.print()

    choice = Prompt.ask(
        "[bold yellow]Select an option[/bold yellow]",
        choices=["1", "2", "3", "4"],
        default="1",
    )
    return choice


async def main() -> None:
    """Main application loop."""
    user_id = get_or_create_user_id()

    while True:
        choice = show_main_menu(user_id)

        if choice == "1":
            # Non-streaming chat
            session_id = select_or_create_session(user_id)
            if session_id:
                chat_non_streaming(user_id, session_id)

        elif choice == "2":
            # Streaming chat
            session_id = select_or_create_session(user_id)
            if session_id:
                await chat_streaming_async(user_id, session_id)

        elif choice == "3":
            # Manage sessions
            manage_sessions(user_id)

        elif choice == "4":
            console.print("\n[bold cyan]Goodbye![/bold cyan]\n")
            break


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        console.print("\n\n[bold cyan]Goodbye![/bold cyan]\n")

```

### File: atomic-examples/fastapi-memory/fastapi_memory/lib/__init__.py

```python
"""Library modules for FastAPI Atomic Agents example."""

```

### File: atomic-examples/fastapi-memory/fastapi_memory/lib/agents/__init__.py

```python
"""Agent implementations for FastAPI example."""

from fastapi_memory.lib.agents.chat_agent import create_async_chat_agent, create_chat_agent

__all__ = ["create_chat_agent", "create_async_chat_agent"]

```

### File: atomic-examples/fastapi-memory/fastapi_memory/lib/agents/chat_agent.py

```python
"""Chat agent configuration and initialization."""

import instructor
import openai
from atomic_agents import AgentConfig, AtomicAgent
from atomic_agents.context import SystemPromptGenerator

from fastapi_memory.lib.config import MODEL_NAME, NUM_SUGGESTED_QUESTIONS, get_api_key
from fastapi_memory.lib.schemas import ChatRequest, ChatResponse


def _create_system_prompt() -> SystemPromptGenerator:
    """Create the system prompt configuration for chat agents.

    Returns:
        SystemPromptGenerator configured for conversational assistance
    """
    return SystemPromptGenerator(
        background=["You are a helpful AI assistant that maintains conversation context."],
        steps=[
            "Understand the user's message",
            "Provide a clear and helpful response",
            f"Generate {NUM_SUGGESTED_QUESTIONS} example questions that the user could type to continue the conversation",
        ],
        output_instructions=[
            "Be concise and friendly",
            "Reference previous context when relevant",
            "Suggested questions must be phrased as if the user is asking them (e.g., 'Tell me more about X', 'How does Y work?', 'What is Z?')",
        ],
    )


def create_chat_agent() -> AtomicAgent[ChatRequest, ChatResponse]:
    """Create a new synchronous chat agent.

    Returns:
        AtomicAgent configured for synchronous chat operations

    Raises:
        ValueError: If OPENAI_API_KEY environment variable is not set
    """
    api_key = get_api_key()
    client = instructor.from_openai(openai.OpenAI(api_key=api_key))
    config = AgentConfig(
        client=client,
        model=MODEL_NAME,
        model_api_parameters={"reasoning_effort": "minimal"},
        system_prompt_generator=_create_system_prompt(),
    )
    return AtomicAgent[ChatRequest, ChatResponse](config=config)


def create_async_chat_agent() -> AtomicAgent[ChatRequest, ChatResponse]:
    """Create a new asynchronous chat agent.

    Returns:
        AtomicAgent configured for asynchronous streaming operations

    Raises:
        ValueError: If OPENAI_API_KEY environment variable is not set
    """
    api_key = get_api_key()
    client = instructor.from_openai(openai.AsyncOpenAI(api_key=api_key))
    config = AgentConfig(
        client=client,
        model=MODEL_NAME,
        model_api_parameters={"reasoning_effort": "minimal"},
        system_prompt_generator=_create_system_prompt(),
    )
    return AtomicAgent[ChatRequest, ChatResponse](config=config)

```

### File: atomic-examples/fastapi-memory/fastapi_memory/lib/config.py

```python
"""Configuration module for FastAPI Atomic Agents example."""

import os


def get_api_key() -> str:
    """Get OpenAI API key from environment variables.

    Returns:
        str: OpenAI API key

    Raises:
        ValueError: If OPENAI_API_KEY environment variable is not set
    """
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise ValueError(
            "OPENAI_API_KEY environment variable is required. "
            "Please set it in your environment before running the application."
        )
    return api_key


# Constants
DEFAULT_SESSION_ID = "default"
MODEL_NAME = "gpt-5-mini"
NUM_SUGGESTED_QUESTIONS = 3

```

### File: atomic-examples/fastapi-memory/fastapi_memory/lib/schemas.py

```python
"""Schema definitions for FastAPI Atomic Agents example."""

from typing import List, Optional

from atomic_agents import BaseIOSchema
from pydantic import Field


class ChatRequest(BaseIOSchema):
    """Request schema for chat endpoint."""

    message: str = Field(..., description="User message")
    user_id: str = Field(..., description="User identifier")
    session_id: Optional[str] = Field(None, description="Session identifier for conversation continuity")


class ChatResponse(BaseIOSchema):
    """Response schema for chat endpoint."""

    response: str = Field(..., description="Agent response")
    session_id: str = Field(..., description="Session identifier")
    suggested_questions: Optional[List[str]] = Field(
        None,
        description="Suggested initial or follow-up questions that the user could ask the assistant",
    )


class SessionCreateRequest(BaseIOSchema):
    """Request schema for creating a new session."""

    user_id: str = Field(..., description="User identifier")


class SessionCreateResponse(BaseIOSchema):
    """Response schema for session creation."""

    session_id: str = Field(..., description="Generated session identifier")
    message: str = Field(..., description="Success message")


class SessionInfo(BaseIOSchema):
    """Information about a single session."""

    session_id: str = Field(..., description="Session identifier")
    created_at: Optional[str] = Field(None, description="Creation timestamp")


class UserSessionsResponse(BaseIOSchema):
    """Response schema for listing user's sessions."""

    user_id: str = Field(..., description="User identifier")
    sessions: List[SessionInfo] = Field(..., description="List of user's sessions")


class SessionDeleteResponse(BaseIOSchema):
    """Response schema for session deletion."""

    message: str = Field(..., description="Status message")


class ConversationMessage(BaseIOSchema):
    """A single message in the conversation history."""

    role: str = Field(..., description="Message role (user or assistant)")
    content: str = Field(..., description="Message content")
    timestamp: str = Field(..., description="Message timestamp")
    suggested_questions: Optional[List[str]] = Field(
        None, description="Suggested follow-up questions (only for assistant messages)"
    )


class ConversationHistory(BaseIOSchema):
    """Conversation history for a session."""

    session_id: str = Field(..., description="Session identifier")
    messages: List[ConversationMessage] = Field(..., description="List of messages in chronological order")

```

### File: atomic-examples/fastapi-memory/fastapi_memory/main.py

```python
"""FastAPI application for conversational AI with session management."""

import json
import uuid
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Dict

from atomic_agents import AtomicAgent
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse

from fastapi_memory.lib.agents.chat_agent import create_async_chat_agent, create_chat_agent
from fastapi_memory.lib.schemas import (
    ChatRequest,
    ChatResponse,
    ConversationHistory,
    ConversationMessage,
    SessionCreateResponse,
    SessionDeleteResponse,
    SessionInfo,
    UserSessionsResponse,
)

# Session storage: user_id -> session_id -> agent
sessions: Dict[str, Dict[str, AtomicAgent[ChatRequest, ChatResponse]]] = {}
async_sessions: Dict[str, Dict[str, AtomicAgent[ChatRequest, ChatResponse]]] = {}

# Session metadata: user_id -> session_id -> creation_timestamp
session_metadata: Dict[str, Dict[str, str]] = {}

# Conversation history: user_id -> session_id -> list of messages
conversation_history: Dict[str, Dict[str, list]] = {}


def _generate_session_id() -> str:
    """Generate a unique session identifier.

    Returns:
        UUID-based session identifier
    """
    return str(uuid.uuid4())


def _ensure_user_exists(user_id: str) -> None:
    """Ensure user exists in all storage dictionaries.

    Args:
        user_id: User identifier
    """
    if user_id not in sessions:
        sessions[user_id] = {}
    if user_id not in async_sessions:
        async_sessions[user_id] = {}
    if user_id not in session_metadata:
        session_metadata[user_id] = {}
    if user_id not in conversation_history:
        conversation_history[user_id] = {}


def _ensure_session_history_exists(user_id: str, session_id: str) -> None:
    """Ensure conversation history exists for a session.

    Args:
        user_id: User identifier
        session_id: Session identifier
    """
    _ensure_user_exists(user_id)
    if session_id not in conversation_history[user_id]:
        conversation_history[user_id][session_id] = []


def _add_message_to_history(
    user_id: str,
    session_id: str,
    role: str,
    content: str,
    suggested_questions: list[str] = None,
) -> None:
    """Add a message to the conversation history.

    Args:
        user_id: User identifier
        session_id: Session identifier
        role: Message role (user or assistant)
        content: Message content
        suggested_questions: Optional list of suggested questions
    """
    _ensure_session_history_exists(user_id, session_id)
    message = {
        "role": role,
        "content": content,
        "timestamp": datetime.now().isoformat(),
        "suggested_questions": suggested_questions,
    }
    conversation_history[user_id][session_id].append(message)


def get_or_create_agent(user_id: str, session_id: str) -> AtomicAgent[ChatRequest, ChatResponse]:
    """Get existing agent or create new synchronous agent for the session.

    Args:
        user_id: User identifier
        session_id: Session identifier

    Returns:
        AtomicAgent configured for synchronous chat operations
    """
    _ensure_user_exists(user_id)
    if session_id not in sessions[user_id]:
        sessions[user_id][session_id] = create_chat_agent()
        if session_id not in session_metadata[user_id]:
            session_metadata[user_id][session_id] = datetime.now().isoformat()
    return sessions[user_id][session_id]


def get_or_create_async_agent(user_id: str, session_id: str) -> AtomicAgent[ChatRequest, ChatResponse]:
    """Get existing agent or create new asynchronous agent for the session.

    Args:
        user_id: User identifier
        session_id: Session identifier

    Returns:
        AtomicAgent configured for asynchronous streaming operations
    """
    _ensure_user_exists(user_id)
    if session_id not in async_sessions[user_id]:
        async_sessions[user_id][session_id] = create_async_chat_agent()
        if session_id not in session_metadata[user_id]:
            session_metadata[user_id][session_id] = datetime.now().isoformat()
    return async_sessions[user_id][session_id]


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager to clean up resources on shutdown.

    Args:
        app: FastAPI application instance

    Yields:
        None
    """
    yield
    sessions.clear()
    async_sessions.clear()
    session_metadata.clear()
    conversation_history.clear()


app = FastAPI(
    title="Atomic Agents FastAPI Example",
    description="Simple example showing FastAPI integration with Atomic Agents",
    version="1.0.0",
    lifespan=lifespan,
)


@app.post("/chat", response_model=ChatResponse, tags=["Chat"])
async def chat(request: ChatRequest) -> ChatResponse:
    """Process a chat message using non-streaming response.

    Args:
        request: Chat request containing message, user_id, and optional session ID

    Returns:
        ChatResponse with agent's reply and suggested questions

    Raises:
        HTTPException: If message processing fails
    """
    try:
        if not request.session_id:
            raise HTTPException(
                status_code=400, detail="session_id is required. Create a session first using POST /users/{user_id}/sessions"
            )

        # Store user message in history
        _add_message_to_history(request.user_id, request.session_id, "user", request.message)

        agent = get_or_create_agent(request.user_id, request.session_id)

        result = agent.run(ChatRequest(message=request.message, user_id=request.user_id))

        # Store assistant response in history
        _add_message_to_history(
            request.user_id,
            request.session_id,
            "assistant",
            result.response,
            getattr(result, "suggested_questions", None),
        )

        return ChatResponse(
            response=result.response,
            session_id=request.session_id,
            suggested_questions=getattr(result, "suggested_questions", None),
        )

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        import traceback

        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Failed to process message: {str(e)}")


@app.post("/chat/stream", tags=["Chat"])
async def chat_stream(request: ChatRequest) -> StreamingResponse:
    """Process a chat message using streaming response.

    Args:
        request: Chat request containing message, user_id, and optional session ID

    Returns:
        StreamingResponse with Server-Sent Events format

    Raises:
        HTTPException: If streaming setup fails
    """
    try:
        if not request.session_id:
            raise HTTPException(
                status_code=400, detail="session_id is required. Create a session first using POST /users/{user_id}/sessions"
            )

        # Store user message in history
        _add_message_to_history(request.user_id, request.session_id, "user", request.message)

        agent = get_or_create_async_agent(request.user_id, request.session_id)

        async def generate():
            """Generate Server-Sent Events stream."""
            full_response = ""
            final_suggested_questions = []
            try:
                async for chunk in agent.run_async_stream(ChatRequest(message=request.message, user_id=request.user_id)):
                    chunk_dict = chunk.model_dump() if hasattr(chunk, "model_dump") else {}
                    response_text = chunk_dict.get("response", "")
                    full_response = response_text  # Keep updating with latest full text

                    if chunk_dict.get("suggested_questions"):
                        final_suggested_questions = chunk_dict.get("suggested_questions")

                    data = {
                        "response": response_text,
                        "session_id": request.session_id,
                        "suggested_questions": chunk_dict.get("suggested_questions"),
                    }
                    yield f"data: {json.dumps(data)}\n\n"

                # Store complete assistant response in history
                if full_response:
                    _add_message_to_history(
                        request.user_id,
                        request.session_id,
                        "assistant",
                        full_response,
                        final_suggested_questions,
                    )

            except Exception as e:
                error_data = {
                    "error": str(e),
                    "session_id": request.session_id,
                }
                yield f"data: {json.dumps(error_data)}\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to setup stream: {str(e)}")


@app.post("/users/{user_id}/sessions", response_model=SessionCreateResponse, tags=["Sessions"])
async def create_session(user_id: str) -> SessionCreateResponse:
    """Create a new chat session for a user.

    Args:
        user_id: User identifier

    Returns:
        SessionCreateResponse with generated session ID

    Raises:
        HTTPException: If session creation fails
    """
    try:
        _ensure_user_exists(user_id)
        session_id = _generate_session_id()
        session_metadata[user_id][session_id] = datetime.now().isoformat()

        return SessionCreateResponse(session_id=session_id, message=f"Session '{session_id}' created successfully")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to create session: {str(e)}")


@app.get("/users/{user_id}/sessions", response_model=UserSessionsResponse, tags=["Sessions"])
async def get_user_sessions(user_id: str) -> UserSessionsResponse:
    """Get all sessions for a specific user.

    Args:
        user_id: User identifier

    Returns:
        UserSessionsResponse with list of user's sessions
    """
    _ensure_user_exists(user_id)

    # Collect all unique session IDs for this user from both dicts
    sync_sessions = set(sessions.get(user_id, {}).keys())
    async_session_ids = set(async_sessions.get(user_id, {}).keys())
    all_session_ids = sync_sessions | async_session_ids

    # Build session info list
    session_list = [
        SessionInfo(session_id=sid, created_at=session_metadata.get(user_id, {}).get(sid)) for sid in sorted(all_session_ids)
    ]

    return UserSessionsResponse(user_id=user_id, sessions=session_list)


@app.get("/users/{user_id}/sessions/{session_id}/history", response_model=ConversationHistory, tags=["Sessions"])
async def get_conversation_history(user_id: str, session_id: str) -> ConversationHistory:
    """Get conversation history for a specific session.

    Args:
        user_id: User identifier
        session_id: Session identifier

    Returns:
        ConversationHistory with all messages in the session

    Raises:
        HTTPException: If session is not found
    """
    _ensure_session_history_exists(user_id, session_id)

    messages = conversation_history.get(user_id, {}).get(session_id, [])

    return ConversationHistory(session_id=session_id, messages=[ConversationMessage(**msg) for msg in messages])


@app.delete("/users/{user_id}/sessions/{session_id}", response_model=SessionDeleteResponse, tags=["Sessions"])
async def delete_session(user_id: str, session_id: str) -> SessionDeleteResponse:
    """Delete a specific session for a user.

    Args:
        user_id: User identifier
        session_id: Session identifier to delete

    Returns:
        SessionDeleteResponse with success message

    Raises:
        HTTPException: If session is not found
    """
    found = False

    if user_id in sessions and session_id in sessions[user_id]:
        del sessions[user_id][session_id]
        found = True

    if user_id in async_sessions and session_id in async_sessions[user_id]:
        del async_sessions[user_id][session_id]
        found = True

    if user_id in session_metadata and session_id in session_metadata[user_id]:
        del session_metadata[user_id][session_id]

    if user_id in conversation_history and session_id in conversation_history[user_id]:
        del conversation_history[user_id][session_id]

    if not found:
        raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found for user '{user_id}'")

    return SessionDeleteResponse(message=f"Session '{session_id}' deleted successfully")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

```

### File: atomic-examples/fastapi-memory/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["fastapi_memory"]

[project]
name = "fastapi-memory"
version = "0.1.0"
description = "Simple FastAPI integration example with Atomic Agents"
readme = "README.md"
authors = [
    { name = "BrainBlend AI" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "fastapi>=0.115.14,<1.0.0",
    "uvicorn>=0.32.1,<1.0.0",
    "instructor==1.14.5",
    "openai>=2.0.0,<3.0.0",
    "pydantic>=2.10.3,<3.0.0",
    "httpx>=0.28.1,<1.0.0",
    "rich>=13.9.4,<14.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```

### File: atomic-examples/fastapi-memory/test_api.py

```python
"""Quick API test script to verify the multi-session architecture."""

import httpx

BASE_URL = "http://localhost:8000"


def test_api():
    """Test the basic API flow."""
    print("Testing FastAPI Memory API...\n")

    # Test user ID
    user_id = "test-user-123"

    # 1. Get user sessions (should be empty initially)
    print(f"1. Fetching sessions for user: {user_id}")
    response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions")
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    # 2. Create a new session
    print("2. Creating new session...")
    response = httpx.post(f"{BASE_URL}/users/{user_id}/sessions")
    print(f"   Status: {response.status_code}")
    data = response.json()
    print(f"   Response: {data}")
    session_id = data["session_id"]
    print(f"   Created session: {session_id}\n")

    # 3. Send a chat message
    print("3. Sending first chat message...")
    response = httpx.post(
        f"{BASE_URL}/chat", json={"message": "Hello, how are you?", "user_id": user_id, "session_id": session_id}
    )
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    # 3b. Send another message to build conversation
    print("3b. Sending second chat message...")
    response = httpx.post(f"{BASE_URL}/chat", json={"message": "Tell me a joke", "user_id": user_id, "session_id": session_id})
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    # 3c. Get conversation history
    print("3c. Fetching conversation history...")
    response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions/{session_id}/history")
    print(f"   Status: {response.status_code}")
    history = response.json()
    print(f"   Number of messages: {len(history.get('messages', []))}")
    for i, msg in enumerate(history.get("messages", []), 1):
        role = msg.get("role")
        content = msg.get("content", "")[:50]  # Truncate for display
        suggested = msg.get("suggested_questions")
        print(f"   Message {i} ({role}): {content}...")
        if role == "assistant":
            print(f"   Suggested questions: {suggested}")
    print()

    # 4. Get user sessions (should have 1 session now)
    print("4. Fetching sessions again...")
    response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions")
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    # 5. Create another session
    print("5. Creating second session...")
    response = httpx.post(f"{BASE_URL}/users/{user_id}/sessions")
    data = response.json()
    session_id_2 = data["session_id"]
    print(f"   Created session: {session_id_2}\n")

    # 6. Get user sessions (should have 2 sessions now)
    print("6. Fetching sessions (should have 2)...")
    response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions")
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    # 7. Delete first session
    print(f"7. Deleting session {session_id}...")
    response = httpx.delete(f"{BASE_URL}/users/{user_id}/sessions/{session_id}")
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    # 8. Get user sessions (should have 1 session now)
    print("8. Fetching sessions (should have 1)...")
    response = httpx.get(f"{BASE_URL}/users/{user_id}/sessions")
    print(f"   Status: {response.status_code}")
    print(f"   Response: {response.json()}\n")

    print("✅ All tests completed!")


if __name__ == "__main__":
    try:
        test_api()
    except httpx.ConnectError:
        print("❌ Could not connect to server. Make sure it's running on http://localhost:8000")
    except Exception as e:
        print(f"❌ Error: {e}")

```


--------------------------------------------------------------------------------
Example: hooks-example
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/hooks-example

## Documentation

# AtomicAgent Hook System Example

This example demonstrates the powerful hook system integration in AtomicAgent, which leverages Instructor's hook system for comprehensive monitoring, error handling, and intelligent retry mechanisms.

## Features Demonstrated

- **🔍 Comprehensive Monitoring**: Track all aspects of agent execution
- **🛡️ Robust Error Handling**: Graceful handling of validation and completion errors  
- **🔄 Intelligent Retry Patterns**: Implement smart retry logic based on error context
- **📊 Performance Metrics**: Monitor response times, success rates, and error patterns
- **🔧 Easy Debugging**: Detailed error information and execution flow visibility
- **⚡ Zero Overhead**: Hooks only execute when registered and enabled

## Getting Started

1. Clone the main Atomic Agents repository:
   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. Navigate to the hooks-example directory:
   ```bash
   cd atomic-agents/atomic-examples/hooks-example
   ```

3. Install the dependencies using uv:
   ```bash
   uv sync
   ```

4. Set up your OpenAI API key:
   ```bash
   export OPENAI_API_KEY="your-api-key-here"
   ```

5. Run the example:
   ```bash
   uv run python hooks_example/main.py
   ```

## What This Example Shows

The example demonstrates several key hook system patterns:

### Basic Hook Registration
- Simple parse error logging
- Completion monitoring and metrics collection

### Advanced Error Handling
- Comprehensive validation error analysis  
- Intelligent retry mechanisms with backoff strategies
- Error isolation to prevent hook failures from disrupting execution

### Performance Monitoring
- Response time tracking
- Success rate calculation
- Error pattern analysis

### Real-World Scenarios
- Handling malformed responses
- Network timeouts and retry logic
- Model switching on repeated failures

## Key Benefits

This hook system implementation provides:

1. **Full Instructor Integration**: All Instructor hook events are supported
2. **Backward Compatibility**: Existing AtomicAgent code works unchanged
3. **Error Context**: Rich error information for intelligent decision making
4. **Performance Insights**: Detailed metrics for optimization
5. **Production Ready**: Robust error handling suitable for production use

## Hook Events Supported

- `parse:error` - Triggered on Pydantic validation failures
- `completion:kwargs` - Before API calls are made
- `completion:response` - After API responses are received
- `completion:error` - On API or network errors

## GitHub Issue Resolution

This example demonstrates the complete resolution of GitHub issue #173, showing how the AtomicAgent hook system enables:

- ✅ Parse error hooks triggering on validation failures
- ✅ Comprehensive error context for retry mechanisms  
- ✅ Full Instructor hook event support
- ✅ 100% backward compatibility
- ✅ Robust error isolation

## Next Steps

After running this example, you can:

1. Experiment with different hook combinations
2. Implement custom retry strategies
3. Add your own monitoring and alerting logic
4. Explore integration with observability platforms

## Source Code


### File: atomic-examples/hooks-example/hooks_example/main.py

```python
#!/usr/bin/env python3
"""
AtomicAgent Hook System Demo

Shows how to monitor agent execution with hooks.
Includes error handling and performance metrics.
"""

import os
import time
import logging

import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from pydantic import Field, ValidationError

from atomic_agents import AtomicAgent, AgentConfig
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from atomic_agents.base.base_io_schema import BaseIOSchema

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
console = Console()
metrics = {
    "total_requests": 0,
    "successful_requests": 0,
    "failed_requests": 0,
    "parse_errors": 0,
    "retry_attempts": 0,
    "total_response_time": 0.0,
    "start_time": time.time(),
}

_request_start_time = None


class UserQuery(BaseIOSchema):
    """Schema for user input containing a chat message."""

    chat_message: str = Field(..., description="User's question or message")


class AgentResponse(BaseIOSchema):
    """Schema for agent response with confidence and reasoning."""

    chat_message: str = Field(..., description="Agent's response to the user")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score (0.0-1.0)")
    reasoning: str = Field(..., description="Brief explanation of the reasoning")


class DetailedResponse(BaseIOSchema):
    """Schema for detailed response with alternatives and confidence level."""

    chat_message: str = Field(..., description="Primary response")
    alternative_suggestions: list[str] = Field(default_factory=list, description="Alternative suggestions")
    confidence_level: str = Field(..., description="Must be 'low', 'medium', or 'high'")
    requires_followup: bool = Field(default=False, description="Whether follow-up is needed")


def setup_api_key() -> str:
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        console.print("[bold red]Error: OPENAI_API_KEY environment variable not set.[/bold red]")
        console.print("Please set it with: export OPENAI_API_KEY='your-api-key-here'")
        exit(1)
    return api_key


def display_metrics():
    runtime = time.time() - metrics["start_time"]
    avg_response_time = metrics["total_response_time"] / metrics["total_requests"] if metrics["total_requests"] > 0 else 0
    success_rate = metrics["successful_requests"] / metrics["total_requests"] * 100 if metrics["total_requests"] > 0 else 0

    table = Table(title="🔍 Hook System Performance Metrics", style="cyan")
    table.add_column("Metric", style="bold")
    table.add_column("Value", style="green")

    table.add_row("Runtime", f"{runtime:.1f}s")
    table.add_row("Total Requests", str(metrics["total_requests"]))
    table.add_row("Successful Requests", str(metrics["successful_requests"]))
    table.add_row("Failed Requests", str(metrics["failed_requests"]))
    table.add_row("Parse Errors", str(metrics["parse_errors"]))
    table.add_row("Retry Attempts", str(metrics["retry_attempts"]))
    table.add_row("Success Rate", f"{success_rate:.1f}%")
    table.add_row("Avg Response Time", f"{avg_response_time:.2f}s")

    console.print(table)


def on_parse_error(error):
    metrics["parse_errors"] += 1
    metrics["failed_requests"] += 1
    logger.error(f"🚨 Parse error occurred: {type(error).__name__}: {error}")

    if isinstance(error, ValidationError):
        console.print("[bold red]❌ Validation Error:[/bold red]")
        for err in error.errors():
            field_path = " -> ".join(str(x) for x in err["loc"])
            console.print(f"  • Field '{field_path}': {err['msg']}")
            logger.error(f"Validation error in field '{field_path}': {err['msg']}")
    else:
        console.print(f"[bold red]❌ Parse Error:[/bold red] {error}")


def on_completion_kwargs(**kwargs):
    global _request_start_time
    metrics["total_requests"] += 1
    model = kwargs.get("model", "unknown")
    messages_count = len(kwargs.get("messages", []))
    logger.info(f"🚀 API call starting - Model: {model}, Messages: {messages_count}")
    _request_start_time = time.time()


def on_completion_response(response, **kwargs):
    global _request_start_time
    if _request_start_time:
        response_time = time.time() - _request_start_time
        metrics["total_response_time"] += response_time
        logger.info(f"✅ API call completed in {response_time:.2f}s")
        _request_start_time = None

    if hasattr(response, "usage"):
        usage = response.usage
        logger.info(
            f"📊 Token usage - Prompt: {usage.prompt_tokens}, "
            f"Completion: {usage.completion_tokens}, "
            f"Total: {usage.total_tokens}"
        )

    metrics["successful_requests"] += 1


def on_completion_error(error, **kwargs):
    global _request_start_time
    metrics["failed_requests"] += 1
    metrics["retry_attempts"] += 1

    if _request_start_time:
        _request_start_time = None

    logger.error(f"🔥 API error: {type(error).__name__}: {error}")
    console.print(f"[bold red]🔥 API Error:[/bold red] {error}")


def create_agent_with_hooks(schema_type: type, system_prompt: str = None) -> AtomicAgent:
    api_key = setup_api_key()
    client = instructor.from_openai(openai.OpenAI(api_key=api_key))

    # Create a system prompt generator if a system prompt is provided
    system_prompt_generator = SystemPromptGenerator(background=[system_prompt]) if system_prompt else None

    config = AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        history=ChatHistory(),
        system_prompt_generator=system_prompt_generator,
    )

    agent = AtomicAgent[UserQuery, schema_type](config)

    agent.register_hook("parse:error", on_parse_error)
    agent.register_hook("completion:kwargs", on_completion_kwargs)
    agent.register_hook("completion:response", on_completion_response)
    agent.register_hook("completion:error", on_completion_error)

    console.print("[bold green]✅ Agent created with comprehensive hook monitoring[/bold green]")
    return agent


def demonstrate_basic_hooks():
    console.print(Panel("🔧 Basic Hook System Demonstration", style="bold blue"))

    agent = create_agent_with_hooks(
        AgentResponse, "You are a helpful assistant. Always provide confident, well-reasoned responses."
    )

    test_queries = [
        "What is the capital of France?",
        "Explain quantum computing in simple terms.",
        "What are the benefits of renewable energy?",
    ]

    for query_text in test_queries:
        console.print(f"\n[bold cyan]Query:[/bold cyan] {query_text}")

        try:
            query = UserQuery(chat_message=query_text)
            response = agent.run(query)

            console.print(f"[bold green]Response:[/bold green] {response.chat_message}")
            console.print(f"[bold yellow]Confidence:[/bold yellow] {response.confidence:.2f}")
            console.print(f"[bold magenta]Reasoning:[/bold magenta] {response.reasoning}")

        except Exception as e:
            console.print(f"[bold red]Error processing query:[/bold red] {e}")

    display_metrics()


def demonstrate_validation_errors():
    console.print(Panel("🚨 Validation Error Handling Demonstration", style="bold red"))

    agent = create_agent_with_hooks(
        DetailedResponse,
        """You are a helpful assistant. INTENTIONALLY use invalid values to test validation:
        - Set confidence_level to something other than 'low', 'medium', or 'high' (like 'very_high' or 'uncertain')
        - This is for testing validation error handling, so please violate the schema constraints intentionally.""",
    )

    validation_test_queries = [
        "Give me a simple yes or no answer about whether the sky is blue.",
        "Provide a complex analysis of climate change with multiple perspectives.",
    ]

    for query_text in validation_test_queries:
        console.print(f"\n[bold cyan]Query:[/bold cyan] {query_text}")

        try:
            query = UserQuery(chat_message=query_text)
            response = agent.run(query)

            console.print(f"[bold green]Main Answer:[/bold green] {response.chat_message}")
            console.print(f"[bold yellow]Confidence Level:[/bold yellow] {response.confidence_level}")
            console.print(f"[bold magenta]Alternatives:[/bold magenta] {response.alternative_suggestions}")
            console.print(f"[bold cyan]Needs Follow-up:[/bold cyan] {response.requires_followup}")

        except Exception as e:
            console.print(f"[bold red]Handled error:[/bold red] {e}")

    display_metrics()


def demonstrate_interactive_mode():
    console.print(Panel("🎮 Interactive Hook System Testing", style="bold magenta"))

    agent = create_agent_with_hooks(
        AgentResponse, "You are a helpful assistant. Provide clear, confident responses with reasoning."
    )

    console.print("[bold green]Welcome to the interactive hook system demo![/bold green]")
    console.print("Type your questions below. Use /metrics to see performance data, /exit to quit.")

    while True:
        try:
            user_input = console.input("\n[bold blue]Your question:[/bold blue] ")

            if user_input.lower() in ["/exit", "/quit"]:
                console.print("Exiting interactive mode...")
                break
            elif user_input.lower() == "/metrics":
                display_metrics()
                continue
            elif user_input.strip() == "":
                continue

            query = UserQuery(chat_message=user_input)
            start_time = time.time()

            response = agent.run(query)

            response_time = time.time() - start_time

            console.print(f"\n[bold green]Answer:[/bold green] {response.chat_message}")
            console.print(f"[bold yellow]Confidence:[/bold yellow] {response.confidence:.2f}")
            console.print(f"[bold magenta]Reasoning:[/bold magenta] {response.reasoning}")
            console.print(f"[dim]Response time: {response_time:.2f}s[/dim]")

        except KeyboardInterrupt:
            console.print("\nExiting on user interrupt...")
            break
        except Exception as e:
            console.print(f"[bold red]Error:[/bold red] {e}")


def main():
    console.print(Panel.fit("🎯 AtomicAgent Hook System Comprehensive Demo", style="bold green"))

    console.print(
        """
[bold cyan]This demonstration showcases:[/bold cyan]
• 🔍 Comprehensive monitoring with hooks
• 🛡️ Robust error handling and validation
• 📊 Real-time performance metrics
• 🔄 Production-ready patterns

[bold yellow]The hook system provides zero-overhead monitoring when hooks aren't registered,
and powerful insights when they are enabled.[/bold yellow]
    """
    )

    try:
        demonstrate_basic_hooks()
        console.print("\n" + "=" * 50)
        demonstrate_validation_errors()
        console.print("\n" + "=" * 50)
        demonstrate_interactive_mode()

    except KeyboardInterrupt:
        console.print("\n[bold yellow]Demo interrupted by user.[/bold yellow]")
    except Exception as e:
        console.print(f"\n[bold red]Demo error:[/bold red] {e}")
        logger.error(f"Demo error: {e}", exc_info=True)
    finally:
        console.print("\n" + "=" * 50)
        console.print(Panel("📊 Final Performance Summary", style="bold green"))
        display_metrics()

        console.print(
            """
[bold green]✅ Hook system demonstration complete![/bold green]

[bold cyan]Key takeaways:[/bold cyan]
• Hooks provide comprehensive monitoring without performance overhead
• Error handling is robust and provides detailed context
• Metrics collection enables performance optimization
• The system is production-ready and scalable

[bold yellow]Next steps:[/bold yellow]
• Implement custom retry logic in hook handlers
• Add monitoring service integration
• Explore advanced error recovery patterns
• Build custom metrics dashboards
        """
        )


if __name__ == "__main__":
    main()

```

### File: atomic-examples/hooks-example/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["hooks_example"]

[project]
name = "hooks-example"
version = "1.0.0"
description = "AtomicAgent hooks system example demonstrating monitoring, error handling, and retry mechanisms"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "instructor==1.14.5",
    "openai>=2.0.0,<3.0.0",
    "python-dotenv>=1.0.1,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: mcp-agent
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/mcp-agent

## Documentation

# MCP Agent Example

This directory contains a complete example of a Model Context Protocol (MCP) implementation, including both client and server components. It demonstrates how to build an intelligent agent that leverages MCP tools via different transport methods.

## Components

This example consists of two main components:

### 1. Example Client (`example-client/`)

An interactive agent that:

- Connects to MCP servers using multiple transport methods (STDIO, SSE, HTTP Stream)
- Dynamically discovers available tools
- Processes natural language queries
- Selects appropriate tools based on user intent
- Executes tools with extracted parameters (sync and async)
- Provides responses in a conversational format

The client features a universal launcher that supports multiple implementations:
- **stdio**: Blocking STDIO CLI client (default)
- **stdio_async**: Async STDIO client
- **sse**: SSE CLI client
- **http_stream**: HTTP Stream CLI client
- **fastapi**: FastAPI HTTP API server

[View Example Client README](example-client/README.md)

### 2. Example MCP Server (`example-mcp-server/`)

A server that:

- Provides MCP tools and resources
- Supports both STDIO and SSE (HTTP) transport methods
- Includes example tools for demonstration
- Can be extended with custom functionality
- Features auto-reload for development

[View Example MCP Server README](example-mcp-server/README.md)

## Understanding the Example

This example shows the flexibility of the MCP architecture with two distinct transport methods:

### STDIO Transport

- The client launches the server as a subprocess
- Communication occurs through standard input/output
- No network connectivity required
- Good for local development and testing

### SSE Transport

- The server runs as a standalone HTTP service
- The client connects via Server-Sent Events (SSE)
- Multiple clients can connect to one server
- Better for production deployments

### HTTP Stream Transport

- The server exposes a single `/mcp` HTTP endpoint for session negotiation, JSON-RPC calls, and termination
- Supports GET (stream/session ID), POST (JSON-RPC payloads), and DELETE (session cancel)
- Useful for HTTP clients that prefer a single transport endpoint

## Getting Started

1. Clone the repository:

   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   cd atomic-agents/atomic-examples/mcp-agent
   ```

2. Set up the server:

   ```bash
   cd example-mcp-server
   uv sync
   ```

3. Set up the client:
   ```bash
   cd ../example-client
   uv sync
   ```

4. Run the example:

   **Using STDIO transport (default):**

   ```bash
   cd example-client
   uv run python -m example_client.main --client stdio
   # or simply:
   uv run python -m example_client.main
   ```

   **Using async STDIO transport:**

   ```bash
   cd example-client
   uv run python -m example_client.main --client stdio_async
   ```

   **Using SSE transport (Deprecated):**

   ```bash
   # First terminal: Start the server
   cd example-mcp-server
   uv run python -m example_mcp_server.server --mode=sse

   # Second terminal: Run the client with SSE transport
   cd example-client
   uv run python -m example_client.main --client sse
   ```

   **Using HTTP Stream transport:**

   ```bash
   # First terminal: Start the server
   cd example-mcp-server
   uv run python -m example_mcp_server.server --mode=http_stream

   # Second terminal: Run the client with HTTP Stream transport
   cd example-client
   uv run python -m example_client.main --client http_stream
   ```

   **Using FastAPI client:**

   ```bash
   # First terminal: Start the MCP server
   cd example-mcp-server
   uv run python -m example_mcp_server.server --mode=http_stream

   # Second terminal: Run the FastAPI client
   cd example-client
   uv run python -m example_client.main --client fastapi
   # Then visit http://localhost:8000 for the API interface
   ```

**Note:** When using SSE, FastAPI or HTTP Stream transport, make sure the server is running before starting the client. The server runs on port 6969 by default.

## Example Queries

The example includes a set of basic arithmetic tools that demonstrate the agent's capability to break down and solve complex mathematical expressions:

### Available Demo Tools

- **AddNumbers**: Adds two numbers together (number1 + number2)
- **SubtractNumbers**: Subtracts the second number from the first (number1 - number2)
- **MultiplyNumbers**: Multiplies two numbers together (number1 * number2)
- **DivideNumbers**: Divides the first number by the second (handles division by zero)

### Conversation Flow

When you interact with the agent, it:

1. Analyzes your input to break it down into sequential operations
2. Selects appropriate tools for each operation
3. Shows its reasoning for each tool selection
4. Executes the tools in sequence
5. Maintains context between operations to build up the final result

For example, when calculating `(5-9)*0.123`:

1. First uses `SubtractNumbers` to compute (5-9) = -4
2. Then uses `MultiplyNumbers` to compute (-4 * 0.123) = -0.492
3. Provides the final result with clear explanation

For more complex expressions like `((4**3)-10)/100)**2`, the agent:

1. Breaks down the expression into multiple steps
2. Uses `MultiplyNumbers` repeatedly for exponentiation (4**3)
3. Uses `SubtractNumbers` for the subtraction operation
4. Uses `DivideNumbers` for division by 100
5. Uses `MultiplyNumbers` again for the final squaring operation

Each step in the conversation shows:

- The tool being executed
- The parameters being used
- The intermediate result
- The agent's reasoning for the next step

Try queries like:

```python
# Simple arithmetic
"What is 2+2?"
# Uses AddNumbers tool directly

# Complex expressions
"(5-9)*0.123"
# Uses SubtractNumbers followed by MultiplyNumbers

# Multi-step calculations
"((4**3)-10)/100)**2"
# Uses multiple tools in sequence to break down the complex expression

# Natural language queries
"Calculate the difference between 50 and 23, then multiply it by 3"
# Understands natural language and breaks it down into appropriate tool calls
```

## Learn More

- [Atomic Agents Documentation](https://github.com/BrainBlend-AI/atomic-agents)
- [Model Context Protocol](https://modelcontextprotocol.io/)


## Source Code


### File: atomic-examples/mcp-agent/example-client/example_client/main.py

```python
# pyright: reportInvalidTypeForm=false
"""
Universal launcher for the MCP examples.
  stdio_async  - runs the async STDIO client
  fastapi      - serves the FastAPI HTTP API
  http_stream  - HTTP-stream CLI client
  sse          - SSE CLI client
  stdio        - blocking STDIO CLI client
"""

import argparse
import asyncio
import importlib
import sys

# Optional import; only used for the FastAPI target
try:
    import uvicorn  # noqa: WPS433 – runtime import is deliberate
except ImportError:  # pragma: no cover
    uvicorn = None


def _run_target(module_name: str, func_name: str | None = "main", *, is_async: bool = False) -> None:
    """
    Import `module_name` and execute `func_name`.

    Args:
        module_name: Python module containing the entry point.
        func_name:   Callable inside that module to execute (skip for FastAPI).
        is_async:    Whether the callable is an async coroutine.
    """
    module = importlib.import_module(module_name)

    if func_name is None:  # fastapi path – start uvicorn directly
        if uvicorn is None:  # pragma: no cover
            sys.exit("uvicorn is not installed - unable to start FastAPI server.")
        # `module_name:app` tells uvicorn where the FastAPI instance lives.
        uvicorn.run(f"{module_name}:app", host="0.0.0.0", port=8000)
        return

    entry = getattr(module, func_name)
    if is_async:
        asyncio.run(entry())
    else:
        entry()


def main() -> None:
    parser = argparse.ArgumentParser(description="MCP Example Launcher")
    parser.add_argument(
        "--client",
        default="stdio",
        choices=[
            "stdio",
            "stdio_async",
            "sse",
            "http_stream",
            "fastapi",
        ],
        help="Which client implementation to start",
    )
    args = parser.parse_args()

    # Map the `--client` value to (module, callable, needs_asyncio)
    dispatch_table: dict[str, tuple[str, str | None, bool]] = {
        "stdio": ("example_client.main_stdio", "main", False),
        "stdio_async": ("example_client.main_stdio_async", "main", True),
        "sse": ("example_client.main_sse", "main", False),
        "http_stream": ("example_client.main_http", "main", False),
        # For FastAPI we hand control to uvicorn – func_name=None signals that.
        "fastapi": ("example_client.main_fastapi", None, False),
    }

    try:
        module_name, func_name, is_async = dispatch_table[args.client]
        _run_target(module_name, func_name, is_async=is_async)
    except KeyError:
        sys.exit(f"Unknown client: {args.client}")
    except (ImportError, AttributeError) as exc:
        sys.exit(f"Failed to load '{args.client}': {exc}")


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-client/example_client/main_fastapi.py

```python
"""FastAPI client example demonstrating async MCP tool usage."""

import os
from typing import Dict, Any, List, Union, Type
from contextlib import asynccontextmanager
from dataclasses import dataclass

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from atomic_agents.connectors.mcp import (
    fetch_mcp_tools_async,
    fetch_mcp_resources_async,
    fetch_mcp_prompts_async,
    MCPTransportType,
)
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
import openai
import instructor


@dataclass
class MCPConfig:
    """Configuration for the MCP Agent system using HTTP Stream transport."""

    mcp_server_url: str = "http://localhost:6969"
    openai_model: str = "gpt-5-mini"
    openai_api_key: str = os.getenv("OPENAI_API_KEY") or ""
    reasoning_effort: str = "low"

    def __post_init__(self):
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is not set")


class NaturalLanguageRequest(BaseModel):
    query: str = Field(..., description="Natural language query for mathematical operations")


class CalculationResponse(BaseModel):
    result: Any
    tools_used: List[str]
    resources_used: List[str]
    prompts_used: List[str]
    query: str


class ResourceResponse(BaseModel):
    content: str
    tools_used: List[str]
    resources_used: List[str]
    prompts_used: List[str]
    query: str


class PromptResponse(BaseModel):
    content: str
    tools_used: List[str]
    resources_used: List[str]
    prompts_fetched: List[str]
    query: str


class MCPOrchestratorInputSchema(BaseIOSchema):
    """Input schema for the MCP orchestrator that processes user queries."""

    query: str = Field(...)


class FinalResponseSchema(BaseIOSchema):
    """Schema for the final response to the user."""

    response_text: str = Field(...)


# Global storage for MCP tools, schema mapping
mcp_tools = {}
mcp_resources = {}
mcp_prompts = {}
tool_schema_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {}
resource_schema_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {}
prompt_schema_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {}
config = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize MCP tools and orchestrator agent on startup."""
    global config

    config = MCPConfig()

    mcp_endpoint = config.mcp_server_url

    try:
        print(f"Attempting to connect to MCP server at {mcp_endpoint}")
        print(f"Using transport type: {MCPTransportType.HTTP_STREAM}")

        import requests

        try:
            response = requests.get(f"{mcp_endpoint}/health", timeout=5)
            print(f"Health check response: {response.status_code}")
        except Exception as health_error:
            print(f"Health check failed: {health_error}")

        tools = await fetch_mcp_tools_async(mcp_endpoint=mcp_endpoint, transport_type=MCPTransportType.HTTP_STREAM)
        resources = await fetch_mcp_resources_async(mcp_endpoint=mcp_endpoint, transport_type=MCPTransportType.HTTP_STREAM)
        prompts = await fetch_mcp_prompts_async(mcp_endpoint=mcp_endpoint, transport_type=MCPTransportType.HTTP_STREAM)

        print(f"fetch_mcp_tools returned {len(tools)} tools")
        print(f"Tools type: {type(tools)}")

        for i, tool in enumerate(tools):
            tool_name = getattr(tool, "mcp_tool_name", tool.__name__)
            mcp_tools[tool_name] = tool
            print(f"Tool {i}: name='{tool_name}', type={type(tool).__name__}")

        print(f"Initialized {len(mcp_tools)} MCP tools: {list(mcp_tools.keys())}")

        # Display resources and prompts if available
        if resources:
            print(f"fetch_mcp_resources returned {len(resources)} resources")
            print(f"Resources type: {type(resources)}")
            for i, resource in enumerate(resources):
                resource_name = getattr(resource, "mcp_resource_name", resource.__name__)
                mcp_resources[resource_name] = resource
                print(f"Resource {i}: name='{resource_name}', type={type(resource).__name__}")
            print(f"Initialized {len(mcp_resources)} MCP resources: {list(mcp_resources.keys())}")
        if prompts:
            print(f"fetch_mcp_prompts returned {len(prompts)} prompts")
            print(f"Prompts type: {type(prompts)}")
            for i, prompt in enumerate(prompts):
                prompt_name = getattr(prompt, "mcp_prompt_name", prompt.__name__)
                mcp_prompts[prompt_name] = prompt
                print(f"Prompt {i}: name='{prompt_name}', type={type(prompt).__name__}")
            print(f"Initialized {len(mcp_prompts)} MCP prompts: {list(mcp_prompts.keys())}")

        tool_schema_map.update(
            {ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")}  # type: ignore
        )

        # Build resource/prompt schema maps and extend available schemas
        resource_schema_map.update(
            {ResourceClass.input_schema: ResourceClass for ResourceClass in resources if hasattr(ResourceClass, "input_schema")}  # type: ignore
        )
        prompt_schema_map.update(
            {PromptClass.input_schema: PromptClass for PromptClass in prompts if hasattr(PromptClass, "input_schema")}  # type: ignore
        )

        available_schemas = (
            tuple(tool_schema_map.keys())
            + tuple(resource_schema_map.keys())
            + tuple(prompt_schema_map.keys())
            + (FinalResponseSchema,)
        )

        client = instructor.from_openai(openai.OpenAI(api_key=config.openai_api_key))
        history = ChatHistory()

        globals()["client"] = client
        globals()["history"] = history
        globals()["available_schemas"] = available_schemas

        print("MCP tools, schema mapping, and agent components initialized successfully")

    except Exception as e:
        print(f"Failed to initialize MCP tools: {e}")
        print(f"Exception type: {type(e).__name__}")
        import traceback

        traceback.print_exc()
        print("\n" + "=" * 60)
        print("ERROR: Could not connect to MCP server!")
        print("Please start the MCP server first:")
        print("  cd /path/to/example-mcp-server")
        print("  uv run python -m example_mcp_server.server --mode=http_stream")
        print("=" * 60)
        raise RuntimeError(f"MCP server connection failed: {e}") from e

    yield

    mcp_tools.clear()
    mcp_resources.clear()
    mcp_prompts.clear()
    tool_schema_map.clear()


app = FastAPI(
    title="MCP FastAPI Client Example",
    description="Demonstrates async MCP tool usage in FastAPI handlers with agent-based architecture",
    lifespan=lifespan,
)


async def execute_with_orchestrator_async(query: str) -> tuple[str, list[str], list[str], list[str]]:
    """Execute using orchestrator agent pattern with async execution."""
    if not config or not tool_schema_map:
        raise HTTPException(status_code=503, detail="Agent components not initialized")

    tools_used = []
    resources_used = []
    prompts_used = []

    try:
        available_schemas = (
            tuple(tool_schema_map.keys())
            + tuple(resource_schema_map.keys())
            + tuple(prompt_schema_map.keys())
            + (FinalResponseSchema,)
        )
        ActionUnion = Union[available_schemas]

        class OrchestratorOutputSchema(BaseIOSchema):
            """Output schema for the MCP orchestrator containing reasoning and selected action."""

            reasoning: str
            action: ActionUnion = Field(
                ...,
                description="The chosen action: either a tool/resource/prompt's input schema instance or a final response schema instance.",
            )

        orchestrator_agent = AtomicAgent[MCPOrchestratorInputSchema, OrchestratorOutputSchema](
            AgentConfig(
                client=globals()["client"],
                model=config.openai_model,
                model_api_parameters={"reasoning_effort": config.reasoning_effort},
                history=ChatHistory(),
                system_prompt_generator=SystemPromptGenerator(
                    background=[
                        "You are an MCP Orchestrator Agent, designed to chat with users and",
                        "determine the best way to handle their queries using the available tools, resources, and prompts.",
                    ],
                    steps=[
                        "1. Use the reasoning field to determine if one or more successive "
                        "tool/resource/prompt calls could be used to handle the user's query.",
                        "2. If so, choose the appropriate tool(s), resource(s), or prompt(s) one "
                        "at a time and extract all necessary parameters from the query.",
                        "3. If a single tool/resource/prompt can not be used to handle the user's query, "
                        "think about how to break down the query into "
                        "smaller tasks and route them to the appropriate tool(s)/resource(s)/prompt(s).",
                        "4. If no sequence of tools/resources/prompts could be used, or if you are "
                        "finished processing the user's query, provide a final response to the user.",
                        "5. If the context is sufficient and no more tools/resources/prompts are needed, provide a final response to the user.",
                    ],
                    output_instructions=[
                        "1. Always provide a detailed explanation of your decision-making process in the 'reasoning' field.",
                        "2. Choose exactly one action schema (either a tool/resource/prompt input or FinalResponseSchema).",
                        "3. Ensure all required parameters for the chosen tool/resource/prompt are properly extracted and validated.",
                        "4. Maintain a professional and helpful tone in all responses.",
                        "5. Break down complex queries into sequential tool/resource/prompt calls "
                        "before giving the final answer via `FinalResponseSchema`.",
                    ],
                ),
            )
        )

        orchestrator_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=query))

        print(f"Debug - orchestrator_output type: {type(orchestrator_output)}, fields: {orchestrator_output.model_dump()}")

        if hasattr(orchestrator_output, "chat_message") and not hasattr(orchestrator_output, "action"):
            action_instance = FinalResponseSchema(response_text=orchestrator_output.chat_message)
            reasoning = "Response generated directly from chat model"
        elif hasattr(orchestrator_output, "action"):
            action_instance = orchestrator_output.action
            reasoning = orchestrator_output.reasoning if hasattr(orchestrator_output, "reasoning") else "No reasoning provided"
        else:
            return "I encountered an unexpected response format. Unable to process.", tools_used, resources_used, prompts_used

        print(f"Debug - Orchestrator reasoning: {reasoning}")
        print(f"Debug - Action instance type: {type(action_instance)}")
        print(f"Debug - Action instance: {action_instance}")

        iteration_count = 0
        max_iterations = 5

        while not isinstance(action_instance, FinalResponseSchema) and iteration_count < max_iterations:
            iteration_count += 1
            print(f"Debug - Iteration {iteration_count}, processing action type: {type(action_instance)}")

            schema_type = type(action_instance)
            schema_type_valid = False

            # Check for tool
            tool_class = tool_schema_map.get(schema_type)
            if tool_class:
                schema_type_valid = True
                tool_name = getattr(tool_class, "mcp_tool_name", "unknown")  # type: ignore
                tools_used.append(tool_name)

                print(f"Debug - Executing {tool_name}...")
                print(f"Debug - Parameters: {action_instance.model_dump()}")
                tool_instance = tool_class()
                try:
                    result = await tool_instance.arun(action_instance)
                    print(f"Debug - Result: {result.result}")

                    next_query = f"Based on the tool result: {result.result}, please provide the final response to the user's original query: {query}"
                    next_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=next_query))

                    print(
                        f"Debug - subsequent orchestrator_output type: {type(next_output)}, fields: {next_output.model_dump()}"
                    )

                    if hasattr(next_output, "action"):
                        action_instance = next_output.action
                        if hasattr(next_output, "reasoning"):
                            print(f"Debug - Orchestrator reasoning: {next_output.reasoning}")
                    else:
                        action_instance = FinalResponseSchema(response_text=next_output.chat_message)

                except Exception as e:
                    print(f"Debug - Error executing tool: {e}")
                    return (
                        f"I encountered an error while executing the tool: {str(e)}",
                        tools_used,
                        resources_used,
                        prompts_used,
                    )

            # Check for resource
            resource_class = globals().get("resource_schema_map", {}).get(schema_type)
            if resource_class:
                schema_type_valid = True
                resource_name = getattr(resource_class, "mcp_resource_name", "unknown")
                resources_used.append(resource_name)

                print(f"Debug - Fetching resource {resource_name}...")
                print(f"Debug - Parameters: {action_instance.model_dump()}")
                resource_instance = resource_class()
                try:
                    result = await resource_instance.aread(action_instance)  # type: ignore
                    print(f"Debug - Result: {result.content}")

                    next_query = (
                        f"Based on the resource content: {result.content}, please provide "
                        f"the final response to the user's original query: {query}"
                    )
                    next_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=next_query))

                    if hasattr(next_output, "action"):
                        action_instance = next_output.action
                        if hasattr(next_output, "reasoning"):
                            print(f"Debug - Orchestrator reasoning: {next_output.reasoning}")
                    else:
                        action_instance = FinalResponseSchema(response_text=getattr(next_output, "chat_message", "No response"))  # type: ignore

                except Exception as e:
                    print(f"Debug - Error fetching resource: {e}")
                    return (
                        f"I encountered an error while fetching the resource: {str(e)}",
                        tools_used,
                        resources_used,
                        prompts_used,
                    )

            # Check for prompt
            prompt_class = globals().get("prompt_schema_map", {}).get(schema_type)  # type: ignore
            if prompt_class:
                schema_type_valid = True
                prompt_name = getattr(prompt_class, "mcp_prompt_name", "unknown")  # type: ignore
                prompts_used.append(prompt_name)

                print(f"Debug - Using prompt {prompt_name}...")
                print(f"Debug - Parameters: {action_instance.model_dump()}")
                prompt_instance = prompt_class()
                try:
                    result = await prompt_instance.agenerate(action_instance)  # type: ignore
                    print(f"Debug - Result: {result.content}")

                    next_query = (
                        f"Based on the prompt content: {result.content}, please provide "
                        f"the final response to the user's original query: {query}"
                    )
                    next_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=next_query))

                    if hasattr(next_output, "action"):
                        action_instance = next_output.action
                        if hasattr(next_output, "reasoning"):
                            print(f"Debug - Orchestrator reasoning: {next_output.reasoning}")
                    else:
                        action_instance = FinalResponseSchema(response_text=getattr(next_output, "chat_message", "No response"))  # type: ignore

                except Exception as e:
                    print(f"Debug - Error using prompt: {e}")
                    return f"I encountered an error while using the prompt: {str(e)}", tools_used, resources_used, prompts_used

            if not schema_type_valid:
                print(f"Debug - Error: No tool/resource/prompt found for schema {schema_type}")
                return (
                    "I encountered an internal error. Could not find the appropriate tool/resource/prompt.",
                    tools_used,
                    resources_used,
                    prompts_used,
                )

        if iteration_count >= max_iterations:
            print(f"Debug - Hit max iterations ({max_iterations}), forcing final response")
            action_instance = FinalResponseSchema(
                response_text="I reached the maximum number of processing steps. Please try rephrasing your query."
            )

        if isinstance(action_instance, FinalResponseSchema):
            return action_instance.response_text, tools_used, resources_used, prompts_used
        else:
            return "Error: Expected final response but got something else", tools_used, resources_used, prompts_used

    except Exception as e:
        print(f"Debug - Orchestrator execution error: {e}")
        import traceback

        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Orchestrator execution failed: {e}")


@app.get("/")
async def root():
    """Root endpoint showing available tools, resources, and prompts, and following the schema structure."""
    return {
        "message": "MCP FastAPI Client Example - Agent-based Architecture",
        "available_tools": list(mcp_tools.keys()),
        "available_resources": list(mcp_resources.keys()),
        "available_prompts": list(mcp_prompts.keys()),
        "tool_schemas": {
            name: tool.input_schema.__name__ if hasattr(tool, "input_schema") else "N/A" for name, tool in mcp_tools.items()
        },
        "endpoints": {
            "calculate": "/calculate - Natural language queries using agent orchestration (e.g., 'multiply 15 by 3')"
        },
        "example_usage": {
            "natural_language": {
                "endpoint": "/calculate",
                "body": {"query": "What is 25 divided by 5?"},
                "description": "Agent will determine the appropriate tool, resource, or prompt",
            }
        },
        "config": {
            "mcp_server_url": config.mcp_server_url if config else "Not initialized",
            "model": config.openai_model if config else "Not initialized",
        },
    }


@app.post("/calculate", response_model=CalculationResponse)
async def calculate_with_agent(request: NaturalLanguageRequest):
    """Calculate using agent-based orchestration with natural language input."""
    try:
        result_text, tools_used, resources_used, prompts_used = await execute_with_orchestrator_async(request.query)
        return CalculationResponse(
            result=result_text,
            tools_used=tools_used,
            resources_used=resources_used,
            prompts_used=prompts_used,
            query=request.query,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Agent calculation failed: {e}")


@app.post("/load_resource", response_model=ResourceResponse)
async def load_resource(request: NaturalLanguageRequest):
    """Calculate using agent-based orchestration with natural language input."""
    try:
        result_text, tools_used, resources_used, prompts_used = await execute_with_orchestrator_async(request.query)
        return ResourceResponse(
            content=result_text,
            tools_used=tools_used,
            resources_used=resources_used,
            prompts_used=prompts_used,
            query=request.query,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Agent resource utilization failed: {e}")


@app.post("/load_prompt", response_model=PromptResponse)
async def load_prompt(request: NaturalLanguageRequest):
    """Calculate using agent-based orchestration with natural language input."""
    try:
        result_text, tools_used, resources_used, prompts_used = await execute_with_orchestrator_async(request.query)
        return PromptResponse(
            content=result_text,
            prompts_fetched=prompts_used,
            tools_used=tools_used,
            resources_used=resources_used,
            query=request.query,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Agent prompt generation failed: {e}")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)


# To test the tool usage:
# curl -X POST http://localhost:8000/calculate -H "Content-Type: application/json" \
#   -d '{"query": "What is 3986733+3375486? Use the tool provided."}' | python -m json.tool
# To test the resource usage:
# curl -X POST http://localhost:8000/load_resource -H "Content-Type: application/json" \
#   -d '{"query": "What is the weather in Dallas?"}' | python -m json.tool
# To test the prompt usage:
# curl -X POST http://localhost:8000/load_prompt -H "Content-Type: application/json" \
#   -d '{"query": "Use the greeting prompt to say hello to Alex."}' | python -m json.tool

```

### File: atomic-examples/mcp-agent/example-client/example_client/main_http.py

```python
"""
HTTP Stream transport client for MCP Agent example.
Communicates with the server_http.py `/mcp` endpoint using HTTP GET/POST/DELETE for JSON-RPC streams.
"""

from atomic_agents.connectors.mcp import (
    fetch_mcp_tools,
    fetch_mcp_resources,
    fetch_mcp_prompts,
    MCPTransportType,
)
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
import sys
from rich.console import Console
from rich.table import Table
from rich.markdown import Markdown
from pydantic import Field
import openai
import os
import instructor
from typing import Union, Type, Dict
from dataclasses import dataclass


@dataclass
class MCPConfig:
    """Configuration for the MCP Agent system using HTTP Stream transport."""

    mcp_server_url: str = "http://localhost:6969"
    openai_model: str = "gpt-5-mini"
    openai_api_key: str = os.getenv("OPENAI_API_KEY")
    reasoning_effort: str = "low"

    def __post_init__(self):
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is not set")


def main():
    # Use default HTTP transport settings from MCPConfig
    config = MCPConfig()
    console = Console()
    client = instructor.from_openai(openai.OpenAI(api_key=config.openai_api_key))

    console.print("[bold green]Initializing MCP Agent System (HTTP Stream mode)...[/bold green]")
    tools = fetch_mcp_tools(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.HTTP_STREAM)
    resources = fetch_mcp_resources(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.HTTP_STREAM)
    prompts = fetch_mcp_prompts(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.HTTP_STREAM)
    if not tools and not resources and not prompts:
        console.print(f"[bold red]No MCP tools or resources or prompts found at {config.mcp_server_url}[/bold red]")
        sys.exit(1)

    # Display available tools
    table = Table(title="Available MCP Tools", box=None)
    table.add_column("Tool Name", style="cyan")
    table.add_column("Input Schema", style="yellow")
    table.add_column("Description", style="magenta")
    for ToolClass in tools:
        schema_name = getattr(ToolClass.input_schema, "__name__", "N/A")
        table.add_row(ToolClass.mcp_tool_name, schema_name, ToolClass.__doc__ or "")
    console.print(table)

    # Display resources and prompts if available
    if resources:
        rtable = Table(title="Available MCP Resources", box=None)
        rtable.add_column("Name", style="cyan")
        rtable.add_column("Description", style="magenta")
        rtable.add_column("Input Schema", style="yellow")
        for ResourceClass in resources:
            schema_name = getattr(ResourceClass.input_schema, "__name__", "N/A")
            rtable.add_row(ResourceClass.mcp_resource_name, schema_name, ResourceClass.__doc__ or "")
        console.print(rtable)
    if prompts:
        ptable = Table(title="Available MCP Prompts", box=None)
        ptable.add_column("Name", style="cyan")
        ptable.add_column("Description", style="magenta")
        ptable.add_column("Input Schema", style="yellow")
        for PromptClass in prompts:
            schema_name = getattr(PromptClass.input_schema, "__name__", "N/A")
            ptable.add_row(PromptClass.mcp_prompt_name, schema_name, PromptClass.__doc__ or "")
        console.print(ptable)

    # Build orchestrator
    class MCPOrchestratorInputSchema(BaseIOSchema):
        """Input schema for the MCP orchestrator that processes user queries."""

        query: str = Field(...)

    class FinalResponseSchema(BaseIOSchema):
        """Schema for the final response to the user."""

        response_text: str = Field(...)

    # Map schemas and define ActionUnion
    tool_schema_map: Dict[Type[BaseIOSchema], Type] = {
        ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")
    }
    resource_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
        ResourceClass.input_schema: ResourceClass for ResourceClass in resources if hasattr(ResourceClass, "input_schema")
    }  # type: ignore
    prompt_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
        PromptClass.input_schema: PromptClass for PromptClass in prompts if hasattr(PromptClass, "input_schema")
    }  # type: ignore
    available_schemas = (
        tuple(tool_schema_map.keys())
        + tuple(resource_schema_to_class_map.keys())
        + tuple(prompt_schema_to_class_map.keys())
        + (FinalResponseSchema,)
    )
    ActionUnion = Union[available_schemas]

    class OrchestratorOutputSchema(BaseIOSchema):
        """Output schema for the MCP orchestrator containing reasoning and selected action."""

        reasoning: str
        action: ActionUnion = Field(  # type: ignore[reportInvalidTypeForm]
            ...,
            description="The chosen action: either a tool/resource/prompt's input schema instance or a final response schema instance.",
        )

    history = ChatHistory()
    orchestrator_agent = AtomicAgent[MCPOrchestratorInputSchema, OrchestratorOutputSchema](
        AgentConfig(
            client=client,
            model=config.openai_model,
            model_api_parameters={"reasoning_effort": config.reasoning_effort},
            history=history,
            system_prompt_generator=SystemPromptGenerator(
                background=[
                    "You are an MCP Orchestrator Agent, designed to chat with users and",
                    "determine the best way to handle their queries using the available tools, resources, and prompts.",
                ],
                steps=[
                    "1. Use the reasoning field to determine if one or more successive "
                    "tool/resource/prompt calls could be used to handle the user's query.",
                    "2. If so, choose the appropriate tool(s), resource(s), or prompt(s) one "
                    "at a time and extract all necessary parameters from the query.",
                    "3. If a single tool/resource/prompt can not be used to handle the user's query, "
                    "think about how to break down the query into "
                    "smaller tasks and route them to the appropriate tool(s)/resource(s)/prompt(s).",
                    "4. If no sequence of tools/resources/prompts could be used, or if you are "
                    "finished processing the user's query, provide a final response to the user.",
                    "5. If the context is sufficient and no more tools/resources/prompts are needed, provide a final response to the user.",
                ],
                output_instructions=[
                    "1. Always provide a detailed explanation of your decision-making process in the 'reasoning' field.",
                    "2. Choose exactly one action schema (either a tool/resource/prompt input or FinalResponseSchema).",
                    "3. Ensure all required parameters for the chosen tool/resource/prompt are properly extracted and validated.",
                    "4. Maintain a professional and helpful tone in all responses.",
                    "5. Break down complex queries into sequential tool/resource/prompt calls "
                    "before giving the final answer via `FinalResponseSchema`.",
                ],
            ),
        )
    )

    console.print("[bold green]HTTP Stream client ready. Type 'exit' to quit.[/bold green]")
    while True:
        query = console.input("[bold yellow]You:[/bold yellow] ").strip()
        if query.lower() in {"exit", "quit"}:
            break
        if not query:
            continue

        try:
            # Initial run with user query
            orchestrator_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=query))

            # Debug output to see what's actually in the output
            console.print(
                f"[dim]Debug - orchestrator_output type: {type(orchestrator_output)}, fields: {orchestrator_output.model_dump()}"
            )

            # Handle the output similar to SSE version
            if hasattr(orchestrator_output, "chat_message") and not hasattr(orchestrator_output, "action"):
                # Convert BasicChatOutputSchema to FinalResponseSchema
                action_instance = FinalResponseSchema(response_text=orchestrator_output.chat_message)
                reasoning = "Response generated directly from chat model"
            elif hasattr(orchestrator_output, "action"):
                action_instance = orchestrator_output.action
                reasoning = (
                    orchestrator_output.reasoning if hasattr(orchestrator_output, "reasoning") else "No reasoning provided"
                )
            else:
                console.print("[yellow]Warning: Unexpected response format. Unable to process.[/yellow]")
                continue

            console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

            # Keep executing until we get a final response
            while not isinstance(action_instance, FinalResponseSchema):
                schema_type = type(action_instance)
                schema_type_valid = False

                try:
                    ToolClass = tool_schema_map.get(schema_type)
                    if ToolClass:
                        schema_type_valid = True
                        tool_name = ToolClass.mcp_tool_name
                        console.print(f"[blue]Executing tool:[/blue] {tool_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        tool_instance = ToolClass()
                        # The persistent session/loop are already part of the ToolClass definition
                        tool_output = tool_instance.run(action_instance)
                        console.print(f"[bold green]Result:[/bold green] {tool_output.result}")

                        # Add tool result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Tool {tool_name} executed with result: " f"{tool_output.result}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    ResourceClass = resource_schema_to_class_map.get(schema_type)
                    if ResourceClass:
                        schema_type_valid = True
                        resource_name = ResourceClass.mcp_resource_name
                        console.print(f"[blue]Reading resource:[/blue] {resource_name}")
                        console.print(f"[dim]Parameters: {action_instance.model_dump()}")

                        resource_instance = ResourceClass()
                        resource_output = resource_instance.read(action_instance)
                        console.print(f"[bold green]Resource content:[/bold green] {resource_output.content}")

                        # Add resource result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Resource {resource_name} read with content: {resource_output.content}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    PromptClass = prompt_schema_to_class_map.get(schema_type)
                    if PromptClass:
                        schema_type_valid = True
                        prompt_name = PromptClass.mcp_prompt_name
                        console.print(f"[blue]Fetching prompt:[/blue] {prompt_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        prompt_instance = PromptClass()
                        prompt_output = prompt_instance.generate(action_instance)
                        console.print(f"[bold green]Prompt content:[/bold green] {prompt_output.content}")

                        # Add prompt result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Prompt {prompt_name} generated content: {prompt_output.content}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    if not schema_type_valid:
                        console.print(f"[red]Error: Unknown schema type {schema_type.__name__}[/red]")
                        action_instance = FinalResponseSchema(
                            response_text="I encountered an internal error. Could not find the appropriate tool/resource/prompt."
                        )
                        break

                    next_output = orchestrator_agent.run()
                    if hasattr(next_output, "action"):
                        action_instance = next_output.action
                        if hasattr(next_output, "reasoning"):
                            console.print(f"[cyan]Orchestrator reasoning:[/cyan] {next_output.reasoning}")
                    else:
                        # If no action, treat as final response
                        action_instance = FinalResponseSchema(response_text=next_output.chat_message)

                except Exception as e:
                    console.print(f"[red]Error executing tool: {e}[/red]")
                    action_instance = FinalResponseSchema(
                        response_text=f"I encountered an error while executing the tool: {str(e)}"
                    )
                    break

            # Display final response
            if isinstance(action_instance, FinalResponseSchema):
                md = Markdown(action_instance.response_text)
                console.print("[bold blue]Agent:[/bold blue]")
                console.print(md)
            else:
                console.print("[red]Error: Expected final response but got something else[/red]")

        except Exception as e:
            console.print(f"[red]Error: {e}[/red]")


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-client/example_client/main_sse.py

```python
# pyright: reportInvalidTypeForm=false
from atomic_agents.connectors.mcp import (
    fetch_mcp_tools,
    fetch_mcp_resources,
    fetch_mcp_prompts,
    MCPTransportType,
)
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from rich.console import Console
from rich.table import Table
from rich.markdown import Markdown
import openai
import os
import instructor
from pydantic import Field
from typing import Union, Type, Dict
from dataclasses import dataclass
import re


# 1. Configuration and environment setup
@dataclass
class MCPConfig:
    """Configuration for the MCP Agent system using SSE transport."""

    mcp_server_url: str = "http://localhost:6969"

    # NOTE: In contrast to other examples, we use gpt-5.1 and not gpt-5-mini here.
    # In my tests, gpt-5-mini was not smart enough to deal with multiple tools like that
    # and at the moment MCP does not yet allow for adding sufficient metadata to
    # clarify tools even more and introduce more constraints.
    openai_model: str = "gpt-5.1"
    openai_api_key: str = os.getenv("OPENAI_API_KEY")
    reasoning_effort: str = "low"

    def __post_init__(self):
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is not set")


config = MCPConfig()
console = Console()
client = instructor.from_openai(openai.OpenAI(api_key=config.openai_api_key))


class FinalResponseSchema(BaseIOSchema):
    """Schema for providing a final text response to the user."""

    response_text: str = Field(..., description="The final text response to the user's query")


# Fetch tools and build ActionUnion statically
tools = fetch_mcp_tools(
    mcp_endpoint=config.mcp_server_url,
    transport_type=MCPTransportType.SSE,
)
resources = fetch_mcp_resources(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.SSE)
prompts = fetch_mcp_prompts(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.SSE)
if not tools and not resources and not prompts:
    raise RuntimeError("No MCP tools/resources/prompts found. Please ensure the MCP server is running and accessible.")

# Build mapping from input_schema to ToolClass
tool_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
    ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")
}
# Collect all tool input schemas
tool_input_schemas = tuple(tool_schema_to_class_map.keys())

resource_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
    ResourceClass.input_schema: ResourceClass for ResourceClass in resources if hasattr(ResourceClass, "input_schema")
}  # type: ignore
prompt_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
    PromptClass.input_schema: PromptClass for PromptClass in prompts if hasattr(PromptClass, "input_schema")
}  # type: ignore
available_schemas = (
    tuple(tool_schema_to_class_map.keys())
    + tuple(resource_schema_to_class_map.keys())
    + tuple(prompt_schema_to_class_map.keys())
    + (FinalResponseSchema,)
)

# Define the Union of all action schemas
ActionUnion = Union[available_schemas]


# 2. Schema and class definitions
class MCPOrchestratorInputSchema(BaseIOSchema):
    """Input schema for the MCP Orchestrator Agent."""

    query: str = Field(..., description="The user's query to analyze.")


class OrchestratorOutputSchema(BaseIOSchema):
    """Output schema for the orchestrator. Contains reasoning and the chosen action."""

    reasoning: str = Field(
        ..., description="Detailed explanation of why this action was chosen and how it will address the user's query."
    )
    action: ActionUnion = Field(  # type: ignore[reportInvalidTypeForm]
        ..., description="The chosen action: either a tool's input schema instance or a final response schema instance."
    )

    model_config = {"arbitrary_types_allowed": True}


# Helper function to format mathematical expressions for better terminal readability
def format_math_expressions(text):
    """
    Format LaTeX-style math expressions for better readability in the terminal.

    Args:
        text (str): Text containing LaTeX-style math expressions

    Returns:
        str: Text with formatted math expressions
    """
    # Replace \( and \) with formatted brackets
    text = re.sub(r"\\[\(\)]", "", text)

    # Replace LaTeX multiplication symbol with a plain x
    text = text.replace("\\times", "×")

    # Format other common LaTeX symbols
    text = text.replace("\\cdot", "·")
    text = text.replace("\\div", "÷")
    text = text.replace("\\sqrt", "√")
    text = text.replace("\\pi", "π")

    return text


# 3. Main logic and script entry point
def main():
    try:
        console.print("[bold green]Initializing MCP Agent System (SSE mode)...[/bold green]")
        resources = fetch_mcp_resources(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.SSE)
        prompts = fetch_mcp_prompts(mcp_endpoint=config.mcp_server_url, transport_type=MCPTransportType.SSE)
        # Display available tools
        table = Table(title="Available MCP Tools", box=None)
        table.add_column("Tool Name", style="cyan")
        table.add_column("Input Schema", style="yellow")
        table.add_column("Description", style="magenta")
        for ToolClass in tools:
            # Fix to handle when input_schema is a property or doesn't have __name__
            if hasattr(ToolClass, "input_schema"):
                if hasattr(ToolClass.input_schema, "__name__"):
                    schema_name = ToolClass.input_schema.__name__
                else:
                    # If it's a property, try to get the type name of the actual class
                    try:
                        schema_instance = ToolClass.input_schema
                        schema_name = schema_instance.__class__.__name__
                    except Exception:
                        schema_name = "Unknown Schema"
            else:
                schema_name = "N/A"
            table.add_row(ToolClass.mcp_tool_name, schema_name, ToolClass.__doc__ or "")
        console.print(table)

        # Display resources and prompts if available
        if resources:
            rtable = Table(title="Available MCP Resources", box=None)
            rtable.add_column("Name", style="cyan")
            rtable.add_column("Description", style="magenta")
            rtable.add_column("Input Schema", style="yellow")
            for ResourceClass in resources:
                schema_name = ResourceClass.input_schema.__name__
                rtable.add_row(ResourceClass.mcp_resource_name, ResourceClass.__doc__ or "", schema_name)
            console.print(rtable)
        if prompts:
            ptable = Table(title="Available MCP Prompts", box=None)
            ptable.add_column("Name", style="cyan")
            ptable.add_column("Description", style="magenta")
            ptable.add_column("Input Schema", style="yellow")
            for PromptClass in prompts:
                schema_name = PromptClass.input_schema.__name__
                ptable.add_row(PromptClass.mcp_prompt_name, PromptClass.__doc__ or "", schema_name)
            console.print(ptable)

        # Create and initialize orchestrator agent
        console.print("[dim]• Creating orchestrator agent...[/dim]")
        history = ChatHistory()
        orchestrator_agent = AtomicAgent[MCPOrchestratorInputSchema, OrchestratorOutputSchema](
            AgentConfig(
                client=client,
                model=config.openai_model,
                model_api_parameters={"reasoning_effort": config.reasoning_effort},
                history=history,
                system_prompt_generator=SystemPromptGenerator(
                    background=[
                        "You are an MCP Orchestrator Agent, designed to chat with users and",
                        "determine the best way to handle their queries using the available tools, resources, and prompts.",
                    ],
                    steps=[
                        "1. Use the reasoning field to determine if one or more successive "
                        "tool/resource/prompt calls could be used to handle the user's query.",
                        "2. If so, choose the appropriate tool(s), resource(s), or prompt(s) one "
                        "at a time and extract all necessary parameters from the query.",
                        "3. If a single tool/resource/prompt can not be used to handle the user's query, "
                        "think about how to break down the query into "
                        "smaller tasks and route them to the appropriate tool(s)/resource(s)/prompt(s).",
                        "4. If no sequence of tools/resources/prompts could be used, or if you are "
                        "finished processing the user's query, provide a final response to the user.",
                        "5. If the context is sufficient and no more tools/resources/prompts are needed, provide a final response to the user.",
                    ],
                    output_instructions=[
                        "1. Always provide a detailed explanation of your decision-making process in the 'reasoning' field.",
                        "2. Choose exactly one action schema (either a tool/resource/prompt input or FinalResponseSchema).",
                        "3. Ensure all required parameters for the chosen tool/resource/prompt are properly extracted and validated.",
                        "4. Maintain a professional and helpful tone in all responses.",
                        "5. Break down complex queries into sequential tool/resource/prompt calls "
                        "before giving the final answer via `FinalResponseSchema`.",
                    ],
                ),
            )
        )
        console.print("[green]Successfully created orchestrator agent.[/green]")
        # Interactive chat loop
        console.print("[bold green]MCP Agent Interactive Chat (SSE mode). Type 'exit' or 'quit' to leave.[/bold green]")
        while True:
            query = console.input("[bold yellow]You:[/bold yellow] ").strip()
            if query.lower() in {"exit", "quit"}:
                console.print("[bold red]Exiting chat. Goodbye![/bold red]")
                break
            if not query:
                continue  # Ignore empty input
            try:
                # Initial run with user query
                orchestrator_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=query))

                # Debug output to see what's actually in the output
                console.print(
                    f"[dim]Debug - orchestrator_output type: {type(orchestrator_output)}, fields: {orchestrator_output.model_dump()}"
                )

                # The model is returning a BasicChatOutputSchema instead of OrchestratorOutputSchema
                # We need to handle this case by creating a FinalResponseSchema directly
                if hasattr(orchestrator_output, "chat_message") and not hasattr(orchestrator_output, "action"):
                    console.print("[yellow]Note: Converting BasicChatOutputSchema to FinalResponseSchema[/yellow]")
                    action_instance = FinalResponseSchema(response_text=orchestrator_output.chat_message)
                    reasoning = "Response generated directly from chat model"
                # Handle the original expected format if it exists
                elif hasattr(orchestrator_output, "action"):
                    action_instance = orchestrator_output.action
                    reasoning = (
                        orchestrator_output.reasoning if hasattr(orchestrator_output, "reasoning") else "No reasoning provided"
                    )
                else:
                    console.print("[yellow]Warning: Unexpected response format. Unable to process.[/yellow]")
                    continue

                console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

                # Keep executing until we get a final response
                while not isinstance(action_instance, FinalResponseSchema):
                    # Handle the case where action_instance is a dictionary
                    if isinstance(action_instance, dict):
                        console.print(
                            "[yellow]Warning: Received dictionary instead of schema object. Attempting to convert...[/yellow]"
                        )
                        console.print(f"[dim]Dictionary contents: {action_instance}[/dim]")

                        # Special handling for function-call format {"recipient_name": "functions.toolname", "parameters": {...}}
                        if "recipient_name" in action_instance and "parameters" in action_instance:
                            console.print("[yellow]Detected function call format with recipient_name and parameters[/yellow]")
                            recipient = action_instance.get("recipient_name", "")
                            parameters = action_instance.get("parameters", {})

                            # Extract tool name from recipient (format might be "functions.toolname")
                            tool_parts = recipient.split(".")
                            if len(tool_parts) > 1:
                                tool_name = tool_parts[-1]  # Take last part after the dot
                                console.print(
                                    f"[yellow]Extracted tool name '{tool_name}' from recipient '{recipient}'[/yellow]"
                                )

                                # Special case for calculator
                                if tool_name.lower() == "calculate":
                                    tool_name = "Calculator"
                                    console.print("[yellow]Mapped 'calculate' to 'Calculator' tool[/yellow]")

                                # Try to find a matching tool class by name
                                matching_tool = next((t for t in tools if t.mcp_tool_name.lower() == tool_name.lower()), None)

                                if matching_tool:
                                    try:
                                        # Create an instance using the parameters
                                        action_instance = matching_tool.input_schema(**parameters)
                                        console.print(
                                            f"[green]Successfully created {matching_tool.input_schema.__name__} from function call format[/green]"
                                        )
                                        continue
                                    except Exception as e:
                                        console.print(f"[red]Error creating schema from function parameters: {e}[/red]")

                        # Try to find a tool_name in the dictionary (original approach)
                        tool_name = action_instance.get("tool_name")

                        # If tool_name is not found, try alternative approaches to identify the tool
                        if not tool_name:
                            # Approach 1: Look for a field that might contain a tool name
                            for key in action_instance.keys():
                                if "tool" in key.lower():
                                    tool_name = action_instance.get(key)
                                    if tool_name:
                                        console.print(
                                            f"[yellow]Found potential tool name '{tool_name}' in field '{key}'[/yellow]"
                                        )

                            # Approach 2: Try to match dictionary fields with tool schemas
                            if not tool_name:
                                console.print("[yellow]Trying to match dictionary fields with available tools...[/yellow]")
                                best_match = None
                                best_match_score = 0

                                for ToolClass in tools:
                                    if not hasattr(ToolClass, "input_schema"):
                                        continue

                                    # Try to create a sample instance to get field names
                                    try:
                                        schema_fields = set(
                                            ToolClass.input_schema.__annotations__.keys()
                                            if hasattr(ToolClass.input_schema, "__annotations__")
                                            else []
                                        )
                                        dict_fields = set(action_instance.keys())

                                        # Count matching fields
                                        matching_fields = len(schema_fields.intersection(dict_fields))
                                        if matching_fields > best_match_score and matching_fields > 0:
                                            best_match_score = matching_fields
                                            best_match = ToolClass
                                            console.print(
                                                f"[dim]Found {matching_fields} matching fields with {ToolClass.mcp_tool_name}[/dim]"
                                            )
                                    except Exception as e:
                                        console.print(
                                            f"[dim]Error checking {getattr(ToolClass, 'mcp_tool_name', 'unknown tool')}: {str(e)}[/dim]"
                                        )

                                if best_match:
                                    tool_name = best_match.mcp_tool_name
                                    console.print(
                                        f"[yellow]Best matching tool: {tool_name} with {best_match_score} matching fields[/yellow]"
                                    )

                        if not tool_name:
                            # Final fallback: Check if this might be a final response
                            if any(
                                key in action_instance for key in ["response_text", "text", "response", "message", "content"]
                            ):
                                response_content = (
                                    action_instance.get("response_text")
                                    or action_instance.get("text")
                                    or action_instance.get("response")
                                    or action_instance.get("message")
                                    or action_instance.get("content")
                                    or "No message content found"
                                )
                                console.print("[yellow]Appears to be a final response. Converting directly.[/yellow]")
                                action_instance = FinalResponseSchema(response_text=response_content)
                                continue

                            console.print("[red]Error: Could not determine tool type from dictionary[/red]")
                            # Create a final response with an error message
                            action_instance = FinalResponseSchema(
                                response_text="I encountered an internal error. The tool could not be determined from the response. "
                                "Please try rephrasing your question."
                            )
                            break

                        # Try to find a matching tool class by name
                        matching_tool = next((t for t in tools if t.mcp_tool_name == tool_name), None)
                        if not matching_tool:
                            console.print(f"[red]Error: No tool found with name {tool_name}[/red]")
                            # Create a final response with an error message
                            action_instance = FinalResponseSchema(
                                response_text=f"I encountered an internal error. Could not find tool named '{tool_name}'."
                            )
                            break

                        # Create an instance of the input schema with the dictionary data
                        try:
                            # Remove tool_name if it's not a field in the schema
                            params = {}
                            has_annotations = hasattr(matching_tool.input_schema, "__annotations__")

                            for k, v in action_instance.items():
                                # Include the key-value pair if it's not "tool_name" or if it's a valid field in the schema
                                if k not in ["tool_name"] or (
                                    has_annotations and k in matching_tool.input_schema.__annotations__.keys()
                                ):
                                    params[k] = v

                            action_instance = matching_tool.input_schema(**params)
                            console.print(
                                f"[green]Successfully converted dictionary to {matching_tool.input_schema.__name__}[/green]"
                            )
                        except Exception as e:
                            console.print(f"[red]Error creating schema instance: {e}[/red]")
                            # Create a final response with an error message
                            action_instance = FinalResponseSchema(
                                response_text=f"I encountered an internal error when trying to use the {tool_name} tool: {str(e)}"
                            )
                            break

                    schema_type = type(action_instance)
                    schema_type_valid = False

                    ToolClass = tool_schema_to_class_map.get(schema_type)
                    if ToolClass:
                        schema_type_valid = True
                        tool_name = ToolClass.mcp_tool_name
                        console.print(f"[blue]Executing tool:[/blue] {tool_name}")
                        console.print(f"[dim]Parameters: {action_instance.model_dump()}")

                        tool_instance = ToolClass()
                        tool_output = tool_instance.run(action_instance)
                        console.print(f"[bold green]Result:[/bold green] {tool_output.result}")

                        # Add tool result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=f"Tool {tool_name} executed with result: {tool_output.result}"
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    ResourceClass = resource_schema_to_class_map.get(schema_type)
                    if ResourceClass:
                        schema_type_valid = True
                        resource_name = ResourceClass.mcp_resource_name  # type: ignore
                        console.print(f"[blue]Fetching resource:[/blue] {resource_name}")
                        console.print(f"[dim]Parameters: {action_instance.model_dump()}")

                        resource_instance = ResourceClass()  # type: ignore
                        resource_output = resource_instance.read(action_instance)  # type: ignore
                        console.print(f"[bold green]Result:[/bold green] {resource_output.content}")

                        # Add resource result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=f"Resource {resource_name} used to fetch content: {resource_output.content}"
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    PromptClass = prompt_schema_to_class_map.get(schema_type)
                    if PromptClass:
                        schema_type_valid = True
                        prompt_name = PromptClass.mcp_prompt_name  # type: ignore
                        console.print(f"[blue]Using prompt:[/blue] {prompt_name}")
                        console.print(f"[dim]Parameters: {action_instance.model_dump()}")

                        prompt_instance = PromptClass()  # type: ignore
                        prompt_output = prompt_instance.generate(action_instance)  # type: ignore
                        console.print(f"[bold green]Result:[/bold green] {prompt_output.content}")

                        # Add prompt result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=f"Prompt {prompt_name} created: {prompt_output.content}"
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    if not schema_type_valid:
                        console.print(f"[red]Unknown schema type '{schema_type.__name__}' returned by orchestrator[/red]")
                        # Create a final response with an error message
                        action_instance = FinalResponseSchema(
                            response_text="I encountered an internal error. The tool/resource/prompt type could not be recognized."
                        )
                        break

                    # Run the agent again without parameters to continue the flow
                    orchestrator_output = orchestrator_agent.run()

                    # Debug output for subsequent responses
                    console.print(
                        f"[dim]Debug - subsequent orchestrator_output type: {type(orchestrator_output)}, fields: {orchestrator_output.model_dump()}"
                    )

                    # Handle different response formats
                    if hasattr(orchestrator_output, "chat_message") and not hasattr(orchestrator_output, "action"):
                        console.print("[yellow]Note: Converting BasicChatOutputSchema to FinalResponseSchema[/yellow]")
                        action_instance = FinalResponseSchema(response_text=orchestrator_output.chat_message)
                        reasoning = "Response generated directly from chat model"
                    elif hasattr(orchestrator_output, "action"):
                        action_instance = orchestrator_output.action
                        reasoning = (
                            orchestrator_output.reasoning
                            if hasattr(orchestrator_output, "reasoning")
                            else "No reasoning provided"
                        )
                    else:
                        console.print("[yellow]Warning: Unexpected response format. Unable to process.[/yellow]")
                        break

                    console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

                # Final response from the agent
                response_text = getattr(
                    action_instance, "response_text", getattr(action_instance, "chat_message", str(action_instance))
                )
                md = Markdown(response_text)
                # Render the response as markdown
                console.print("[bold blue]Agent: [/bold blue]")
                console.print(md)

            except Exception as e:
                console.print(f"[red]Error processing query:[/red] {str(e)}")
                console.print_exception()
    except Exception as e:
        console.print(f"[bold red]Fatal error:[/bold red] {str(e)}")
        console.print_exception()


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-client/example_client/main_stdio.py

```python
# pyright: reportInvalidTypeForm=false
from atomic_agents.connectors.mcp import (
    fetch_mcp_tools,
    fetch_mcp_resources,
    fetch_mcp_prompts,
    MCPTransportType,
)
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from rich.console import Console
from rich.table import Table
import openai
import os
import instructor
import asyncio
import shlex
from contextlib import AsyncExitStack
from pydantic import Field
from typing import Union, Type, Dict, Optional
from dataclasses import dataclass
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


# 1. Configuration and environment setup
@dataclass
class MCPConfig:
    """Configuration for the MCP Agent system using STDIO transport."""

    # NOTE: In contrast to other examples, we use gpt-5-mini and not gpt-5-mini here.
    # In my tests, gpt-5-mini was not smart enough to deal with multiple tools like that
    # and at the moment MCP does not yet allow for adding sufficient metadata to
    # clarify tools even more and introduce more constraints.
    openai_model: str = "gpt-5-mini"
    openai_api_key: str = os.getenv("OPENAI_API_KEY")
    reasoning_effort: str = "low"

    # Command to run the STDIO server.
    # In practice, this could be something like "pipx some-other-persons-server or npx some-other-persons-server
    # if working with a server you did not write yourself.
    mcp_stdio_server_command: str = "uv run example-mcp-server --mode stdio"

    def __post_init__(self):
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is not set")


config = MCPConfig()
console = Console()
client = instructor.from_openai(openai.OpenAI(api_key=config.openai_api_key))


class FinalResponseSchema(BaseIOSchema):
    """Schema for providing a final text response to the user."""

    response_text: str = Field(..., description="The final text response to the user's query")


# --- Bootstrap persistent STDIO session ---
stdio_session: Optional[ClientSession] = None
stdio_loop: Optional[asyncio.AbstractEventLoop] = None
stdio_exit_stack: Optional[AsyncExitStack] = None

# Initialize STDIO session
stdio_loop = asyncio.new_event_loop()


async def _bootstrap_stdio():
    global stdio_exit_stack  # Allow modification of the global variable
    stdio_exit_stack = AsyncExitStack()
    command_parts = shlex.split(config.mcp_stdio_server_command)
    server_params = StdioServerParameters(command=command_parts[0], args=command_parts[1:], env=None)
    read_stream, write_stream = await stdio_exit_stack.enter_async_context(stdio_client(server_params))
    session = await stdio_exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
    await session.initialize()
    return session


stdio_session = stdio_loop.run_until_complete(_bootstrap_stdio())
# The stdio_exit_stack is kept to clean up later

# Fetch tools and build ActionUnion statically
tools = fetch_mcp_tools(
    mcp_endpoint=None,
    transport_type=MCPTransportType.STDIO,
    client_session=stdio_session,  # Pass persistent session
    event_loop=stdio_loop,  # Pass corresponding loop
)
resources = fetch_mcp_resources(
    mcp_endpoint=None, transport_type=MCPTransportType.STDIO, client_session=stdio_session, event_loop=stdio_loop
)
prompts = fetch_mcp_prompts(
    mcp_endpoint=None, transport_type=MCPTransportType.STDIO, client_session=stdio_session, event_loop=stdio_loop
)
if not tools and not resources and not prompts:
    raise RuntimeError("No MCP tools or resources or prompts found. Please ensure the MCP server is running and accessible.")

# Build mapping from input_schema to ToolClass
tool_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
    ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")
}
# Collect all tool input schemas
tool_input_schemas = tuple(tool_schema_to_class_map.keys())

# Build mapping for resources and prompts
resource_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
    ResourceClass.input_schema: ResourceClass for ResourceClass in resources if hasattr(ResourceClass, "input_schema")
}  # type: ignore
resource_input_schemas = tuple(resource_schema_to_class_map.keys())

prompt_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
    PromptClass.input_schema: PromptClass for PromptClass in prompts if hasattr(PromptClass, "input_schema")
}  # type: ignore
prompt_input_schemas = tuple(prompt_schema_to_class_map.keys())

# Available schemas include all tool input schemas, resource schemas, prompts and the final response schema
available_schemas = tool_input_schemas + resource_input_schemas + prompt_input_schemas + (FinalResponseSchema,)

# Define the Union of all action schemas
ActionUnion = Union[available_schemas]


# 2. Schema and class definitions
class MCPOrchestratorInputSchema(BaseIOSchema):
    """Input schema for the MCP Orchestrator Agent."""

    query: str = Field(..., description="The user's query to analyze.")


class OrchestratorOutputSchema(BaseIOSchema):
    """Output schema for the orchestrator. Contains reasoning and the chosen action."""

    reasoning: str = Field(
        ..., description="Detailed explanation of why this action was chosen and how it will address the user's query."
    )
    action: ActionUnion = Field(  # type: ignore[reportInvalidTypeForm]
        ...,
        description="The chosen action: either a tool/resource/prompt's input schema instance or a final response schema instance.",
    )

    model_config = {"arbitrary_types_allowed": True}


# 3. Main logic and script entry point
def main():
    try:
        console.print("[bold green]Initializing MCP Agent System (STDIO mode)...[/bold green]")
        # Display available tools
        table = Table(title="Available MCP Tools", box=None)
        table.add_column("Tool Name", style="cyan")
        table.add_column("Input Schema", style="yellow")
        table.add_column("Description", style="magenta")
        for ToolClass in tools:
            schema_name = ToolClass.input_schema.__name__ if hasattr(ToolClass, "input_schema") else "N/A"
            table.add_row(ToolClass.mcp_tool_name, schema_name, ToolClass.__doc__ or "")
        console.print(table)

        # Display resources and prompts if available
        if resources:
            rtable = Table(title="Available MCP Resources", box=None)
            rtable.add_column("Name", style="cyan")
            rtable.add_column("Description", style="magenta")
            rtable.add_column("Input Schema", style="yellow")
            for ResourceClass in resources:
                schema_name = ResourceClass.input_schema.__name__ if hasattr(ResourceClass, "input_schema") else "N/A"
                rtable.add_row(ResourceClass.mcp_resource_name, ResourceClass.__doc__ or "", schema_name)
            console.print(rtable)
        if prompts:
            ptable = Table(title="Available MCP Prompts", box=None)
            ptable.add_column("Name", style="cyan")
            ptable.add_column("Description", style="magenta")
            ptable.add_column("Input Schema", style="yellow")
            for PromptClass in prompts:
                schema_name = PromptClass.input_schema.__name__ if hasattr(PromptClass, "input_schema") else "N/A"
                ptable.add_row(PromptClass.mcp_prompt_name, PromptClass.__doc__ or "", schema_name)
            console.print(ptable)

        # Create and initialize orchestrator agent
        console.print("[dim]• Creating orchestrator agent...[/dim]")
        history = ChatHistory()
        orchestrator_agent = AtomicAgent[MCPOrchestratorInputSchema, OrchestratorOutputSchema](
            AgentConfig(
                client=client,
                model=config.openai_model,
                model_api_parameters={"reasoning_effort": config.reasoning_effort},
                history=history,
                system_prompt_generator=SystemPromptGenerator(
                    background=[
                        "You are an MCP Orchestrator Agent, designed to chat with users and",
                        "determine the best way to handle their queries using the available tools, resources, and prompts.",
                    ],
                    steps=[
                        "1. Use the reasoning field to determine if one or more successive "
                        "tool/resource/prompt calls could be used to handle the user's query.",
                        "2. If so, choose the appropriate tool(s), resource(s), or prompt(s) one "
                        "at a time and extract all necessary parameters from the query.",
                        "3. If a single tool/resource/prompt can not be used to handle the user's query, "
                        "think about how to break down the query into "
                        "smaller tasks and route them to the appropriate tool(s)/resource(s)/prompt(s).",
                        "4. If no sequence of tools/resources/prompts could be used, or if you are "
                        "finished processing the user's query, provide a final response to the user.",
                        "5. If the context is sufficient and no more tools/resources/prompts are needed, provide a final response to the user.",
                    ],
                    output_instructions=[
                        "1. Always provide a detailed explanation of your decision-making process in the 'reasoning' field.",
                        "2. Choose exactly one action schema (either a tool/resource/prompt input or FinalResponseSchema).",
                        "3. Ensure all required parameters for the chosen tool/resource/prompt are properly extracted and validated.",
                        "4. Maintain a professional and helpful tone in all responses.",
                        "5. Break down complex queries into sequential tool/resource/prompt calls "
                        "before giving the final answer via `FinalResponseSchema`.",
                    ],
                ),
            )
        )
        console.print("[green]Successfully created orchestrator agent.[/green]")
        console.print("[bold green]MCP Agent Interactive Chat (STDIO mode). Type '/exit' or '/quit' to leave.[/bold green]")
        while True:
            query = console.input("[bold yellow]You:[/bold yellow] ").strip()
            if query.lower() in {"/exit", "/quit"}:
                console.print("[bold red]Exiting chat. Goodbye![/bold red]")
                break
            if not query:
                continue  # Ignore empty input
            try:
                # Initial run with user query
                orchestrator_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=query))
                action_instance = orchestrator_output.action
                reasoning = orchestrator_output.reasoning
                console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

                # Keep executing until we get a final response
                while not isinstance(action_instance, FinalResponseSchema):
                    schema_type = type(action_instance)
                    schema_type_valid = False

                    ToolClass = tool_schema_to_class_map.get(schema_type)
                    if ToolClass:
                        schema_type_valid = True
                        tool_name = ToolClass.mcp_tool_name
                        console.print(f"[blue]Executing tool:[/blue] {tool_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        tool_instance = ToolClass()
                        # The persistent session/loop are already part of the ToolClass definition
                        tool_output = tool_instance.run(action_instance)
                        console.print(f"[bold green]Result:[/bold green] {tool_output.result}")

                        # Add tool result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Tool {tool_name} executed with result: " f"{tool_output.result}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    ResourceClass = resource_schema_to_class_map.get(schema_type)
                    if ResourceClass:
                        schema_type_valid = True
                        resource_name = ResourceClass.mcp_resource_name
                        console.print(f"[blue]Reading resource:[/blue] {resource_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        resource_instance = ResourceClass()
                        resource_output = resource_instance.read(action_instance)
                        console.print(f"[bold green]Resource content:[/bold green] {resource_output.content}")

                        # Add resource result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Resource {resource_name} read with content: {resource_output.content}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    PromptClass = prompt_schema_to_class_map.get(schema_type)
                    if PromptClass:
                        schema_type_valid = True
                        prompt_name = PromptClass.mcp_prompt_name
                        console.print(f"[blue]Fetching prompt:[/blue] {prompt_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        prompt_instance = PromptClass()
                        prompt_output = prompt_instance.generate(action_instance)
                        console.print(f"[bold green]Prompt content:[/bold green] {prompt_output.content}")

                        # Add prompt result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f'Prompt {prompt_name} generated successfully. Content: "{prompt_output.content}"')
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    if not schema_type_valid:
                        raise ValueError(f"Unknown schema type '" f"{schema_type.__name__}" f"' returned by orchestrator")

                    # Run the agent again without parameters to continue the flow
                    orchestrator_output = orchestrator_agent.run()
                    action_instance = orchestrator_output.action
                    reasoning = orchestrator_output.reasoning
                    console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

                # Final response from the agent
                console.print(f"[bold blue]Agent:[/bold blue] {action_instance.response_text}")

            except Exception as e:
                console.print(f"[red]Error processing query:[/red] {str(e)}")
                console.print_exception()
    except Exception as e:
        console.print(f"[bold red]Fatal error:[/bold red] {str(e)}")
        console.print_exception()
        return
    finally:
        # Cleanup persistent STDIO resources
        if stdio_loop and stdio_exit_stack:
            console.print("\n[dim]Cleaning up STDIO resources...[/dim]")
            try:
                stdio_loop.run_until_complete(stdio_exit_stack.aclose())
            except Exception as cleanup_err:
                console.print(f"[red]Error during STDIO cleanup:[/red] {cleanup_err}")
            finally:
                stdio_loop.close()


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-client/example_client/main_stdio_async.py

```python
# pyright: reportInvalidTypeForm=false
from atomic_agents.connectors.mcp import (
    fetch_mcp_tools_async,
    fetch_mcp_resources_async,
    fetch_mcp_prompts_async,
    MCPToolOutputSchema,
    MCPTransportType,
)
from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from rich.console import Console
from rich.table import Table
import openai
import os
import instructor
import asyncio
import shlex
from contextlib import AsyncExitStack
from pydantic import Field
from typing import Union, Type, Dict, Any
from dataclasses import dataclass
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


# 1. Configuration and environment setup
@dataclass
class MCPConfig:
    """Configuration for the MCP Agent system using STDIO transport."""

    # NOTE: In contrast to other examples, we use gpt-5.1 and not gpt-5-mini here.
    # In my tests, gpt-5-mini was not smart enough to deal with multiple tools like that
    # and at the moment MCP does not yet allow for adding sufficient metadata to
    # clarify tools even more and introduce more constraints.
    openai_model: str = "gpt-5.1"
    openai_api_key: str = os.getenv("OPENAI_API_KEY")
    reasoning_effort: str = "low"

    # Command to run the STDIO server.
    # In practice, this could be something like "pipx some-other-persons-server or npx some-other-persons-server
    # if working with a server you did not write yourself.
    mcp_stdio_server_command: str = "uv run example-mcp-server --mode stdio"

    def __post_init__(self):
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is not set")


config = MCPConfig()
console = Console()
client = instructor.from_openai(openai.OpenAI(api_key=config.openai_api_key))


class FinalResponseSchema(BaseIOSchema):
    """Schema for providing a final text response to the user."""

    response_text: str = Field(..., description="The final text response to the user's query")


async def main():
    async with AsyncExitStack() as stack:
        # Start MCP server
        cmd, *args = shlex.split(config.mcp_stdio_server_command)
        read_stream, write_stream = await stack.enter_async_context(
            stdio_client(StdioServerParameters(command=cmd, args=args))
        )
        session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
        await session.initialize()

        # Fetch tools, resources and prompts - factory sees running loop
        tools = await fetch_mcp_tools_async(
            transport_type=MCPTransportType.STDIO,
            client_session=session,  # factory sees running loop
        )
        resources = await fetch_mcp_resources_async(
            transport_type=MCPTransportType.STDIO,
            client_session=session,
        )
        prompts = await fetch_mcp_prompts_async(
            transport_type=MCPTransportType.STDIO,
            client_session=session,
        )

        if not tools and not resources and not prompts:
            raise RuntimeError(
                "No MCP tools or resources or prompts found. Please ensure the MCP server is running and accessible."
            )

        # Build mapping from input_schema to ToolClass
        tool_schema_to_class_map: Dict[Type[BaseIOSchema], Type[AtomicAgent]] = {
            ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")
        }
        # Collect all tool input schemas
        tool_input_schemas = tuple(tool_schema_to_class_map.keys())

        # Build mapping for resources and prompts
        resource_schema_to_class_map: Dict[Type[BaseIOSchema], Any] = {  # type: ignore
            ResourceClass.input_schema: ResourceClass for ResourceClass in resources if hasattr(ResourceClass, "input_schema")
        }
        resource_input_schemas = tuple(resource_schema_to_class_map.keys())

        prompt_schema_to_class_map: Dict[Type[BaseIOSchema], Any] = {  # type: ignore
            PromptClass.input_schema: PromptClass for PromptClass in prompts if hasattr(PromptClass, "input_schema")
        }
        prompt_input_schemas = tuple(prompt_schema_to_class_map.keys())

        # Available schemas include all tool input schemas, resource schemas, prompts and the final response schema
        available_schemas = tool_input_schemas + resource_input_schemas + prompt_input_schemas + (FinalResponseSchema,)

        # Define the Union of all action schemas
        ActionUnion = Union[available_schemas]

        # 2. Schema and class definitions
        class MCPOrchestratorInputSchema(BaseIOSchema):
            """Input schema for the MCP Orchestrator Agent."""

            query: str = Field(..., description="The user's query to analyze.")

        class OrchestratorOutputSchema(BaseIOSchema):
            """Output schema for the orchestrator. Contains reasoning and the chosen action."""

            reasoning: str = Field(
                ..., description="Detailed explanation of why this action was chosen and how it will address the user's query."
            )
            action: ActionUnion = Field(  # type: ignore
                ...,
                description="The chosen action: either a tool/resource/prompt's input schema instance or a final response schema instance.",
            )

            model_config = {"arbitrary_types_allowed": True}

        # 3. Main logic
        console.print("[bold green]Initializing MCP Agent System (STDIO mode - Async)...[/bold green]")

        # Display available tools
        table = Table(title="Available MCP Tools", box=None)
        table.add_column("Tool Name", style="cyan")
        table.add_column("Input Schema", style="yellow")
        table.add_column("Description", style="magenta")
        for ToolClass in tools:
            schema_name = ToolClass.input_schema.__name__ if hasattr(ToolClass, "input_schema") else "N/A"
            table.add_row(ToolClass.mcp_tool_name, schema_name, ToolClass.__doc__ or "")
        console.print(table)

        # Display resources and prompts if available
        if resources:
            rtable = Table(title="Available MCP Resources", box=None)
            rtable.add_column("Name", style="cyan")
            rtable.add_column("Description", style="magenta")
            rtable.add_column("Input Schema", style="yellow")
            for ResourceClass in resources:
                schema_name = ResourceClass.input_schema.__name__
                rtable.add_row(ResourceClass.mcp_resource_name, ResourceClass.__doc__ or "", schema_name)
            console.print(rtable)
        if prompts:
            ptable = Table(title="Available MCP Prompts", box=None)
            ptable.add_column("Name", style="cyan")
            ptable.add_column("Description", style="magenta")
            ptable.add_column("Input Schema", style="yellow")
            for PromptClass in prompts:
                schema_name = PromptClass.input_schema.__name__
                ptable.add_row(PromptClass.mcp_prompt_name, PromptClass.__doc__ or "", schema_name)
            console.print(ptable)

        # Create and initialize orchestrator agent
        console.print("[dim]• Creating orchestrator agent...[/dim]")
        history = ChatHistory()
        orchestrator_agent = AtomicAgent[MCPOrchestratorInputSchema, OrchestratorOutputSchema](
            AgentConfig(
                client=client,
                model=config.openai_model,
                model_api_parameters={"reasoning_effort": config.reasoning_effort},
                history=history,
                system_prompt_generator=SystemPromptGenerator(
                    background=[
                        "You are an MCP Orchestrator Agent, designed to chat with users and",
                        "determine the best way to handle their queries using the available tools, resources, and prompts.",
                    ],
                    steps=[
                        "1. Use the reasoning field to determine if one or more successive "
                        "tool/resource/prompt calls could be used to handle the user's query.",
                        "2. If so, choose the appropriate tool(s), resource(s), or prompt(s) one "
                        "at a time and extract all necessary parameters from the query.",
                        "3. If a single tool/resource/prompt can not be used to handle the user's query, "
                        "think about how to break down the query into "
                        "smaller tasks and route them to the appropriate tool(s)/resource(s)/prompt(s).",
                        "4. If no sequence of tools/resources/prompts could be used, or if you are "
                        "finished processing the user's query, provide a final response to the user.",
                        "5. If the context is sufficient and no more tools/resources/prompts are needed, provide a final response to the user.",
                    ],
                    output_instructions=[
                        "1. Always provide a detailed explanation of your decision-making process in the 'reasoning' field.",
                        "2. Choose exactly one action schema (either a tool/resource/prompt input or FinalResponseSchema).",
                        "3. Ensure all required parameters for the chosen tool/resource/prompt are properly extracted and validated.",
                        "4. Maintain a professional and helpful tone in all responses.",
                        "5. Break down complex queries into sequential tool/resource/prompt calls "
                        "before giving the final answer via `FinalResponseSchema`.",
                    ],
                ),
            )
        )
        console.print("[green]Successfully created orchestrator agent.[/green]")

        # Interactive chat loop
        console.print(
            "[bold green]MCP Agent Interactive Chat (STDIO mode - Async). Type '/exit' or '/quit' to leave.[/bold green]"
        )
        while True:
            query = console.input("[bold yellow]You:[/bold yellow] ").strip()
            if query.lower() in {"/exit", "/quit"}:
                console.print("[bold red]Exiting chat. Goodbye![/bold red]")
                break
            if not query:
                continue  # Ignore empty input

            try:
                # Initial run with user query
                orchestrator_output = orchestrator_agent.run(MCPOrchestratorInputSchema(query=query))
                action_instance = orchestrator_output.action
                reasoning = orchestrator_output.reasoning
                console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

                # Keep executing until we get a final response
                while not isinstance(action_instance, FinalResponseSchema):
                    schema_type = type(action_instance)
                    schema_type_valid = False

                    ToolClass = tool_schema_to_class_map.get(schema_type)
                    if ToolClass:
                        schema_type_valid = True
                        tool_name = ToolClass.mcp_tool_name
                        console.print(f"[blue]Executing tool:[/blue] {tool_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        # Execute the MCP tool using the session directly to avoid event loop conflicts
                        arguments = action_instance.model_dump(exclude={"tool_name"}, exclude_none=True)
                        tool_result = await session.call_tool(name=tool_name, arguments=arguments)

                        # Process the result similar to how the factory does it
                        if hasattr(tool_result, "content"):
                            actual_result_content = tool_result.content
                        elif isinstance(tool_result, dict) and "content" in tool_result:
                            actual_result_content = tool_result["content"]
                        else:
                            actual_result_content = tool_result

                        # Create output schema instance
                        OutputSchema = type(
                            f"{tool_name}OutputSchema", (MCPToolOutputSchema,), {"__doc__": f"Output schema for {tool_name}"}
                        )
                        tool_output = OutputSchema(result=actual_result_content)
                        console.print(f"[bold green]Result:[/bold green] {tool_output.result}")

                        # Add tool result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Tool {tool_name} executed with result: " f"{tool_output.result}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    ResourceClass = resource_schema_to_class_map.get(schema_type)
                    if ResourceClass:
                        schema_type_valid = True
                        resource_name = ResourceClass.mcp_resource_name
                        console.print(f"[blue]Reading resource:[/blue] {resource_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        resource_instance = ResourceClass()
                        resource_output = await resource_instance.aread(action_instance)
                        console.print(f"[bold green]Resource content:[/bold green] {resource_output.content}")

                        # Add resource result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Resource {resource_name} read with content: {resource_output.content}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    PromptClass = prompt_schema_to_class_map.get(schema_type)
                    if PromptClass:
                        schema_type_valid = True
                        prompt_name = PromptClass.mcp_prompt_name
                        console.print(f"[blue]Fetching prompt:[/blue] {prompt_name}")
                        console.print(f"[dim]Parameters:[/dim] " f"{action_instance.model_dump()}")

                        prompt_instance = PromptClass()
                        prompt_output = await prompt_instance.agenerate(action_instance)
                        console.print(f"[bold green]Prompt content:[/bold green] {prompt_output.content}")

                        # Add prompt result to agent history
                        result_message = MCPOrchestratorInputSchema(
                            query=(f"Prompt {prompt_name} generated content: {prompt_output.content}")
                        )
                        orchestrator_agent.add_tool_result(result_message)

                    if not schema_type_valid:
                        raise ValueError(f"Unknown schema type '" f"{schema_type.__name__}" f"' returned by orchestrator")

                    # Run the agent again without parameters to continue the flow
                    orchestrator_output = orchestrator_agent.run()
                    action_instance = orchestrator_output.action
                    reasoning = orchestrator_output.reasoning
                    console.print(f"[cyan]Orchestrator reasoning:[/cyan] {reasoning}")

                # Final response from the agent
                console.print(f"[bold blue]Agent:[/bold blue] {action_instance.response_text}")

            except Exception as e:
                console.print(f"[red]Error processing query:[/red] {str(e)}")
                console.print_exception()


if __name__ == "__main__":
    asyncio.run(main())

```

### File: atomic-examples/mcp-agent/example-client/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["example_client"]

[project]
name = "example-client"
version = "0.1.0"
description = "Example: Choosing the right MCP tool for a user query using the MCP Tool Factory."
authors = [
    { name = "Your Name", email = "you@example.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "example-mcp-server",
    "pydantic>=2.10.3,<3.0.0",
    "rich>=13.0.0",
    "openai>=2.0.0,<3.0.0",
    "mcp[cli]>=1.9.4",
    "fastapi>=0.115.14,<1.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }
example-mcp-server = { workspace = true }

```

### File: atomic-examples/mcp-agent/example-mcp-server/demo_tools.py

```python
#!/usr/bin/env python3
"""
Demo script to list available tools from MCP servers.

This script demonstrates how to:
1. Connect to an MCP server using STDIO transport
2. Connect to an MCP server using SSE transport
3. List available tools from both transports
4. Call each available tool with appropriate input
"""

import asyncio
import random
import json
import datetime
from contextlib import AsyncExitStack
from typing import Dict, Any

# Import MCP client libraries
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client

# Rich library for pretty output
from rich.console import Console
from rich.table import Table
from rich.syntax import Syntax


class MCPClient:
    """A simple client that can connect to MCP servers using either STDIO or SSE transport."""

    def __init__(self):
        self.session = None
        self.exit_stack = AsyncExitStack()
        self.transport_type = None  # Will be set to 'stdio' or 'sse'

    async def connect_to_stdio_server(self, server_script_path: str):
        """Connect to an MCP server via STDIO transport.

        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        try:
            # Determine script type (Python or JavaScript)
            is_python = server_script_path.endswith(".py")
            is_js = server_script_path.endswith(".js")

            if not (is_python or is_js):
                raise ValueError("Server script must be a .py or .js file")

            command = "python" if is_python else "node"

            # Set up STDIO transport
            server_params = StdioServerParameters(command=command, args=[server_script_path], env=None)

            # Connect to the server
            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
            read_stream, write_stream = stdio_transport

            # Initialize the session
            self.session = await self.exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
            await self.session.initialize()
            self.transport_type = "stdio"

        except Exception as e:
            await self.cleanup()
            raise e

    async def connect_to_sse_server(self, server_url: str):
        """Connect to an MCP server via SSE transport.

        Args:
            server_url: URL of the SSE server (e.g., http://localhost:6969)
        """
        try:
            # Initialize SSE transport with the correct endpoint
            sse_transport = await self.exit_stack.enter_async_context(sse_client(f"{server_url}/sse"))
            read_stream, write_stream = sse_transport

            # Initialize the session
            self.session = await self.exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
            await self.session.initialize()
            self.transport_type = "sse"

        except Exception as e:
            await self.cleanup()
            raise e

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]):
        """Call a tool with the given arguments.

        Args:
            tool_name: Name of the tool to call
            arguments: Arguments to pass to the tool

        Returns:
            The result of the tool call
        """
        if not self.session:
            raise RuntimeError("Session not initialized")

        return await self.session.call_tool(name=tool_name, arguments=arguments)

    async def cleanup(self):
        """Clean up resources."""
        if self.session:
            await self.exit_stack.aclose()
            self.session = None
            self.transport_type = None


def generate_input_for_tool(tool_name: str, input_schema: Dict[str, Any]) -> Dict[str, Any]:
    """Generate appropriate input based on the tool name and input schema.

    This function creates sensible inputs for different tool types.

    Args:
        tool_name: The name of the tool
        input_schema: The JSON schema of the tool input

    Returns:
        A dictionary with values matching the schema
    """
    result = {}

    # Special handling for known tool types
    if tool_name == "AddNumbers":
        result = {"number1": random.randint(1, 100), "number2": random.randint(1, 100)}
    elif tool_name == "DateDifference":
        # Generate two dates with a reasonable difference
        today = datetime.date.today()
        days_diff = random.randint(1, 30)
        date1 = today - datetime.timedelta(days=days_diff)
        date2 = today
        result = {"date1": date1.isoformat(), "date2": date2.isoformat()}
    elif tool_name == "ReverseString":
        words = ["hello", "world", "testing", "reverse", "string", "tool"]
        result = {"text_to_reverse": random.choice(words)}
    elif tool_name == "RandomNumber":
        min_val = random.randint(0, 50)
        max_val = random.randint(min_val + 10, min_val + 100)
        result = {"min_value": min_val, "max_value": max_val}
    elif tool_name == "CurrentTime":
        # This tool doesn't need any input
        result = {}
    else:
        # Generic handling for unknown tools
        if "properties" in input_schema:
            for prop_name, prop_schema in input_schema["properties"].items():
                prop_type = prop_schema.get("type")

                if prop_type == "string":
                    result[prop_name] = f"random_string_{random.randint(1, 1000)}"
                elif prop_type == "number" or prop_type == "integer":
                    result[prop_name] = random.randint(1, 100)
                elif prop_type == "boolean":
                    result[prop_name] = random.choice([True, False])
                elif prop_type == "array":
                    result[prop_name] = []
                    if random.choice([True, False]):
                        item_type = prop_schema.get("items", {}).get("type", "string")
                        if item_type == "string":
                            result[prop_name].append(f"item_{random.randint(1, 100)}")
                        elif item_type == "number" or item_type == "integer":
                            result[prop_name].append(random.randint(1, 100))
                elif prop_type == "object":
                    result[prop_name] = {}

    return result


def format_parameter_info(schema: Dict[str, Any]) -> str:
    """Format parameter information including descriptions.

    Args:
        schema: The JSON schema of a tool input

    Returns:
        A formatted string with parameter information
    """
    result = []

    if "properties" in schema:
        for prop_name, prop_schema in schema["properties"].items():
            prop_type = prop_schema.get("type", "unknown")
            description = prop_schema.get("description", "No description")
            default = prop_schema.get("default", "required")

            param_info = f"{prop_name} ({prop_type})"
            if default != "required":
                param_info += f" = {default}"
            param_info += f": {description}"

            result.append(param_info)

    return "\n".join(result) if result else "No parameters"


async def test_tools_with_client(client: MCPClient, console: Console, connection_info: str):
    """Test all tools with the provided client.

    Args:
        client: The initialized MCP client
        console: Rich console for output
        connection_info: Info about the connection for display
    """
    # List available tools from the server
    console.print(f"\n[bold green]Available Tools ({connection_info}):[/bold green]")
    response = await client.session.list_tools()

    # Create a table to display the tools
    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("Tool Name")
    table.add_column("Description")
    table.add_column("Parameters")

    # Add each tool to the table
    for tool in response.tools:
        parameters = format_parameter_info(tool.inputSchema)

        table.add_row(tool.name, tool.description or "No description available", parameters)

    console.print(table)

    # Call each available tool with appropriate input
    for tool in response.tools:
        console.print(f"\n[bold yellow]Calling tool ({connection_info}): {tool.name}[/bold yellow]")

        # Generate appropriate input based on the tool
        input_args = generate_input_for_tool(tool.name, tool.inputSchema)

        # Display the input we're using
        console.print("[bold cyan]Input arguments:[/bold cyan]")
        syntax = Syntax(json.dumps(input_args, indent=2), "json")
        console.print(syntax)

        # Call the tool
        result = await client.call_tool(tool.name, input_args)

        # Display the result
        console.print("[bold green]Result:[/bold green]")
        if hasattr(result, "content"):
            for content_item in result.content:
                if content_item.type == "text":
                    console.print(content_item.text)
                else:
                    console.print(f"Content type: {content_item.type}")
        else:
            # Try to format as JSON if possible
            try:
                if isinstance(result, dict) or isinstance(result, list):
                    console.print(Syntax(json.dumps(result, indent=2), "json"))
                else:
                    console.print(str(result))
            except Exception:
                console.print(str(result))


async def list_server_tools():
    """Connect to MCP servers using both STDIO and SSE in sequence and list available tools."""
    console = Console()
    client = MCPClient()

    # Define the paths/URLs for both types of servers
    stdio_server_path = "example_mcp_server/server_stdio.py"  # Path to STDIO server
    sse_server_url = "http://localhost:6969"  # SSE server URL (default port)

    try:
        # 1. First test STDIO transport
        console.print("\n[bold blue]===== Testing STDIO Transport =====")
        console.print("[bold blue]Connecting to MCP server via STDIO...[/bold blue]")

        # Connect to the STDIO server
        await client.connect_to_stdio_server(stdio_server_path)

        # Test the tools available through STDIO
        await test_tools_with_client(client, console, "STDIO transport")

        # Clean up STDIO connection before moving to SSE
        await client.cleanup()

        # 2. Then test SSE transport
        console.print("\n[bold blue]===== Testing SSE Transport =====")
        console.print("[bold blue]Connecting to MCP server via SSE...[/bold blue]")

        # Connect to the SSE server
        await client.connect_to_sse_server(sse_server_url)

        # Test the tools available through SSE
        await test_tools_with_client(client, console, "SSE transport")

    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] {str(e)}")
    finally:
        # Clean up resources
        await client.cleanup()


if __name__ == "__main__":
    try:
        asyncio.run(list_server_tools())
    except KeyboardInterrupt:
        print("\nExiting...")
    except Exception as e:
        print(f"Fatal error: {str(e)}")

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/__init__.py

```python
"""example-mcp-server package."""

__version__ = "0.1.0"

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/interfaces/__init__.py

```python
"""Interface definitions for the application."""

from .tool import Tool, BaseToolInput, ToolResponse, ToolContent
from .resource import Resource, BaseResourceInput, ResourceContent, ResourceResponse
from .prompt import Prompt, BasePromptInput, PromptContent, PromptResponse

__all__ = [
    "Tool",
    "BaseToolInput",
    "ToolResponse",
    "ToolContent",
    "Resource",
    "BaseResourceInput",
    "ResourceContent",
    "ResourceResponse",
    "Prompt",
    "BasePromptInput",
    "PromptContent",
    "PromptResponse",
]

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/interfaces/prompt.py

```python
"""Interfaces for prompt abstractions."""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, ClassVar, Type, TypeVar
from pydantic import BaseModel, Field

# Define a type variable for generic model support
T = TypeVar("T", bound=BaseModel)


class BasePromptInput(BaseModel):
    """Base class for prompt input models."""

    model_config = {"extra": "forbid"}  # Equivalent to additionalProperties: false


class PromptContent(BaseModel):
    """Model for content in prompt responses."""

    type: str = Field(default="text", description="Content type identifier")

    # Common fields for all content types
    content_id: Optional[str] = Field(None, description="Optional content identifier")

    # Type-specific fields (using discriminated unions pattern)
    # Text content
    text: Optional[str] = Field(None, description="Text content when type='text'")

    # JSON content (for structured data)
    json_data: Optional[Dict[str, Any]] = Field(None, description="JSON data when type='json'")

    # Model content (will be converted to json_data during serialization)
    model: Optional[Any] = Field(None, exclude=True, description="Pydantic model instance")

    def model_post_init(self, __context: Any) -> None:
        """Post-initialization hook to handle model conversion."""
        if self.model and not self.json_data:
            # Convert model to json_data
            if isinstance(self.model, BaseModel):
                self.json_data = self.model.model_dump()
                if not self.type or self.type == "text":
                    self.type = "json"


class PromptResponse(BaseModel):
    """Model for prompt responses."""

    content: List[PromptContent]

    @classmethod
    def from_model(cls, model: BaseModel) -> "PromptResponse":
        """Create a PromptResponse from a Pydantic model.

        This makes it easier to return structured data directly.

        Args:
            model: A Pydantic model instance to convert

        Returns:
            A PromptResponse with the model data in JSON format
        """
        return cls(content=[PromptContent(type="json", json_data=model.model_dump(), model=model)])

    @classmethod
    def from_text(cls, text: str) -> "PromptResponse":
        """Create a PromptResponse from plain text.

        Args:
            text: The text content

        Returns:
            A PromptResponse with text content
        """
        return cls(content=[PromptContent(type="text", text=text)])


class Prompt(ABC):
    """Abstract base class for all prompts."""

    name: ClassVar[str]
    description: ClassVar[str]
    input_model: ClassVar[Type[BasePromptInput]]
    output_model: ClassVar[Optional[Type[BaseModel]]] = None

    @abstractmethod
    async def generate(self, input_data: BasePromptInput) -> PromptResponse:
        """Generate the prompt with given arguments."""
        pass

    def get_schema(self) -> Dict[str, Any]:
        """Get JSON schema for the prompt."""
        schema = {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
        }

        if self.output_model:
            schema["output"] = self.output_model.model_json_schema()

        return schema

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/interfaces/resource.py

```python
"""Interfaces for resource abstractions."""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, ClassVar, Type, TypeVar
from pydantic import BaseModel, Field

# Define a type variable for generic model support
T = TypeVar("T", bound=BaseModel)


class BaseResourceInput(BaseModel):
    """Base class for resource input models."""

    model_config = {"extra": "forbid"}  # Equivalent to additionalProperties: false


class ResourceContent(BaseModel):
    """Model for content in resource responses."""

    type: str = Field(default="text", description="Content type identifier")

    # Common fields for all content types
    content_id: Optional[str] = Field(None, description="Optional content identifier")

    # Type-specific fields (using discriminated unions pattern)
    # Text content
    text: Optional[str] = Field(None, description="Text content when type='text'")

    # JSON content (for structured data)
    json_data: Optional[Dict[str, Any]] = Field(None, description="JSON data when type='json'")

    # Model content (will be converted to json_data during serialization)
    model: Optional[Any] = Field(None, exclude=True, description="Pydantic model instance")

    # Resource-specific fields
    uri: Optional[str] = Field(None, description="URI of the resource")
    mime_type: Optional[str] = Field(None, description="MIME type of the resource")

    # Add more content types as needed (e.g., binary, image, etc.)

    def model_post_init(self, __context: Any) -> None:
        """Post-initialization hook to handle model conversion."""
        if self.model and not self.json_data:
            # Convert model to json_data
            if isinstance(self.model, BaseModel):
                self.json_data = self.model.model_dump()
                if not self.type or self.type == "text":
                    self.type = "json"


class ResourceResponse(BaseModel):
    """Model for resource responses."""

    content: List[ResourceContent]

    @classmethod
    def from_model(cls, model: BaseModel) -> "ResourceResponse":
        """Create a ResourceResponse from a Pydantic model.

        This makes it easier to return structured data directly.

        Args:
            model: A Pydantic model instance to convert

        Returns:
            A ResourceResponse with the model data in JSON format
        """
        return cls(content=[ResourceContent(type="json", json_data=model.model_dump(), model=model)])

    @classmethod
    def from_text(cls, text: str, uri: Optional[str] = None, mime_type: Optional[str] = None) -> "ResourceResponse":
        """Create a ResourceResponse from plain text.

        Args:
            text: The text content
            uri: Optional URI of the resource
            mime_type: Optional MIME type

        Returns:
            A ResourceResponse with text content
        """
        return cls(content=[ResourceContent(type="text", text=text, uri=uri, mime_type=mime_type)])


class Resource(ABC):
    """Abstract base class for all resources."""

    name: ClassVar[str]
    description: ClassVar[str]
    uri: ClassVar[str]
    mime_type: ClassVar[Optional[str]] = None
    input_model: ClassVar[Optional[Type[BaseResourceInput]]] = None
    output_model: ClassVar[Optional[Type[BaseModel]]] = None

    @abstractmethod
    async def read(self, input_data: BaseResourceInput) -> ResourceResponse:
        """Execute the resource with given arguments."""
        pass

    def get_schema(self) -> Dict[str, Any]:
        """Get JSON schema for the resource."""
        schema = {
            "name": self.name,
            "description": self.description,
            "uri": self.uri,
        }

        if self.mime_type:
            schema["mime_type"] = self.mime_type

        if self.input_model:
            schema["input"] = self.input_model.model_json_schema()

        if self.output_model:
            schema["output"] = self.output_model.model_json_schema()

        return schema

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/interfaces/tool.py

```python
"""Interfaces for tool abstractions."""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, ClassVar, Type, TypeVar
from pydantic import BaseModel, Field

# Define a type variable for generic model support
T = TypeVar("T", bound=BaseModel)


class BaseToolInput(BaseModel):
    """Base class for tool input models."""

    model_config = {"extra": "forbid"}  # Equivalent to additionalProperties: false


class ToolContent(BaseModel):
    """Model for content in tool responses."""

    type: str = Field(default="text", description="Content type identifier")

    # Common fields for all content types
    content_id: Optional[str] = Field(None, description="Optional content identifier")

    # Type-specific fields (using discriminated unions pattern)
    # Text content
    text: Optional[str] = Field(None, description="Text content when type='text'")

    # JSON content (for structured data)
    json_data: Optional[Dict[str, Any]] = Field(None, description="JSON data when type='json'")

    # Model content (will be converted to json_data during serialization)
    model: Optional[Any] = Field(None, exclude=True, description="Pydantic model instance")

    # Add more content types as needed (e.g., binary, image, etc.)

    def model_post_init(self, __context: Any) -> None:
        """Post-initialization hook to handle model conversion."""
        if self.model and not self.json_data:
            # Convert model to json_data
            if isinstance(self.model, BaseModel):
                self.json_data = self.model.model_dump()
                if not self.type or self.type == "text":
                    self.type = "json"


class ToolResponse(BaseModel):
    """Model for tool responses."""

    content: List[ToolContent]

    @classmethod
    def from_model(cls, model: BaseModel) -> "ToolResponse":
        """Create a ToolResponse from a Pydantic model.

        This makes it easier to return structured data directly.

        Args:
            model: A Pydantic model instance to convert

        Returns:
            A ToolResponse with the model data in JSON format
        """
        return cls(content=[ToolContent(type="json", json_data=model.model_dump(), model=model)])

    @classmethod
    def from_text(cls, text: str) -> "ToolResponse":
        """Create a ToolResponse from plain text.

        Args:
            text: The text content

        Returns:
            A ToolResponse with text content
        """
        return cls(content=[ToolContent(type="text", text=text)])


class Tool(ABC):
    """Abstract base class for all tools."""

    name: ClassVar[str]
    description: ClassVar[str]
    input_model: ClassVar[Type[BaseToolInput]]
    output_model: ClassVar[Optional[Type[BaseModel]]] = None

    @abstractmethod
    async def execute(self, input_data: BaseToolInput) -> ToolResponse:
        """Execute the tool with given arguments."""
        pass

    def get_schema(self) -> Dict[str, Any]:
        """Get JSON schema for the tool."""
        schema = {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
        }

        if self.output_model:
            schema["output"] = self.output_model.model_json_schema()

        return schema

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/prompts/sample_prompts.py

```python
"""Sample prompt implementations."""

from typing import Dict, Any, Union
from pydantic import Field, BaseModel, ConfigDict

from ..interfaces.prompt import Prompt, BasePromptInput, PromptResponse


class GreetingInput(BasePromptInput):
    """Input schema for the GreetingPrompt."""

    model_config = ConfigDict(json_schema_extra={"examples": [{"name": "Alice"}, {"name": "Bob"}]})

    name: str = Field(description="The name of the person to greet", examples=["Alice", "Bob"])


class GreetingOutput(BaseModel):
    """Output schema for the GreetingPrompt."""

    model_config = ConfigDict(
        json_schema_extra={
            "examples": [
                {"content": "Hello Alice, welcome!"},
                {"content": "Hello Bob, welcome!"},
            ]
        }
    )

    content: str = Field(description="The generated greeting message")
    error: Union[str, None] = Field(default=None, description="An error message if the operation failed.")


class GreetingPrompt(Prompt):
    """A prompt that greets the user by name."""

    name = "GreetingPrompt"
    description = "Generate a prompt that greets the user by name"
    input_model = GreetingInput
    output_model = GreetingOutput

    def get_schema(self) -> Dict[str, Any]:
        """Get the JSON schema for this prompt."""
        schema = {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
        }

        if self.output_model:
            schema["output"] = self.output_model.model_json_schema()

        return schema

    async def generate(self, input_data: GreetingInput, **kwargs) -> PromptResponse:
        """Execute the greeting prompt.

        Args:
            input_data: The validated input for the prompt

        Returns:
            A response containing the greeting message
        """
        greeting_input = GreetingInput.model_validate(input_data.model_dump())
        content = f"Hello {greeting_input.name.title()}, welcome to the project!"
        output = GreetingOutput(content=content, error=None)
        return PromptResponse.from_model(output)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/resources/__init__.py

```python
"""Resource exports."""

__all__ = []

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/resources/sample_resources.py

```python
"""Sample text resource."""

from typing import Dict, Any, Union

from pydantic import Field, BaseModel, ConfigDict

from ..interfaces.resource import Resource, BaseResourceInput, ResourceResponse
from urllib.parse import unquote as decode_uri


class TestWeatherInput(BaseResourceInput):
    """Input schema for the TestWeatherResource."""

    model_config = ConfigDict(
        json_schema_extra={"examples": [{"country": "USA", "city": "New York"}, {"country": "Canada", "city": "Toronto"}]}
    )

    country: str = Field(description="The country name", examples=["USA", "Canada"])
    city: str = Field(description="The city name", examples=["New York", "Toronto"])


class TestWeatherOutput(BaseModel):
    """Output schema for the TestWeatherResource."""

    model_config = ConfigDict(json_schema_extra={"examples": [{"weather": "72 F and pleasant", "error": None}]})

    weather: str = Field(description="The weather information")
    error: Union[str, None] = Field(default=None, description="An error message if the operation failed.")


class TestWeatherResource(Resource):
    """A sample weather resource that returns static weather content."""

    name = "TestWeatherService"
    description = "Fetch weather based on country and city name."
    uri = "resource://weather/{country}/{city}"
    mime_type = "text/plain"
    input_model = TestWeatherInput
    output_model = TestWeatherOutput

    def get_schema(self) -> Dict[str, Any]:
        """Get the JSON schema for this resource."""
        schema = {
            "name": self.name,
            "description": self.description,
            "uri": self.uri,
            "mime_type": self.mime_type,
            "input": self.input_model.model_json_schema(),
        }

        if self.output_model:
            schema["output"] = self.output_model.model_json_schema()

        return schema

    async def read(self, input_data: TestWeatherInput) -> ResourceResponse:
        """Execute the weather resource.

        Args:
            input_data: The validated input for the resource

        Returns:
            A response containing the weather information
        """
        city = decode_uri(input_data.city.title())
        country = decode_uri(input_data.country)
        weather_info = f"Temperature in {city}, {country} is 72 F and pleasant."
        output = TestWeatherOutput(weather=weather_info, error=None)
        return ResourceResponse.from_model(output)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/server.py

```python
"""example-mcp-server MCP Server unified entry point."""

import argparse
import sys


def main():
    """Entry point for the server."""
    parser = argparse.ArgumentParser(description="example-mcp-server MCP Server")
    parser.add_argument(
        "--mode",
        type=str,
        required=True,
        choices=["stdio", "sse", "http_stream"],
        help="Server mode: stdio for standard I/O, sse for Server-Sent Events, or http_stream for HTTP Stream Transport",
    )

    # HTTP Stream specific arguments
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to (sse/http_stream mode only)")
    parser.add_argument("--port", type=int, default=6969, help="Port to listen on (sse/http_stream mode only)")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development (sse/http_stream mode only)")

    args = parser.parse_args()

    if args.mode == "stdio":
        # Import and run the stdio server
        from example_mcp_server.server_stdio import main as stdio_main

        stdio_main()
    elif args.mode == "sse":
        # Import and run the SSE server with appropriate arguments
        from example_mcp_server.server_sse import main as sse_main

        sys.argv = [sys.argv[0], "--host", args.host, "--port", str(args.port)]
        if args.reload:
            sys.argv.append("--reload")
        sse_main()
    elif args.mode == "http_stream":
        # Import and run the HTTP Stream Transport server
        from example_mcp_server.server_http import main as http_main

        sys.argv = [sys.argv[0], "--host", args.host, "--port", str(args.port)]
        if args.reload:
            sys.argv.append("--reload")
        http_main()
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/server_http.py

```python
"""example-mcp-server MCP Server HTTP Stream Transport."""

from typing import List
import argparse

import uvicorn
from starlette.middleware.cors import CORSMiddleware

from mcp.server.fastmcp import FastMCP

from example_mcp_server.services.tool_service import ToolService
from example_mcp_server.services.resource_service import ResourceService
from example_mcp_server.services.prompt_service import PromptService
from example_mcp_server.interfaces.tool import Tool
from example_mcp_server.interfaces.resource import Resource
from example_mcp_server.interfaces.prompt import Prompt
from example_mcp_server.tools import (
    AddNumbersTool,
    SubtractNumbersTool,
    MultiplyNumbersTool,
    DivideNumbersTool,
    BatchCalculatorTool,
)
from example_mcp_server.resources.sample_resources import TestWeatherResource
from example_mcp_server.prompts.sample_prompts import GreetingPrompt


def get_available_tools() -> List[Tool]:
    """Get list of all available tools."""
    return [
        AddNumbersTool(),
        SubtractNumbersTool(),
        MultiplyNumbersTool(),
        DivideNumbersTool(),
        BatchCalculatorTool(),
    ]


def get_available_resources() -> List[Resource]:
    """Get list of all available resources."""
    return [
        TestWeatherResource(),
        # Add more resources here as you create them
    ]


def get_available_prompts() -> List[Prompt]:
    """Get list of all available prompts."""
    return [
        GreetingPrompt(),
        # Add more prompts here as you create them
    ]


def create_mcp_server() -> FastMCP:
    """Create and configure the MCP server."""
    mcp = FastMCP("example-mcp-server")
    tool_service = ToolService()
    resource_service = ResourceService()
    prompt_service = PromptService()

    # Register all tools and their MCP handlers
    tool_service.register_tools(get_available_tools())
    tool_service.register_mcp_handlers(mcp)

    # Register all resources and their MCP handlers
    resource_service.register_resources(get_available_resources())
    resource_service.register_mcp_handlers(mcp)

    # Register all prompts and their MCP handlers
    prompt_service.register_prompts(get_available_prompts())
    prompt_service.register_mcp_handlers(mcp)

    return mcp


def create_http_app():
    """Create a FastMCP HTTP app with CORS middleware."""
    mcp_server = create_mcp_server()

    # Use FastMCP directly as the app instead of mounting it
    # This avoids the task group initialization issue
    # See: https://github.com/modelcontextprotocol/python-sdk/issues/732
    app = mcp_server.streamable_http_app()  # type: ignore[attr-defined]

    # Apply CORS middleware manually
    app = CORSMiddleware(
        app,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
        allow_credentials=True,
    )

    return app


def main():
    """Entry point for the HTTP Stream Transport server."""
    parser = argparse.ArgumentParser(description="Run MCP HTTP Stream server")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=6969, help="Port to listen on")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development")
    args = parser.parse_args()

    app = create_http_app()
    print(f"MCP HTTP Stream Server starting on {args.host}:{args.port}")
    uvicorn.run(
        app,
        host=args.host,
        port=args.port,
        reload=args.reload,
    )


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/server_sse.py

```python
"""example-mcp-server MCP Server implementation with SSE transport."""

from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
from mcp.server.sse import SseServerTransport
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Mount, Route
from mcp.server import Server
import uvicorn
from typing import List
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware

from example_mcp_server.services.tool_service import ToolService
from example_mcp_server.services.resource_service import ResourceService
from example_mcp_server.services.prompt_service import PromptService
from example_mcp_server.interfaces.tool import Tool
from example_mcp_server.interfaces.resource import Resource
from example_mcp_server.interfaces.prompt import Prompt
from example_mcp_server.tools import AddNumbersTool, SubtractNumbersTool, MultiplyNumbersTool, DivideNumbersTool
from example_mcp_server.resources.sample_resources import TestWeatherResource
from example_mcp_server.prompts.sample_prompts import GreetingPrompt


def get_available_tools() -> List[Tool]:
    """Get list of all available tools."""
    return [
        AddNumbersTool(),
        SubtractNumbersTool(),
        MultiplyNumbersTool(),
        DivideNumbersTool(),
    ]


def get_available_resources() -> List[Resource]:
    """Get list of all available resources."""
    return [
        TestWeatherResource(),
        # Add more resources here as you create them
    ]


def get_available_prompts() -> List[Prompt]:
    """Get list of all available prompts."""
    return [
        GreetingPrompt(),
        # Add more prompts here as you create them
    ]


def create_starlette_app(mcp_server: Server) -> Starlette:
    """Create a Starlette application that can serve the provided mcp server with SSE."""
    sse = SseServerTransport("/messages/")

    async def handle_sse(request: Request) -> Response:
        async with sse.connect_sse(
            request.scope,
            request.receive,
            request._send,  # noqa: SLF001
        ) as (read_stream, write_stream):
            await mcp_server.run(
                read_stream,
                write_stream,
                mcp_server.create_initialization_options(),
            )
        return Response("SSE connection closed", status_code=200)

    middleware = [
        Middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_methods=["*"],
            allow_headers=["*"],
            allow_credentials=True,
        )
    ]

    return Starlette(
        routes=[
            Route("/sse", endpoint=handle_sse),
            Mount("/messages/", app=sse.handle_post_message),
        ],
        middleware=middleware,
    )


# Initialize FastMCP server with SSE
mcp = FastMCP("example-mcp-server")
tool_service = ToolService()
resource_service = ResourceService()
prompt_service = PromptService()

# Register all tools and their MCP handlers
tool_service.register_tools(get_available_tools())
tool_service.register_mcp_handlers(mcp)

# Register all resources and their MCP handlers
resource_service.register_resources(get_available_resources())
resource_service.register_mcp_handlers(mcp)

# Register all prompts and their MCP handlers
prompt_service.register_prompts(get_available_prompts())
prompt_service.register_mcp_handlers(mcp)

# Get the MCP server
mcp_server = mcp._mcp_server  # noqa: WPS437

# Create the Starlette app
app = create_starlette_app(mcp_server)

# Export the app
__all__ = ["app"]


def main():
    """Entry point for the server."""
    import argparse

    parser = argparse.ArgumentParser(description="Run MCP SSE-based server")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=6969, help="Port to listen on")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development")
    args = parser.parse_args()

    # Run the server with auto-reload if enabled
    uvicorn.run(
        "example_mcp_server.server_sse:app",  # Use the app from server_sse.py directly
        host=args.host,
        port=args.port,
        reload=args.reload,
        reload_dirs=["example_mcp_server"],  # Watch this directory for changes
        timeout_graceful_shutdown=5,  # Add timeout
    )


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/server_stdio.py

```python
"""example-mcp-server MCP Server implementation."""

from mcp.server.fastmcp import FastMCP
from typing import List

from example_mcp_server.services.tool_service import ToolService
from example_mcp_server.services.resource_service import ResourceService
from example_mcp_server.services.prompt_service import PromptService
from example_mcp_server.interfaces.tool import Tool
from example_mcp_server.interfaces.resource import Resource
from example_mcp_server.interfaces.prompt import Prompt

from example_mcp_server.tools import (
    AddNumbersTool,
    SubtractNumbersTool,
    MultiplyNumbersTool,
    DivideNumbersTool,
)
from example_mcp_server.resources.sample_resources import TestWeatherResource
from example_mcp_server.prompts.sample_prompts import GreetingPrompt


def get_available_tools() -> List[Tool]:
    """Get list of all available tools."""
    return [
        # HelloWorldTool(), # Removed
        AddNumbersTool(),
        SubtractNumbersTool(),
        MultiplyNumbersTool(),
        DivideNumbersTool(),
        # Add more tools here as you create them
    ]


def get_available_resources() -> List[Resource]:
    """Get list of all available resources."""
    return [
        TestWeatherResource(),
        # Add more resources here as you create them
    ]


def get_available_prompts() -> List[Prompt]:
    """Get list of all available prompts."""
    return [
        GreetingPrompt(),
        # Add more prompts here as you create them
    ]


def main():
    """Entry point for the server."""
    mcp = FastMCP("example-mcp-server")
    tool_service = ToolService()
    resource_service = ResourceService()
    prompt_service = PromptService()

    # Register all tools and their MCP handlers
    tool_service.register_tools(get_available_tools())
    tool_service.register_mcp_handlers(mcp)

    # Register all resources and their MCP handlers
    resource_service.register_resources(get_available_resources())
    resource_service.register_mcp_handlers(mcp)

    # Register all prompts and their MCP handlers
    prompt_service.register_prompts(get_available_prompts())
    prompt_service.register_mcp_handlers(mcp)

    mcp.run()


if __name__ == "__main__":
    main()

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/services/__init__.py

```python
"""Service layer for the application."""

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/services/prompt_service.py

```python
"""Service layer for managing prompts."""

from typing import Dict, List, Any
import logging
import inspect
from mcp.server.fastmcp import FastMCP
from example_mcp_server.interfaces.prompt import Prompt, PromptResponse, PromptContent


class PromptService:
    """Service for managing and executing prompts."""

    def __init__(self):
        self._prompts: Dict[str, Prompt] = {}

    def register_prompt(self, prompt: Prompt) -> None:
        """Register a new prompt."""
        self._prompts[prompt.name] = prompt

    def register_prompts(self, prompts: List[Prompt]) -> None:
        """Register multiple prompts."""
        for prompt in prompts:
            self.register_prompt(prompt)

    def get_prompt(self, prompt_name: str) -> Prompt:
        """Get a prompt by name."""
        if prompt_name not in self._prompts:
            raise ValueError(f"Prompt not found: {prompt_name}")
        return self._prompts[prompt_name]

    async def generate_prompt(self, prompt_name: str, input_data: Dict[str, Any]) -> PromptResponse:
        """Execute a prompt by name with given arguments.

        This validates the input against the prompt's input model and calls
        the prompt's async generate method.
        """
        prompt = self.get_prompt(prompt_name)

        # Validate input using Pydantic model_validate to support nested models
        input_model = prompt.input_model.model_validate(input_data)

        return await prompt.generate(input_model)

    def _process_prompt_content(self, content: PromptContent) -> str | Dict[str, Any] | None:
        """Process a PromptContent object into a serializable form."""
        if content.type == "text":
            return content.text
        elif content.type == "json" and content.json_data is not None:
            return content.json_data
        else:
            return content.text or content.json_data or {}

    def _serialize_response(self, response: PromptResponse) -> Any:
        """Serialize a PromptResponse to return to clients.

        If there's a single content item, return it directly; otherwise return a list.
        """
        if not response.content:
            return {}

        if len(response.content) == 1:  # Not a list
            return self._process_prompt_content(response.content[0])

        return [self._process_prompt_content(content) for content in response.content]

    def register_mcp_handlers(self, mcp: FastMCP) -> None:
        """Register all prompts as MCP handlers."""
        for prompt in self._prompts.values():
            # Create a handler that uses the prompt's Pydantic input model directly for schema generation
            def create_handler(prompt: Prompt):
                # Get the fields of the input_model
                input_fields = prompt.input_model.model_fields

                sig = inspect.Signature(
                    [
                        inspect.Parameter(
                            field_name,
                            inspect.Parameter.POSITIONAL_OR_KEYWORD,
                            annotation=field_info.annotation,
                        )
                        for field_name, field_info in input_fields.items()
                    ]
                )

                # Create the handler function
                async def handler(*args, **kwargs):
                    """Execute the prompt with the given input data."""
                    # Bind the arguments to the signature
                    bound_args = sig.bind(*args, **kwargs)
                    bound_args.apply_defaults()

                    input_data = dict(bound_args.arguments)
                    logger = logging.getLogger("example_mcp_server.prompt_service")
                    logger.debug("Received input_data for prompt '%s': %s", prompt.name, input_data)

                    # Validate the input using the Pydantic model
                    input_model = prompt.input_model.model_validate(input_data)
                    result = await self.generate_prompt(prompt.name, input_model.model_dump())
                    return self._serialize_response(result)

                # Set the signature and metadata on the handler
                handler.__signature__ = sig
                handler.__name__ = prompt.name
                handler.__doc__ = prompt.description or ""

                # Set annotations
                handler.__annotations__ = {
                    field_name: field_info.annotation for field_name, field_info in input_fields.items()
                }
                handler.__annotations__["return"] = Any

                return handler

            handler = create_handler(prompt)

            # Register the prompt with FastMCP. Use the prompt name as the handler name.
            mcp.prompt(name=prompt.name, description=prompt.description)(handler)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/services/resource_service.py

```python
"""Service layer for managing resources."""

from typing import Dict, List
import re
import inspect
from mcp.server.fastmcp import FastMCP
from example_mcp_server.interfaces.resource import Resource, ResourceResponse


class ResourceService:
    """Service for managing and executing resources."""

    def __init__(self):
        self._resources: Dict[str, Resource] = {}
        self._uri_patterns: Dict[str, Resource] = {}

    def register_resource(self, resource: Resource) -> None:
        """Register a new resource."""
        # Store the resource by its URI pattern for handler registration
        self._uri_patterns[resource.uri] = resource

        # If the URI doesn't have parameters, also store by exact URI
        if "{" not in resource.uri:
            self._resources[resource.uri] = resource

    def register_resources(self, resources: List[Resource]) -> None:
        """Register multiple resources."""
        for resource in resources:
            self.register_resource(resource)

    def get_resource_by_pattern(self, uri_pattern: str) -> Resource:
        """Get a resource by its URI pattern."""
        if uri_pattern not in self._uri_patterns:
            raise ValueError(f"Resource not found for pattern: {uri_pattern}")
        return self._uri_patterns[uri_pattern]

    def get_resource(self, uri: str) -> Resource:
        """Get a resource by exact URI."""
        # First check if there's an exact match for the URI
        if uri in self._resources:
            return self._resources[uri]

        # If not, try to find a pattern that matches
        for pattern, resource in self._uri_patterns.items():
            # Convert the pattern to a regex by replacing {param} with (?P<param>[^/]+)
            regex_pattern = re.sub(r"\{([^}]+)\}", r"(?P<\1>[^/]+)", pattern)
            # Ensure we match the whole URI by adding anchors
            regex_pattern = f"^{regex_pattern}$"

            match = re.match(regex_pattern, uri)
            if match:
                # Found a matching pattern, extract parameters
                # Cache the resource with the specific URI for future lookups
                self._resources[uri] = resource
                return resource

        raise ValueError(f"Resource not found: {uri}")

    def extract_params_from_uri(self, pattern: str, uri: str) -> Dict[str, str]:
        """Extract parameters from a URI based on a pattern."""
        # Convert the pattern to a regex by replacing {param} with (?P<param>[^/]+)
        regex_pattern = re.sub(r"\{([^}]+)\}", r"(?P<\1>[^/]+)", pattern)
        # Ensure we match the whole URI by adding anchors
        regex_pattern = f"^{regex_pattern}$"

        match = re.match(regex_pattern, uri)
        if match:
            return match.groupdict()
        return {}

    def create_handler(self, resource: Resource, uri_pattern: str):
        """Create a handler function for a resource with the correct parameters."""
        # Extract parameters from URI pattern
        uri_params = set(re.findall(r"\{([^}]+)\}", uri_pattern))

        if not uri_params:
            # For static resources with no parameters
            async def static_handler() -> ResourceResponse:
                """Handle static resource request."""
                # Create empty input for resources without parameters
                input_data = resource.input_model()
                return await resource.read(input_data)

            # Set metadata for the handler
            static_handler.__name__ = resource.name
            static_handler.__doc__ = resource.description
            return static_handler
        else:
            # For resources with parameters
            # Create parameters for the signature
            uri_params_list = list(uri_params)
            sig = inspect.Signature(
                [
                    inspect.Parameter(param, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=str)
                    for param in uri_params_list
                ]
            )

            # Create the handler function
            async def param_handler(*args, **kwargs):
                """Handle parameterized resource request."""
                # Bind the arguments to the signature
                bound_args = sig.bind(*args, **kwargs)
                bound_args.apply_defaults()

                # Create input data from bound arguments
                input_data = resource.input_model(**bound_args.arguments)
                return await resource.read(input_data)

            # Set the signature and metadata on the handler
            param_handler.__signature__ = sig
            param_handler.__name__ = resource.name
            param_handler.__doc__ = resource.description

            # Set annotations
            param_handler.__annotations__ = {param: str for param in uri_params_list}
            param_handler.__annotations__["return"] = ResourceResponse

            return param_handler

    def register_mcp_handlers(self, mcp: FastMCP) -> None:
        """Register all resources as MCP handlers."""
        for uri_pattern, resource in self._uri_patterns.items():
            handler = self.create_handler(resource, uri_pattern)

            # Register the resource with the full metadata
            wrapped_handler = mcp.resource(
                uri=uri_pattern, name=resource.name, description=resource.description, mime_type=resource.mime_type
            )(handler)

            # Ensure the handler's metadata is preserved
            wrapped_handler.__name__ = resource.name
            wrapped_handler.__doc__ = resource.description

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/services/tool_service.py

```python
"""Service layer for managing tools."""

from typing import Dict, List, Any
from mcp.server.fastmcp import FastMCP
from example_mcp_server.interfaces.tool import Tool, ToolResponse, ToolContent


class ToolService:
    """Service for managing and executing tools."""

    def __init__(self):
        self._tools: Dict[str, Tool] = {}

    def register_tool(self, tool: Tool) -> None:
        """Register a new tool."""
        self._tools[tool.name] = tool

    def register_tools(self, tools: List[Tool]) -> None:
        """Register multiple tools."""
        for tool in tools:
            self.register_tool(tool)

    def get_tool(self, tool_name: str) -> Tool:
        """Get a tool by name."""
        if tool_name not in self._tools:
            raise ValueError(f"Tool not found: {tool_name}")
        return self._tools[tool_name]

    async def execute_tool(self, tool_name: str, input_data: Dict[str, Any]) -> ToolResponse:
        """Execute a tool by name with given arguments.

        Args:
            tool_name: The name of the tool to execute
            input_data: Dictionary of input arguments for the tool

        Returns:
            The tool's response containing the execution results

        Raises:
            ValueError: If the tool is not found
            ValidationError: If the input data is invalid
        """
        tool = self.get_tool(tool_name)

        # Use model_validate to handle complex nested objects properly
        input_model = tool.input_model.model_validate(input_data)

        # Execute the tool with validated input
        return await tool.execute(input_model)

    def _process_tool_content(self, content: ToolContent) -> Any:
        """Process a ToolContent object based on its type.

        Args:
            content: The ToolContent to process

        Returns:
            The appropriate representation of the content based on its type
        """
        if content.type == "text":
            return content.text
        elif content.type == "json" and content.json_data is not None:
            return content.json_data
        else:
            # Default to returning whatever is available
            return content.text or content.json_data or {}

    def _serialize_response(self, response: ToolResponse) -> Any:
        """Serialize a ToolResponse to return to the client.

        This handles the actual response serialization based on content types.

        Args:
            response: The ToolResponse to serialize

        Returns:
            The serialized response
        """
        if not response.content:
            return {}

        # If there's only one content item, return it directly
        if len(response.content) == 1:
            return self._process_tool_content(response.content[0])

        # If there are multiple content items, return them as a list
        return [self._process_tool_content(content) for content in response.content]

    def register_mcp_handlers(self, mcp: FastMCP) -> None:
        """Register all tools as MCP handlers."""
        for tool in self._tools.values():
            # Create a handler that uses the tool's input model directly for schema generation
            def create_handler(tool_instance):
                # Use the actual Pydantic model as the function parameter
                # This ensures FastMCP gets the complete schema including nested objects
                async def handler(input_data: tool_instance.input_model):
                    f'"""{tool_instance.description}"""'
                    result = await self.execute_tool(tool_instance.name, input_data.model_dump())
                    return self._serialize_response(result)

                return handler

            # Create the handler
            handler = create_handler(tool)

            # Register with FastMCP - it should auto-detect the schema from the type annotation
            mcp.tool(name=tool.name, description=tool.description)(handler)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/tools/__init__.py

```python
"""Tool exports."""

from .add_numbers import AddNumbersTool
from .subtract_numbers import SubtractNumbersTool
from .multiply_numbers import MultiplyNumbersTool
from .divide_numbers import DivideNumbersTool
from .batch_operations import BatchCalculatorTool

__all__ = [
    "AddNumbersTool",
    "SubtractNumbersTool",
    "MultiplyNumbersTool",
    "DivideNumbersTool",
    "BatchCalculatorTool",
    # Add additional tools to the __all__ list as you create them
]

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/tools/add_numbers.py

```python
"""Tool for adding two numbers."""

from typing import Dict, Any, Union

from pydantic import Field, BaseModel, ConfigDict

from ..interfaces.tool import Tool, BaseToolInput, ToolResponse


class AddNumbersInput(BaseToolInput):
    """Input schema for the AddNumbers tool."""

    model_config = ConfigDict(
        json_schema_extra={"examples": [{"number1": 5, "number2": 3}, {"number1": -2.5, "number2": 1.5}]}
    )

    number1: float = Field(description="The first number to add", examples=[5, -2.5])
    number2: float = Field(description="The second number to add", examples=[3, 1.5])


class AddNumbersOutput(BaseModel):
    """Output schema for the AddNumbers tool."""

    model_config = ConfigDict(json_schema_extra={"examples": [{"sum": 8, "error": None}, {"sum": -1.0, "error": None}]})

    sum: float = Field(description="The sum of the two numbers")
    error: Union[str, None] = Field(default=None, description="An error message if the operation failed.")


class AddNumbersTool(Tool):
    """Tool that adds two numbers together."""

    name = "AddNumbers"
    description = "Adds two numbers (number1 + number2) and returns the sum"
    input_model = AddNumbersInput
    output_model = AddNumbersOutput

    def get_schema(self) -> Dict[str, Any]:
        """Get the JSON schema for this tool."""
        return {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
            "output": self.output_model.model_json_schema(),
        }

    async def execute(self, input_data: AddNumbersInput) -> ToolResponse:
        """Execute the add numbers tool.

        Args:
            input_data: The validated input for the tool

        Returns:
            A response containing the sum
        """
        result = input_data.number1 + input_data.number2
        output = AddNumbersOutput(sum=result, error=None)
        return ToolResponse.from_model(output)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/tools/batch_operations.py

```python
# Tool: BatchCalculatorTool
from typing import List, Union, Literal, Annotated, Dict, Any
from pydantic import BaseModel, Field, ConfigDict
from ..interfaces.tool import Tool, BaseToolInput, ToolResponse


# ---- ops (discriminated union) ----
class Add(BaseModel):
    op: Literal["add"]
    nums: List[float] = Field(min_items=1)


class Mul(BaseModel):
    op: Literal["mul"]
    nums: List[float] = Field(min_items=1)


Op = Annotated[Union[Add, Mul], Field(discriminator="op")]


# ---- IO ----
class BatchInput(BaseToolInput):
    model_config = ConfigDict(
        title="BatchInput",
        json_schema_extra={
            "examples": [{"mode": "sum", "tasks": [{"op": "add", "nums": [1, 2, 3]}, {"op": "mul", "nums": [2, 3]}]}]
        },
    )
    tasks: List[Op] = Field(description="List of operations to run (add|mul)")
    mode: Literal["sum", "avg"] = Field(default="sum", description="Combine per-task results by sum or average")
    explain: bool = False


class BatchOutput(BaseModel):
    results: List[float]
    combined: float
    mode_used: Literal["sum", "avg"]
    summary: str | None = None


# ---- Tool ----
class BatchCalculatorTool(Tool):
    name = "BatchCalculator"
    description = (
        "Run a batch of simple ops. \nExamples:\n"
        '- {"tasks":[{"op":"add","nums":[1,2,3]}, {"op":"mul","nums":[4,5]}], "mode":"sum"}\n'
        '- {"tasks":[{"op":"mul","nums":[2,3,4]}], "mode":"avg"}\n'
        '- {"tasks":[{"op":"add","nums":[10,20]}, {"op":"add","nums":[30,40]}], "mode":"avg"}'
    )
    input_model = BatchInput
    output_model = BatchOutput

    def get_schema(self) -> Dict[str, Any]:
        inp = self.input_model.model_json_schema()
        return {
            "name": self.name,
            "description": self.description,
            "input": inp,
            "output": self.output_model.model_json_schema(),
            "examples": inp.get("examples", []),
        }

    async def execute(self, data: BatchInput) -> ToolResponse:
        def run(op: Op) -> float:
            if op.op == "add":
                return float(sum(op.nums))
            prod = 1.0
            for x in op.nums:
                prod *= float(x)
            return prod

        results = [run(t) for t in data.tasks]
        combined = float(sum(results)) if data.mode == "sum" else (float(sum(results)) / len(results) if results else 0.0)
        summary = (f"tasks={len(results)}, results={results}, combined={combined} ({data.mode})") if data.explain else None

        return ToolResponse.from_model(BatchOutput(results=results, combined=combined, mode_used=data.mode, summary=summary))

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/tools/divide_numbers.py

```python
"""Tool for dividing two numbers."""

from typing import Dict, Any, Union

from pydantic import Field, BaseModel, ConfigDict

from ..interfaces.tool import Tool, BaseToolInput, ToolResponse


class DivideNumbersInput(BaseToolInput):
    """Input schema for the DivideNumbers tool."""

    model_config = ConfigDict(
        json_schema_extra={
            "examples": [{"dividend": 10, "divisor": 2}, {"dividend": 5, "divisor": 0}, {"dividend": 7.5, "divisor": 2.5}]
        }
    )

    dividend: float = Field(description="The number to be divided", examples=[10, 5, 7.5])
    divisor: float = Field(description="The number to divide by", examples=[2, 0, 2.5])


class DivideNumbersOutput(BaseModel):
    """Output schema for the DivideNumbers tool."""

    model_config = ConfigDict(
        json_schema_extra={"examples": [{"quotient": 5.0}, {"error": "Division by zero is not allowed."}, {"quotient": 3.0}]}
    )

    quotient: Union[float, None] = Field(
        default=None, description="The result of the division (dividend / divisor). None if division by zero occurred."
    )
    error: Union[str, None] = Field(
        default=None, description="An error message if the operation failed (e.g., division by zero)."
    )


class DivideNumbersTool(Tool):
    """Tool that divides one number by another."""

    name = "DivideNumbers"
    description = "Divides the first number (dividend) by the second number (divisor) and returns the quotient. Handles division by zero."
    input_model = DivideNumbersInput
    output_model = DivideNumbersOutput

    def get_schema(self) -> Dict[str, Any]:
        """Get the JSON schema for this tool."""
        return {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
            "output": self.output_model.model_json_schema(),
        }

    async def execute(self, input_data: DivideNumbersInput) -> ToolResponse:
        """Execute the divide numbers tool.

        Args:
            input_data: The validated input for the tool

        Returns:
            A response containing the quotient or an error message
        """
        if input_data.divisor == 0:
            output = DivideNumbersOutput(error="Division by zero is not allowed.")
            # Optionally set a specific status code if your ToolResponse supports it
            # return ToolResponse(status_code=400, content=ToolContent.from_model(output))
            return ToolResponse.from_model(output)
        else:
            result = input_data.dividend / input_data.divisor
            output = DivideNumbersOutput(quotient=result)
            return ToolResponse.from_model(output)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/tools/multiply_numbers.py

```python
"""Tool for multiplying two numbers."""

from typing import Dict, Any, Union

from pydantic import Field, BaseModel, ConfigDict

from ..interfaces.tool import Tool, BaseToolInput, ToolResponse


class MultiplyNumbersInput(BaseToolInput):
    """Input schema for the MultiplyNumbers tool."""

    model_config = ConfigDict(json_schema_extra={"examples": [{"number1": 5, "number2": 3}, {"number1": -2.5, "number2": 4}]})

    number1: float = Field(description="The first number to multiply", examples=[5, -2.5])
    number2: float = Field(description="The second number to multiply", examples=[3, 4])


class MultiplyNumbersOutput(BaseModel):
    """Output schema for the MultiplyNumbers tool."""

    model_config = ConfigDict(
        json_schema_extra={"examples": [{"product": 15, "error": None}, {"product": -10.0, "error": None}]}
    )

    product: float = Field(description="The product of the two numbers (number1 * number2)")
    error: Union[str, None] = Field(default=None, description="An error message if the operation failed.")


class MultiplyNumbersTool(Tool):
    """Tool that multiplies two numbers together."""

    name = "MultiplyNumbers"
    description = "Multiplies two numbers (number1 * number2) and returns the product"
    input_model = MultiplyNumbersInput
    output_model = MultiplyNumbersOutput

    def get_schema(self) -> Dict[str, Any]:
        """Get the JSON schema for this tool."""
        return {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
            "output": self.output_model.model_json_schema(),
        }

    async def execute(self, input_data: MultiplyNumbersInput) -> ToolResponse:
        """Execute the multiply numbers tool.

        Args:
            input_data: The validated input for the tool

        Returns:
            A response containing the product
        """
        result = input_data.number1 * input_data.number2
        output = MultiplyNumbersOutput(product=result, error=None)
        return ToolResponse.from_model(output)

```

### File: atomic-examples/mcp-agent/example-mcp-server/example_mcp_server/tools/subtract_numbers.py

```python
"""Tool for subtracting two numbers."""

from typing import Dict, Any, Union

from pydantic import Field, BaseModel, ConfigDict

from ..interfaces.tool import Tool, BaseToolInput, ToolResponse


class SubtractNumbersInput(BaseToolInput):
    """Input schema for the SubtractNumbers tool."""

    model_config = ConfigDict(json_schema_extra={"examples": [{"number1": 5, "number2": 3}, {"number1": 1.5, "number2": 2.5}]})

    number1: float = Field(description="The number to subtract from", examples=[5, 1.5])
    number2: float = Field(description="The number to subtract", examples=[3, 2.5])


class SubtractNumbersOutput(BaseModel):
    """Output schema for the SubtractNumbers tool."""

    model_config = ConfigDict(
        json_schema_extra={"examples": [{"difference": 2, "error": None}, {"difference": -1.0, "error": None}]}
    )

    difference: float = Field(description="The difference between the two numbers (number1 - number2)")
    error: Union[str, None] = Field(default=None, description="An error message if the operation failed.")


class SubtractNumbersTool(Tool):
    """Tool that subtracts one number from another."""

    name = "SubtractNumbers"
    description = "Subtracts the second number from the first number (number1 - number2) and returns the difference"
    input_model = SubtractNumbersInput
    output_model = SubtractNumbersOutput

    def get_schema(self) -> Dict[str, Any]:
        """Get the JSON schema for this tool."""
        return {
            "name": self.name,
            "description": self.description,
            "input": self.input_model.model_json_schema(),
            "output": self.output_model.model_json_schema(),
        }

    async def execute(self, input_data: SubtractNumbersInput) -> ToolResponse:
        """Execute the subtract numbers tool.

        Args:
            input_data: The validated input for the tool

        Returns:
            A response containing the difference
        """
        result = input_data.number1 - input_data.number2
        output = SubtractNumbersOutput(difference=result, error=None)
        return ToolResponse.from_model(output)

```

### File: atomic-examples/mcp-agent/example-mcp-server/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["example_mcp_server"]

[project]
name = "example-mcp-server"
version = "0.1.0"
description = "example-mcp-server MCP server"
authors = []
requires-python = ">=3.12"
dependencies = [
    "mcp[cli]>=1.9.4",
    "rich>=13.0.0",
    "pydantic>=2.0.0",
    "uvicorn>=0.15.0",
]

[project.scripts]
example-mcp-server = "example_mcp_server.server:main"

```


--------------------------------------------------------------------------------
Example: nested-multimodal
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/nested-multimodal

## Documentation

# Nested Multimodal Example

This example demonstrates how to use the Atomic Agents framework with **nested multimodal content** — images and PDFs inside nested Pydantic schemas, not just at the top level.

This showcases the fixes for:
- [#208](https://github.com/BrainBlend-AI/atomic-agents/issues/208): ChatHistory crashes with `TypeError` when schemas have both multimodal fields and nested Pydantic models
- [#141](https://github.com/BrainBlend-AI/atomic-agents/issues/141): AgentMemory doesn't support multimodal data inside nested schemas

## Features

1. **Nested Multimodal Schemas**: Images embedded inside nested Pydantic models (e.g., `Document.image`)
2. **Mixed Content**: Top-level multimodal fields combined with nested Pydantic context objects
3. **End-to-End Verification**: Verifies the chat history format is correct before making the LLM call

## Getting Started

1. Navigate to the nested-multimodal directory:

   ```bash
   cd atomic-agents/atomic-examples/nested-multimodal
   ```

2. Install dependencies using uv:

   ```bash
   uv sync
   ```

3. Set up environment variables:

   Create a `.env` file with:
   ```env
   OPENAI_API_KEY=your_openai_api_key
   ```

4. Run the example:

   ```bash
   uv run python nested_multimodal/main.py
   ```

## Schema Design

The example uses nested schemas that would have previously caused errors:

```python
class AnalysisContext(BaseIOSchema):
    """Nested context — a plain Pydantic model alongside multimodal fields."""
    focus_area: str
    detail_level: str

class ImageWithContext(BaseIOSchema):
    """Image wrapped in a nested schema with metadata."""
    image: instructor.Image
    label: str

class AnalysisInput(BaseIOSchema):
    """Top-level input combining nested multimodal + nested context."""
    documents: List[ImageWithContext]   # Images nested inside schemas
    context: AnalysisContext            # Nested Pydantic model
    instruction: str
```

The framework recursively extracts `Image` objects from any nesting depth and serializes the remaining fields using Pydantic's `model_dump_json(exclude=...)`.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.


## Source Code


### File: atomic-examples/nested-multimodal/nested_multimodal/main.py

```python
"""
Nested Multimodal Example
=========================

Demonstrates that Atomic Agents correctly handles multimodal content (images,
PDFs) inside nested Pydantic schemas — not just at the top level.

This example exercises the fixes for:
  - GitHub #208: nested Pydantic model + top-level multimodal → TypeError
  - GitHub #141: multimodal inside nested schemas invisible to ChatHistory
"""

import json
import os
from typing import List

import instructor
import openai
from dotenv import load_dotenv
from pydantic import Field

from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator

load_dotenv()

# ---------------------------------------------------------------------------
# API key
# ---------------------------------------------------------------------------
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )


# ---------------------------------------------------------------------------
# Schemas — nested multimodal content
# ---------------------------------------------------------------------------
class AnalysisContext(BaseIOSchema):
    """Additional context for the analysis request."""

    focus_area: str = Field(..., description="What aspect to focus the analysis on")
    detail_level: str = Field(..., description="How detailed the analysis should be (brief / detailed)")


class ImageWithContext(BaseIOSchema):
    """An image wrapped in a nested schema together with metadata."""

    image: instructor.Image = Field(..., description="The image to analyze")
    label: str = Field(..., description="A short human-readable label for this image")


class AnalysisInput(BaseIOSchema):
    """Input schema that combines nested multimodal content with a nested Pydantic context object."""

    documents: List[ImageWithContext] = Field(..., description="Images to analyze, each with a label")
    context: AnalysisContext = Field(..., description="Analysis context and preferences")
    instruction: str = Field(..., description="What the agent should do with the images")


class AnalysisOutput(BaseIOSchema):
    """Structured output from the image analysis."""

    summary: str = Field(..., description="Overall summary of all analyzed images")
    per_image: List[str] = Field(..., description="One description per image, in the same order as the input")


# ---------------------------------------------------------------------------
# Agent
# ---------------------------------------------------------------------------
agent = AtomicAgent[AnalysisInput, AnalysisOutput](
    config=AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=API_KEY)),
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are an image analysis assistant.",
                "You receive images wrapped inside document objects, each with a label.",
                "You also receive a context object that tells you what to focus on.",
            ],
            steps=[
                "1. Look at each image and its label.",
                "2. Analyze according to the focus_area and detail_level in the context.",
                "3. Write a per-image description and an overall summary.",
            ],
            output_instructions=[
                "Return a summary covering all images and a list of per-image descriptions.",
            ],
        ),
    )
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def verify_history_format(agent_instance: AtomicAgent) -> None:
    """Print the serialized chat history so we can confirm the fix works."""
    history = agent_instance.history.get_history()
    print("\n--- Chat history entries ---")
    for i, entry in enumerate(history):
        role = entry["role"]
        content = entry["content"]
        if isinstance(content, list):
            text_parts = [json.loads(c) if isinstance(c, str) else type(c).__name__ for c in content]
            print(f"  [{i}] role={role}  content (list with {len(content)} items):")
            for j, part in enumerate(text_parts):
                print(f"        [{j}] {part}")
        else:
            print(f"  [{i}] role={role}  content={content[:120]}...")
    print("--- end ---\n")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    print("=== Nested Multimodal Example ===\n")

    # Build the input — images nested inside ImageWithContext schemas
    script_dir = os.path.dirname(os.path.abspath(__file__))
    test_images_dir = os.path.join(os.path.dirname(script_dir), "test_images")
    image_path = os.path.join(test_images_dir, "nutrition_label_1.png")

    analysis_input = AnalysisInput(
        documents=[
            ImageWithContext(
                image=instructor.Image.from_path(image_path),
                label="Nutrition label photo",
            ),
        ],
        context=AnalysisContext(
            focus_area="nutritional content",
            detail_level="brief",
        ),
        instruction="Describe what you see in each image, paying attention to the focus area.",
    )

    # --- Verify the history format (no LLM call yet) -----------------------
    print("Step 1: Adding message to history and verifying serialization...\n")
    agent.history.add_message("user", analysis_input)
    verify_history_format(agent)

    # Confirm the nested Image was extracted and the nested AnalysisContext
    # was serialized properly (this is what Issues #208 / #141 broke).
    history = agent.history.get_history()
    assert isinstance(history[0]["content"], list), "Content should be a multimodal list"
    json_part = json.loads(history[0]["content"][0])
    assert "context" in json_part, "Nested AnalysisContext should be in the JSON"
    assert json_part["context"]["focus_area"] == "nutritional content"
    assert any(
        isinstance(item, instructor.Image) for item in history[0]["content"]
    ), "Image should be extracted into the content list"
    print("Serialization OK — nested context preserved, nested image extracted.\n")

    # Reset history before the real run (the agent adds messages internally)
    agent.reset_history()

    # --- End-to-end LLM call ------------------------------------------------
    print("Step 2: Running the agent end-to-end...\n")
    result = agent.run(analysis_input)

    print("Agent response:")
    print(f"  Summary : {result.summary}")
    for i, desc in enumerate(result.per_image, 1):
        print(f"  Image {i}: {desc}")

    # Show the full history after the run
    verify_history_format(agent)
    print("Done — nested multimodal schemas work end-to-end!")


if __name__ == "__main__":
    main()

```

### File: atomic-examples/nested-multimodal/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["nested_multimodal"]

[project]
name = "nested-multimodal"
version = "1.0.0"
description = "Nested multimodal example demonstrating images/PDFs inside nested Pydantic schemas"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "instructor==1.14.5",
    "openai>=2.0.0,<3.0.0",
    "python-dotenv>=1.0.0,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: orchestration-agent
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/orchestration-agent

## Documentation

# Orchestration Agent Example

This example demonstrates how to create an Orchestrator Agent that intelligently decides between using a search tool or a calculator tool based on user input.

## Features
- Intelligent tool selection between search and calculator tools
- Dynamic input/output schema handling
- Real-time date context provider
- Rich console output formatting
- Final answer generation based on tool outputs

## Getting Started

1. Clone the Atomic Agents repository:
   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. Navigate to the orchestration-agent directory:
   ```bash
   cd atomic-agents/atomic-examples/orchestration-agent
   ```

3. Install dependencies using uv:
   ```bash
   uv sync
   ```

4. Set up environment variables:
   Create a `.env` file in the `orchestration-agent` directory with:
   ```env
   OPENAI_API_KEY=your_openai_api_key
   ```

5. Install SearXNG (See: https://github.com/searxng/searxng)

6. Run the example:
   ```bash
   uv run python orchestration_agent/orchestrator.py
   ```

## Components

### Input/Output Schemas

- **OrchestratorInputSchema**: Handles user input messages
- **OrchestratorOutputSchema**: Specifies tool selection and parameters
- **FinalAnswerSchema**: Formats the final response

### Tools
These tools were installed using the Atomic Assembler CLI (See the main README [here](../../README.md) for more info)
The agent orchestrates between two tools:
- **SearXNG Search Tool**: For queries requiring factual information
- **Calculator Tool**: For mathematical calculations

### Context Providers

- **CurrentDateProvider**: Provides the current date in YYYY-MM-DD format


## Source Code


### File: atomic-examples/orchestration-agent/orchestration_agent/orchestrator.py

```python
from typing import Union
import openai
from pydantic import Field
from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator, BaseDynamicContextProvider

from orchestration_agent.tools.searxng_search import (
    SearXNGSearchTool,
    SearXNGSearchToolConfig,
    SearXNGSearchToolInputSchema,
    SearXNGSearchToolOutputSchema,
)
from orchestration_agent.tools.calculator import (
    CalculatorTool,
    CalculatorToolConfig,
    CalculatorToolInputSchema,
    CalculatorToolOutputSchema,
)

import instructor
from datetime import datetime


########################
# INPUT/OUTPUT SCHEMAS #
########################
class OrchestratorInputSchema(BaseIOSchema):
    """Input schema for the Orchestrator Agent. Contains the user's message to be processed."""

    chat_message: str = Field(..., description="The user's input message to be analyzed and responded to.")


class OrchestratorOutputSchema(BaseIOSchema):
    """Combined output schema for the Orchestrator Agent. Contains the tool parameters."""

    tool_parameters: Union[SearXNGSearchToolInputSchema, CalculatorToolInputSchema] = Field(
        ..., description="The parameters for the selected tool"
    )


class FinalAnswerSchema(BaseIOSchema):
    """Schema for the final answer generated by the Orchestrator Agent."""

    final_answer: str = Field(..., description="The final answer generated based on the tool output and user query.")


#######################
# AGENT CONFIGURATION #
#######################
class OrchestratorAgentConfig(AgentConfig):
    """Configuration for the Orchestrator Agent."""

    searxng_config: SearXNGSearchToolConfig
    calculator_config: CalculatorToolConfig


#####################
# CONTEXT PROVIDERS #
#####################
class CurrentDateProvider(BaseDynamicContextProvider):
    def __init__(self, title):
        super().__init__(title)
        self.date = datetime.now().strftime("%Y-%m-%d")

    def get_info(self) -> str:
        return f"Current date in format YYYY-MM-DD: {self.date}"


######################
# ORCHESTRATOR AGENT #
######################
orchestrator_agent_config = AgentConfig(
    client=instructor.from_openai(openai.OpenAI()),
    model="gpt-5-mini",
    model_api_parameters={"reasoning_effort": "low"},
    system_prompt_generator=SystemPromptGenerator(
        background=[
            "You are an Orchestrator Agent that decides between using a search tool or a calculator tool based on user input.",
            "Use the search tool for queries requiring factual information, current events, or specific data.",
            "Use the calculator tool for mathematical calculations and expressions.",
        ],
        output_instructions=[
            "Analyze the input to determine whether it requires a web search or a calculation.",
            "For search queries, use the 'search' tool and provide 1-3 relevant search queries.",
            "For calculations, use the 'calculator' tool and provide the mathematical expression to evaluate.",
            "When uncertain, prefer using the search tool.",
            "Format the output using the appropriate schema.",
        ],
    ),
)
orchestrator_agent = AtomicAgent[OrchestratorInputSchema, OrchestratorOutputSchema](config=orchestrator_agent_config)
orchestrator_agent_final = AtomicAgent[OrchestratorInputSchema, FinalAnswerSchema](config=orchestrator_agent_config)

# Register the current date provider
orchestrator_agent.register_context_provider("current_date", CurrentDateProvider("Current Date"))
orchestrator_agent_final.register_context_provider("current_date", CurrentDateProvider("Current Date"))


def execute_tool(
    searxng_tool: SearXNGSearchTool, calculator_tool: CalculatorTool, orchestrator_output: OrchestratorOutputSchema
) -> Union[SearXNGSearchToolOutputSchema, CalculatorToolOutputSchema]:
    if isinstance(orchestrator_output.tool_parameters, SearXNGSearchToolInputSchema):
        return searxng_tool.run(orchestrator_output.tool_parameters)
    elif isinstance(orchestrator_output.tool_parameters, CalculatorToolInputSchema):
        return calculator_tool.run(orchestrator_output.tool_parameters)
    else:
        raise ValueError(f"Unknown tool parameters type: {type(orchestrator_output.tool_parameters)}")


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    import os
    from dotenv import load_dotenv
    from rich.console import Console
    from rich.panel import Panel
    from rich.syntax import Syntax

    load_dotenv()

    # Set up the OpenAI client
    client = instructor.from_openai(openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")))

    # Initialize the tools
    searxng_tool = SearXNGSearchTool(SearXNGSearchToolConfig(base_url="http://localhost:8080", max_results=5))
    calculator_tool = CalculatorTool(CalculatorToolConfig())

    # Initialize Rich console
    console = Console()

    # Print the full system prompt
    console.print(Panel(orchestrator_agent.system_prompt_generator.generate_prompt(), title="System Prompt", expand=False))
    console.print("\n")

    # Example inputs
    inputs = [
        "Who won the Nobel Prize in Physics in 2024?",
        "Please calculate the sine of pi/3 to the third power",
    ]

    for user_input in inputs:
        console.print(Panel(f"[bold cyan]User Input:[/bold cyan] {user_input}", expand=False))

        # Create the input schema
        input_schema = OrchestratorInputSchema(chat_message=user_input)

        # Print the input schema
        console.print("\n[bold yellow]Generated Input Schema:[/bold yellow]")
        input_syntax = Syntax(str(input_schema.model_dump_json(indent=2)), "json", theme="monokai", line_numbers=True)
        console.print(input_syntax)

        # Run the orchestrator to get the tool selection and input
        orchestrator_output = orchestrator_agent.run(input_schema)

        # Print the orchestrator output
        console.print("\n[bold magenta]Orchestrator Output:[/bold magenta]")
        orchestrator_syntax = Syntax(
            str(orchestrator_output.model_dump_json(indent=2)), "json", theme="monokai", line_numbers=True
        )
        console.print(orchestrator_syntax)

        # Run the selected tool
        response = execute_tool(searxng_tool, calculator_tool, orchestrator_output)

        # Print the tool output
        console.print("\n[bold green]Tool Output:[/bold green]")
        output_syntax = Syntax(str(response.model_dump_json(indent=2)), "json", theme="monokai", line_numbers=True)
        console.print(output_syntax)

        console.print("\n" + "-" * 80 + "\n")

        # Switch agent
        history = orchestrator_agent.history
        orchestrator_agent = orchestrator_agent_final
        orchestrator_agent.history = history
        orchestrator_agent.add_tool_result(response)
        final_answer = orchestrator_agent.run(input_schema)
        console.print(f"\n[bold blue]Final Answer:[/bold blue] {final_answer.final_answer}")
        # Reset the agent to the original
        orchestrator_agent = AtomicAgent[OrchestratorInputSchema, OrchestratorOutputSchema](config=orchestrator_agent_config)

```

### File: atomic-examples/orchestration-agent/orchestration_agent/tools/calculator.py

```python
from pydantic import Field
from sympy import sympify

from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class CalculatorToolInputSchema(BaseIOSchema):
    """
    Tool for performing calculations. Supports basic arithmetic operations
    like addition, subtraction, multiplication, and division, as well as more
    complex operations like exponentiation and trigonometric functions.
    Use this tool to evaluate mathematical expressions.
    """

    expression: str = Field(..., description="Mathematical expression to evaluate. For example, '2 + 2'.")


#################
# OUTPUT SCHEMA #
#################
class CalculatorToolOutputSchema(BaseIOSchema):
    """
    Schema for the output of the CalculatorTool.
    """

    result: str = Field(..., description="Result of the calculation.")


#################
# CONFIGURATION #
#################
class CalculatorToolConfig(BaseToolConfig):
    """
    Configuration for the CalculatorTool.
    """

    pass


#####################
# MAIN TOOL & LOGIC #
#####################
class CalculatorTool(BaseTool[CalculatorToolInputSchema, CalculatorToolOutputSchema]):
    """
    Tool for performing calculations based on the provided mathematical expression.

    Attributes:
        input_schema (CalculatorToolInputSchema): The schema for the input data.
        output_schema (CalculatorToolOutputSchema): The schema for the output data.
    """

    input_schema = CalculatorToolInputSchema
    output_schema = CalculatorToolOutputSchema

    def __init__(self, config: CalculatorToolConfig = CalculatorToolConfig()):
        """
        Initializes the CalculatorTool.

        Args:
            config (CalculatorToolConfig): Configuration for the tool.
        """
        super().__init__(config)

    def run(self, params: CalculatorToolInputSchema) -> CalculatorToolOutputSchema:
        """
        Executes the CalculatorTool with the given parameters.

        Args:
            params (CalculatorToolInputSchema): The input parameters for the tool.

        Returns:
            CalculatorToolOutputSchema: The result of the calculation.
        """
        # Convert the expression string to a symbolic expression
        parsed_expr = sympify(str(params.expression))

        # Evaluate the expression numerically
        result = parsed_expr.evalf()
        return CalculatorToolOutputSchema(result=str(result))


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    calculator = CalculatorTool()
    result = calculator.run(CalculatorToolInputSchema(expression="sin(pi/2) + cos(pi/4)"))
    print(result)  # Expected output: {"result":"1.70710678118655"}

```

### File: atomic-examples/orchestration-agent/orchestration_agent/tools/searxng_search.py

```python
from typing import List, Literal, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

import aiohttp
from pydantic import Field

from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class SearXNGSearchToolInputSchema(BaseIOSchema):
    """
    Schema for input to a tool for searching for information, news, references, and other content using SearXNG.
    Returns a list of search results with a short description or content snippet and URLs for further exploration
    """

    queries: List[str] = Field(..., description="List of search queries.")
    category: Optional[Literal["general", "news", "social_media"]] = Field(
        "general", description="Category of the search queries."
    )


####################
# OUTPUT SCHEMA(S) #
####################
class SearXNGSearchResultItemSchema(BaseIOSchema):
    """This schema represents a single search result item"""

    url: str = Field(..., description="The URL of the search result")
    title: str = Field(..., description="The title of the search result")
    content: Optional[str] = Field(None, description="The content snippet of the search result")
    query: str = Field(..., description="The query used to obtain this search result")


class SearXNGSearchToolOutputSchema(BaseIOSchema):
    """This schema represents the output of the SearXNG search tool."""

    results: List[SearXNGSearchResultItemSchema] = Field(..., description="List of search result items")
    category: Optional[str] = Field(None, description="The category of the search results")


##############
# TOOL LOGIC #
##############
class SearXNGSearchToolConfig(BaseToolConfig):
    base_url: str = ""
    max_results: int = 10


class SearXNGSearchTool(BaseTool[SearXNGSearchToolInputSchema, SearXNGSearchToolOutputSchema]):
    """
    Tool for performing searches on SearXNG based on the provided queries and category.

    Attributes:
        input_schema (SearXNGSearchToolInputSchema): The schema for the input data.
        output_schema (SearXNGSearchToolOutputSchema): The schema for the output data.
        max_results (int): The maximum number of search results to return.
        base_url (str): The base URL for the SearXNG instance to use.
    """

    input_schema = SearXNGSearchToolInputSchema
    output_schema = SearXNGSearchToolOutputSchema

    def __init__(self, config: SearXNGSearchToolConfig = SearXNGSearchToolConfig()):
        """
        Initializes the SearXNGTool.

        Args:
            config (SearXNGSearchToolConfig):
                Configuration for the tool, including base URL, max results, and optional title and description overrides.
        """
        super().__init__(config)
        self.base_url = config.base_url
        self.max_results = config.max_results

    async def _fetch_search_results(self, session: aiohttp.ClientSession, query: str, category: Optional[str]) -> List[dict]:
        """
        Fetches search results for a single query asynchronously.

        Args:
            session (aiohttp.ClientSession): The aiohttp session to use for the request.
            query (str): The search query.
            category (Optional[str]): The category of the search query.

        Returns:
            List[dict]: A list of search result dictionaries.

        Raises:
            Exception: If the request to SearXNG fails.
        """
        query_params = {
            "q": query,
            "safesearch": "0",
            "format": "json",
            "language": "en",
            "engines": "bing,duckduckgo,google,startpage,yandex",
        }

        if category:
            query_params["categories"] = category

        async with session.get(f"{self.base_url}/search", params=query_params) as response:
            if response.status != 200:
                raise Exception(f"Failed to fetch search results for query '{query}': {response.status} {response.reason}")
            data = await response.json()
            results = data.get("results", [])

            # Add the query to each result
            for result in results:
                result["query"] = query

            return results

    async def run_async(
        self, params: SearXNGSearchToolInputSchema, max_results: Optional[int] = None
    ) -> SearXNGSearchToolOutputSchema:
        """
        Runs the SearXNGTool asynchronously with the given parameters.

        Args:
            params (SearXNGSearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
            max_results (Optional[int]): The maximum number of search results to return.

        Returns:
            SearXNGSearchToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            ValueError: If the base URL is not provided.
            Exception: If the request to SearXNG fails.
        """
        async with aiohttp.ClientSession() as session:
            tasks = [self._fetch_search_results(session, query, params.category) for query in params.queries]
            results = await asyncio.gather(*tasks)

        all_results = [item for sublist in results for item in sublist]

        # Sort the combined results by score in descending order
        sorted_results = sorted(all_results, key=lambda x: x.get("score", 0), reverse=True)

        # Remove duplicates while preserving order
        seen_urls = set()
        unique_results = []
        for result in sorted_results:
            if "content" not in result or "title" not in result or "url" not in result or "query" not in result:
                continue
            if result["url"] not in seen_urls:
                unique_results.append(result)
                if "metadata" in result:
                    result["title"] = f"{result['title']} - (Published {result['metadata']})"
                if "publishedDate" in result and result["publishedDate"]:
                    result["title"] = f"{result['title']} - (Published {result['publishedDate']})"
                seen_urls.add(result["url"])

        # Filter results to include only those with the correct category if it is set
        if params.category:
            filtered_results = [result for result in unique_results if result.get("category") == params.category]
        else:
            filtered_results = unique_results

        filtered_results = filtered_results[: max_results or self.max_results]

        return SearXNGSearchToolOutputSchema(
            results=[
                SearXNGSearchResultItemSchema(
                    url=result["url"], title=result["title"], content=result.get("content"), query=result["query"]
                )
                for result in filtered_results
            ],
            category=params.category,
        )

    def run(self, params: SearXNGSearchToolInputSchema, max_results: Optional[int] = None) -> SearXNGSearchToolOutputSchema:
        """
        Runs the SearXNGTool synchronously with the given parameters.

        This method creates an event loop in a separate thread to run the asynchronous operations.

        Args:
            params (SearXNGSearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
            max_results (Optional[int]): The maximum number of search results to return.

        Returns:
            SearXNGSearchToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            ValueError: If the base URL is not provided.
            Exception: If the request to SearXNG fails.
        """
        with ThreadPoolExecutor() as executor:
            return executor.submit(asyncio.run, self.run_async(params, max_results)).result()


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console
    from dotenv import load_dotenv

    load_dotenv()
    rich_console = Console()

    search_tool_instance = SearXNGSearchTool(config=SearXNGSearchToolConfig(base_url="http://localhost:8080", max_results=5))

    search_input = SearXNGSearchTool.input_schema(
        queries=["Python programming", "Machine learning", "Artificial intelligence"],
        category="news",
    )

    output = search_tool_instance.run(search_input)

    rich_console.print(output)

```

### File: atomic-examples/orchestration-agent/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["orchestration_agent"]

[project]
name = "orchestration-agent"
version = "0.1.0"
description = "Orchestration agent example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "KennyVaneetvelde", email = "kenny@inosta.be" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "instructor==1.14.5",
    "pydantic>=2.10.3,<3.0.0",
    "sympy>=1.13.3,<2.0.0",
    "python-dotenv>=1.0.1,<2.0.0",
    "openai>=2.0.0,<3.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```


--------------------------------------------------------------------------------
Example: progressive-disclosure
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/progressive-disclosure

## Documentation

# Progressive Disclosure Example

This example demonstrates **Anthropic's "progressive disclosure" pattern** for efficient MCP tool loading using the Atomic Agents framework with **three MCP servers** and **24 total tools**.

## The Problem

As documented by [Anthropic's Engineering Blog](https://www.anthropic.com/engineering/code-execution-with-mcp):

- **Context window bloat**: Loading all tool definitions upfront consumes massive context space
- **Performance degradation**: Agents connecting to 2-3+ MCP servers see significant accuracy drops
- **Cost inefficiency**: Traditional approach for multi-server setup: ~25,000+ tokens just for tool schemas

## The Solution: Progressive Disclosure

Instead of loading all 24 tool definitions upfront, a **sub-agent discovers relevant tools on-demand**:

```
┌─────────────────────────────────────────────────────────────────┐
│ WITHOUT Progressive Disclosure                                  │
│                                                                 │
│ Agent Context Window:                                          │
│ ┌─────────────────────────────────────────────────────────────┐│
│ │ math-server: 8 tools × ~500 tokens = 4,000 tokens          ││
│ │ text-server: 8 tools × ~500 tokens = 4,000 tokens          ││
│ │ data-server: 8 tools × ~500 tokens = 4,000 tokens          ││
│ │ ─────────────────────────────────────────────────           ││
│ │ Total: ~12,000 tokens just for tool definitions!           ││
│ └─────────────────────────────────────────────────────────────┘│
└─────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────┐
│ WITH Progressive Disclosure                                     │
│                                                                 │
│ Agent Context Window:                                          │
│ ┌─────────────────────────────────────────────────────────────┐│
│ │ add_numbers (500 tokens)                                    ││
│ │ multiply_numbers (500 tokens)                               ││
│ │ ─────────────────────────────────────────────────           ││
│ │ Total: ~1,000 tokens (92% reduction!)                      ││
│ └─────────────────────────────────────────────────────────────┘│
└─────────────────────────────────────────────────────────────────┘
```

## Project Structure

```
progressive-disclosure/
├── pyproject.toml
├── README.md
├── servers/                           # Three MCP servers
│   ├── math_server/                   # 8 arithmetic tools
│   │   ├── pyproject.toml
│   │   └── math_server/
│   │       ├── __init__.py
│   │       └── server.py              # FastMCP server
│   ├── text_server/                   # 8 text manipulation tools
│   │   ├── pyproject.toml
│   │   └── text_server/
│   │       ├── __init__.py
│   │       └── server.py
│   └── data_server/                   # 8 list/data tools
│       ├── pyproject.toml
│       └── data_server/
│           ├── __init__.py
│           └── server.py
└── progressive_disclosure/            # Client with progressive disclosure
    ├── __init__.py
    ├── main.py                        # Entry point
    ├── registry/
    │   └── tool_registry.py           # Lightweight tool metadata
    ├── tools/
    │   └── search_tools.py            # Tool search functionality
    └── agents/
        ├── tool_finder_agent.py       # Sub-agent for discovery
        └── orchestrator_agent.py      # Dynamic orchestrator factory
```

## Available Tools (24 Total)

### math-server (8 tools)
| Tool | Description |
|------|-------------|
| `add_numbers` | Add two numbers (a + b) |
| `subtract_numbers` | Subtract b from a (a - b) |
| `multiply_numbers` | Multiply two numbers (a * b) |
| `divide_numbers` | Divide a by b (a / b) |
| `power` | Raise base to exponent |
| `square_root` | Calculate square root |
| `modulo` | Calculate remainder (a % b) |
| `absolute_value` | Get absolute value |

### text-server (8 tools)
| Tool | Description |
|------|-------------|
| `uppercase` | Convert to UPPERCASE |
| `lowercase` | Convert to lowercase |
| `reverse_text` | Reverse character order |
| `word_count` | Count words in text |
| `char_count` | Count characters |
| `concatenate` | Join two strings |
| `replace_text` | Find and replace |
| `split_text` | Split by delimiter |

### data-server (8 tools)
| Tool | Description |
|------|-------------|
| `sort_list` | Sort numbers in a list |
| `filter_greater_than` | Filter values > threshold |
| `filter_less_than` | Filter values < threshold |
| `sum_list` | Sum all values |
| `average_list` | Calculate average |
| `min_value` | Find minimum |
| `max_value` | Find maximum |
| `unique_values` | Remove duplicates |

## Architecture

```
User Query: "Calculate (5 + 3) * 2 and reverse 'hello'"
    │
    ▼
┌─────────────────────────────────────────────────────┐
│     Phase 1: Tool Discovery                         │
│     ─────────────────────────                       │
│     Tool Finder Agent (gpt-5-mini)                │
│     - Searches lightweight registry                 │
│     - Registry has 24 tool names + descriptions    │
│     - Returns: ["add_numbers", "multiply_numbers",  │
│                 "reverse_text"]                     │
└─────────────┬───────────────────────────────────────┘
              │
              ▼
┌─────────────────────────────────────────────────────┐
│     Phase 2: Dynamic Orchestrator Creation          │
│     ─────────────────────────                       │
│     OrchestratorFactory                            │
│     - Loads ONLY 3 tool schemas (not 24!)          │
│     - Creates Union type dynamically               │
│     - 92% context reduction achieved               │
└─────────────┬───────────────────────────────────────┘
              │
              ▼
┌─────────────────────────────────────────────────────┐
│     Phase 3: Query Execution                        │
│     ─────────────────────────                       │
│     Main Orchestrator Agent (gpt-4o)               │
│     - Executes add_numbers(5, 3) → 8               │
│     - Executes multiply_numbers(8, 2) → 16         │
│     - Executes reverse_text("hello") → "olleh"     │
│     - Returns final response                        │
└─────────────────────────────────────────────────────┘
```

## Getting Started

### Prerequisites

- Python 3.12+
- OpenAI API key
- uv package manager

### Installation

```bash
# Clone the repository
git clone https://github.com/BrainBlend-AI/atomic-agents
cd atomic-agents/atomic-examples/progressive-disclosure

# Install dependencies
uv sync
```

### Configuration

Create a `.env` file:
```bash
OPENAI_API_KEY=your-api-key-here
```

### Running the Demo

```bash
uv run python -m progressive_disclosure.main
```

## Example Session

```
╭──────────────────────────────────────────────────────╮
│ Progressive Disclosure Demo                          │
│ Demonstrating Anthropic's pattern with 3 MCP servers │
╰──────────────────────────────────────────────────────╯

Connecting to MCP servers...
Connecting to math-server...
  Connected: 8 tools
Connecting to text-server...
  Connected: 8 tools
Connecting to data-server...
  Connected: 8 tools

Total: 24 tools across 3 servers

Ready! Type '/exit' to quit, '/stats' for statistics.
Example queries:
  - 'Calculate (5 + 3) * 2'              (math tools)
  - 'Convert HELLO WORLD to lowercase'  (text tools)
  - 'Find the average of [1,2,3,4,5]'   (data tools)
  - 'Reverse the text ABC and add 10+5' (multi-server!)

You: Calculate (5 + 3) * 2

Phase 1: Tool Discovery
Sub-agent searching 24 tools across 3 servers...
Selected 2 tools: ['add_numbers', 'multiply_numbers']
Reasoning: The query requires addition and multiplication operations

Phase 2: Creating Focused Orchestrator
Orchestrator context: 2 tools (filtered 92% = saved ~11000 tokens)

Phase 3: Query Execution
Executing: add_numbers
Parameters: {'a': 5, 'b': 3}
Executing: multiply_numbers
Parameters: {'a': 8, 'b': 2}

Response: The result of (5 + 3) * 2 is 16.

╭──────────────────────────────────────────────────────╮
│ Progressive Disclosure: 2/24 tools loaded (92%)     │
╰──────────────────────────────────────────────────────╯
```

## Key Benefits

| Metric | Without PD | With PD | Improvement |
|--------|-----------|---------|-------------|
| Tools in context | 24 | 2-5 | 90%+ reduction |
| Token usage | ~12,000 | ~1,000 | 92% savings |
| Tool accuracy | Lower | Higher | Better focus |
| Scalability | Limited | Excellent | Many servers |

## How Atomic Agents Enables This

This example demonstrates several Atomic Agents patterns:

1. **Sub-Agent Pattern**: Tool Finder as specialized discovery agent
2. **Dynamic Schema Creation**: `Union` types built at runtime from selected tools
3. **Multi-Server MCP**: Connecting to multiple MCP servers simultaneously
4. **Tool Registry**: Lightweight metadata storage without full schemas
5. **Context Efficiency**: Only relevant information loaded

## The FastMCP Servers

Each server is a simple FastMCP application:

```python
from fastmcp import FastMCP

mcp = FastMCP("math-server")

@mcp.tool
def add_numbers(a: float, b: float) -> float:
    """Add two numbers together (a + b)."""
    return a + b

# ... more tools ...

if __name__ == "__main__":
    mcp.run()
```

## References

- [Anthropic: Code Execution with MCP](https://www.anthropic.com/engineering/code-execution-with-mcp)
- [FastMCP Documentation](https://gofastmcp.com)
- [Model Context Protocol](https://modelcontextprotocol.io/)
- [Atomic Agents Documentation](https://github.com/BrainBlend-AI/atomic-agents)

## See Also

- [MCP Agent Example](../mcp-agent/) - Basic single-server MCP integration
- [Orchestration Agent Example](../orchestration-agent/) - Tool orchestration patterns
- [Deep Research Example](../deep-research/) - Multi-agent pipelines


## Source Code


### File: atomic-examples/progressive-disclosure/progressive_disclosure/__init__.py

```python
"""Progressive Disclosure example for Atomic Agents.

This module demonstrates Anthropic's "progressive disclosure" pattern where
MCP tools are discovered on-demand rather than loaded all at once, significantly
reducing context window usage and improving tool selection accuracy.
"""

__version__ = "0.1.0"

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/agents/__init__.py

```python
"""Agents module for progressive disclosure."""

from progressive_disclosure.agents.tool_finder_agent import (
    ToolFinderInputSchema,
    ToolFinderOutputSchema,
    create_tool_finder_agent,
)
from progressive_disclosure.agents.orchestrator_agent import (
    OrchestratorFactory,
    OrchestratorInputSchema,
    FinalResponseSchema,
)

__all__ = [
    "ToolFinderInputSchema",
    "ToolFinderOutputSchema",
    "create_tool_finder_agent",
    "OrchestratorFactory",
    "OrchestratorInputSchema",
    "FinalResponseSchema",
]

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/agents/orchestrator_agent.py

```python
# pyright: reportInvalidTypeForm=false
"""Dynamic Orchestrator Factory for progressive disclosure.

This module provides a factory for creating orchestrator agents with
dynamically filtered tool sets. Instead of loading all available MCP tools,
the orchestrator is created with only the tools selected by the Tool Finder Agent.

This is the key component that achieves context window efficiency through
progressive disclosure.

Supports both sequential and parallel tool execution modes.
"""

from typing import List, Type, Dict, Union, Optional, Any, Callable
from pydantic import Field
import instructor
import asyncio
from concurrent.futures import ThreadPoolExecutor, as_completed

from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import ChatHistory, SystemPromptGenerator
from atomic_agents.base.base_tool import BaseTool
from atomic_agents.connectors.mcp import (
    fetch_mcp_tools,
    MCPTransportType,
)


########################
# INPUT/OUTPUT SCHEMAS #
########################
class OrchestratorInputSchema(BaseIOSchema):
    """Input schema for the orchestrator agent."""

    query: str = Field(
        ...,
        description="The user's query to process using the available tools.",
    )


class FinalResponseSchema(BaseIOSchema):
    """Schema for the final response to the user."""

    response_text: str = Field(
        ...,
        description="The final text response to the user's query.",
    )


class MCPToolOutputSchema(BaseIOSchema):
    """Generic output schema for MCP tool execution."""

    result: Any = Field(..., description="The result from the tool execution.")


#######################
# ORCHESTRATOR OUTPUT #
#######################
def create_orchestrator_output_schema(
    tool_schemas: tuple[Type[BaseIOSchema], ...],
    parallel: bool = False,
) -> Type[BaseIOSchema]:
    """Dynamically create an orchestrator output schema with the given tools.

    Args:
        tool_schemas: Tuple of tool input schema classes.
        parallel: If True, creates schema supporting multiple parallel actions.

    Returns:
        A new BaseIOSchema class with the dynamic action field(s).
    """
    # Create the union of all schemas
    all_schemas = tool_schemas + (FinalResponseSchema,)
    ActionUnion = Union[all_schemas]  # type: ignore[valid-type]

    if parallel:

        class ParallelOrchestratorOutputSchema(BaseIOSchema):
            """Orchestrator output schema supporting parallel tool execution."""

            reasoning: str = Field(
                ...,
                description="Explanation of why these tools are needed and how they work together.",
            )
            actions: List[ActionUnion] = Field(  # type: ignore[valid-type]
                ...,
                description="List of tool executions. Independent tools will run in parallel. Include FinalResponseSchema when done.",
            )

            model_config = {"arbitrary_types_allowed": True}

        return ParallelOrchestratorOutputSchema
    else:

        class DynamicOrchestratorOutputSchema(BaseIOSchema):
            """Dynamically generated orchestrator output schema."""

            reasoning: str = Field(
                ...,
                description="Detailed explanation of why this action was chosen and how it addresses the user's query.",
            )
            action: ActionUnion = Field(  # type: ignore[valid-type]
                ...,
                description="The chosen action: either a tool's input schema instance or a final response.",
            )

            model_config = {"arbitrary_types_allowed": True}

        return DynamicOrchestratorOutputSchema


######################
# ORCHESTRATOR CLASS #
######################
class OrchestratorFactory:
    """Factory for creating orchestrator agents with filtered tool sets.

    This factory creates orchestrator agents that only have access to
    the specific tools selected by the Tool Finder Agent, implementing
    the progressive disclosure pattern.

    Supports both sequential (one tool at a time) and parallel execution modes.

    Example:
        >>> factory = OrchestratorFactory(
        ...     mcp_endpoint="http://localhost:6969",
        ...     transport_type=MCPTransportType.HTTP_STREAM,
        ...     client=instructor.from_openai(openai.OpenAI()),
        ...     parallel_execution=True,  # Enable parallel mode
        ... )
        >>> orchestrator, tool_map = factory.create_with_tools(
        ...     ["AddNumbers", "SubtractNumbers"],
        ...     all_tools=all_mcp_tools,
        ... )
    """

    def __init__(
        self,
        mcp_endpoint: Optional[str],
        transport_type: MCPTransportType,
        client: instructor.Instructor,
        model: str = "gpt-5.1",
        client_session: Optional[Any] = None,
        event_loop: Optional[asyncio.AbstractEventLoop] = None,
        parallel_execution: bool = True,
    ):
        """Initialize the orchestrator factory.

        Args:
            mcp_endpoint: MCP server endpoint URL (None for STDIO).
            transport_type: MCP transport type (HTTP_STREAM, SSE, STDIO).
            client: Instructor-wrapped LLM client.
            model: Model to use for orchestration.
            client_session: Optional MCP client session for STDIO transport.
            event_loop: Optional event loop for STDIO transport.
            parallel_execution: If True, enables parallel tool execution mode.
        """
        self.mcp_endpoint = mcp_endpoint
        self.transport_type = transport_type
        self.client = client
        self.model = model
        self.client_session = client_session
        self.event_loop = event_loop
        self.parallel_execution = parallel_execution

    def create_with_tools(
        self,
        tool_names: List[str],
        all_tools: Optional[List[Type[BaseTool]]] = None,
    ) -> tuple[AtomicAgent, Dict[Type[BaseIOSchema], Type[BaseTool]]]:
        """Create an orchestrator with only the specified tools.

        This is the core method that achieves progressive disclosure:
        only the selected tools are included in the orchestrator's schema,
        keeping the context window lean and focused.

        Args:
            tool_names: Names of tools to include (from Tool Finder Agent).
            all_tools: Optional pre-fetched list of all MCP tools. If not provided,
                      tools will be fetched from the MCP server.

        Returns:
            Tuple of (orchestrator_agent, tool_schema_to_class_map).

        Raises:
            ValueError: If no matching tools are found.
        """
        # Get all tools if not provided
        if all_tools is None:
            all_tools = fetch_mcp_tools(
                mcp_endpoint=self.mcp_endpoint,
                transport_type=self.transport_type,
                client_session=self.client_session,
                event_loop=self.event_loop,
            )

        # Filter to only the requested tools
        filtered_tools = [tool for tool in all_tools if getattr(tool, "mcp_tool_name", None) in tool_names]

        if not filtered_tools:
            # If no MCP tools match, create a minimal orchestrator
            return self._create_minimal_orchestrator(), {}

        # Build schema-to-class mapping for execution
        tool_schema_to_class: Dict[Type[BaseIOSchema], Type[BaseTool]] = {tool.input_schema: tool for tool in filtered_tools}

        # Create the dynamic output schema with only filtered tools
        tool_input_schemas = tuple(tool.input_schema for tool in filtered_tools)
        output_schema = create_orchestrator_output_schema(tool_input_schemas, parallel=self.parallel_execution)

        # Build tool descriptions for the system prompt
        tool_descriptions = []
        for tool in filtered_tools:
            tool_name = getattr(tool, "mcp_tool_name", tool.__name__)
            tool_desc = tool.__doc__ or "No description available"
            tool_descriptions.append(f"- {tool_name}: {tool_desc}")

        # Create system prompt based on execution mode
        if self.parallel_execution:
            background = [
                "You are an Orchestrator Agent that MUST use the provided tools.",
                "You have a FOCUSED set of tools for this task.",
                "",
                "Available tools:",
                *tool_descriptions,
                "",
                "CRITICAL: You MUST call tools - never compute results yourself!",
                "PARALLEL MODE: Batch independent tool calls together for speed.",
            ]
            steps = [
                "1. Identify ALL tool calls needed for the query",
                "2. Batch 1: Call ALL tools whose inputs are already known",
                "3. Wait for results, then Batch 2: Call tools using those results",
                "4. Only return FinalResponseSchema AFTER all tools have been called",
            ]
            output_instructions = [
                "MANDATORY: Use tools for ALL calculations - never compute in your head",
                "BATCH independent calls: char_count('a'), char_count('b') → 2 actions together",
                "NEVER skip tools - even for simple math like sqrt or counting",
                "FinalResponseSchema: Only after ALL required tools have returned results",
            ]
        else:
            background = [
                "You are an Orchestrator Agent that processes user queries using available tools.",
                "You have been given a FOCUSED set of tools relevant to the current task.",
                "",
                "Available tools:",
                *tool_descriptions,
                "",
                "SEQUENTIAL MODE: Execute ONE tool per turn.",
                "You will be called multiple times, receiving tool results after each execution.",
            ]
            steps = [
                "1. Analyze what needs to be done next (considering previous results if any)",
                "2. Choose exactly ONE tool to execute, or provide the final response",
                "3. Fill in the tool's parameters directly in the action field",
                "4. After receiving results, continue with the next tool or finalize",
            ]
            output_instructions = [
                "Execute exactly ONE tool per turn",
                "The 'action' field must contain a SINGLE tool's input schema directly",
                "When all tools have been executed, use FinalResponseSchema with the complete answer",
            ]

        # Create the orchestrator agent
        orchestrator = AtomicAgent[OrchestratorInputSchema, output_schema](
            config=AgentConfig(
                client=self.client,
                model=self.model,
                history=ChatHistory(),
                system_prompt_generator=SystemPromptGenerator(
                    background=background,
                    steps=steps,
                    output_instructions=output_instructions,
                ),
            )
        )

        return orchestrator, tool_schema_to_class

    def _create_minimal_orchestrator(self) -> AtomicAgent:
        """Create a minimal orchestrator with no tools (for conversation only)."""
        output_schema = create_orchestrator_output_schema(tuple(), parallel=self.parallel_execution)

        if self.parallel_execution:
            output_instructions = [
                "Provide clear, helpful responses",
                "Use FinalResponseSchema in the actions list for your response",
            ]
        else:
            output_instructions = [
                "Provide clear, helpful responses",
                "Use FinalResponseSchema for your response",
            ]

        return AtomicAgent[OrchestratorInputSchema, output_schema](
            config=AgentConfig(
                client=self.client,
                model=self.model,
                history=ChatHistory(),
                system_prompt_generator=SystemPromptGenerator(
                    background=[
                        "You are an assistant that responds to user queries.",
                        "No tools are currently available for this query.",
                    ],
                    steps=[
                        "1. Analyze the user's query",
                        "2. Provide a helpful response based on your knowledge",
                    ],
                    output_instructions=output_instructions,
                ),
            )
        )


##################################
# SEQUENTIAL EXECUTION (LEGACY)  #
##################################
def execute_orchestrator_loop(
    orchestrator: AtomicAgent,
    tool_schema_to_class: Dict[Type[BaseIOSchema], Type[BaseTool]],
    initial_query: str,
    max_iterations: int = 10,
    on_tool_execution: Optional[Callable] = None,
) -> str:
    """Execute the orchestrator loop sequentially (one tool at a time).

    This function handles the multi-turn interaction where the orchestrator
    selects and executes tools until it produces a final response.

    Args:
        orchestrator: The orchestrator agent.
        tool_schema_to_class: Mapping from input schemas to tool classes.
        initial_query: The user's initial query.
        max_iterations: Maximum number of tool executions.
        on_tool_execution: Optional callback for tool execution events.

    Returns:
        The final response text.
    """
    # Initial run with user query
    output = orchestrator.run(OrchestratorInputSchema(query=initial_query))
    action = output.action

    iteration = 0
    while not isinstance(action, FinalResponseSchema) and iteration < max_iterations:
        iteration += 1
        schema_type = type(action)

        # Find and execute the matching tool
        tool_class = tool_schema_to_class.get(schema_type)
        if tool_class is None:
            raise ValueError(f"Unknown action schema: {schema_type.__name__}")

        # Execute the tool
        tool_instance = tool_class()
        tool_name = getattr(tool_class, "mcp_tool_name", tool_class.__name__)

        if on_tool_execution:
            on_tool_execution(tool_name, action.model_dump())

        tool_output = tool_instance.run(action)

        # Add result to history
        result_message = OrchestratorInputSchema(query=f"Tool '{tool_name}' executed. Result: {tool_output.result}")
        orchestrator.add_tool_result(result_message)

        # Continue the loop
        output = orchestrator.run()
        action = output.action

    if isinstance(action, FinalResponseSchema):
        return action.response_text
    else:
        return "Maximum iterations reached. Please try a simpler query."


##################################
# PARALLEL EXECUTION             #
##################################
def execute_orchestrator_loop_parallel(
    orchestrator: AtomicAgent,
    tool_schema_to_class: Dict[Type[BaseIOSchema], Type[BaseTool]],
    initial_query: str,
    max_iterations: int = 10,
    on_tool_execution: Optional[Callable] = None,
    on_parallel_batch: Optional[Callable] = None,
    max_parallel_workers: int = 5,
) -> str:
    """Execute the orchestrator loop with parallel tool execution.

    When the orchestrator returns multiple independent tools in its 'actions' list,
    they are executed concurrently using a thread pool for maximum efficiency.

    Args:
        orchestrator: The orchestrator agent (must be created with parallel_execution=True).
        tool_schema_to_class: Mapping from input schemas to tool classes.
        initial_query: The user's initial query.
        max_iterations: Maximum number of execution rounds (not individual tools).
        on_tool_execution: Optional callback for each tool execution.
        on_parallel_batch: Optional callback when a parallel batch starts, receives count.
        max_parallel_workers: Maximum concurrent tool executions.

    Returns:
        The final response text.
    """
    # Initial run with user query
    output = orchestrator.run(OrchestratorInputSchema(query=initial_query))
    actions = output.actions  # List of actions in parallel mode

    # Track executed tool calls to prevent duplicates
    executed_calls: set[str] = set()

    def get_call_signature(action) -> str:
        """Create a unique signature for a tool call."""
        tool_class = tool_schema_to_class.get(type(action))
        if tool_class is None:
            return ""
        tool_name = getattr(tool_class, "mcp_tool_name", tool_class.__name__)
        # Create signature from tool name + sorted params
        params = action.model_dump()
        params.pop("tool_name", None)  # Remove tool_name from params
        param_str = str(sorted(params.items()))
        return f"{tool_name}:{param_str}"

    iteration = 0
    while iteration < max_iterations:
        iteration += 1

        # Separate final response from tool actions
        final_responses = [a for a in actions if isinstance(a, FinalResponseSchema)]
        tool_actions = [a for a in actions if not isinstance(a, FinalResponseSchema)]

        # Filter out duplicate tool calls
        unique_tool_actions = []
        skipped_duplicates = 0
        for action in tool_actions:
            sig = get_call_signature(action)
            if sig and sig not in executed_calls:
                unique_tool_actions.append(action)
                executed_calls.add(sig)
            else:
                skipped_duplicates += 1

        tool_actions = unique_tool_actions

        # If no tool actions, we're done - return final response or error
        if not tool_actions:
            if final_responses:
                return final_responses[0].response_text
            # If we skipped duplicates, prompt model for final answer
            if skipped_duplicates > 0:
                prompt_msg = OrchestratorInputSchema(
                    query="All tool results are now available. Please provide your final answer using FinalResponseSchema."
                )
                orchestrator.add_tool_result(prompt_msg)
                output = orchestrator.run()
                actions = output.actions
                continue  # Re-check for FinalResponseSchema
            return "No actions returned by orchestrator."

        # Notify about parallel batch
        if on_parallel_batch and len(tool_actions) > 1:
            on_parallel_batch(len(tool_actions))

        # Execute tools in parallel using ThreadPoolExecutor
        def execute_single_tool(action):
            schema_type = type(action)
            tool_class = tool_schema_to_class.get(schema_type)
            if tool_class is None:
                return {"error": f"Unknown action schema: {schema_type.__name__}"}

            tool_instance = tool_class()
            tool_name = getattr(tool_class, "mcp_tool_name", tool_class.__name__)

            if on_tool_execution:
                on_tool_execution(tool_name, action.model_dump())

            try:
                # Use sync run method - it handles async internally for MCP tools
                import warnings

                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", RuntimeWarning)
                    tool_output = tool_instance.run(action)
                return {
                    "tool_name": tool_name,
                    "result": tool_output.result,
                    "success": True,
                }
            except Exception as e:
                return {
                    "tool_name": tool_name,
                    "error": str(e),
                    "success": False,
                }

        # Execute all tools in parallel
        results = []
        with ThreadPoolExecutor(max_workers=max_parallel_workers) as executor:
            future_to_action = {executor.submit(execute_single_tool, action): action for action in tool_actions}
            for future in as_completed(future_to_action):
                results.append(future.result())

        # Build result message for history
        if len(results) == 1:
            r = results[0]
            if r.get("success"):
                result_text = f"Tool '{r['tool_name']}' executed. Result: {r['result']}"
            else:
                result_text = f"Tool '{r['tool_name']}' failed. Error: {r.get('error')}"
        else:
            result_lines = ["Tools executed in parallel:"]
            for r in results:
                if r.get("success"):
                    result_lines.append(f"  - {r['tool_name']}: {r['result']}")
                else:
                    result_lines.append(f"  - {r['tool_name']}: ERROR - {r.get('error')}")
            result_text = "\n".join(result_lines)

        # Add results to history
        result_message = OrchestratorInputSchema(query=result_text)
        orchestrator.add_tool_result(result_message)

        # Continue the loop
        output = orchestrator.run()
        actions = output.actions

    return "Maximum iterations reached. Please try a simpler query."


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console

    console = Console()
    console.print("[bold]Orchestrator Factory Demo[/bold]")
    console.print("This module is typically used via main.py")
    console.print("See main.py for a complete example of progressive disclosure in action.")
    console.print("")
    console.print("[cyan]Parallel Execution Mode:[/cyan]")
    console.print("  - Multiple independent tools execute concurrently")
    console.print("  - Example: sqrt(14) + sqrt(10) runs both sqrt calls in parallel")
    console.print("  - Reduces latency by ~50% for independent operations")

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/agents/tool_finder_agent.py

```python
"""Tool Finder Agent for progressive disclosure.

This agent is responsible for discovering relevant tools for a given user query.
It analyzes the lightweight tool registry to find the most appropriate tools,
allowing the main orchestrator to be created with only the necessary tools
loaded into its context window.

This implements the "search_tools" pattern from Anthropic's progressive disclosure.
"""

from typing import List, Optional
from pydantic import Field
import instructor

from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator, BaseDynamicContextProvider

from progressive_disclosure.registry.tool_registry import ToolRegistry


########################
# INPUT/OUTPUT SCHEMAS #
########################
class ToolFinderInputSchema(BaseIOSchema):
    """Input for the tool finder agent."""

    user_query: str = Field(
        ...,
        description="The user's original query that needs to be analyzed to determine required tools.",
    )
    task_context: Optional[str] = Field(
        default=None,
        description="Additional context about the task that might help with tool selection.",
    )


class ToolFinderOutputSchema(BaseIOSchema):
    """Output containing selected tools for the main orchestrator."""

    reasoning: str = Field(
        ...,
        description="Detailed explanation of why these specific tools were selected and how they relate to the user's query.",
    )
    selected_tools: List[str] = Field(
        ...,
        description="Names of tools that should be loaded for the main orchestrator. Keep this list minimal.",
    )
    search_queries_used: List[str] = Field(
        default_factory=list,
        description="Keywords or concepts used to identify these tools.",
    )
    confidence: str = Field(
        default="high",
        description="Confidence level in tool selection: 'high', 'medium', or 'low'.",
    )


#####################
# CONTEXT PROVIDERS #
#####################
class ToolRegistryProvider(BaseDynamicContextProvider):
    """Provides the full tool registry to the finder agent."""

    def __init__(self, registry: ToolRegistry, title: str = "Available Tools"):
        super().__init__(title)
        self._registry = registry

    def get_info(self) -> str:
        """Get all available tools with descriptions."""
        tools = self._registry.get_all_tools()
        if not tools:
            return "No tools available in registry."

        lines = ["The following tools are available:\n"]
        for tool in tools:
            lines.append(f"- **{tool.name}**: {tool.description}")

        lines.append("\nSelect ONLY the tools needed to complete the user's query.")
        return "\n".join(lines)


#############################
# TOOL FINDER AGENT FACTORY #
#############################
def create_tool_finder_agent(
    registry: ToolRegistry,
    client: instructor.Instructor,
    model: str = "gpt-5-mini",
) -> tuple[AtomicAgent, None, None]:
    """Create a tool finder agent with access to tool metadata.

    The tool finder agent uses a lightweight model to analyze user queries
    and determine which MCP tools should be loaded for the main orchestrator.

    Args:
        registry: Tool registry containing metadata about available tools.
        client: Instructor-wrapped LLM client.
        model: Model to use for the finder agent. Default is gpt-5-mini
               for cost efficiency since this is a discovery task.

    Returns:
        Tuple of (agent, None, None) - the None values maintain API compatibility.

    Example:
        >>> registry = ToolRegistry()
        >>> registry.register_from_mcp(mcp_definitions)
        >>> client = instructor.from_openai(openai.OpenAI())
        >>> agent, _, _ = create_tool_finder_agent(registry, client)
        >>> result = run_tool_finder(agent, None, None, "Calculate 2+2")
        >>> print(result.selected_tools)
        ['add_numbers']
    """
    # Create the agent
    agent = AtomicAgent[ToolFinderInputSchema, ToolFinderOutputSchema](
        config=AgentConfig(
            client=client,
            model=model,
            system_prompt_generator=SystemPromptGenerator(
                background=[
                    "You are a Tool Finder Agent specialized in discovering relevant tools for user queries.",
                    "Your role is to analyze user queries and find the MINIMUM set of tools needed to accomplish the task.",
                    "You have access to a list of available MCP tools with their descriptions.",
                    "",
                    "IMPORTANT: Your goal is CONTEXT EFFICIENCY - select only the tools that are directly needed.",
                    "The tools you select will be loaded into another agent's context window.",
                    "Loading unnecessary tools wastes context space and reduces accuracy.",
                ],
                steps=[
                    "1. Analyze the user's query to understand what capabilities are needed",
                    "2. Review the available tools list provided in your context",
                    "3. Select ONLY the tools that are necessary for this specific query",
                    "4. Provide your selection with clear reasoning",
                ],
                output_instructions=[
                    "Select the MINIMUM number of tools needed - prefer fewer tools over more",
                    "Only include tools that are directly relevant to accomplishing the user's task",
                    "If no tools are needed (e.g., general conversation), return an empty list",
                    "Include clear reasoning for each selected tool",
                    "Rate your confidence: 'high' if certain, 'medium' if tools might work, 'low' if unsure",
                    "Use the exact tool names as they appear in the available tools list",
                ],
            ),
        )
    )

    # Register context provider with full tool list
    agent.register_context_provider(
        "tool_registry",
        ToolRegistryProvider(registry, "Available Tools"),
    )

    return agent, None, None


def run_tool_finder(
    agent: AtomicAgent,
    search_tool,  # Not used, kept for API compatibility
    list_tool,  # Not used, kept for API compatibility
    user_query: str,
    task_context: Optional[str] = None,
    max_iterations: int = 5,  # Not used, kept for API compatibility
) -> ToolFinderOutputSchema:
    """Run the tool finder agent to discover relevant tools.

    This is a single-pass approach - the agent sees all tool metadata
    and selects the relevant tools in one call.

    Args:
        agent: The tool finder agent.
        search_tool: Not used (kept for API compatibility).
        list_tool: Not used (kept for API compatibility).
        user_query: The user's query to analyze.
        task_context: Optional additional context.
        max_iterations: Not used (kept for API compatibility).

    Returns:
        ToolFinderOutputSchema with the selected tools.
    """
    input_schema = ToolFinderInputSchema(
        user_query=user_query,
        task_context=task_context,
    )

    # Single-pass tool selection
    result = agent.run(input_schema)
    return result


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    import os
    from dotenv import load_dotenv
    from rich.console import Console
    import openai
    from progressive_disclosure.registry.tool_registry import ToolMetadata

    load_dotenv()
    console = Console()

    # Create a test registry
    registry = ToolRegistry()
    registry.register(
        ToolMetadata(
            name="add_numbers",
            description="Add two numbers together",
            keywords=["add", "sum", "plus", "arithmetic"],
            category="math",
        )
    )
    registry.register(
        ToolMetadata(
            name="subtract_numbers",
            description="Subtract one number from another",
            keywords=["subtract", "minus", "difference", "arithmetic"],
            category="math",
        )
    )
    registry.register(
        ToolMetadata(
            name="multiply_numbers",
            description="Multiply two numbers together",
            keywords=["multiply", "times", "product", "arithmetic"],
            category="math",
        )
    )
    registry.register(
        ToolMetadata(
            name="divide_numbers",
            description="Divide one number by another",
            keywords=["divide", "quotient", "arithmetic"],
            category="math",
        )
    )
    registry.register(
        ToolMetadata(
            name="uppercase",
            description="Convert text to uppercase",
            keywords=["upper", "capitalize", "text"],
            category="text",
        )
    )
    registry.register(
        ToolMetadata(
            name="reverse_text",
            description="Reverse the characters in text",
            keywords=["reverse", "backwards", "text"],
            category="text",
        )
    )

    # Create client
    client = instructor.from_openai(openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")))

    # Create finder agent
    agent, _, _ = create_tool_finder_agent(registry, client)

    # Test queries
    test_queries = [
        "What is 5 plus 3?",
        "Calculate (10 - 4) * 2",
        "Reverse the word HELLO and convert ABC to uppercase",
    ]

    for query in test_queries:
        console.print(f"\n[bold cyan]Query:[/bold cyan] {query}")
        result = run_tool_finder(agent, None, None, query)
        console.print(f"[bold green]Selected tools:[/bold green] {result.selected_tools}")
        console.print(f"[dim]Reasoning: {result.reasoning}[/dim]")
        console.print(f"[dim]Confidence: {result.confidence}[/dim]")

        # Reset history for next query
        agent.history.history = []
        agent.history.current_turn_id = None

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/main.py

```python
# pyright: reportInvalidTypeForm=false
"""Progressive Disclosure Demo with Multiple MCP Servers.

This script demonstrates Anthropic's "progressive disclosure" pattern where
MCP tools are discovered on-demand rather than loaded all at once.

We have THREE MCP servers:
- math-server: 8 arithmetic tools (add, subtract, multiply, divide, power, sqrt, modulo, abs)
- text-server: 8 text manipulation tools (uppercase, lowercase, reverse, word_count, etc.)
- data-server: 8 list/data tools (sort, filter, sum, average, min, max, unique)

Total: 24 tools across 3 servers.

The progressive disclosure pattern:
1. Tool Finder Agent searches for relevant tools based on user query
2. Only selected tools (typically 2-5) are loaded into the Main Orchestrator
3. Result: ~90% reduction in context window usage

Without progressive disclosure: All 24 tool schemas in context (~12,000 tokens)
With progressive disclosure: Only 2-5 relevant tools (~1,000 tokens)
"""

import asyncio
import os
import shlex
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from typing import List, Type, Dict

import instructor
import openai
from dotenv import load_dotenv
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from atomic_agents.connectors.mcp import (
    fetch_mcp_tools,
    MCPTransportType,
)
from atomic_agents.base.base_tool import BaseTool

from progressive_disclosure.registry.tool_registry import ToolRegistry, MCPToolDefinition
from progressive_disclosure.agents.tool_finder_agent import (
    create_tool_finder_agent,
    run_tool_finder,
)
from progressive_disclosure.agents.orchestrator_agent import (
    OrchestratorFactory,
    execute_orchestrator_loop,
    execute_orchestrator_loop_parallel,
)


########################
# CONFIGURATION        #
########################
@dataclass
class ServerConfig:
    """Configuration for an MCP server."""

    name: str
    command: str
    category: str  # For tool categorization


@dataclass
class ProgressiveDisclosureConfig:
    """Configuration for the Progressive Disclosure demo."""

    openai_api_key: str = field(default_factory=lambda: os.getenv("OPENAI_API_KEY", ""))
    finder_model: str = "gpt-5-mini"  # Lightweight model for tool discovery
    orchestrator_model: str = "gpt-5.1"  # More capable model for execution
    parallel_execution: bool = True  # Enable parallel tool execution

    # Three MCP servers demonstrating multi-server progressive disclosure
    servers: List[ServerConfig] = field(
        default_factory=lambda: [
            ServerConfig(name="math-server", command="uv run pd-math-server", category="math"),
            ServerConfig(name="text-server", command="uv run pd-text-server", category="text"),
            ServerConfig(name="data-server", command="uv run pd-data-server", category="data"),
        ]
    )

    def __post_init__(self):
        if not self.openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is not set")


########################
# SERVER SESSION MGR   #
########################
class MCPServerManager:
    """Manages connections to multiple MCP servers."""

    def __init__(self, server_configs: List[ServerConfig]):
        self.server_configs = server_configs
        self.sessions: Dict[str, ClientSession] = {}
        self.loops: Dict[str, asyncio.AbstractEventLoop] = {}
        self.exit_stacks: Dict[str, AsyncExitStack] = {}
        self.tools_by_server: Dict[str, List[Type[BaseTool]]] = {}
        self.all_tools: List[Type[BaseTool]] = []

    async def _connect_server(self, config: ServerConfig) -> ClientSession:
        """Connect to a single MCP server."""
        exit_stack = AsyncExitStack()
        self.exit_stacks[config.name] = exit_stack

        command_parts = shlex.split(config.command)
        server_params = StdioServerParameters(command=command_parts[0], args=command_parts[1:], env=None)

        read_stream, write_stream = await exit_stack.enter_async_context(stdio_client(server_params))
        session = await exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
        await session.initialize()
        return session

    def connect_all(self, console: Console) -> None:
        """Connect to all configured MCP servers."""
        for config in self.server_configs:
            console.print(f"[dim]Connecting to {config.name}...[/dim]")

            # Create event loop for this server
            loop = asyncio.new_event_loop()
            self.loops[config.name] = loop

            # Connect
            session = loop.run_until_complete(self._connect_server(config))
            self.sessions[config.name] = session

            # Fetch tools
            tools = fetch_mcp_tools(
                mcp_endpoint=None,
                transport_type=MCPTransportType.STDIO,
                client_session=session,
                event_loop=loop,
            )

            self.tools_by_server[config.name] = tools
            self.all_tools.extend(tools)

            console.print(f"[green]  Connected: {len(tools)} tools[/green]")

    def close_all(self, console: Console) -> None:
        """Close all server connections."""
        for name in list(self.sessions.keys()):
            console.print(f"[dim]Closing {name}...[/dim]")
            try:
                loop = self.loops.get(name)
                exit_stack = self.exit_stacks.get(name)
                if loop and exit_stack:
                    loop.run_until_complete(exit_stack.aclose())
                    loop.close()
            except Exception as e:
                console.print(f"[red]Error closing {name}: {e}[/red]")


########################
# STATISTICS TRACKING  #
########################
@dataclass
class DisclosureStats:
    """Track statistics to demonstrate progressive disclosure benefits."""

    total_tools_available: int = 0
    tools_selected: int = 0
    servers_with_selected_tools: int = 0
    search_queries_made: int = 0
    tool_executions: int = 0
    parallel_batches: int = 0
    tools_in_parallel: int = 0

    @property
    def tools_filtered_percentage(self) -> float:
        """Percentage of tools that were NOT loaded."""
        if self.total_tools_available == 0:
            return 0.0
        return ((self.total_tools_available - self.tools_selected) / self.total_tools_available) * 100

    def display(self, console: Console) -> None:
        """Display statistics."""
        table = Table(title="Progressive Disclosure Statistics", box=None)
        table.add_column("Metric", style="cyan")
        table.add_column("Value", style="green")

        table.add_row("Total tools (3 servers)", str(self.total_tools_available))
        table.add_row("Tools selected for query", str(self.tools_selected))
        table.add_row("Context reduction", f"{self.tools_filtered_percentage:.1f}%")
        table.add_row("Search queries made", str(self.search_queries_made))
        table.add_row("Tool executions", str(self.tool_executions))
        if self.parallel_batches > 0:
            table.add_row("Parallel batches", str(self.parallel_batches))
            table.add_row("Tools run in parallel", str(self.tools_in_parallel))

        console.print(table)


########################
# MAIN DEMO FUNCTION   #
########################
def main():
    """Run the progressive disclosure demonstration with multiple MCP servers."""
    load_dotenv()
    console = Console()
    config = ProgressiveDisclosureConfig()

    console.print(
        Panel.fit(
            "[bold cyan]Progressive Disclosure Demo[/bold cyan]\n"
            "[dim]Demonstrating Anthropic's pattern with 3 MCP servers (24 total tools)[/dim]",
            border_style="cyan",
        )
    )

    # Initialize instructor client
    client = instructor.from_openai(openai.OpenAI(api_key=config.openai_api_key))

    # Initialize server manager
    server_manager = MCPServerManager(config.servers)

    try:
        # Connect to all servers
        console.print("\n[bold]Connecting to MCP servers...[/bold]")
        server_manager.connect_all(console)

        all_tools = server_manager.all_tools
        if not all_tools:
            console.print("[red]No tools found across any server.[/red]")
            return

        # Display all available tools by server
        for server_config in config.servers:
            server_tools = server_manager.tools_by_server.get(server_config.name, [])
            table = Table(title=f"{server_config.name} Tools", box=None)
            table.add_column("Tool", style="cyan")
            table.add_column("Description", style="dim", max_width=50)

            for tool in server_tools:
                name = getattr(tool, "mcp_tool_name", tool.__name__)
                desc = (tool.__doc__ or "")[:50]
                table.add_row(name, desc)
            console.print(table)

        console.print(f"\n[bold green]Total: {len(all_tools)} tools across {len(config.servers)} servers[/bold green]")

        # Create lightweight tool registry
        console.print("\n[dim]Building lightweight tool registry (metadata only)...[/dim]")
        registry = ToolRegistry()

        mcp_definitions = []
        for server_config in config.servers:
            for tool in server_manager.tools_by_server.get(server_config.name, []):
                name = getattr(tool, "mcp_tool_name", tool.__name__)
                description = tool.__doc__ or ""
                mcp_definitions.append(
                    MCPToolDefinition(
                        name=name,
                        description=description,
                        input_schema={},
                    )
                )

        registry.register_from_mcp(mcp_definitions)

        # Create Tool Finder Agent
        console.print("[dim]Creating Tool Finder Agent (sub-agent)...[/dim]")
        finder_agent, search_tool, list_tool = create_tool_finder_agent(
            registry=registry,
            client=client,
            model=config.finder_model,
        )
        console.print(f"[green]Tool Finder ready (using {config.finder_model})[/green]")

        # Create Orchestrator Factory
        # We'll pass all tools and let the factory filter
        orchestrator_factory = OrchestratorFactory(
            mcp_endpoint=None,
            transport_type=MCPTransportType.STDIO,
            client=client,
            model=config.orchestrator_model,
            parallel_execution=config.parallel_execution,
            # We don't pass session/loop since tools already have them bound
        )

        # Interactive loop
        console.print("\n[bold green]Ready! Type '/exit' to quit, '/stats' for statistics.[/bold green]")
        console.print("[dim]Example queries:[/dim]")
        console.print("[dim]  - 'Calculate (5 + 3) * 2'              (math tools)[/dim]")
        console.print("[dim]  - 'Convert HELLO WORLD to lowercase'  (text tools)[/dim]")
        console.print("[dim]  - 'Find the average of [1,2,3,4,5]'   (data tools)[/dim]")
        console.print("[dim]  - 'Reverse the text ABC and add 10+5' (multi-server!)[/dim]\n")

        stats = DisclosureStats(total_tools_available=len(all_tools))

        while True:
            query = console.input("[bold yellow]You:[/bold yellow] ").strip()

            if query.lower() in {"/exit", "/quit"}:
                console.print("[bold red]Exiting. Goodbye![/bold red]")
                break

            if query.lower() == "/stats":
                stats.display(console)
                continue

            if not query:
                continue

            try:
                # Phase 1: Tool Discovery (Progressive Disclosure)
                console.print("\n[bold cyan]Phase 1: Tool Discovery[/bold cyan]")
                console.print(f"[dim]Sub-agent searching {len(all_tools)} tools across {len(config.servers)} servers...[/dim]")

                finder_result = run_tool_finder(
                    agent=finder_agent,
                    search_tool=search_tool,
                    list_tool=list_tool,
                    user_query=query,
                )

                stats.search_queries_made += len(finder_result.search_queries_used)
                stats.tools_selected = len(finder_result.selected_tools)

                console.print(
                    f"[green]Selected {len(finder_result.selected_tools)} tools:[/green] {finder_result.selected_tools}"
                )
                console.print(f"[dim]Reasoning: {finder_result.reasoning}[/dim]")

                # Phase 2: Dynamic Orchestrator Creation
                console.print("\n[bold cyan]Phase 2: Creating Focused Orchestrator[/bold cyan]")

                orchestrator, tool_map = orchestrator_factory.create_with_tools(
                    tool_names=finder_result.selected_tools,
                    all_tools=all_tools,
                )

                if finder_result.selected_tools:
                    tools_count = len(finder_result.selected_tools)
                    tokens_saved = (len(all_tools) - tools_count) * 500
                    console.print(
                        f"[green]Orchestrator context: {tools_count} tools "
                        f"(filtered {stats.tools_filtered_percentage:.0f}% = "
                        f"saved ~{tokens_saved} tokens)[/green]"
                    )
                else:
                    console.print("[yellow]No tools needed - conversational response[/yellow]")

                # Phase 3: Query Execution
                console.print("\n[bold cyan]Phase 3: Query Execution[/bold cyan]")

                def on_tool_execution(tool_name: str, params: dict):
                    stats.tool_executions += 1
                    console.print(f"[blue]Executing:[/blue] {tool_name}")
                    console.print(f"[dim]Parameters: {params}[/dim]")

                def on_parallel_batch(count: int):
                    stats.parallel_batches += 1
                    stats.tools_in_parallel += count
                    console.print(f"[magenta]⚡ Parallel batch:[/magenta] {count} tools executing simultaneously")

                if config.parallel_execution:
                    response = execute_orchestrator_loop_parallel(
                        orchestrator=orchestrator,
                        tool_schema_to_class=tool_map,
                        initial_query=query,
                        on_tool_execution=on_tool_execution,
                        on_parallel_batch=on_parallel_batch,
                    )
                else:
                    response = execute_orchestrator_loop(
                        orchestrator=orchestrator,
                        tool_schema_to_class=tool_map,
                        initial_query=query,
                        on_tool_execution=on_tool_execution,
                    )

                console.print(f"\n[bold green]Response:[/bold green] {response}")

                # Show savings summary
                savings_pct = stats.tools_filtered_percentage
                parallel_info = " | ⚡ Parallel mode" if config.parallel_execution else ""
                console.print(
                    Panel(
                        f"[dim]Progressive Disclosure: {len(finder_result.selected_tools)}/{len(all_tools)} tools loaded "
                        f"({savings_pct:.0f}% context reduction){parallel_info}[/dim]",
                        border_style="dim",
                    )
                )

                # Reset histories for next query
                finder_agent.history.history = []
                finder_agent.history.current_turn_id = None
                orchestrator.history.history = []
                orchestrator.history.current_turn_id = None

            except Exception as e:
                console.print(f"[red]Error:[/red] {str(e)}")
                import traceback

                console.print(f"[dim]{traceback.format_exc()}[/dim]")

    finally:
        # Cleanup all servers
        console.print("\n[dim]Cleaning up server connections...[/dim]")
        server_manager.close_all(console)


if __name__ == "__main__":
    main()

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/registry/__init__.py

```python
"""Tool registry module for progressive disclosure."""

from progressive_disclosure.registry.tool_registry import ToolRegistry, ToolMetadata

__all__ = ["ToolRegistry", "ToolMetadata"]

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/registry/tool_registry.py

```python
"""Lightweight tool registry for progressive disclosure.

This module provides a registry that holds tool metadata (name, description, keywords)
without loading the full schema definitions. This enables efficient tool discovery
where the sub-agent can search through available tools without incurring the context
window cost of full schema definitions.
"""

import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, NamedTuple


class MCPToolDefinition(NamedTuple):
    """Definition of an MCP tool (matching atomic-agents structure)."""

    name: str
    description: Optional[str]
    input_schema: Dict[str, Any]


@dataclass
class ToolMetadata:
    """Lightweight tool representation for search.

    This stores only the essential metadata needed for tool discovery,
    avoiding the full JSON schema that would bloat the context window.
    """

    name: str
    description: str
    keywords: List[str] = field(default_factory=list)
    category: Optional[str] = None

    def to_search_string(self) -> str:
        """Create a searchable string representation."""
        parts = [self.name, self.description]
        if self.keywords:
            parts.extend(self.keywords)
        if self.category:
            parts.append(self.category)
        return " ".join(parts).lower()


class ToolRegistry:
    """Registry that holds tool metadata for progressive discovery.

    The registry stores lightweight metadata about available tools,
    enabling efficient search without loading full schema definitions.
    This is a key component of the progressive disclosure pattern.

    Example:
        >>> registry = ToolRegistry()
        >>> registry.register_from_mcp(mcp_definitions)
        >>> results = registry.search("calculate numbers", max_results=3)
        >>> for tool in results:
        ...     print(f"{tool.name}: {tool.description}")
    """

    def __init__(self):
        self._tools: Dict[str, ToolMetadata] = {}

    def register(self, metadata: ToolMetadata) -> None:
        """Register a single tool's metadata."""
        self._tools[metadata.name] = metadata

    def register_from_mcp(self, mcp_definitions: List[MCPToolDefinition]) -> None:
        """Register tools from MCP definitions (metadata only, no schemas).

        Args:
            mcp_definitions: List of MCP tool definitions to register.
                            Only name and description are stored.
        """
        for defn in mcp_definitions:
            keywords = self._extract_keywords(defn.description)
            category = self._infer_category(defn.name, defn.description)

            self._tools[defn.name] = ToolMetadata(
                name=defn.name,
                description=defn.description or "",
                keywords=keywords,
                category=category,
            )

    def _extract_keywords(self, description: Optional[str]) -> List[str]:
        """Extract keywords from a tool description.

        Uses simple heuristics to identify important terms:
        - Words longer than 3 characters
        - Words not in common stop words
        - Capitalized words (likely proper nouns or technical terms)
        """
        if not description:
            return []

        # Common stop words to filter out
        stop_words = {
            "the",
            "a",
            "an",
            "and",
            "or",
            "but",
            "in",
            "on",
            "at",
            "to",
            "for",
            "of",
            "with",
            "by",
            "from",
            "as",
            "is",
            "was",
            "are",
            "were",
            "been",
            "be",
            "have",
            "has",
            "had",
            "do",
            "does",
            "did",
            "will",
            "would",
            "could",
            "should",
            "may",
            "might",
            "must",
            "shall",
            "can",
            "this",
            "that",
            "these",
            "those",
            "it",
            "its",
            "they",
            "them",
            "their",
            "we",
            "us",
            "our",
            "you",
            "your",
            "i",
            "me",
            "my",
            "he",
            "she",
            "his",
            "her",
            "which",
            "who",
            "whom",
            "what",
            "when",
            "where",
            "why",
            "how",
            "all",
            "each",
            "every",
            "both",
            "few",
            "more",
            "most",
            "other",
            "some",
            "such",
            "no",
            "not",
            "only",
            "own",
            "same",
            "so",
            "than",
            "too",
            "very",
            "just",
            "also",
            "now",
        }

        # Extract words
        words = re.findall(r"\b[a-zA-Z]+\b", description.lower())

        # Filter and deduplicate
        keywords = []
        seen = set()
        for word in words:
            if len(word) > 3 and word not in stop_words and word not in seen:
                keywords.append(word)
                seen.add(word)

        return keywords[:10]  # Limit to top 10 keywords

    def _infer_category(self, name: str, description: Optional[str]) -> Optional[str]:
        """Infer a category for the tool based on name and description.

        Categories help with broad filtering before detailed search.
        """
        text = f"{name} {description or ''}".lower()

        categories = {
            "math": ["add", "subtract", "multiply", "divide", "calculate", "math", "number", "arithmetic"],
            "search": ["search", "find", "query", "lookup", "fetch"],
            "file": ["file", "read", "write", "save", "load", "open", "close"],
            "data": ["data", "database", "sql", "json", "xml", "csv"],
            "web": ["http", "api", "request", "url", "web", "download", "upload"],
            "text": ["text", "string", "parse", "format", "convert"],
        }

        for category, keywords in categories.items():
            if any(kw in text for kw in keywords):
                return category

        return None

    def search(self, query: str, max_results: int = 5, category: Optional[str] = None) -> List[ToolMetadata]:
        """Search for tools matching the query.

        Uses a simple relevance scoring based on:
        - Exact name match (highest weight)
        - Name contains query terms
        - Description contains query terms
        - Keyword matches

        Args:
            query: Search query string
            max_results: Maximum number of results to return
            category: Optional category filter

        Returns:
            List of ToolMetadata sorted by relevance
        """
        query_terms = set(query.lower().split())
        scored_results: List[tuple[float, ToolMetadata]] = []

        for metadata in self._tools.values():
            # Apply category filter if specified
            if category and metadata.category != category:
                continue

            score = self._calculate_relevance(metadata, query_terms)
            if score > 0:
                scored_results.append((score, metadata))

        # Sort by score descending
        scored_results.sort(key=lambda x: x[0], reverse=True)

        return [metadata for _, metadata in scored_results[:max_results]]

    def _calculate_relevance(self, metadata: ToolMetadata, query_terms: set[str]) -> float:
        """Calculate relevance score for a tool against query terms."""
        score = 0.0
        name_lower = metadata.name.lower()
        search_string = metadata.to_search_string()

        for term in query_terms:
            # Exact name match - highest weight
            if term == name_lower:
                score += 10.0

            # Name contains term
            elif term in name_lower:
                score += 5.0

            # Term in description/keywords
            if term in search_string:
                score += 2.0

            # Partial match in keywords
            for keyword in metadata.keywords:
                if term in keyword or keyword in term:
                    score += 1.0

        return score

    def get_all_metadata(self) -> List[ToolMetadata]:
        """Get all tool metadata (for context provider listing)."""
        return list(self._tools.values())

    def get_all_tools(self) -> List[ToolMetadata]:
        """Get all tool metadata (alias for get_all_metadata)."""
        return self.get_all_metadata()

    def get_tool(self, name: str) -> Optional[ToolMetadata]:
        """Get metadata for a specific tool by name."""
        return self._tools.get(name)

    def get_tool_names(self) -> List[str]:
        """Get list of all registered tool names."""
        return list(self._tools.keys())

    def __len__(self) -> int:
        """Return the number of registered tools."""
        return len(self._tools)

    def __contains__(self, name: str) -> bool:
        """Check if a tool is registered."""
        return name in self._tools

    def get_summary(self) -> str:
        """Get a summary string of all tools for context injection.

        This provides a lightweight overview suitable for the tool finder agent's
        context, listing all available tools without full schema definitions.
        """
        lines = ["Available tools:"]
        for metadata in self._tools.values():
            category_str = f" [{metadata.category}]" if metadata.category else ""
            lines.append(f"- {metadata.name}{category_str}: {metadata.description}")
        return "\n".join(lines)

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/tools/__init__.py

```python
"""Tools module for progressive disclosure."""

from progressive_disclosure.tools.search_tools import (
    SearchToolsTool,
    SearchToolsInputSchema,
    SearchToolsOutputSchema,
    SearchToolsConfig,
    ListAllToolsTool,
    ListAllToolsInputSchema,
    ListAllToolsOutputSchema,
)

__all__ = [
    "SearchToolsTool",
    "SearchToolsInputSchema",
    "SearchToolsOutputSchema",
    "SearchToolsConfig",
    "ListAllToolsTool",
    "ListAllToolsInputSchema",
    "ListAllToolsOutputSchema",
]

```

### File: atomic-examples/progressive-disclosure/progressive_disclosure/tools/search_tools.py

```python
"""Tool for searching available MCP tools.

This tool enables the Tool Finder Agent to search through the registry
of available tools without loading their full schemas into context.
"""

from typing import Dict, List, Optional
from pydantic import Field

from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig
from progressive_disclosure.registry.tool_registry import ToolRegistry


################
# INPUT SCHEMA #
################
class SearchToolsInputSchema(BaseIOSchema):
    """Search for available tools that match a query.

    Use this tool to find relevant MCP tools for a given task.
    The search looks at tool names, descriptions, and keywords.
    """

    search_query: str = Field(
        ...,
        description="Search query to find relevant tools. Can include keywords describing the desired functionality.",
    )
    max_results: int = Field(
        default=5,
        description="Maximum number of results to return. Use fewer for focused tasks, more for exploratory searches.",
        ge=1,
        le=20,
    )
    category: Optional[str] = Field(
        default=None,
        description="Optional category filter (e.g., 'math', 'search', 'file', 'data', 'web', 'text').",
    )


#################
# OUTPUT SCHEMA #
#################
class SearchToolsOutputSchema(BaseIOSchema):
    """Results from searching available tools."""

    matched_tools: List[str] = Field(
        ...,
        description="List of tool names that matched the search query, ordered by relevance.",
    )
    tool_descriptions: Dict[str, str] = Field(
        ...,
        description="Mapping of tool name to description for each matched tool.",
    )
    total_tools_available: int = Field(
        ...,
        description="Total number of tools available in the registry.",
    )
    search_query_used: str = Field(
        ...,
        description="The search query that was used.",
    )


#################
# CONFIGURATION #
#################
class SearchToolsConfig(BaseToolConfig):
    """Configuration for the SearchToolsTool."""

    registry: Optional[ToolRegistry] = None

    model_config = {"arbitrary_types_allowed": True}


#####################
# MAIN TOOL & LOGIC #
#####################
class SearchToolsTool(BaseTool[SearchToolsInputSchema, SearchToolsOutputSchema]):
    """Tool for searching available MCP tools by query.

    This is a key component of the progressive disclosure pattern,
    allowing the Tool Finder Agent to discover relevant tools without
    having all tool schemas in its context window.

    Example:
        >>> registry = ToolRegistry()
        >>> registry.register_from_mcp(mcp_definitions)
        >>> tool = SearchToolsTool(SearchToolsConfig(registry=registry))
        >>> result = tool.run(SearchToolsInputSchema(search_query="calculate math"))
        >>> print(result.matched_tools)
        ['AddNumbers', 'SubtractNumbers', 'MultiplyNumbers']
    """

    input_schema = SearchToolsInputSchema
    output_schema = SearchToolsOutputSchema

    def __init__(self, config: SearchToolsConfig = SearchToolsConfig()):
        """Initialize the SearchToolsTool.

        Args:
            config: Configuration containing the tool registry.
        """
        super().__init__(config)
        self._registry = config.registry

    @property
    def registry(self) -> ToolRegistry:
        """Get the tool registry."""
        if self._registry is None:
            raise ValueError("Tool registry not configured. Pass a registry via SearchToolsConfig.")
        return self._registry

    def run(self, params: SearchToolsInputSchema) -> SearchToolsOutputSchema:
        """Execute the search and return matching tools.

        Args:
            params: Search parameters including query and optional filters.

        Returns:
            SearchToolsOutputSchema containing matched tools and their descriptions.
        """
        results = self.registry.search(
            query=params.search_query,
            max_results=params.max_results,
            category=params.category,
        )

        return SearchToolsOutputSchema(
            matched_tools=[tool.name for tool in results],
            tool_descriptions={tool.name: tool.description for tool in results},
            total_tools_available=len(self.registry),
            search_query_used=params.search_query,
        )


class ListAllToolsInputSchema(BaseIOSchema):
    """List all available tools in the registry.

    Use this to get an overview of all tools when you need to understand
    the full capabilities available.
    """

    include_categories: bool = Field(
        default=True,
        description="Whether to include category information for each tool.",
    )


class ListAllToolsOutputSchema(BaseIOSchema):
    """List of all available tools."""

    tools: List[Dict[str, str]] = Field(
        ...,
        description="List of tools with their name, description, and optionally category.",
    )
    total_count: int = Field(
        ...,
        description="Total number of tools available.",
    )
    categories_found: List[str] = Field(
        ...,
        description="List of unique categories found among the tools.",
    )


class ListAllToolsTool(BaseTool[ListAllToolsInputSchema, ListAllToolsOutputSchema]):
    """Tool for listing all available tools.

    Useful when the Tool Finder Agent needs to see the complete
    set of available capabilities.
    """

    input_schema = ListAllToolsInputSchema
    output_schema = ListAllToolsOutputSchema

    def __init__(self, config: SearchToolsConfig = SearchToolsConfig()):
        """Initialize the ListAllToolsTool.

        Args:
            config: Configuration containing the tool registry.
        """
        super().__init__(config)
        self._registry = config.registry

    @property
    def registry(self) -> ToolRegistry:
        """Get the tool registry."""
        if self._registry is None:
            raise ValueError("Tool registry not configured. Pass a registry via SearchToolsConfig.")
        return self._registry

    def run(self, params: ListAllToolsInputSchema) -> ListAllToolsOutputSchema:
        """List all available tools.

        Args:
            params: Parameters for listing tools.

        Returns:
            ListAllToolsOutputSchema containing all tools.
        """
        all_tools = self.registry.get_all_metadata()
        categories = set()

        tools_list = []
        for tool in all_tools:
            tool_info = {
                "name": tool.name,
                "description": tool.description,
            }
            if params.include_categories and tool.category:
                tool_info["category"] = tool.category
                categories.add(tool.category)
            tools_list.append(tool_info)

        return ListAllToolsOutputSchema(
            tools=tools_list,
            total_count=len(all_tools),
            categories_found=sorted(list(categories)),
        )


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from progressive_disclosure.registry.tool_registry import ToolMetadata

    # Create a test registry
    registry = ToolRegistry()
    registry.register(
        ToolMetadata(
            name="AddNumbers",
            description="Add two numbers together",
            keywords=["add", "sum", "plus", "arithmetic"],
            category="math",
        )
    )
    registry.register(
        ToolMetadata(
            name="SubtractNumbers",
            description="Subtract one number from another",
            keywords=["subtract", "minus", "difference", "arithmetic"],
            category="math",
        )
    )
    registry.register(
        ToolMetadata(
            name="SearchWeb",
            description="Search the web for information",
            keywords=["search", "web", "query", "find"],
            category="search",
        )
    )

    # Test search
    search_tool = SearchToolsTool(SearchToolsConfig(registry=registry))
    result = search_tool.run(SearchToolsInputSchema(search_query="add numbers math"))
    print("Search results:")
    print(f"  Matched: {result.matched_tools}")
    print(f"  Descriptions: {result.tool_descriptions}")

    # Test list all
    list_tool = ListAllToolsTool(SearchToolsConfig(registry=registry))
    all_result = list_tool.run(ListAllToolsInputSchema())
    print(f"\nAll tools ({all_result.total_count}):")
    for tool in all_result.tools:
        print(f"  - {tool['name']}: {tool['description']}")

```

### File: atomic-examples/progressive-disclosure/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["progressive_disclosure"]

[project]
name = "progressive-disclosure"
version = "0.1.0"
description = "Progressive Disclosure example for Atomic Agents - demonstrating Anthropic's pattern for efficient MCP tool loading"
readme = "README.md"
authors = [
    { name = "KennyVaneetvelde", email = "kenny@inosta.be" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "pd-math-server",
    "pd-text-server",
    "pd-data-server",
    "instructor==1.14.5",
    "pydantic>=2.10.3,<3.0.0",
    "rich>=13.0.0",
    "openai>=2.0.0,<3.0.0",
    "mcp[cli]>=1.9.4",
    "fastmcp>=2.0.0",
    "python-dotenv>=1.0.1,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }
pd-math-server = { path = "servers/math_server" }
pd-text-server = { path = "servers/text_server" }
pd-data-server = { path = "servers/data_server" }

```

### File: atomic-examples/progressive-disclosure/servers/data_server/data_server/__init__.py

```python
"""Data MCP Server - list/data operations for progressive disclosure demo."""

```

### File: atomic-examples/progressive-disclosure/servers/data_server/data_server/server.py

```python
"""Data MCP Server with list/data manipulation tools.

This server provides 8 data/list operations to demonstrate
progressive disclosure - when combined with other servers,
the agent will select only the relevant data tools.
"""

from typing import List
from fastmcp import FastMCP

mcp = FastMCP("data-server")


@mcp.tool
def sort_list(items: List[float], descending: bool = False) -> List[float]:
    """Sort a list of numbers. Use ascending=True for descending order."""
    return sorted(items, reverse=descending)


@mcp.tool
def filter_greater_than(items: List[float], threshold: float) -> List[float]:
    """Filter list to only include items greater than the threshold."""
    return [x for x in items if x > threshold]


@mcp.tool
def filter_less_than(items: List[float], threshold: float) -> List[float]:
    """Filter list to only include items less than the threshold."""
    return [x for x in items if x < threshold]


@mcp.tool
def sum_list(items: List[float]) -> float:
    """Calculate the sum of all numbers in a list. Use for totaling values."""
    return sum(items)


@mcp.tool
def average_list(items: List[float]) -> float:
    """Calculate the average (mean) of all numbers in a list."""
    if not items:
        return 0.0
    return sum(items) / len(items)


@mcp.tool
def min_value(items: List[float]) -> float:
    """Find the minimum value in a list. Use to find smallest number."""
    if not items:
        raise ValueError("Cannot find minimum of empty list")
    return min(items)


@mcp.tool
def max_value(items: List[float]) -> float:
    """Find the maximum value in a list. Use to find largest number."""
    if not items:
        raise ValueError("Cannot find maximum of empty list")
    return max(items)


@mcp.tool
def unique_values(items: List[float]) -> List[float]:
    """Remove duplicate values from a list, preserving order."""
    seen = set()
    result = []
    for item in items:
        if item not in seen:
            seen.add(item)
            result.append(item)
    return result


def main():
    """Run the data server."""
    mcp.run()


if __name__ == "__main__":
    main()

```

### File: atomic-examples/progressive-disclosure/servers/data_server/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["data_server"]

[project]
name = "pd-data-server"
version = "0.1.0"
description = "MCP server with data/list tools for progressive disclosure demo"
authors = [
    { name = "KennyVaneetvelde", email = "kenny@inosta.be" }
]
requires-python = ">=3.12"
dependencies = [
    "fastmcp>=2.0.0",
]

[project.scripts]
pd-data-server = "data_server.server:main"

```

### File: atomic-examples/progressive-disclosure/servers/math_server/math_server/__init__.py

```python
"""Math MCP Server - arithmetic operations for progressive disclosure demo."""

```

### File: atomic-examples/progressive-disclosure/servers/math_server/math_server/server.py

```python
"""Math MCP Server with arithmetic tools.

This server provides 8 arithmetic operations to demonstrate
progressive disclosure - when combined with other servers,
the agent will select only the relevant math tools.
"""

import math
from fastmcp import FastMCP

mcp = FastMCP("math-server")


@mcp.tool
def add_numbers(a: float, b: float) -> float:
    """Add two numbers together (a + b). Use for addition operations."""
    return a + b


@mcp.tool
def subtract_numbers(a: float, b: float) -> float:
    """Subtract b from a (a - b). Use for subtraction operations."""
    return a - b


@mcp.tool
def multiply_numbers(a: float, b: float) -> float:
    """Multiply two numbers (a * b). Use for multiplication operations."""
    return a * b


@mcp.tool
def divide_numbers(a: float, b: float) -> float:
    """Divide a by b (a / b). Use for division operations. Returns error message if b is 0."""
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b


@mcp.tool
def power(base: float, exponent: float) -> float:
    """Raise base to the power of exponent (base ** exponent). Use for exponentiation."""
    return base**exponent


@mcp.tool
def square_root(number: float) -> float:
    """Calculate the square root of a number. Use for sqrt operations."""
    if number < 0:
        raise ValueError("Cannot calculate square root of negative number")
    return math.sqrt(number)


@mcp.tool
def modulo(a: float, b: float) -> float:
    """Calculate the remainder of a divided by b (a % b). Use for modulo operations."""
    return a % b


@mcp.tool
def absolute_value(number: float) -> float:
    """Return the absolute value of a number. Use to remove negative signs."""
    return abs(number)


def main():
    """Run the math server."""
    mcp.run()


if __name__ == "__main__":
    main()

```

### File: atomic-examples/progressive-disclosure/servers/math_server/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["math_server"]

[project]
name = "pd-math-server"
version = "0.1.0"
description = "MCP server with arithmetic tools for progressive disclosure demo"
authors = [
    { name = "KennyVaneetvelde", email = "kenny@inosta.be" }
]
requires-python = ">=3.12"
dependencies = [
    "fastmcp>=2.0.0",
]

[project.scripts]
pd-math-server = "math_server.server:main"

```

### File: atomic-examples/progressive-disclosure/servers/text_server/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["text_server"]

[project]
name = "pd-text-server"
version = "0.1.0"
description = "MCP server with text manipulation tools for progressive disclosure demo"
authors = [
    { name = "KennyVaneetvelde", email = "kenny@inosta.be" }
]
requires-python = ">=3.12"
dependencies = [
    "fastmcp>=2.0.0",
]

[project.scripts]
pd-text-server = "text_server.server:main"

```

### File: atomic-examples/progressive-disclosure/servers/text_server/text_server/__init__.py

```python
"""Text MCP Server - text manipulation operations for progressive disclosure demo."""

```

### File: atomic-examples/progressive-disclosure/servers/text_server/text_server/server.py

```python
"""Text MCP Server with text manipulation tools.

This server provides 8 text operations to demonstrate
progressive disclosure - when combined with other servers,
the agent will select only the relevant text tools.
"""

from typing import List
from fastmcp import FastMCP

mcp = FastMCP("text-server")


@mcp.tool
def uppercase(text: str) -> str:
    """Convert text to all uppercase letters. Use for capitalizing text."""
    return text.upper()


@mcp.tool
def lowercase(text: str) -> str:
    """Convert text to all lowercase letters. Use for lowercasing text."""
    return text.lower()


@mcp.tool
def reverse_text(text: str) -> str:
    """Reverse the order of characters in text. Use to flip text backwards."""
    return text[::-1]


@mcp.tool
def word_count(text: str) -> int:
    """Count the number of words in text. Use to count words."""
    return len(text.split())


@mcp.tool
def char_count(text: str, include_spaces: bool = True) -> int:
    """Count the number of characters in text. Can optionally exclude spaces."""
    if not include_spaces:
        text = text.replace(" ", "")
    return len(text)


@mcp.tool
def concatenate(text1: str, text2: str, separator: str = "") -> str:
    """Join two texts together with an optional separator. Use for combining strings."""
    return text1 + separator + text2


@mcp.tool
def replace_text(text: str, search: str, replacement: str) -> str:
    """Replace all occurrences of search string with replacement. Use for find-and-replace."""
    return text.replace(search, replacement)


@mcp.tool
def split_text(text: str, delimiter: str = " ") -> List[str]:
    """Split text into parts using a delimiter. Use to break text into pieces."""
    return text.split(delimiter)


def main():
    """Run the text server."""
    mcp.run()


if __name__ == "__main__":
    main()

```


--------------------------------------------------------------------------------
Example: quickstart
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/quickstart

## Documentation

# Atomic Agents Quickstart Examples

This directory contains quickstart examples for the Atomic Agents project. These examples demonstrate various features and capabilities of the Atomic Agents framework.

## Getting Started

To run these examples:

1. Clone the main Atomic Agents repository:

   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. Navigate to the quickstart directory:

   ```bash
   cd atomic-agents/atomic-examples/quickstart
   ```

3. Install the dependencies using uv:

   ```bash
   uv sync
   ```

4. Run the examples using uv:

   ```bash
   uv run python quickstart/1_0_basic_chatbot.py
   ```

## Example Files

### 1_0. Basic Chatbot (1_0_basic_chatbot.py)

This example demonstrates a simple chatbot using the Atomic Agents framework. It includes:

- Setting up the OpenAI API client
- Initializing a basic agent with default configurations
- Running a chat loop where the user can interact with the agent

### 1_1. Basic Streaming Chatbot (1_1_basic_chatbot_streaming.py)

This example is similar to 1_0 but it uses `run_stream` method.

### 1_2. Basic Async Streaming Chatbot (1_2_basic_chatbot_async_streaming.py)

This example is similar to 1_0 but it uses an async client and `run_async_stream` method.

### 2. Custom Chatbot (2_basic_custom_chatbot.py)

This example shows how to create a custom chatbot with:

- A custom system prompt
- Customized agent configuration
- A chat loop with rhyming responses

### 3_0. Custom Chatbot with Custom Schema (3_0_basic_custom_chatbot_with_custom_schema.py)

This example demonstrates:

- Creating a custom output schema for the agent
- Implementing suggested follow-up questions in the agent's responses
- Using a custom system prompt and agent configuration

### 3_1. Custom Streaming Chatbot with Custom Schema

This example is similar to 3_0 but uses an async client and `run_async_stream` method.

### 4. Chatbot with Different Providers (4_basic_chatbot_different_providers.py)

This example showcases:

- How to use different AI providers (OpenAI, Groq, Ollama)
- Dynamically selecting a provider at runtime
- Adapting the agent configuration based on the chosen provider

### 5. Custom System Role (5_custom_system_role_for_reasoning_models.py)

This example showcases a usage of `system_role` parameter for a reasoning model.

### 6_0. Asynchronous Processing (6_0_asynchronous_processing.py)

This example showcases a utilization of `run_async` method for a concurrent processing of multiple data.

### 6_1. Asynchronous Streaming Processing

This example adds streaming to 6_0.

## Running the Examples

To run any of the examples, use the following command:

```bash
uv run python quickstart/<example_file_name>.py
```

Replace `<example_file_name>` with the name of the example you want to run (e.g., `1_basic_chatbot.py`).

These examples provide a great starting point for understanding and working with the Atomic Agents framework. Feel free to modify and experiment with them to learn more about the capabilities of Atomic Agents.


## Source Code


### File: atomic-examples/quickstart/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["quickstart"]

[project]
name = "quickstart"
version = "1.0.0"
description = "Quickstart example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "instructor[anthropic,groq,google-genai]==1.14.5",
    "openai>=2.0.0,<3.0.0",
    "python-dotenv>=1.0.1,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```

### File: atomic-examples/quickstart/quickstart/1_0_basic_chatbot.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from atomic_agents.context import ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()

# Initialize history with an initial message from the assistant
initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?")
history.add_message("assistant", initial_message)

# OpenAI client setup using the Instructor library
client = instructor.from_openai(openai.OpenAI(api_key=API_KEY))

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        history=history,
    )
)

# Generate the default system prompt for the agent
default_system_prompt = agent.system_prompt_generator.generate_prompt()
# Display the system prompt in a styled panel
console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan")

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="bold green"))

# Start an infinite loop to handle user inputs and agent responses
while True:
    # Prompt the user for input with a styled prompt
    user_input = console.input("[bold blue]You:[/bold blue] ")
    # Check if the user wants to exit the chat
    if user_input.lower() in ["/exit", "/quit"]:
        console.print("Exiting chat...")
        break
    # Check if the user wants to see token count
    if user_input.lower() == "/tokens":
        token_info = agent.get_context_token_count()
        console.print("[bold magenta]Token Usage:[/bold magenta]")
        console.print(f"  Total: {token_info.total} tokens")
        console.print(f"  System prompt: {token_info.system_prompt} tokens")
        console.print(f"  History: {token_info.history} tokens")
        if token_info.utilization:
            console.print(f"  Context utilization: {token_info.utilization:.1%}")
        continue

    # Process the user's input through the agent and get the response
    input_schema = BasicChatInputSchema(chat_message=user_input)
    response = agent.run(input_schema)

    agent_message = Text(response.chat_message, style="bold green")
    console.print(Text("Agent:", style="bold green"), end=" ")
    console.print(agent_message)

```

### File: atomic-examples/quickstart/quickstart/1_1_basic_chatbot_streaming.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from atomic_agents.context import ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()

# Initialize history with an initial message from the assistant
initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?")
history.add_message("assistant", initial_message)

# OpenAI client setup using the Instructor library for synchronous operations
client = instructor.from_openai(openai.OpenAI(api_key=API_KEY))

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        history=history,
    )
)

# Generate the default system prompt for the agent
default_system_prompt = agent.system_prompt_generator.generate_prompt()
# Display the system prompt in a styled panel
console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan")

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="green"))


def main():
    """
    Main function to handle the chat loop using synchronous streaming.
    This demonstrates how to use AtomicAgent.run_stream() instead of the async version.
    """
    # Start an infinite loop to handle user inputs and agent responses
    while True:
        # Prompt the user for input with a styled prompt
        user_input = console.input("\n[bold blue]You:[/bold blue] ")
        # Check if the user wants to exit the chat
        if user_input.lower() in ["/exit", "/quit"]:
            console.print("Exiting chat...")
            break

        # Process the user's input through the agent
        input_schema = BasicChatInputSchema(chat_message=user_input)
        console.print()  # Add newline before response
        console.print(Text("Agent: ", style="bold green"), end="")

        # Current display string to avoid repeating output
        current_display = ""

        # Use run_stream for synchronous streaming responses
        for partial_response in agent.run_stream(input_schema):
            if hasattr(partial_response, "chat_message") and partial_response.chat_message:
                # Only output the incremental part of the message
                new_content = partial_response.chat_message
                if new_content != current_display:
                    # Only print the new part since the last update
                    if new_content.startswith(current_display):
                        incremental_text = new_content[len(current_display) :]
                        console.print(Text(incremental_text, style="green"), end="")
                        current_display = new_content
                    else:
                        # If there's a mismatch, print the full message
                        # (this should rarely happen with most LLMs)
                        console.print(Text(new_content, style="green"), end="")
                        current_display = new_content

                    # Flush to ensure output is displayed immediately
                    console.file.flush()

        console.print()  # Add a newline after the response is complete


if __name__ == "__main__":
    main()

```

### File: atomic-examples/quickstart/quickstart/1_2_basic_chatbot_async_streaming.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.live import Live
from atomic_agents.context import ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()

# Initialize history with an initial message from the assistant
initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?")
history.add_message("assistant", initial_message)

# OpenAI client setup using the Instructor library for async operations
client = instructor.from_openai(openai.AsyncOpenAI(api_key=API_KEY))

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        history=history,
    )
)

# Generate the default system prompt for the agent
default_system_prompt = agent.system_prompt_generator.generate_prompt()
# Display the system prompt in a styled panel
console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan")

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="green"))


async def main():
    # Start an infinite loop to handle user inputs and agent responses
    while True:
        # Prompt the user for input with a styled prompt
        user_input = console.input("\n[bold blue]You:[/bold blue] ")
        # Check if the user wants to exit the chat
        if user_input.lower() in ["/exit", "/quit"]:
            console.print("Exiting chat...")
            break

        # Process the user's input through the agent and get the streaming response
        input_schema = BasicChatInputSchema(chat_message=user_input)
        console.print()  # Add newline before response

        # Use Live display to show streaming response
        with Live("", refresh_per_second=10, auto_refresh=True) as live:
            current_response = ""
            # Use run_async_stream instead of run_async for streaming functionality
            async for partial_response in agent.run_async_stream(input_schema):
                if hasattr(partial_response, "chat_message") and partial_response.chat_message:
                    # Only update if we have new content
                    if partial_response.chat_message != current_response:
                        current_response = partial_response.chat_message
                        # Combine the label and response in the live display
                        display_text = Text.assemble(("Agent: ", "bold green"), (current_response, "green"))
                        live.update(display_text)


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

```

### File: atomic-examples/quickstart/quickstart/2_basic_custom_chatbot.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from atomic_agents.context import SystemPromptGenerator, ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()

# Initialize history with an initial message from the assistant
initial_message = BasicChatOutputSchema(
    chat_message="How do you do? What can I do for you? Tell me, pray, what is your need today?"
)
history.add_message("assistant", initial_message)

# OpenAI client setup using the Instructor library
# Note, you can also set up a client using any other LLM provider, such as Anthropic, Cohere, etc.
# See the Instructor library for more information: https://github.com/instructor-ai/instructor
client = instructor.from_openai(openai.OpenAI(api_key=API_KEY))

# Instead of the default system prompt, we can set a custom system prompt
system_prompt_generator = SystemPromptGenerator(
    background=[
        "This assistant is a general-purpose AI designed to be helpful and friendly.",
    ],
    steps=["Understand the user's input and provide a relevant response.", "Respond to the user."],
    output_instructions=[
        "Provide helpful and relevant information to assist the user.",
        "Be friendly and respectful in all interactions.",
        "Always answer in rhyming verse.",
    ],
)
console.print(Panel(system_prompt_generator.generate_prompt(), width=console.width, style="bold cyan"), style="bold cyan")

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=system_prompt_generator,
        history=history,
    )
)

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="bold green"))

# Start an infinite loop to handle user inputs and agent responses
while True:
    # Prompt the user for input with a styled prompt
    user_input = console.input("[bold blue]You:[/bold blue] ")
    # Check if the user wants to exit the chat
    if user_input.lower() in ["/exit", "/quit"]:
        console.print("Exiting chat...")
        break

    # Process the user's input through the agent and get the response and display it
    response = agent.run(agent.input_schema(chat_message=user_input))
    agent_message = Text(response.chat_message, style="bold green")
    console.print(Text("Agent:", style="bold green"), end=" ")
    console.print(agent_message)

```

### File: atomic-examples/quickstart/quickstart/3_0_basic_custom_chatbot_with_custom_schema.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from typing import List
from pydantic import Field
from atomic_agents.context import SystemPromptGenerator, ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BaseIOSchema

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()


# Custom output schema
class CustomOutputSchema(BaseIOSchema):
    """This schema represents the response generated by the chat agent, including suggested follow-up questions."""

    chat_message: str = Field(
        ...,
        description="The chat message exchanged between the user and the chat agent.",
    )
    suggested_user_questions: List[str] = Field(
        ...,
        description="A list of suggested follow-up questions the user could ask the agent.",
    )


# Initialize history with an initial message from the assistant
initial_message = CustomOutputSchema(
    chat_message="Hello! How can I assist you today?",
    suggested_user_questions=["What can you do?", "Tell me a joke", "Tell me about how you were made"],
)
history.add_message("assistant", initial_message)

# OpenAI client setup using the Instructor library
client = instructor.from_openai(openai.OpenAI(api_key=API_KEY))

# Custom system prompt
system_prompt_generator = SystemPromptGenerator(
    background=[
        "This assistant is a knowledgeable AI designed to be helpful, friendly, and informative.",
        "It has a wide range of knowledge on various topics and can engage in diverse conversations.",
    ],
    steps=[
        "Analyze the user's input to understand the context and intent.",
        "Formulate a relevant and informative response based on the assistant's knowledge.",
        "Generate 3 suggested follow-up questions for the user to explore the topic further.",
        "When you get a simple number from the user, choose the corresponding question from the last list of "
        "suggested questions and answer it. Note that the first question is 1, the second is 2, and so on.",
    ],
    output_instructions=[
        "Provide clear, concise, and accurate information in response to user queries.",
        "Maintain a friendly and professional tone throughout the conversation.",
        "Conclude each response with 3 relevant suggested questions for the user.",
    ],
)
console.print(Panel(system_prompt_generator.generate_prompt(), width=console.width, style="bold cyan"), style="bold cyan")

# Agent setup with specified configuration and custom output schema
agent = AtomicAgent[BasicChatInputSchema, CustomOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=system_prompt_generator,
        history=history,
    )
)

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="bold green"))

# Display initial suggested questions
console.print("\n[bold cyan]Suggested questions you could ask:[/bold cyan]")
for i, question in enumerate(initial_message.suggested_user_questions, 1):
    console.print(f"[cyan]{i}. {question}[/cyan]")
console.print()  # Add an empty line for better readability

# Start an infinite loop to handle user inputs and agent responses
while True:
    # Prompt the user for input with a styled prompt
    user_input = console.input("[bold blue]You:[/bold blue] ")
    # Check if the user wants to exit the chat
    if user_input.lower() in ["/exit", "/quit"]:
        console.print("Exiting chat...")
        break

    # Process the user's input through the agent and get the response
    response = agent.run(BasicChatInputSchema(chat_message=user_input))

    # Display the agent's response
    agent_message = Text(response.chat_message, style="bold green")
    console.print(Text("Agent:", style="bold green"), end=" ")
    console.print(agent_message)

    # Display follow-up questions
    console.print("\n[bold cyan]Suggested questions you could ask:[/bold cyan]")
    for i, question in enumerate(response.suggested_user_questions, 1):
        console.print(f"[cyan]{i}. {question}[/cyan]")
    console.print()  # Add an empty line for better readability

```

### File: atomic-examples/quickstart/quickstart/3_1_basic_custom_chatbot_with_custom_schema_streaming.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.live import Live
from typing import List
from pydantic import Field
from atomic_agents.context import SystemPromptGenerator, ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BaseIOSchema

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()


# Custom output schema
class CustomOutputSchema(BaseIOSchema):
    """This schema represents the response generated by the chat agent, including suggested follow-up questions."""

    chat_message: str = Field(
        ...,
        description="The chat message exchanged between the user and the chat agent.",
    )
    suggested_user_questions: List[str] = Field(
        ...,
        description="A list of suggested follow-up questions the user could ask the agent.",
    )


# Initialize history with an initial message from the assistant
initial_message = CustomOutputSchema(
    chat_message="Hello! How can I assist you today?",
    suggested_user_questions=["What can you do?", "Tell me a joke", "Tell me about how you were made"],
)
history.add_message("assistant", initial_message)

# OpenAI client setup using the Instructor library for async operations
client = instructor.from_openai(openai.AsyncOpenAI(api_key=API_KEY))

# Custom system prompt
system_prompt_generator = SystemPromptGenerator(
    background=[
        "This assistant is a knowledgeable AI designed to be helpful, friendly, and informative.",
        "It has a wide range of knowledge on various topics and can engage in diverse conversations.",
    ],
    steps=[
        "Analyze the user's input to understand the context and intent.",
        "Formulate a relevant and informative response based on the assistant's knowledge.",
        "Generate 3 suggested follow-up questions for the user to explore the topic further.",
        "When you get a simple number from the user,"
        "choose the corresponding question from the last list of suggested questions and answer it."
        "Note that the first question is 1, the second is 2, and so on.",
    ],
    output_instructions=[
        "Provide clear, concise, and accurate information in response to user queries.",
        "Maintain a friendly and professional tone throughout the conversation.",
        "Conclude each response with 3 relevant suggested questions for the user.",
    ],
)
console.print(Panel(system_prompt_generator.generate_prompt(), width=console.width, style="bold cyan"), style="bold cyan")

# Agent setup with specified configuration and custom output schema
agent = AtomicAgent[BasicChatInputSchema, CustomOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=system_prompt_generator,
        history=history,
    )
)

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="green"))

# Display initial suggested questions
console.print("\n[bold cyan]Suggested questions you could ask:[/bold cyan]")
for i, question in enumerate(initial_message.suggested_user_questions, 1):
    console.print(f"[cyan]{i}. {question}[/cyan]")
console.print()  # Add an empty line for better readability


async def main():
    # Start an infinite loop to handle user inputs and agent responses
    while True:
        # Prompt the user for input with a styled prompt
        user_input = console.input("[bold blue]You:[/bold blue] ")
        # Check if the user wants to exit the chat
        if user_input.lower() in ["/exit", "/quit"]:
            console.print("Exiting chat...")
            break

        # Process the user's input through the agent and get the streaming response
        input_schema = BasicChatInputSchema(chat_message=user_input)
        console.print()  # Add newline before response

        # Use Live display to show streaming response
        with Live("", refresh_per_second=10, auto_refresh=True) as live:
            current_response = ""
            current_questions: List[str] = []

            async for partial_response in agent.run_async_stream(input_schema):
                if hasattr(partial_response, "chat_message") and partial_response.chat_message:
                    # Update the message part
                    if partial_response.chat_message != current_response:
                        current_response = partial_response.chat_message

                    # Update questions if available
                    if hasattr(partial_response, "suggested_user_questions"):
                        current_questions = partial_response.suggested_user_questions

                    # Combine all elements for display
                    display_text = Text.assemble(("Agent: ", "bold green"), (current_response, "green"))

                    # Add questions if we have them
                    if current_questions:
                        display_text.append("\n\n")
                        display_text.append("Suggested questions you could ask:\n", style="bold cyan")
                        for i, question in enumerate(current_questions, 1):
                            display_text.append(f"{i}. {question}\n", style="cyan")

                    live.update(display_text)

        console.print()  # Add an empty line for better readability


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

```

### File: atomic-examples/quickstart/quickstart/4_basic_chatbot_different_providers.py

```python
import os
import instructor
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from atomic_agents.context import ChatHistory
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema
from dotenv import load_dotenv

load_dotenv()

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()


# Function to set up the client based on the chosen provider
def setup_client(provider):
    console.log(f"provider: {provider}")
    if provider == "1" or provider == "openai":
        from openai import OpenAI

        api_key = os.getenv("OPENAI_API_KEY")
        client = instructor.from_openai(OpenAI(api_key=api_key))
        model = "gpt-5-mini"
        model_api_parameters = {"reasoning_effort": "low", "max_tokens": 2048}
        assistant_role = "assistant"
    elif provider == "2" or provider == "anthropic":
        from anthropic import Anthropic

        api_key = os.getenv("ANTHROPIC_API_KEY")
        client = instructor.from_anthropic(Anthropic(api_key=api_key))
        model = "claude-3-5-haiku-20241022"
        model_api_parameters = {"max_tokens": 2048}
        assistant_role = "assistant"
    elif provider == "3" or provider == "groq":
        from groq import Groq

        api_key = os.getenv("GROQ_API_KEY")
        client = instructor.from_groq(Groq(api_key=api_key), mode=instructor.Mode.JSON)
        model = "mixtral-8x7b-32768"
        model_api_parameters = {"max_tokens": 2048}
        assistant_role = "assistant"
    elif provider == "4" or provider == "ollama":
        from openai import OpenAI as OllamaClient

        client = instructor.from_openai(
            OllamaClient(base_url="http://localhost:11434/v1", api_key="ollama"), mode=instructor.Mode.JSON
        )
        model = "llama3"
        model_api_parameters = {"max_tokens": 2048}
        assistant_role = "assistant"
    elif provider == "5" or provider == "gemini":
        import google.genai

        api_key = os.getenv("GEMINI_API_KEY")
        client = instructor.from_genai(
            google.genai.Client(api_key=api_key),
            mode=instructor.Mode.GENAI_TOOLS,
        )
        model = "gemini-2.5-flash"
        model_api_parameters = {}
        assistant_role = "model"
    elif provider == "6" or provider == "openrouter":
        from openai import OpenAI as OpenRouterClient

        api_key = os.getenv("OPENROUTER_API_KEY")
        client = instructor.from_openai(OpenRouterClient(base_url="https://openrouter.ai/api/v1", api_key=api_key))
        model = "mistral/ministral-8b"
        model_api_parameters = {"max_tokens": 2048}
        assistant_role = "assistant"
    elif provider == "7" or provider == "minimax":
        from openai import OpenAI as MiniMaxClient

        api_key = os.getenv("MINIMAX_API_KEY")
        client = instructor.from_openai(
            MiniMaxClient(base_url="https://api.minimax.io/v1", api_key=api_key),
            mode=instructor.Mode.JSON,
        )
        model = "MiniMax-M3"
        model_api_parameters = {"max_tokens": 2048}
        assistant_role = "assistant"
    else:
        raise ValueError(f"Unsupported provider: {provider}")

    return client, model, model_api_parameters, assistant_role


# Prompt the user to choose a provider from one in the list below.
providers_list = ["openai", "anthropic", "groq", "ollama", "gemini", "openrouter", "minimax"]
y = "bold yellow"
b = "bold blue"
g = "bold green"
provider_inner_str = (
    f"{' / '.join(f'[[{g}]{i + 1}[/{g}]]. [{b}]{provider}[/{b}]' for i, provider in enumerate(providers_list))}"
)
providers_str = f"[{y}]Choose a provider ({provider_inner_str}): [/{y}]"

provider = console.input(providers_str).lower()

# Set up the client and model based on the chosen provider
client, model, model_api_parameters, assistant_role = setup_client(provider)

# Initialize history with an initial message from the assistant
initial_message = BasicChatOutputSchema(chat_message="Hello! How can I assist you today?")
history.add_message(assistant_role, initial_message)

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](
    config=AgentConfig(
        client=client,
        model=model,
        history=history,
        assistant_role=assistant_role,
        model_api_parameters=model_api_parameters,
    )
)

# Generate the default system prompt for the agent
default_system_prompt = agent.system_prompt_generator.generate_prompt()
# Display the system prompt in a styled panel
console.print(Panel(default_system_prompt, width=console.width, style="bold cyan"), style="bold cyan")

# Display the initial message from the assistant
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(Text(initial_message.chat_message, style="bold green"))

# Start an infinite loop to handle user inputs and agent responses
while True:
    # Prompt the user for input with a styled prompt
    user_input = console.input("[bold blue]You:[/bold blue] ")
    # Check if the user wants to exit the chat
    if user_input.lower() in ["/exit", "/quit"]:
        console.print("Exiting chat...")
        break
    # Check if the user wants to see token count (works with any provider!)
    if user_input.lower() == "/tokens":
        token_info = agent.get_context_token_count()
        console.print(f"[bold magenta]Token Usage ({model}):[/bold magenta]")
        console.print(f"  Total: {token_info.total} tokens")
        console.print(f"  System prompt: {token_info.system_prompt} tokens")
        console.print(f"  History: {token_info.history} tokens")
        if token_info.max_tokens:
            console.print(f"  Max context: {token_info.max_tokens} tokens")
        if token_info.utilization:
            console.print(f"  Context utilization: {token_info.utilization:.1%}")
        continue

    # Process the user's input through the agent and get the response
    input_schema = BasicChatInputSchema(chat_message=user_input)
    response = agent.run(input_schema)

    agent_message = Text(response.chat_message, style="bold green")
    console.print(Text("Agent:", style="bold green"), end=" ")
    console.print(agent_message)

```

### File: atomic-examples/quickstart/quickstart/5_custom_system_role_for_reasoning_models.py

```python
import os
import instructor
import openai
from rich.console import Console
from rich.text import Text
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BasicChatOutputSchema
from atomic_agents.context import SystemPromptGenerator

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# OpenAI client setup using the Instructor library
client = instructor.from_openai(openai.OpenAI(api_key=API_KEY))

# System prompt generator setup
system_prompt_generator = SystemPromptGenerator(
    background=["You are a math genius."],
    steps=["Think logically step by step and solve a math problem."],
    output_instructions=["Answer in plain English plus formulas."],
)
# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema](
    config=AgentConfig(
        client=client,
        model="o3-mini",
        system_prompt_generator=system_prompt_generator,
        # It is a convention to use "developer" as the system role for reasoning models from OpenAI such as o1, o3-mini.
        # Also these models are often used without a system prompt, which you can do by setting system_role=None
        system_role="developer",
    )
)

# Prompt the user for input with a styled prompt
user_input = "Decompose this number to prime factors: 1234567890"
console.print(Text("User:", style="bold green"), end=" ")
console.print(user_input)

# Process the user's input through the agent and get the response
input_schema = BasicChatInputSchema(chat_message=user_input)
response = agent.run(input_schema)

agent_message = Text(response.chat_message, style="bold green")
console.print(Text("Agent:", style="bold green"), end=" ")
console.print(agent_message)

```

### File: atomic-examples/quickstart/quickstart/6_0_asynchronous_processing.py

```python
import os
import asyncio
import instructor
import openai
from rich.console import Console
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig, BasicChatInputSchema
from atomic_agents.context import SystemPromptGenerator

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# OpenAI client setup using the Instructor library
client = instructor.from_openai(openai.AsyncOpenAI(api_key=API_KEY))


# Define a schema for the output data
class PersonSchema(BaseIOSchema):
    """Schema for person information."""

    name: str
    age: int
    pronouns: list[str]
    profession: str


# System prompt generator setup
system_prompt_generator = SystemPromptGenerator(
    background=["You parse a sentence and extract elements."],
    steps=[],
    output_instructions=[],
)

dataset = [
    "My name is Mike, I am 30 years old, my pronouns are he/him, and I am a software engineer.",
    "My name is Sarah, I am 25 years old, my pronouns are she/her, and I am a data scientist.",
    "My name is John, I am 40 years old, my pronouns are he/him, and I am a product manager.",
    "My name is Emily, I am 35 years old, my pronouns are she/her, and I am a UX designer.",
    "My name is David, I am 28 years old, my pronouns are he/him, and I am a web developer.",
    "My name is Anna, I am 32 years old, my pronouns are she/her, and I am a graphic designer.",
]

sem = asyncio.Semaphore(2)

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, PersonSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=system_prompt_generator,
    )
)


async def exec_agent(message: str):
    """Execute the agent with the provided message."""
    user_input = BasicChatInputSchema(chat_message=message)
    agent.reset_history()
    response = await agent.run_async(user_input)
    return response


async def process(dataset: list[str]):
    """Process the dataset asynchronously."""
    async with sem:
        # Run the agent asynchronously for each message in the dataset
        # and collect the responses
        responses = await asyncio.gather(*(exec_agent(message) for message in dataset))

    return responses


responses = asyncio.run(process(dataset))
console.print(responses)

```

### File: atomic-examples/quickstart/quickstart/6_1_asynchronous_processing_streaming.py

```python
import os
import asyncio
import instructor
import openai
from rich.console import Console
from rich.live import Live
from rich.table import Table
from rich.text import Text
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig, BasicChatInputSchema
from atomic_agents.context import SystemPromptGenerator

# API Key setup
API_KEY = ""
if not API_KEY:
    API_KEY = os.getenv("OPENAI_API_KEY")

if not API_KEY:
    raise ValueError(
        "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY."
    )

# Initialize a Rich Console for pretty console outputs
console = Console()

# OpenAI client setup using the Instructor library
client = instructor.from_openai(openai.AsyncOpenAI(api_key=API_KEY))


# Define a schema for the output data
class PersonSchema(BaseIOSchema):
    """Schema for person information."""

    name: str
    age: int
    pronouns: list[str]
    profession: str


# System prompt generator setup
system_prompt_generator = SystemPromptGenerator(
    background=["You parse a sentence and extract elements."],
    steps=[],
    output_instructions=[],
)

dataset = [
    "My name is Mike, I am 30 years old, my pronouns are he/him, and I am a software engineer.",
    "My name is Sarah, I am 25 years old, my pronouns are she/her, and I am a data scientist.",
    "My name is John, I am 40 years old, my pronouns are he/him, and I am a product manager.",
    "My name is Emily, I am 35 years old, my pronouns are she/her, and I am a UX designer.",
    "My name is David, I am 28 years old, my pronouns are he/him, and I am a web developer.",
    "My name is Anna, I am 32 years old, my pronouns are she/her, and I am a graphic designer.",
]

# Max concurrent requests - adjust this to see performance differences
MAX_CONCURRENT = 3
sem = asyncio.Semaphore(MAX_CONCURRENT)

# Agent setup with specified configuration
agent = AtomicAgent[BasicChatInputSchema, PersonSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=system_prompt_generator,
    )
)


async def exec_agent(message: str, idx: int, progress_dict: dict):
    """Execute the agent with the provided message and update progress in real-time."""
    # Acquire the semaphore to limit concurrent executions
    async with sem:
        user_input = BasicChatInputSchema(chat_message=message)
        agent.reset_history()

        # Track streaming progress
        partial_data = {}
        progress_dict[idx] = {"status": "Processing", "data": partial_data, "message": message}

        partial_response = None
        # Actually demonstrate streaming by processing each partial response
        async for partial_response in agent.run_async_stream(user_input):
            if partial_response:
                # Extract any available fields from the partial response
                response_dict = partial_response.model_dump()
                for field in ["name", "age", "pronouns", "profession"]:
                    if field in response_dict and response_dict[field]:
                        partial_data[field] = response_dict[field]

                # Update progress dictionary to display changes in real-time
                progress_dict[idx]["data"] = partial_data.copy()
                # Small sleep to simulate processing and make streaming more visible
                await asyncio.sleep(0.05)

        assert partial_response
        # Final response with complete data
        response = PersonSchema(**partial_response.model_dump())
        progress_dict[idx]["status"] = "Complete"
        progress_dict[idx]["data"] = response.model_dump()

        return response


def generate_status_table(progress_dict: dict) -> Table:
    """Generate a rich table showing the current processing status."""
    table = Table(title="Asynchronous Stream Processing Demo")

    table.add_column("ID", justify="center")
    table.add_column("Status", justify="center")
    table.add_column("Input", style="cyan")
    table.add_column("Current Data", style="green")

    for idx, info in progress_dict.items():
        # Format the partial data nicely
        data_str = ""
        if info["data"]:
            for k, v in info["data"].items():
                data_str += f"{k}: {v}\n"

        status_style = "yellow" if info["status"] == "Processing" else "green"

        # Add row with current processing information
        table.add_row(
            f"{idx + 1}",
            f"[{status_style}]{info['status']}[/{status_style}]",
            Text(info["message"][:30] + "..." if len(info["message"]) > 30 else info["message"]),
            data_str or "Waiting...",
        )

    return table


async def process_all(dataset: list[str]):
    """Process all items in dataset with visual progress tracking."""
    progress_dict = {}  # Track processing status for visualization

    # Create tasks for each message processing
    tasks = []
    for idx, message in enumerate(dataset):
        # Initialize entry in progress dictionary
        progress_dict[idx] = {"status": "Waiting", "data": {}, "message": message}
        # Create task without awaiting it
        task = asyncio.create_task(exec_agent(message, idx, progress_dict))
        tasks.append(task)

    # Display live updating status while tasks run
    with Live(generate_status_table(progress_dict), refresh_per_second=10) as live:
        while not all(task.done() for task in tasks):
            # Update the live display with current progress
            live.update(generate_status_table(progress_dict))
            await asyncio.sleep(0.1)

        # Final update after all tasks complete
        live.update(generate_status_table(progress_dict))

    # Gather all results when complete
    responses = await asyncio.gather(*tasks)
    return responses


if __name__ == "__main__":
    console.print("[bold blue]Starting Asynchronous Stream Processing Demo[/bold blue]")
    console.print(f"Processing {len(dataset)} items with max {MAX_CONCURRENT} concurrent requests\n")

    responses = asyncio.run(process_all(dataset))

    # Display final results in a structured table
    results_table = Table(title="Processing Results")
    results_table.add_column("Name", style="cyan")
    results_table.add_column("Age", justify="center")
    results_table.add_column("Pronouns")
    results_table.add_column("Profession")

    for resp in responses:
        results_table.add_row(resp.name, str(resp.age), "/".join(resp.pronouns), resp.profession)

    console.print(results_table)

```


--------------------------------------------------------------------------------
Example: rag-chatbot
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/rag-chatbot

## Documentation

# RAG Chatbot

This directory contains the RAG (Retrieval-Augmented Generation) Chatbot example for the Atomic Agents project. This example demonstrates how to build an intelligent chatbot that uses document retrieval to provide context-aware responses using the Atomic Agents framework.

## Features

1. Document Chunking: Automatically splits documents into manageable chunks with configurable overlap
2. Vector Storage: Supports both [ChromaDB](https://www.trychroma.com/) and [Qdrant](https://qdrant.tech/) for efficient storage and retrieval of document chunks
3. Semantic Search: Generates and executes semantic search queries to find relevant context
4. Context-Aware Responses: Provides detailed answers based on retrieved document chunks
5. Interactive UI: Rich console interface with progress indicators and formatted output

## Getting Started

To get started with the RAG Chatbot:

1. **Clone the main Atomic Agents repository:**

   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. **Navigate to the RAG Chatbot directory:**

   ```bash
   cd atomic-agents/atomic-examples/rag-chatbot
   ```

3. **Install the dependencies using uv:**

   ```bash
   uv sync
   ```

4. **Set up environment variables:**
   Create a `.env` file in the `rag-chatbot` directory with the following content:

   ```env
   OPENAI_API_KEY=your_openai_api_key
   VECTOR_DB_TYPE=chroma  # or 'qdrant'
   ```

   Replace `your_openai_api_key` with your actual OpenAI API key.

5. **Run the RAG Chatbot:**

   ```bash
   uv run python rag_chatbot/main.py
   ```

## Vector Database Configuration

The RAG Chatbot supports two vector databases:

### ChromaDB (Default)

- **Local storage**: Data is stored locally in the `chroma_db/` directory
- **Configuration**: Set `VECTOR_DB_TYPE=chroma` in your `.env` file

### Qdrant

- **Local storage**: Data is stored locally in the `qdrant_db/` directory
- **Configuration**: Set `VECTOR_DB_TYPE=qdrant` in your `.env` file

## Usage

### Using ChromaDB (Default)

```bash
export VECTOR_DB_TYPE=chroma
uv run python rag_chatbot/main.py
```

### Using Qdrant (Local)

```bash
export VECTOR_DB_TYPE=qdrant
uv run python rag_chatbot/main.py
```

## Components

### 1. Query Agent (`agents/query_agent.py`)

Generates semantic search queries based on user questions to find relevant document chunks.

### 2. QA Agent (`agents/qa_agent.py`)

Analyzes retrieved chunks and generates comprehensive answers to user questions.

### 3. Vector Database Services (`services/`)

- **Base Service** (`services/base.py`): Abstract interface for vector database operations
- **ChromaDB Service** (`services/chroma_db.py`): ChromaDB implementation
- **Qdrant Service** (`services/qdrant_db.py`): Qdrant implementation
- **Factory** (`services/factory.py`): Creates the appropriate service based on configuration

### 4. Context Provider (`context_providers.py`)

Provides retrieved document chunks as context to the agents.

### 5. Main Script (`main.py`)

Orchestrates the entire process, from document processing to user interaction.

## How It Works

1. The system initializes by:
   - Downloading a sample document (State of the Union address)
   - Splitting it into chunks with configurable overlap
   - Storing chunks in the selected vector database with vector embeddings

2. For each user question:
   - The Query Agent generates an optimized semantic search query
   - Relevant chunks are retrieved from the vector database
   - The QA Agent analyzes the chunks and generates a detailed answer
   - The system displays the thought process and final answer

## Customization

You can customize the RAG Chatbot by:

- Modifying chunk size and overlap in `config.py`
- Adjusting the number of chunks to retrieve for each query
- Using different documents as the knowledge base
- Customizing the system prompts for both agents
- Switching between ChromaDB and Qdrant by changing the `VECTOR_DB_TYPE` environment variable

## Example Usage

The chatbot can answer questions about the loaded document, such as:

- "What were the main points about the economy?"
- "What did the president say about healthcare?"
- "How did he address foreign policy?"

## Contributing

Contributions are welcome! Please fork the repository and submit a pull request with your enhancements or bug fixes.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.


## Source Code


### File: atomic-examples/rag-chatbot/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["rag_chatbot"]

[project]
name = "rag-chatbot"
version = "0.1.0"
description = "A RAG chatbot example using Atomic Agents and ChromaDB/Qdrant"
readme = "README.md"
authors = [
    { name = "Your Name", email = "your.email@example.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "chromadb>=1.0.20,<2.0.0",
    "qdrant-client>=1.15.1,<2.0.0",
    "numpy>=2.3.2,<3.0.0",
    "python-dotenv>=1.0.1,<2.0.0",
    "openai>=2.0.0,<3.0.0",
    "pulsar-client>=3.8.0,<4.0.0",
    "rich>=13.7.0,<14.0.0",
    "wget>=3.2,<4.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```

### File: atomic-examples/rag-chatbot/rag_chatbot/agents/qa_agent.py

```python
import instructor
import openai
from pydantic import Field
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
from atomic_agents.context import SystemPromptGenerator

from rag_chatbot.config import ChatConfig


class RAGQuestionAnsweringAgentInputSchema(BaseIOSchema):
    """Input schema for the RAG QA agent."""

    question: str = Field(..., description="The user's question to answer")


class RAGQuestionAnsweringAgentOutputSchema(BaseIOSchema):
    """Output schema for the RAG QA agent."""

    reasoning: str = Field(..., description="The reasoning process leading up to the final answer")
    answer: str = Field(..., description="The answer to the user's question based on the retrieved context")


qa_agent = AtomicAgent[RAGQuestionAnsweringAgentInputSchema, RAGQuestionAnsweringAgentOutputSchema](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are an expert at answering questions using retrieved context chunks from a RAG system.",
                "Your role is to synthesize information from the chunks to provide accurate, well-supported answers.",
                "You must explain your reasoning process before providing the answer.",
            ],
            steps=[
                "1. Analyze the question and available context chunks",
                "2. Identify the most relevant information in the chunks",
                "3. Explain how you'll use this information to answer the question",
                "4. Synthesize information into a coherent answer",
            ],
            output_instructions=[
                "First explain your reasoning process clearly",
                "Then provide a clear, direct answer based on the context",
                "If context is insufficient, state this in your reasoning",
                "Never make up information not present in the chunks",
                "Focus on being accurate and concise",
            ],
        ),
    )
)

```

### File: atomic-examples/rag-chatbot/rag_chatbot/agents/query_agent.py

```python
import instructor
import openai
from pydantic import Field
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
from atomic_agents.context import SystemPromptGenerator

from rag_chatbot.config import ChatConfig


class RAGQueryAgentInputSchema(BaseIOSchema):
    """Input schema for the RAG query agent."""

    user_message: str = Field(..., description="The user's question or message to generate a semantic search query for")


class RAGQueryAgentOutputSchema(BaseIOSchema):
    """Output schema for the RAG query agent."""

    reasoning: str = Field(..., description="The reasoning process leading up to the final query")
    query: str = Field(..., description="The semantic search query to use for retrieving relevant chunks")


query_agent = AtomicAgent[RAGQueryAgentInputSchema, RAGQueryAgentOutputSchema](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
        model=ChatConfig.model,
        model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are an expert at formulating semantic search queries for RAG systems.",
                "Your role is to convert user questions into effective semantic search queries that will retrieve the most relevant text chunks.",
            ],
            steps=[
                "1. Analyze the user's question to identify key concepts and information needs",
                "2. Reformulate the question into a semantic search query that will match relevant content",
                "3. Ensure the query captures the core meaning while being general enough to match similar content",
            ],
            output_instructions=[
                "Generate a clear, concise semantic search query",
                "Focus on key concepts and entities from the user's question",
                "Avoid overly specific details that might miss relevant matches",
                "Include synonyms or related terms when appropriate",
                "Explain your reasoning for the query formulation",
            ],
        ),
    )
)

```

### File: atomic-examples/rag-chatbot/rag_chatbot/config.py

```python
import os
from dataclasses import dataclass
from enum import Enum


class VectorDBType(Enum):
    CHROMA = "chroma"
    QDRANT = "qdrant"


def get_api_key() -> str:
    """Retrieve API key from environment or raise error"""
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("API key not found. Please set the OPENAI_API_KEY environment variable.")
    return api_key


def get_vector_db_type() -> VectorDBType:
    """Get the vector database type from environment variable"""
    db_type = os.getenv("VECTOR_DB_TYPE", "chroma").lower()
    try:
        return VectorDBType(db_type)
    except ValueError:
        raise ValueError(f"Invalid VECTOR_DB_TYPE: {db_type}. Must be 'chroma' or 'qdrant'")


@dataclass
class ChatConfig:
    """Configuration for the chat application"""

    api_key: str = get_api_key()
    model: str = "gpt-5-mini"
    reasoning_effort: str = "low"
    exit_commands: set[str] = frozenset({"/exit", "exit", "quit", "/quit"})

    def __init__(self):
        # Prevent instantiation
        raise TypeError("ChatConfig is not meant to be instantiated")


# Model Configuration
EMBEDDING_MODEL = "text-embedding-3-small"  # OpenAI's latest embedding model
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200

# Vector Search Configuration
NUM_CHUNKS_TO_RETRIEVE = 3
SIMILARITY_METRIC = "cosine"

# Vector Database Configuration
VECTOR_DB_TYPE = get_vector_db_type()

# ChromaDB Configuration
CHROMA_PERSIST_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "chroma_db")

# Qdrant Configuration
QDRANT_PERSIST_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "qdrant_db")

# History Configuration
HISTORY_SIZE = 10  # Number of messages to keep in conversation history
MAX_CONTEXT_LENGTH = 4000  # Maximum length of combined context to send to the model

```

### File: atomic-examples/rag-chatbot/rag_chatbot/context_providers.py

```python
from dataclasses import dataclass
from typing import List
from atomic_agents.context import BaseDynamicContextProvider


@dataclass
class ChunkItem:
    content: str
    metadata: dict


class RAGContextProvider(BaseDynamicContextProvider):
    def __init__(self, title: str):
        super().__init__(title=title)
        self.chunks: List[ChunkItem] = []

    def get_info(self) -> str:
        return "\n\n".join(
            [
                f"Chunk {idx}:\nMetadata: {item.metadata}\nContent:\n{item.content}\n{'-' * 80}"
                for idx, item in enumerate(self.chunks, 1)
            ]
        )

```

### File: atomic-examples/rag-chatbot/rag_chatbot/main.py

```python
import os
from typing import List
import wget
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.table import Table
from rich import box
from rich.progress import Progress, SpinnerColumn, TextColumn

from rag_chatbot.agents.query_agent import query_agent, RAGQueryAgentInputSchema, RAGQueryAgentOutputSchema
from rag_chatbot.agents.qa_agent import qa_agent, RAGQuestionAnsweringAgentInputSchema, RAGQuestionAnsweringAgentOutputSchema
from rag_chatbot.context_providers import RAGContextProvider, ChunkItem
from rag_chatbot.services.factory import create_vector_db_service
from rag_chatbot.services.base import BaseVectorDBService
from rag_chatbot.config import CHUNK_SIZE, CHUNK_OVERLAP, NUM_CHUNKS_TO_RETRIEVE, VECTOR_DB_TYPE


console = Console()

WELCOME_MESSAGE = """
Welcome to the RAG Chatbot! I can help you find information from the State of the Union address.
Ask me any questions about the speech and I'll use my knowledge base to provide accurate answers.

I'll show you my thought process:
1. First, I'll generate a semantic search query from your question
2. Then, I'll retrieve relevant chunks of text from the speech
3. Finally, I'll analyze these chunks to provide you with an answer

Using vector database: {db_type}
"""

STARTER_QUESTIONS = [
    "What were the main points about the economy?",
    "What did the president say about healthcare?",
    "How did he address foreign policy?",
]


def download_document() -> str:
    """Download the sample document if it doesn't exist."""
    url = "https://raw.githubusercontent.com/IBM/watson-machine-learning-samples/master/cloud/data/foundation_models/state_of_the_union.txt"
    output_path = "downloads/state_of_the_union.txt"

    if not os.path.exists("downloads"):
        os.makedirs("downloads")

    if not os.path.exists(output_path):
        console.print("\n[bold yellow]📥 Downloading sample document...[/bold yellow]")
        wget.download(url, output_path)
        console.print("\n[bold green]✓ Download complete![/bold green]")

    return output_path


def chunk_document(file_path: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[str]:
    """Split the document into chunks with overlap."""
    with open(file_path, "r", encoding="utf-8") as file:
        text = file.read()

    # Split into paragraphs first
    paragraphs = text.split("\n\n")
    chunks = []
    current_chunk = ""
    current_size = 0

    for i, paragraph in enumerate(paragraphs):
        if current_size + len(paragraph) > chunk_size:
            if current_chunk:
                chunks.append(current_chunk.strip())
            # Include some overlap from the previous chunk
            if overlap > 0 and chunks:
                last_chunk = chunks[-1]
                overlap_text = " ".join(last_chunk.split()[-overlap:])
                current_chunk = overlap_text + "\n\n" + paragraph
            else:
                current_chunk = paragraph
            current_size = len(current_chunk)
        else:
            current_chunk += "\n\n" + paragraph if current_chunk else paragraph
            current_size += len(paragraph)

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks


def initialize_system() -> tuple[BaseVectorDBService, RAGContextProvider]:
    """Initialize the RAG system components."""
    console.print("\n[bold magenta]🚀 Initializing RAG Chatbot System...[/bold magenta]")

    try:
        # Download and chunk document
        doc_path = download_document()
        chunks = chunk_document(doc_path)
        console.print(f"[dim]• Created {len(chunks)} document chunks[/dim]")

        # Initialize vector database
        console.print(f"[dim]• Initializing {VECTOR_DB_TYPE.value} vector database...[/dim]")
        vector_db = create_vector_db_service(collection_name="state_of_union", recreate_collection=True)

        # Add chunks to vector database
        console.print("[dim]• Adding document chunks to vector database...[/dim]")
        chunk_ids = vector_db.add_documents(
            documents=chunks, metadatas=[{"source": "state_of_union", "chunk_index": i} for i in range(len(chunks))]
        )
        console.print(f"[dim]• Added {len(chunk_ids)} chunks to vector database[/dim]")

        # Initialize context provider
        console.print("[dim]• Creating context provider...[/dim]")
        rag_context = RAGContextProvider("RAG Context")

        # Register context provider with agents
        console.print("[dim]• Registering context provider with agents...[/dim]")
        query_agent.register_context_provider("rag_context", rag_context)
        qa_agent.register_context_provider("rag_context", rag_context)

        console.print("[bold green]✨ System initialized successfully![/bold green]\n")
        return vector_db, rag_context

    except Exception as e:
        console.print(f"\n[bold red]Error during initialization:[/bold red] {str(e)}")
        raise


def display_welcome() -> None:
    """Display welcome message and starter questions."""
    welcome_panel = Panel(
        WELCOME_MESSAGE.format(db_type=VECTOR_DB_TYPE.value.upper()),
        title="[bold blue]RAG Chatbot[/bold blue]",
        border_style="blue",
        padding=(1, 2),
    )
    console.print("\n")
    console.print(welcome_panel)

    table = Table(
        show_header=True, header_style="bold cyan", box=box.ROUNDED, title="[bold]Example Questions to Get Started[/bold]"
    )
    table.add_column("№", style="dim", width=4)
    table.add_column("Question", style="green")

    for i, question in enumerate(STARTER_QUESTIONS, 1):
        table.add_row(str(i), question)

    console.print("\n")
    console.print(table)
    console.print("\n" + "─" * 80 + "\n")


def display_chunks(chunks: List[ChunkItem]) -> None:
    """Display the retrieved chunks in a formatted way."""
    console.print("\n[bold cyan]📚 Retrieved Text Chunks:[/bold cyan]")

    for i, chunk in enumerate(chunks, 1):
        chunk_panel = Panel(
            Markdown(chunk.content),
            title=f"[bold]Chunk {i} (Distance: {chunk.metadata['distance']:.4f})[/bold]",
            border_style="blue",
            padding=(1, 2),
        )
        console.print(chunk_panel)
        console.print()


def display_query_info(query_output: RAGQueryAgentOutputSchema) -> None:
    """Display information about the generated query."""
    query_panel = Panel(
        f"[yellow]Generated Query:[/yellow] {query_output.query}\n\n" f"[yellow]Reasoning:[/yellow] {query_output.reasoning}",
        title="[bold]🔍 Semantic Search Strategy[/bold]",
        border_style="yellow",
        padding=(1, 2),
    )
    console.print("\n")
    console.print(query_panel)


def display_answer(qa_output: RAGQuestionAnsweringAgentOutputSchema) -> None:
    """Display the reasoning and answer from the QA agent."""
    # Display reasoning
    reasoning_panel = Panel(
        Markdown(qa_output.reasoning),
        title="[bold]🤔 Analysis & Reasoning[/bold]",
        border_style="green",
        padding=(1, 2),
    )
    console.print("\n")
    console.print(reasoning_panel)

    # Display answer
    answer_panel = Panel(
        Markdown(qa_output.answer),
        title="[bold]💡 Answer[/bold]",
        border_style="blue",
        padding=(1, 2),
    )
    console.print("\n")
    console.print(answer_panel)


def chat_loop(vector_db: BaseVectorDBService, rag_context: RAGContextProvider) -> None:
    """Main chat loop."""
    display_welcome()

    while True:
        try:
            user_message = console.input("\n[bold blue]Your question:[/bold blue] ").strip()

            if user_message.lower() in ["/exit", "/quit"]:
                console.print("\n[bold]👋 Goodbye! Thanks for using the RAG Chatbot.[/bold]")
                break

            try:
                i_question = int(user_message) - 1
                if 0 <= i_question < len(STARTER_QUESTIONS):
                    user_message = STARTER_QUESTIONS[i_question]
            except ValueError:
                pass

            console.print("\n" + "─" * 80)
            console.print("\n[bold magenta]🔄 Processing your question...[/bold magenta]")

            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                console=console,
            ) as progress:
                # Generate search query
                task = progress.add_task("[cyan]Generating semantic search query...", total=None)
                query_output = query_agent.run(RAGQueryAgentInputSchema(user_message=user_message))
                progress.remove_task(task)

                # Display query information
                display_query_info(query_output)

                # Perform vector search
                task = progress.add_task("[cyan]Searching knowledge base...", total=None)
                search_results = vector_db.query(query_text=query_output.query, n_results=NUM_CHUNKS_TO_RETRIEVE)

                # Update context with retrieved chunks
                rag_context.chunks = [
                    ChunkItem(content=doc, metadata={"chunk_id": id, "distance": dist})
                    for doc, id, dist in zip(search_results["documents"], search_results["ids"], search_results["distances"])
                ]
                progress.remove_task(task)

                # Display retrieved chunks
                display_chunks(rag_context.chunks)

                # Generate answer
                task = progress.add_task("[cyan]Analyzing chunks and generating answer...", total=None)
                qa_output = qa_agent.run(RAGQuestionAnsweringAgentInputSchema(question=user_message))
                progress.remove_task(task)

                # Display answer
                display_answer(qa_output)

            console.print("\n" + "─" * 80)

        except Exception as e:
            console.print(f"\n[bold red]Error:[/bold red] {str(e)}")
            console.print("[dim]Please try again or type 'exit' to quit.[/dim]")


if __name__ == "__main__":
    try:
        vector_db, rag_context = initialize_system()
        chat_loop(vector_db, rag_context)
    except KeyboardInterrupt:
        console.print("\n[bold]👋 Goodbye! Thanks for using the RAG Chatbot.[/bold]")
    except Exception as e:
        console.print(f"\n[bold red]Fatal error:[/bold red] {str(e)}")

```

### File: atomic-examples/rag-chatbot/rag_chatbot/services/__init__.py

```python

```

### File: atomic-examples/rag-chatbot/rag_chatbot/services/base.py

```python
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, TypedDict


class QueryResult(TypedDict):
    documents: List[str]
    metadatas: List[Dict[str, str]]
    distances: List[float]
    ids: List[str]


class BaseVectorDBService(ABC):
    """Abstract base class for vector database services."""

    @abstractmethod
    def add_documents(
        self,
        documents: List[str],
        metadatas: Optional[List[Dict[str, str]]] = None,
        ids: Optional[List[str]] = None,
    ) -> List[str]:
        """Add documents to the collection.

        Args:
            documents: List of text documents to add
            metadatas: Optional list of metadata dicts for each document
            ids: Optional list of IDs for each document. If not provided, UUIDs will be generated.

        Returns:
            List[str]: The IDs of the added documents
        """
        pass

    @abstractmethod
    def query(
        self,
        query_text: str,
        n_results: int = 5,
        where: Optional[Dict[str, str]] = None,
    ) -> QueryResult:
        """Query the collection for similar documents.

        Args:
            query_text: Text to find similar documents for
            n_results: Number of results to return
            where: Optional filter criteria

        Returns:
            QueryResult containing documents, metadata, distances and IDs
        """
        pass

    @abstractmethod
    def delete_collection(self, collection_name: Optional[str] = None) -> None:
        """Delete a collection by name.

        Args:
            collection_name: Name of the collection to delete. If None, deletes the current collection.
        """
        pass

    @abstractmethod
    def delete_by_ids(self, ids: List[str]) -> None:
        """Delete documents from the collection by their IDs.

        Args:
            ids: List of IDs to delete
        """
        pass

```

### File: atomic-examples/rag-chatbot/rag_chatbot/services/chroma_db.py

```python
import os
import shutil
import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
from typing import Dict, List, Optional
import uuid

from .base import BaseVectorDBService, QueryResult


class ChromaDBService(BaseVectorDBService):
    """Service for interacting with ChromaDB using OpenAI embeddings."""

    def __init__(
        self,
        collection_name: str,
        persist_directory: str = "./chroma_db",
        recreate_collection: bool = False,
    ) -> None:
        """Initialize ChromaDB service with OpenAI embeddings.

        Args:
            collection_name: Name of the collection to use
            persist_directory: Directory to persist ChromaDB data
            recreate_collection: If True, deletes the collection if it exists before creating
        """
        # Initialize embedding function with OpenAI
        self.embedding_function = OpenAIEmbeddingFunction(
            api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
        )

        # If recreating, delete the entire persist directory
        if recreate_collection and os.path.exists(persist_directory):
            shutil.rmtree(persist_directory)
            os.makedirs(persist_directory)

        # Initialize persistent client
        self.client = chromadb.PersistentClient(path=persist_directory)

        # Get or create collection
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embedding_function,
            metadata={"hnsw:space": "cosine"},  # Explicitly set distance metric
        )

    def add_documents(
        self,
        documents: List[str],
        metadatas: Optional[List[Dict[str, str]]] = None,
        ids: Optional[List[str]] = None,
    ) -> List[str]:
        """Add documents to the collection.

        Args:
            documents: List of text documents to add
            metadatas: Optional list of metadata dicts for each document
            ids: Optional list of IDs for each document. If not provided, UUIDs will be generated.

        Returns:
            List[str]: The IDs of the added documents
        """
        if ids is None:
            ids = [str(uuid.uuid4()) for _ in documents]

        self.collection.add(documents=documents, metadatas=metadatas, ids=ids)
        return ids

    def query(
        self,
        query_text: str,
        n_results: int = 5,
        where: Optional[Dict[str, str]] = None,
    ) -> QueryResult:
        """Query the collection for similar documents.

        Args:
            query_text: Text to find similar documents for
            n_results: Number of results to return
            where: Optional filter criteria

        Returns:
            QueryResult containing documents, metadata, distances and IDs
        """
        results = self.collection.query(
            query_texts=[query_text],
            n_results=n_results,
            where=where,
            include=["documents", "metadatas", "distances"],
        )

        return {
            "documents": results["documents"][0],
            "metadatas": results["metadatas"][0],
            "distances": results["distances"][0],
            "ids": results["ids"][0],
        }

    def delete_collection(self, collection_name: Optional[str] = None) -> None:
        """Delete a collection by name.

        Args:
            collection_name: Name of the collection to delete. If None, deletes the current collection.
        """
        name_to_delete = collection_name if collection_name is not None else self.collection.name
        self.client.delete_collection(name_to_delete)

    def delete_by_ids(self, ids: List[str]) -> None:
        """Delete documents from the collection by their IDs.

        Args:
            ids: List of IDs to delete
        """
        self.collection.delete(ids=ids)


if __name__ == "__main__":
    chroma_db_service = ChromaDBService(collection_name="test", recreate_collection=True)
    added_ids = chroma_db_service.add_documents(
        documents=["Hello, world!", "This is a test document."],
        metadatas=[{"source": "test"}, {"source": "test"}],
    )
    print("Added documents with IDs:", added_ids)

    results = chroma_db_service.query(query_text="Hello, world!")
    print("Query results:", results)

    chroma_db_service.delete_by_ids([added_ids[0]])
    print("Deleted document with ID:", added_ids[0])

    updated_results = chroma_db_service.query(query_text="Hello, world!")
    print("Updated results after deletion:", updated_results)

```

### File: atomic-examples/rag-chatbot/rag_chatbot/services/factory.py

```python
from .base import BaseVectorDBService
from .chroma_db import ChromaDBService
from .qdrant_db import QdrantDBService
from ..config import VECTOR_DB_TYPE, CHROMA_PERSIST_DIR, QDRANT_PERSIST_DIR


def create_vector_db_service(
    collection_name: str,
    recreate_collection: bool = False,
) -> BaseVectorDBService:
    """Create a vector database service based on configuration.

    Args:
        collection_name: Name of the collection to use
        recreate_collection: If True, deletes the collection if it exists before creating

    Returns:
        BaseVectorDBService: The appropriate vector database service instance
    """

    if VECTOR_DB_TYPE == VECTOR_DB_TYPE.CHROMA:
        return ChromaDBService(
            collection_name=collection_name,
            persist_directory=CHROMA_PERSIST_DIR,
            recreate_collection=recreate_collection,
        )
    elif VECTOR_DB_TYPE == VECTOR_DB_TYPE.QDRANT:
        return QdrantDBService(
            collection_name=collection_name,
            persist_directory=QDRANT_PERSIST_DIR,
            recreate_collection=recreate_collection,
        )
    else:
        raise ValueError(f"Unsupported database type: {VECTOR_DB_TYPE}")

```

### File: atomic-examples/rag-chatbot/rag_chatbot/services/qdrant_db.py

```python
import os
import shutil
import uuid
from typing import Dict, List, Optional
from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance,
    VectorParams,
    PointStruct,
    Filter,
    FieldCondition,
    MatchValue,
)
import openai

from .base import BaseVectorDBService, QueryResult


class QdrantDBService(BaseVectorDBService):
    """Service for interacting with Qdrant using OpenAI embeddings."""

    def __init__(
        self,
        collection_name: str,
        persist_directory: str = "./qdrant_db",
        recreate_collection: bool = False,
    ) -> None:
        """Initialize Qdrant service with OpenAI embeddings.

        Args:
            collection_name: Name of the collection to use
            persist_directory: Directory to persist Qdrant data
            recreate_collection: If True, deletes the collection if it exists before creating
        """
        self.openai_client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.embedding_model = "text-embedding-3-small"

        if recreate_collection and os.path.exists(persist_directory):
            shutil.rmtree(persist_directory)
            os.makedirs(persist_directory)

        self.client = QdrantClient(path=persist_directory)
        self.collection_name = collection_name
        self._ensure_collection_exists(recreate_collection)

    def _ensure_collection_exists(self, recreate_collection: bool = False) -> None:
        collection_exists = self.client.collection_exists(self.collection_name)

        if recreate_collection and collection_exists:
            self.client.delete_collection(self.collection_name)
            collection_exists = False

        if not collection_exists:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=1536,  # OpenAI text-embedding-3-small dimension
                    distance=Distance.COSINE,
                ),
            )

    def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
        response = self.openai_client.embeddings.create(model=self.embedding_model, input=texts)

        return [embedding.embedding for embedding in response.data]

    def add_documents(
        self,
        documents: List[str],
        metadatas: Optional[List[Dict[str, str]]] = None,
        ids: Optional[List[str]] = None,
    ) -> List[str]:
        ids = ids or [str(uuid.uuid4()) for _ in documents]
        metadatas = metadatas or [{} for _ in documents]

        embeddings = self._get_embeddings(documents)

        points = []
        for doc_id, doc, embedding, metadata in zip(ids, documents, embeddings, metadatas):
            point = PointStruct(id=doc_id, vector=embedding, payload={"text": doc, "metadata": metadata})
            points.append(point)

        self.client.upsert(collection_name=self.collection_name, points=points)

        return ids

    def query(
        self,
        query_text: str,
        n_results: int = 5,
        where: Optional[Dict[str, str]] = None,
    ) -> QueryResult:
        query_embedding = self._get_embeddings([query_text])[0]

        filter_condition = None
        if where:
            conditions = []
            for key, value in where.items():
                conditions.append(FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value)))
            if conditions:
                filter_condition = Filter(must=conditions)

        search_results = self.client.query_points(
            collection_name=self.collection_name,
            query=query_embedding,
            limit=n_results,
            query_filter=filter_condition,
            with_payload=True,
        ).points

        # Extract results
        documents = []
        metadatas = []
        distances = []
        ids = []

        for result in search_results:
            documents.append(result.payload["text"])
            metadatas.append(result.payload["metadata"])
            distances.append(result.score)
            ids.append(result.id)

        return {
            "documents": documents,
            "metadatas": metadatas,
            "distances": distances,
            "ids": ids,
        }

    def delete_collection(self, collection_name: Optional[str] = None) -> None:
        name_to_delete = collection_name if collection_name is not None else self.collection_name
        self.client.delete_collection(name_to_delete)

    def delete_by_ids(self, ids: List[str]) -> None:
        self.client.delete(collection_name=self.collection_name, points_selector=ids)


if __name__ == "__main__":
    qdrant_db_service = QdrantDBService(collection_name="test", recreate_collection=True)
    added_ids = qdrant_db_service.add_documents(
        documents=["Hello, world!", "This is a test document."],
        metadatas=[{"source": "test"}, {"source": "test"}],
    )
    print("Added documents with IDs:", added_ids)

    results = qdrant_db_service.query(query_text="Hello, world!")
    print("Query results:", results)

    qdrant_db_service.delete_by_ids([added_ids[0]])
    print("Deleted document with ID:", added_ids[0])

    updated_results = qdrant_db_service.query(query_text="Hello, world!")
    print("Updated results after deletion:", updated_results)

```


--------------------------------------------------------------------------------
Example: web-search-agent
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/web-search-agent

## Documentation

# Web Search Agent

This project demonstrates an intelligent web search agent built using the Atomic Agents framework. The agent can perform web searches, generate relevant queries, and provide detailed answers to user questions based on the search results.

## Features

1. Query Generation: Automatically generates relevant search queries based on user input.
2. Web Search: Utilizes SearXNG to perform web searches across multiple search engines.
3. Question Answering: Provides detailed answers to user questions based on search results.
4. Follow-up Questions: Suggests related questions to encourage further exploration of the topic.

## Components

The Web Search Agent consists of several key components:

1. Query Agent (`query_agent.py`): Generates diverse and relevant search queries based on user input.
2. SearXNG Search Tool (`searxng_search.py`): Performs web searches using the SearXNG meta-search engine.
3. Question Answering Agent (`question_answering_agent.py`): Analyzes search results and provides detailed answers to user questions.
4. Main Script (`main.py`): Orchestrates the entire process, from query generation to final answer presentation.

## Getting Started

To run the Web Search Agent:

1. Setting up SearXNG server if you haven't:

Make sure to add these lines to `settings.tml`:

   ```yaml
   search:
   formats:
      - html
      - json
   ```

1. Clone the Atomic Agents repository:

   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

1. Navigate to the web-search-agent directory:

   ```bash
   cd atomic-agents/atomic-examples/web-search-agent
   ```

1. Install dependencies using uv:

   ```bash
   uv sync
   ```

1. Set up environment variables:
   Create a `.env` file in the `web-search-agent` directory with the following content:

   ```bash
   OPENAI_API_KEY=your_openai_api_key
   SEARXNG_BASE_URL=your_searxng_instance_url
   ```

   Replace `your_openai_api_key` with your actual OpenAI API key and `your_searxng_instance_url` with the URL of your SearXNG instance.  
   If you do not have a SearxNG instance, see the instructions below to set up one locally with docker.

2. Run the Web Search Agent:

   ```bash
   uv run python web_search_agent/main.py
   ```

## How It Works

1. The user provides an initial question or topic for research.
2. The Query Agent generates multiple relevant search queries based on the user's input.
3. The SearXNG Search Tool performs web searches using the generated queries.
4. The Question Answering Agent analyzes the search results and formulates a detailed answer.
5. The main script presents the answer, along with references and follow-up questions.

## SearxNG Setup with docker

From the [official instructions](https://docs.searxng.org/admin/installation-docker.html):

```shell
mkdir my-instance
cd my-instance
export PORT=8080
docker pull searxng/searxng
docker run --rm \
           -d -p ${PORT}:8080 \
           -v "${PWD}/searxng:/etc/searxng" \
           -e "BASE_URL=http://localhost:$PORT/" \
           -e "INSTANCE_NAME=my-instance" \
           searxng/searxng
```

Set the `SEARXNG_BASE_URL` environment variable to `http://localhost:8080/` in your `.env` file.


Note: for the agent to communicate with SearxNG, the instance must enable the JSON engine, which is disabled by default.
Edit `/etc/searxng/settings.yml` and add `- json` in the `search.formats` section, then restart the container.


## Customization

You can customize the Web Search Agent by modifying the following:

- Adjust the number of generated queries in `main.py`.
- Modify the search categories or parameters in `searxng_search.py`.
- Customize the system prompts for the Query Agent and Question Answering Agent in their respective files.

## Contributing

Contributions to the Web Search Agent project are welcome! Please fork the repository and submit a pull request with your enhancements or bug fixes.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.



## Source Code


### File: atomic-examples/web-search-agent/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["web_search_agent"]

[project]
name = "web-search-agent"
version = "1.0.0"
description = "Web search agent example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12"
dependencies = [
    "atomic-agents",
    "openai>=2.0.0,<3.0.0",
    "pydantic>=2.9.2,<3.0.0",
    "instructor==1.14.5",
    "python-dotenv>=1.0.1,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```

### File: atomic-examples/web-search-agent/web_search_agent/agents/query_agent.py

```python
import instructor
import openai
from pydantic import Field
from typing import List
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
from atomic_agents.context import SystemPromptGenerator


class QueryAgentInputSchema(BaseIOSchema):
    """This is the input schema for the QueryAgent."""

    instruction: str = Field(..., description="A detailed instruction or request to generate deep research queries for.")
    num_queries: int = Field(..., description="The number of queries to generate.")


class QueryAgentOutputSchema(BaseIOSchema):
    """This is the output schema for the QueryAgent."""

    queries: List[str] = Field(..., description="A list of search queries.")


query_agent = AtomicAgent[QueryAgentInputSchema, QueryAgentOutputSchema](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI()),
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are an advanced search query generator.",
                "Your task is to convert user questions into multiple effective search queries.",
            ],
            steps=[
                "Analyze the user's question to understand the core information need.",
                "Generate multiple search queries that capture the question's essence from different angles.",
                "Ensure each query is optimized for search engines (compact, focused, and unambiguous).",
            ],
            output_instructions=[
                "Generate 3-5 different search queries.",
                "Do not include special search operators or syntax.",
                "Each query should be concise and focused on retrieving relevant information.",
            ],
        ),
    )
)

```

### File: atomic-examples/web-search-agent/web_search_agent/agents/question_answering_agent.py

```python
import instructor
import openai
from pydantic import Field, HttpUrl
from typing import List
from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig
from atomic_agents.context import SystemPromptGenerator


class QuestionAnsweringAgentInputSchema(BaseIOSchema):
    """This schema defines the input schema for the QuestionAnsweringAgent."""

    question: str = Field(..., description="A question that needs to be answered based on the provided context.")


class QuestionAnsweringAgentOutputSchema(BaseIOSchema):
    """This schema defines the output schema for the QuestionAnsweringAgent."""

    markdown_output: str = Field(..., description="The answer to the question in markdown format.")
    references: List[HttpUrl] = Field(
        ..., max_items=3, description="A list of up to 3 HTTP URLs used as references for the answer."
    )
    followup_questions: List[str] = Field(
        ..., max_items=3, description="A list of up to 3 follow-up questions related to the answer."
    )


# Create the question answering agent
question_answering_agent = AtomicAgent[QuestionAnsweringAgentInputSchema, QuestionAnsweringAgentOutputSchema](
    AgentConfig(
        client=instructor.from_openai(openai.OpenAI()),
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "You are an intelligent question answering expert.",
                "Your task is to provide accurate and detailed answers to user questions based on the given context.",
            ],
            steps=[
                "You will receive a question and the context information.",
                "Provide up to 3 relevant references (HTTP URLs) used in formulating the answer.",
                "Generate up to 3 follow-up questions related to the answer.",
            ],
            output_instructions=[
                "Ensure clarity and conciseness in each answer.",
                "Ensure the answer is directly relevant to the question and context provided.",
                "Include up to 3 relevant HTTP URLs as references.",
                "Provide up to 3 follow-up questions to encourage further exploration of the topic.",
            ],
        ),
    )
)

```

### File: atomic-examples/web-search-agent/web_search_agent/main.py

```python
import os
from dotenv import load_dotenv
from rich.console import Console
from rich.markdown import Markdown
from pydantic import Field

from atomic_agents import BaseIOSchema
from atomic_agents.context import ChatHistory, BaseDynamicContextProvider

from web_search_agent.tools.searxng_search import (
    SearXNGSearchTool,
    SearXNGSearchToolConfig,
    SearXNGSearchToolInputSchema,
    SearXNGSearchToolOutputSchema,
)

from web_search_agent.agents.query_agent import QueryAgentInputSchema, query_agent
from web_search_agent.agents.question_answering_agent import question_answering_agent, QuestionAnsweringAgentInputSchema


load_dotenv()

# Initialize a Rich Console for pretty console outputs
console = Console()

# History setup
history = ChatHistory()

# Initialize the SearXNGSearchTool
search_tool = SearXNGSearchTool(config=SearXNGSearchToolConfig(base_url=os.getenv("SEARXNG_BASE_URL"), max_results=5))


class SearchResultsProvider(BaseDynamicContextProvider):
    def __init__(self, title: str, search_results: SearXNGSearchToolOutputSchema | Exception):
        super().__init__(title=title)
        self.search_results = search_results

    def get_info(self) -> str:
        return f"{self.title}: {self.search_results}"


# Define input/output schemas for the main agent
class MainAgentInputSchema(BaseIOSchema):
    """Input schema for the main agent."""

    chat_message: str = Field(..., description="Chat message from the user.")


class MainAgentOutputSchema(BaseIOSchema):
    """Output schema for the main agent."""

    chat_message: str = Field(..., description="Response to the user's message.")


# Example usage
instruction = "Tell me about the Atomic Agents AI agent framework."
num_queries = 3
console.print(f"[bold blue]Instruction:[/bold blue] {instruction}")

while True:
    # Generate queries using the query agent
    query_input = QueryAgentInputSchema(instruction=instruction, num_queries=num_queries)
    generated_queries = query_agent.run(query_input)

    console.print("[bold blue]Generated Queries:[/bold blue]")
    for query in generated_queries.queries:
        console.print(f"- {query}")

    # Perform searches using the generated queries
    search_input = SearXNGSearchToolInputSchema(queries=generated_queries.queries, category="general")

    try:
        search_results = search_tool.run(search_input)
        search_results_provider = SearchResultsProvider("Search Results", search_results)
    except Exception as e:
        search_results_provider = SearchResultsProvider("Search Failed", e)

    question_answering_agent.register_context_provider("search results", search_results_provider)

    answer = question_answering_agent.run(QuestionAnsweringAgentInputSchema(question=instruction))

    # Create a Rich Console instance
    console = Console()

    # Print the answer using Rich's Markdown rendering
    console.print("\n[bold blue]Answer:[/bold blue]")
    console.print(Markdown(answer.markdown_output))

    # Print references
    console.print("\n[bold blue]References:[/bold blue]")
    for ref in answer.references:
        console.print(f"- {ref}")

    # Print follow-up questions
    console.print("\n[bold blue]Follow-up Questions:[/bold blue]")
    for i, question in enumerate(answer.followup_questions, 1):
        console.print(f"[cyan]{i}. {question}[/cyan]")

    console.print()  # Add an empty line for better readability
    instruction = console.input("[bold blue]You:[/bold blue] ")
    if instruction.lower() in ["/exit", "/quit"]:
        console.print("Exiting chat...")
        break

    try:
        followup_question_id = int(instruction.strip())
        if 1 <= followup_question_id <= len(answer.followup_questions):
            instruction = answer.followup_questions[followup_question_id - 1]
            console.print(f"[bold blue]Follow-up Question:[/bold blue] {instruction}")
    except ValueError:
        pass

```

### File: atomic-examples/web-search-agent/web_search_agent/tools/searxng_search.py

```python
import os
from typing import List, Literal, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

import aiohttp
from pydantic import Field

from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class SearXNGSearchToolInputSchema(BaseIOSchema):
    """
    Schema for input to a tool for searching for information, news, references, and other content using SearXNG.
    Returns a list of search results with a short description or content snippet and URLs for further exploration
    """

    queries: List[str] = Field(..., description="List of search queries.")
    category: Optional[Literal["general", "news", "social_media"]] = Field(
        "general", description="Category of the search queries."
    )


####################
# OUTPUT SCHEMA(S) #
####################
class SearXNGSearchResultItemSchema(BaseIOSchema):
    """This schema represents a single search result item"""

    url: str = Field(..., description="The URL of the search result")
    title: str = Field(..., description="The title of the search result")
    content: Optional[str] = Field(None, description="The content snippet of the search result")
    query: str = Field(..., description="The query used to obtain this search result")


class SearXNGSearchToolOutputSchema(BaseIOSchema):
    """This schema represents the output of the SearXNG search tool."""

    results: List[SearXNGSearchResultItemSchema] = Field(..., description="List of search result items")
    category: Optional[str] = Field(None, description="The category of the search results")


##############
# TOOL LOGIC #
##############
class SearXNGSearchToolConfig(BaseToolConfig):
    base_url: str = ""
    max_results: int = 10


class SearXNGSearchTool(BaseTool[SearXNGSearchToolInputSchema, SearXNGSearchToolOutputSchema]):
    """
    Tool for performing searches on SearXNG based on the provided queries and category.

    Attributes:
        input_schema (SearXNGSearchToolInputSchema): The schema for the input data.
        output_schema (SearXNGSearchToolOutputSchema): The schema for the output data.
        max_results (int): The maximum number of search results to return.
        base_url (str): The base URL for the SearXNG instance to use.
    """

    def __init__(self, config: SearXNGSearchToolConfig = SearXNGSearchToolConfig()):
        """
        Initializes the SearXNGTool.

        Args:
            config (SearXNGSearchToolConfig):
                Configuration for the tool, including base URL, max results, and optional title and description overrides.
        """
        super().__init__(config)
        self.base_url = config.base_url
        self.max_results = config.max_results

    async def _fetch_search_results(self, session: aiohttp.ClientSession, query: str, category: Optional[str]) -> List[dict]:
        """
        Fetches search results for a single query asynchronously.

        Args:
            session (aiohttp.ClientSession): The aiohttp session to use for the request.
            query (str): The search query.
            category (Optional[str]): The category of the search query.

        Returns:
            List[dict]: A list of search result dictionaries.

        Raises:
            Exception: If the request to SearXNG fails.
        """
        query_params = {
            "q": query,
            "safesearch": "0",
            "format": "json",
            "language": "en",
            "engines": "bing,duckduckgo,google,startpage,yandex",
        }

        if category:
            query_params["categories"] = category

        async with session.get(f"{self.base_url}/search", params=query_params) as response:
            if response.status != 200:
                raise Exception(f"Failed to fetch search results for query '{query}': {response.status} {response.reason}")
            data = await response.json()
            results = data.get("results", [])

            # Add the query to each result
            for result in results:
                result["query"] = query

            return results

    async def run_async(
        self, params: SearXNGSearchToolInputSchema, max_results: Optional[int] = None
    ) -> SearXNGSearchToolOutputSchema:
        """
        Runs the SearXNGTool asynchronously with the given parameters.

        Args:
            params (SearXNGSearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
            max_results (Optional[int]): The maximum number of search results to return.

        Returns:
            SearXNGSearchToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            ValueError: If the base URL is not provided.
            Exception: If the request to SearXNG fails.
        """
        async with aiohttp.ClientSession() as session:
            tasks = [self._fetch_search_results(session, query, params.category) for query in params.queries]
            results = await asyncio.gather(*tasks)

        all_results = [item for sublist in results for item in sublist]

        # Sort the combined results by score in descending order
        sorted_results = sorted(all_results, key=lambda x: x.get("score", 0), reverse=True)

        # Remove duplicates while preserving order
        seen_urls = set()
        unique_results = []
        for result in sorted_results:
            if "content" not in result or "title" not in result or "url" not in result or "query" not in result:
                continue
            if result["url"] not in seen_urls:
                unique_results.append(result)
                if "metadata" in result:
                    result["title"] = f"{result['title']} - (Published {result['metadata']})"
                if "publishedDate" in result and result["publishedDate"]:
                    result["title"] = f"{result['title']} - (Published {result['publishedDate']})"
                seen_urls.add(result["url"])

        # Filter results to include only those with the correct category if it is set
        if params.category:
            filtered_results = [result for result in unique_results if result.get("category") == params.category]
        else:
            filtered_results = unique_results

        filtered_results = filtered_results[: max_results or self.max_results]

        return SearXNGSearchToolOutputSchema(
            results=[
                SearXNGSearchResultItemSchema(
                    url=result["url"], title=result["title"], content=result.get("content"), query=result["query"]
                )
                for result in filtered_results
            ],
            category=params.category,
        )

    def run(self, params: SearXNGSearchToolInputSchema, max_results: Optional[int] = None) -> SearXNGSearchToolOutputSchema:
        """
        Runs the SearXNGTool synchronously with the given parameters.

        This method creates an event loop in a separate thread to run the asynchronous operations.

        Args:
            params (SearXNGSearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
            max_results (Optional[int]): The maximum number of search results to return.

        Returns:
            SearXNGSearchToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            ValueError: If the base URL is not provided.
            Exception: If the request to SearXNG fails.
        """
        with ThreadPoolExecutor() as executor:
            return executor.submit(asyncio.run, self.run_async(params, max_results)).result()


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console
    from dotenv import load_dotenv

    load_dotenv()
    rich_console = Console()

    search_tool_instance = SearXNGSearchTool(
        config=SearXNGSearchToolConfig(base_url=os.getenv("SEARXNG_BASE_URL"), max_results=5)
    )

    search_input = SearXNGSearchTool.input_schema(
        queries=["Python programming", "Machine learning", "Artificial intelligence"],
        category="news",
    )

    output = search_tool_instance.run(search_input)

    rich_console.print(output)

```


--------------------------------------------------------------------------------
Example: youtube-summarizer
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/youtube-summarizer

## Documentation

# YouTube Summarizer

This directory contains the YouTube Summarizer example for the Atomic Agents project. This example demonstrates how to extract and summarize knowledge from YouTube videos using the Atomic Agents framework.

## Getting Started

To get started with the YouTube Summarizer:

1. **Clone the main Atomic Agents repository:**
   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. **Navigate to the YouTube Summarizer directory:**
   ```bash
   cd atomic-agents/atomic-examples/youtube-summarizer
   ```

3. **Install the dependencies using uv:**
   ```bash
   uv sync
   ```

4. **Set up environment variables:**

   Create a `.env` file in the `youtube-summarizer` directory with the following content:
   ```env
   OPENAI_API_KEY=your_openai_api_key
   YOUTUBE_API_KEY=your_youtube_api_key
   ```

   To get your YouTube API key, follow the instructions in the [YouTube Scraper README](/atomic-forge/tools/youtube_transcript_scraper/README.md).

   Replace `your_openai_api_key` and `your_youtube_api_key` with your actual API keys.

5. **Run the YouTube Summarizer:**
   ```bash
   uv run python youtube_summarizer/main.py
   ```
   or

   ```bash
   uv run python -m youtube_summarizer.main
   ```

## File Explanation

### 1. Agent (`agent.py`)

This module defines the `YouTubeKnowledgeExtractionAgent`, responsible for extracting summaries, insights, quotes, and more from YouTube video transcripts.


### 2. YouTube Transcript Scraper (`tools/youtube_transcript_scraper.py`)

This tool comes from the [Atomic Forge](/atomic-forge/README.md) and handles fetching transcripts and metadata from YouTube videos.

### 3. Main (`main.py`)

The entry point for the YouTube Summarizer application. It orchestrates fetching transcripts, processing them through the agent, and displaying the results.

## Customization

You can modify the `video_url` variable in `main.py` to analyze different YouTube videos. Additionally, you can adjust the agent's configuration in `agent.py` to tailor the summaries and insights according to your requirements.

## Contributing

Contributions are welcome! Please fork the repository and submit a pull request with your enhancements or bug fixes.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.


## Source Code


### File: atomic-examples/youtube-summarizer/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["youtube_summarizer"]

[project]
name = "youtube-summarizer"
version = "1.0.0"
description = "Youtube Summarizer example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12,<3.14"
dependencies = [
    "atomic-agents",
    "openai>=2.0.0,<3.0.0",
    "pydantic>=2.10.3,<3.0.0",
    "google-api-python-client>=2.101.0,<3.0.0",
    "youtube-transcript-api>=1.1.1,<2.0.0",
    "instructor==1.14.5",
    "python-dotenv>=1.0.1,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```

### File: atomic-examples/youtube-summarizer/youtube_summarizer/agent.py

```python
import instructor
import openai
from pydantic import Field
from typing import List, Optional

from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import BaseDynamicContextProvider, SystemPromptGenerator


class YtTranscriptProvider(BaseDynamicContextProvider):
    def __init__(self, title):
        super().__init__(title)
        self.transcript = None
        self.duration = None
        self.metadata = None

    def get_info(self) -> str:
        return f'VIDEO TRANSCRIPT: "{self.transcript}"\n\nDURATION: {self.duration}\n\nMETADATA: {self.metadata}'


class YouTubeKnowledgeExtractionInputSchema(BaseIOSchema):
    """This schema defines the input schema for the YouTubeKnowledgeExtractionAgent."""

    video_url: str = Field(..., description="The URL of the YouTube video to analyze")


class YouTubeKnowledgeExtractionOutputSchema(BaseIOSchema):
    """This schema defines an elaborate set of insights about the contentof the video."""

    summary: str = Field(
        ..., description="A short summary of the content, including who is presenting and the content being discussed."
    )
    insights: List[str] = Field(
        ..., min_items=5, max_items=5, description="exactly 5 of the best insights and ideas from the input."
    )
    quotes: List[str] = Field(
        ...,
        min_items=5,
        max_items=5,
        description="exactly 5 of the most surprising, insightful, and/or interesting quotes from the input.",
    )
    habits: Optional[List[str]] = Field(
        None,
        min_items=5,
        max_items=5,
        description="exactly 5 of the most practical and useful personal habits mentioned by the speakers.",
    )
    facts: List[str] = Field(
        ...,
        min_items=5,
        max_items=5,
        description="exactly 5 of the most surprising, insightful, and/or interesting valid facts about the greater world mentioned in the content.",
    )
    recommendations: List[str] = Field(
        ...,
        min_items=5,
        max_items=5,
        description="exactly 5 of the most surprising, insightful, and/or interesting recommendations from the content.",
    )
    references: List[str] = Field(
        ...,
        description="All mentions of writing, art, tools, projects, and other sources of inspiration mentioned by the speakers.",
    )
    one_sentence_takeaway: str = Field(
        ..., description="The most potent takeaways and recommendations condensed into a single 20-word sentence."
    )


transcript_provider = YtTranscriptProvider(title="YouTube Transcript")

youtube_knowledge_extraction_agent = AtomicAgent[
    YouTubeKnowledgeExtractionInputSchema, YouTubeKnowledgeExtractionOutputSchema
](
    config=AgentConfig(
        client=instructor.from_openai(openai.OpenAI()),
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "This Assistant is an expert at extracting knowledge and other insightful and interesting information from YouTube transcripts."
            ],
            steps=[
                "Analyse the YouTube transcript thoroughly to extract the most valuable insights, facts, and recommendations.",
                "Adhere strictly to the provided schema when extracting information from the input content.",
                "Ensure that the output matches the field descriptions, types and constraints exactly.",
            ],
            output_instructions=[
                "Only output Markdown-compatible strings.",
                "Ensure you follow ALL these instructions when creating your output.",
            ],
            context_providers={"yt_transcript": transcript_provider},
        ),
    )
)

```

### File: atomic-examples/youtube-summarizer/youtube_summarizer/main.py

```python
import os
from dotenv import load_dotenv
from rich.console import Console

from youtube_summarizer.tools.youtube_transcript_scraper import (
    YouTubeTranscriptTool,
    YouTubeTranscriptToolConfig,
    YouTubeTranscriptToolInputSchema,
)

from youtube_summarizer.agent import (
    YouTubeKnowledgeExtractionInputSchema,
    youtube_knowledge_extraction_agent,
    transcript_provider,
)

load_dotenv()

# Initialize a Rich Console for pretty console outputs
console = Console()

# Initialize the YouTubeTranscriptTool
transcript_tool = YouTubeTranscriptTool(config=YouTubeTranscriptToolConfig(api_key=os.getenv("YOUTUBE_API_KEY")))

# Remove the infinite loop and perform a one-time transcript extraction
video_url = "https://www.youtube.com/watch?v=Sp30YsjGUW0"

transcript_input = YouTubeTranscriptToolInputSchema(video_url=video_url, language="en")
try:
    transcript_output = transcript_tool.run(transcript_input)
    console.print(f"[bold green]Transcript:[/bold green] {transcript_output.transcript}")
    console.print(f"[bold green]Duration:[/bold green] {transcript_output.duration} seconds")

    # Update transcript_provider with the scraped transcript data
    transcript_provider.transcript = transcript_output.transcript
    transcript_provider.duration = transcript_output.duration
    transcript_provider.metadata = transcript_output.metadata  # Assuming metadata is available in transcript_output

    # Run the transcript through the agent
    transcript_input_schema = YouTubeKnowledgeExtractionInputSchema(video_url=video_url)
    agent_response = youtube_knowledge_extraction_agent.run(transcript_input_schema)

    # Print the output schema in a formatted way
    console.print("[bold blue]Agent Output Schema:[/bold blue]")
    console.print(agent_response)
except Exception as e:
    console.print(f"[bold red]Error:[/bold red] {str(e)}")

```

### File: atomic-examples/youtube-summarizer/youtube_summarizer/tools/youtube_transcript_scraper.py

```python
import os
from typing import List, Optional

from pydantic import Field, BaseModel
from datetime import datetime

from googleapiclient.discovery import build
from youtube_transcript_api import (
    NoTranscriptFound,
    TranscriptsDisabled,
    YouTubeTranscriptApi,
)
from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class YouTubeTranscriptToolInputSchema(BaseIOSchema):
    """
    Tool for fetching the transcript of a YouTube video using the YouTube Transcript API.
    Returns the transcript with text, start time, and duration.
    """

    video_url: str = Field(..., description="URL of the YouTube video to fetch the transcript for.")
    language: Optional[str] = Field(None, description="Language code for the transcript (e.g., 'en' for English).")


#################
# OUTPUT SCHEMA #
#################
class VideoMetadata(BaseModel):
    """Schema for YouTube video metadata."""

    id: str = Field(..., description="The YouTube video ID.")
    title: str = Field(..., description="The title of the YouTube video.")
    channel: str = Field(..., description="The name of the YouTube channel.")
    published_at: datetime = Field(..., description="The publication date and time of the video.")


class YouTubeTranscriptToolOutputSchema(BaseIOSchema):
    """
    Output schema for the YouTubeTranscriptTool. Contains the transcript text, duration, comments, and metadata.
    """

    transcript: str = Field(..., description="Transcript of the YouTube video.")
    duration: float = Field(..., description="Duration of the YouTube video in seconds.")
    comments: List[str] = Field(default_factory=list, description="Comments on the YouTube video.")
    metadata: VideoMetadata = Field(..., description="Metadata of the YouTube video.")


#################
# CONFIGURATION #
#################
class YouTubeTranscriptToolConfig(BaseToolConfig):
    """
    Configuration for the YouTubeTranscriptTool.

    Attributes:
        languages (List[str]): List of language codes to try when fetching transcripts.
    """

    languages: List[str] = ["en", "en-US", "en-GB"]


#####################
# MAIN TOOL & LOGIC #
#####################
class YouTubeTranscriptTool(BaseTool[YouTubeTranscriptToolInputSchema, YouTubeTranscriptToolOutputSchema]):
    """
    Tool for extracting transcripts from YouTube videos.

    Attributes:
        input_schema (YouTubeTranscriptToolInputSchema): The schema for the input data.
        output_schema (YouTubeTranscriptToolOutputSchema): The schema for the output data.
        languages (List[str]): List of language codes to try when fetching transcripts.
    """

    input_schema = YouTubeTranscriptToolInputSchema
    output_schema = YouTubeTranscriptToolOutputSchema

    def __init__(self, config: YouTubeTranscriptToolConfig = YouTubeTranscriptToolConfig()):
        """
        Initializes the YouTubeTranscriptTool.

        Args:
            config (YouTubeTranscriptToolConfig): Configuration for the tool.
        """
        super().__init__(config)
        self.languages = config.languages

    def run(self, params: YouTubeTranscriptToolInputSchema) -> YouTubeTranscriptToolOutputSchema:
        """
        Runs the YouTubeTranscriptTool with the given parameters.

        Args:
            params (YouTubeTranscriptToolInputSchema): The input parameters for the tool, adhering to the input schema.

        Returns:
            YouTubeTranscriptToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            Exception: If fetching the transcript fails.
        """
        video_id = self.extract_video_id(params.video_url)
        try:
            if params.language:
                transcripts = YouTubeTranscriptApi.get_transcript(video_id, languages=[params.language])
            else:
                transcripts = YouTubeTranscriptApi.get_transcript(video_id)
        except (NoTranscriptFound, TranscriptsDisabled) as e:
            raise Exception(f"Failed to fetch transcript for video '{video_id}': {str(e)}")

        transcript_text = " ".join([transcript["text"] for transcript in transcripts])
        total_duration = sum([transcript["duration"] for transcript in transcripts])

        metadata = self.fetch_video_metadata(video_id)

        return YouTubeTranscriptToolOutputSchema(
            transcript=transcript_text,
            duration=total_duration,
            comments=[],
            metadata=metadata,
        )

    @staticmethod
    def extract_video_id(url: str) -> str:
        """
        Extracts the video ID from a YouTube URL.

        Args:
            url (str): The YouTube video URL.

        Returns:
            str: The extracted video ID.
        """
        return url.split("v=")[-1].split("&")[0]

    def fetch_video_metadata(self, video_id: str) -> VideoMetadata:
        """
        Fetches metadata for a YouTube video.

        Args:
            video_id (str): The YouTube video ID.

        Returns:
            VideoMetadata: The metadata of the video.

        Raises:
            Exception: If no metadata is found for the video.
        """
        api_key = os.getenv("YOUTUBE_API_KEY")
        youtube = build("youtube", "v3", developerKey=api_key)
        request = youtube.videos().list(part="snippet", id=video_id)
        response = request.execute()

        if not response["items"]:
            raise Exception(f"No metadata found for video '{video_id}'")

        video_info = response["items"][0]["snippet"]
        return VideoMetadata(
            id=video_id,
            title=video_info["title"],
            channel=video_info["channelTitle"],
            published_at=datetime.fromisoformat(video_info["publishedAt"].rstrip("Z")),
        )


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console
    from dotenv import load_dotenv

    load_dotenv()

    rich_console = Console()
    search_tool_instance = YouTubeTranscriptTool(config=YouTubeTranscriptToolConfig())

    search_input = YouTubeTranscriptTool.input_schema(video_url="https://www.youtube.com/watch?v=t1e8gqXLbsU", language="en")

    output = search_tool_instance.run(search_input)
    rich_console.print(output)

```


--------------------------------------------------------------------------------
Example: youtube-to-recipe
--------------------------------------------------------------------------------

**View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/youtube-to-recipe

## Documentation

# YouTube Recipe Extractor

This directory contains the YouTube Recipe Extractor example for the Atomic Agents project. This example demonstrates how to extract structured recipe information from cooking videos using the Atomic Agents framework.

## Getting Started

To get started with the YouTube Recipe Extractor:

1. **Clone the main Atomic Agents repository:**

   ```bash
   git clone https://github.com/BrainBlend-AI/atomic-agents
   ```

2. **Navigate to the YouTube Recipe Extractor directory:**

   ```bash
   cd atomic-agents/atomic-examples/youtube-to-recipe
   ```

3. **Install the dependencies using uv:**

   ```bash
   uv sync
   ```

4. **Set up environment variables:**

   Create a `.env` file in the `youtube-to-recipe` directory with the following content:

   ```env
   OPENAI_API_KEY=your_openai_api_key
   YOUTUBE_API_KEY=your_youtube_api_key
   ```

   To get your YouTube API key, follow the instructions in the [YouTube Scraper README](/atomic-forge/tools/youtube_transcript_scraper/README.md).

   Replace `your_openai_api_key` and `your_youtube_api_key` with your actual API keys.

5. **Run the YouTube Recipe Extractor:**

   ```bash
   uv run python youtube_to_recipe/main.py
   ```

## File Explanation

### 1. Agent (`agent.py`)

This module defines the `YouTubeRecipeExtractionAgent`, responsible for extracting structured recipe information from cooking video transcripts. It extracts:

- Recipe name and description
- Ingredients with quantities and units
- Step-by-step cooking instructions
- Required equipment
- Cooking times and temperatures
- Tips and dietary information

### 2. YouTube Transcript Scraper (`tools/youtube_transcript_scraper.py`)

This tool comes from the [Atomic Forge](/atomic-forge/README.md) and handles fetching transcripts and metadata from YouTube cooking videos.

### 3. Main (`main.py`)

The entry point for the YouTube Recipe Extractor application. It orchestrates fetching transcripts, processing them through the agent, and outputting structured recipe information.

## Example Output

The agent extracts recipe information in a structured format including:

- Detailed ingredient lists with measurements
- Step-by-step cooking instructions with timing and temperature
- Required kitchen equipment
- Cooking tips and tricks
- Dietary information and cuisine type
- Preparation and cooking times

## Customization

You can modify the `video_url` variable in `main.py` to extract recipes from different cooking videos. Additionally, you can adjust the agent's configuration in `agent.py` to customize the recipe extraction format or add additional fields to capture more recipe details.

## Contributing

Contributions are welcome! Please fork the repository and submit a pull request with your enhancements or bug fixes.

## License

This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details.


## Source Code


### File: atomic-examples/youtube-to-recipe/pyproject.toml

```toml
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["youtube_to_recipe"]

[project]
name = "youtube-to-recipe"
version = "1.0.0"
description = "Youtube Recipe Extractor example for Atomic Agents"
readme = "README.md"
authors = [
    { name = "Kenny Vaneetvelde", email = "kenny.vaneetvelde@gmail.com" }
]
requires-python = ">=3.12,<3.14"
dependencies = [
    "atomic-agents",
    "openai>=2.0.0,<3.0.0",
    "pydantic>=2.10.3,<3.0.0",
    "google-api-python-client>=2.101.0,<3.0.0",
    "youtube-transcript-api>=1.1.1,<2.0.0",
    "instructor==1.14.5",
    "python-dotenv>=1.0.1,<2.0.0",
]

[tool.uv.sources]
atomic-agents = { workspace = true }

```

### File: atomic-examples/youtube-to-recipe/youtube_to_recipe/agent.py

```python
import instructor
import openai
from pydantic import BaseModel, Field
from typing import List, Optional

from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema
from atomic_agents.context import BaseDynamicContextProvider, SystemPromptGenerator


class YtTranscriptProvider(BaseDynamicContextProvider):
    def __init__(self, title):
        super().__init__(title)
        self.transcript = None
        self.duration = None
        self.metadata = None

    def get_info(self) -> str:
        return f'VIDEO TRANSCRIPT: "{self.transcript}"\n\nDURATION: {self.duration}\n\nMETADATA: {self.metadata}'


class YouTubeRecipeExtractionInputSchema(BaseIOSchema):
    """This schema defines the input schema for the YouTubeRecipeExtractionAgent."""

    video_url: str = Field(..., description="The URL of the YouTube cooking video to analyze")


class Ingredient(BaseModel):
    """Model for recipe ingredients"""

    item: str = Field(..., description="The ingredient name")
    amount: str = Field(..., description="The quantity of the ingredient")
    unit: Optional[str] = Field(None, description="The unit of measurement, if applicable")
    notes: Optional[str] = Field(None, description="Any special notes about the ingredient")


class Step(BaseModel):
    """Model for recipe steps"""

    instruction: str = Field(..., description="The cooking instruction")
    duration: Optional[str] = Field(None, description="Time required for this step, if mentioned")
    temperature: Optional[str] = Field(None, description="Cooking temperature, if applicable")
    tips: Optional[str] = Field(None, description="Any tips or warnings for this step")


class YouTubeRecipeExtractionOutputSchema(BaseIOSchema):
    """This schema defines the structured recipe information extracted from the video."""

    recipe_name: str = Field(..., description="The name of the recipe being prepared")
    chef: Optional[str] = Field(None, description="The name of the chef/cook presenting the recipe")
    description: str = Field(..., description="A brief description of the dish and its characteristics")
    prep_time: Optional[str] = Field(None, description="Total preparation time mentioned in the video")
    cook_time: Optional[str] = Field(None, description="Total cooking time mentioned in the video")
    servings: Optional[int] = Field(None, description="Number of servings this recipe makes")
    ingredients: List[Ingredient] = Field(..., description="List of ingredients with their quantities and units")
    steps: List[Step] = Field(..., description="Detailed step-by-step cooking instructions")
    equipment: List[str] = Field(..., description="List of kitchen equipment and tools needed")
    tips: List[str] = Field(..., description="General cooking tips and tricks mentioned in the video")
    difficulty_level: Optional[str] = Field(None, description="Difficulty level of the recipe (e.g., Easy, Medium, Hard)")
    cuisine_type: Optional[str] = Field(None, description="Type of cuisine (e.g., Italian, Mexican, Japanese)")
    dietary_info: List[str] = Field(
        default_factory=list, description="Dietary information (e.g., Vegetarian, Vegan, Gluten-free)"
    )


transcript_provider = YtTranscriptProvider(title="YouTube Recipe Transcript")

youtube_recipe_extraction_agent = AtomicAgent[YouTubeRecipeExtractionInputSchema, YouTubeRecipeExtractionOutputSchema](
    config=AgentConfig(
        client=instructor.from_openai(openai.OpenAI()),
        model="gpt-5-mini",
        model_api_parameters={"reasoning_effort": "low"},
        system_prompt_generator=SystemPromptGenerator(
            background=[
                "This Assistant is an expert at extracting detailed recipe information from cooking video transcripts.",
                "It understands cooking terminology, measurements, and techniques.",
            ],
            steps=[
                "Analyze the cooking video transcript thoroughly to extract recipe details.",
                "Convert approximate measurements and instructions into precise recipe format.",
                "Identify all ingredients, steps, equipment, and cooking tips mentioned.",
                "Ensure all critical recipe information is captured accurately.",
            ],
            output_instructions=[
                "Only output Markdown-compatible strings.",
                "Maintain proper units and measurements in recipe format.",
                "Include all safety warnings and important cooking notes.",
            ],
            context_providers={"yt_transcript": transcript_provider},
        ),
    )
)

```

### File: atomic-examples/youtube-to-recipe/youtube_to_recipe/main.py

```python
import os
from dotenv import load_dotenv
from rich.console import Console

from youtube_to_recipe.tools.youtube_transcript_scraper import (
    YouTubeTranscriptTool,
    YouTubeTranscriptToolConfig,
    YouTubeTranscriptToolInputSchema,
)

from youtube_to_recipe.agent import YouTubeRecipeExtractionInputSchema, youtube_recipe_extraction_agent, transcript_provider

load_dotenv()

# Initialize a Rich Console for pretty console outputs
console = Console()

# Initialize the YouTubeTranscriptTool
transcript_tool = YouTubeTranscriptTool(config=YouTubeTranscriptToolConfig(api_key=os.getenv("YOUTUBE_API_KEY")))

# Remove the infinite loop and perform a one-time transcript extraction
video_url = "https://www.youtube.com/watch?v=kUymAc9Oldk"

transcript_input = YouTubeTranscriptToolInputSchema(video_url=video_url, language="en")
try:
    transcript_output = transcript_tool.run(transcript_input)
    console.print(f"[bold green]Transcript:[/bold green] {transcript_output.transcript}")
    console.print(f"[bold green]Duration:[/bold green] {transcript_output.duration} seconds")

    # Update transcript_provider with the scraped transcript data
    transcript_provider.transcript = transcript_output.transcript
    transcript_provider.duration = transcript_output.duration
    transcript_provider.metadata = transcript_output.metadata  # Assuming metadata is available in transcript_output

    # Run the transcript through the agent
    transcript_input_schema = YouTubeRecipeExtractionInputSchema(video_url=video_url)
    agent_response = youtube_recipe_extraction_agent.run(transcript_input_schema)

    # Print the output schema in a formatted way
    console.print("[bold blue]Agent Output Schema:[/bold blue]")
    console.print(agent_response)
except Exception as e:
    console.print(f"[bold red]Error:[/bold red] {str(e)}")

```

### File: atomic-examples/youtube-to-recipe/youtube_to_recipe/tools/youtube_transcript_scraper.py

```python
import os
from typing import List, Optional

from pydantic import Field, BaseModel
from datetime import datetime

from googleapiclient.discovery import build
from youtube_transcript_api import (
    NoTranscriptFound,
    TranscriptsDisabled,
    YouTubeTranscriptApi,
)
from atomic_agents import BaseIOSchema, BaseTool, BaseToolConfig


################
# INPUT SCHEMA #
################
class YouTubeTranscriptToolInputSchema(BaseIOSchema):
    """
    Tool for fetching the transcript of a YouTube video using the YouTube Transcript API.
    Returns the transcript with text, start time, and duration.
    """

    video_url: str = Field(..., description="URL of the YouTube video to fetch the transcript for.")
    language: Optional[str] = Field(None, description="Language code for the transcript (e.g., 'en' for English).")


#################
# OUTPUT SCHEMA #
#################
class VideoMetadata(BaseModel):
    """Schema for YouTube video metadata."""

    id: str = Field(..., description="The YouTube video ID.")
    title: str = Field(..., description="The title of the YouTube video.")
    channel: str = Field(..., description="The name of the YouTube channel.")
    published_at: datetime = Field(..., description="The publication date and time of the video.")


class YouTubeTranscriptToolOutputSchema(BaseIOSchema):
    """
    Output schema for the YouTubeTranscriptTool. Contains the transcript text, duration, comments, and metadata.
    """

    transcript: str = Field(..., description="Transcript of the YouTube video.")
    duration: float = Field(..., description="Duration of the YouTube video in seconds.")
    comments: List[str] = Field(default_factory=list, description="Comments on the YouTube video.")
    metadata: VideoMetadata = Field(..., description="Metadata of the YouTube video.")


#################
# CONFIGURATION #
#################
class YouTubeTranscriptToolConfig(BaseToolConfig):
    """
    Configuration for the YouTubeTranscriptTool.

    Attributes:
        languages (List[str]): List of language codes to try when fetching transcripts.
    """

    languages: List[str] = ["en", "en-US", "en-GB"]


#####################
# MAIN TOOL & LOGIC #
#####################
class YouTubeTranscriptTool(BaseTool[YouTubeTranscriptToolInputSchema, YouTubeTranscriptToolOutputSchema]):
    """
    Tool for extracting transcripts from YouTube videos.

    Attributes:
        input_schema (YouTubeTranscriptToolInputSchema): The schema for the input data.
        output_schema (YouTubeTranscriptToolOutputSchema): The schema for the output data.
        languages (List[str]): List of language codes to try when fetching transcripts.
    """

    def __init__(self, config: YouTubeTranscriptToolConfig = YouTubeTranscriptToolConfig()):
        """
        Initializes the YouTubeTranscriptTool.

        Args:
            config (YouTubeTranscriptToolConfig): Configuration for the tool.
        """
        super().__init__(config)
        self.languages = config.languages

    def run(self, params: YouTubeTranscriptToolInputSchema) -> YouTubeTranscriptToolOutputSchema:
        """
        Runs the YouTubeTranscriptTool with the given parameters.

        Args:
            params (YouTubeTranscriptToolInputSchema): The input parameters for the tool, adhering to the input schema.

        Returns:
            YouTubeTranscriptToolOutputSchema: The output of the tool, adhering to the output schema.

        Raises:
            Exception: If fetching the transcript fails.
        """
        video_id = self.extract_video_id(params.video_url)
        try:
            if params.language:
                transcripts = YouTubeTranscriptApi.get_transcript(video_id, languages=[params.language])
            else:
                transcripts = YouTubeTranscriptApi.get_transcript(video_id)
        except (NoTranscriptFound, TranscriptsDisabled) as e:
            raise Exception(f"Failed to fetch transcript for video '{video_id}': {str(e)}")

        transcript_text = " ".join([transcript["text"] for transcript in transcripts])
        total_duration = sum([transcript["duration"] for transcript in transcripts])

        metadata = self.fetch_video_metadata(video_id)

        return YouTubeTranscriptToolOutputSchema(
            transcript=transcript_text,
            duration=total_duration,
            comments=[],
            metadata=metadata,
        )

    @staticmethod
    def extract_video_id(url: str) -> str:
        """
        Extracts the video ID from a YouTube URL.

        Args:
            url (str): The YouTube video URL.

        Returns:
            str: The extracted video ID.
        """
        return url.split("v=")[-1].split("&")[0]

    def fetch_video_metadata(self, video_id: str) -> VideoMetadata:
        """
        Fetches metadata for a YouTube video.

        Args:
            video_id (str): The YouTube video ID.

        Returns:
            VideoMetadata: The metadata of the video.

        Raises:
            Exception: If no metadata is found for the video.
        """
        api_key = os.getenv("YOUTUBE_API_KEY")
        youtube = build("youtube", "v3", developerKey=api_key)
        request = youtube.videos().list(part="snippet", id=video_id)
        response = request.execute()

        if not response["items"]:
            raise Exception(f"No metadata found for video '{video_id}'")

        video_info = response["items"][0]["snippet"]
        return VideoMetadata(
            id=video_id,
            title=video_info["title"],
            channel=video_info["channelTitle"],
            published_at=datetime.fromisoformat(video_info["publishedAt"].rstrip("Z")),
        )


#################
# EXAMPLE USAGE #
#################
if __name__ == "__main__":
    from rich.console import Console
    from dotenv import load_dotenv

    load_dotenv()

    rich_console = Console()
    search_tool_instance = YouTubeTranscriptTool(config=YouTubeTranscriptToolConfig())

    search_input = YouTubeTranscriptTool.input_schema(video_url="https://www.youtube.com/watch?v=t1e8gqXLbsU", language="en")

    output = search_tool_instance.run(search_input)
    rich_console.print(output)

```


================================================================================
END OF DOCUMENT
================================================================================