================================================================================ 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 Poetry: ```bash poetry install ``` 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 poetry 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 [tool.poetry] name = "basic-multimodal" version = "1.0.0" description = "Basic Multimodal Quickstart example for Atomic Agents" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} instructor = "==1.9.2" openai = ">=1.35.12,<2.0.0" groq = ">=0.11.0,<1.0.0" mistralai = ">=1.1.0,<2.0.0" anthropic = ">=0.39.0,<1.0.0" [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" ``` -------------------------------------------------------------------------------- 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 Poetry: ```bash poetry install ``` 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 poetry 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.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 [tool.poetry] name = "basic-pdf-analysis" version = "1.0.0" description = "Basic PDF analysis Quickstart example for Atomic Agents" authors = ["Renaud Dufour "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<3.14" atomic-agents = {path = "../..", develop = true} instructor = "==1.9.2" google-genai = ">=1.18.0,<2.0.0" jsonref = ">=1.1.0,<2.0.0" [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" ``` -------------------------------------------------------------------------------- Example: deep-research -------------------------------------------------------------------------------- **View on GitHub:** https://github.com/BrainBlend-AI/atomic-agents/tree/main/atomic-examples/deep-research ## Documentation # Deep Research Agent This directory contains the Deep Research Agent example for the Atomic Agents project. This example demonstrates how to create an intelligent research assistant that performs web searches and provides detailed answers with relevant follow-up questions. ## 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 Poetry:** ```bash poetry install ``` 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 ``` **Important**: To find your SearXNG secret key, check the `secret_key` value in your SearXNG `settings.yml` file, typically located at `/etc/searxng/settings.yml` or in the SearXNG installation directory. 5. **Set up SearXNG:** - Install SearXNG from the [official repository](https://github.com/searxng/searxng) - Default configuration expects SearXNG at `http://localhost:8080` 6. **Run the Deep Research Agent:** ```bash poetry run python deep_research/main.py ``` ## File Explanation ### 1. Agents (`agents/`) - **ChoiceAgent** (`choice_agent.py`): Determines when new searches are needed - **QueryAgent** (`query_agent.py`): Generates optimized search queries - **QuestionAnsweringAgent** (`qa_agent.py`): Processes information and generates responses ### 2. Tools (`tools/`) - **SearXNG Search** (`searxng_search.py`): Performs web searches across multiple engines - **Webpage Scraper** (`webpage_scraper.py`): Extracts and processes web content ### 3. Main (`main.py`) The entry point for the Deep Research Agent application. It orchestrates the research process: - Deciding when to perform new searches - Generating and executing search queries - Processing information and providing answers - Suggesting relevant follow-up questions ## Example Usage The agent can handle various research queries and provides: - Detailed answers based on current web information - Relevant citations and sources - Specific follow-up questions for deeper exploration - Context-aware responses that build on previous interactions ## Contributing Contributions are welcome! Please fork the repository and submit a pull request with your improvements. ## License This project is licensed under the MIT License. See the [LICENSE](../../LICENSE) file for details. ## Source Code ### File: atomic-examples/deep-research/deep_research/agents/choice_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 deep_research.config import ChatConfig class ChoiceAgentInputSchema(BaseIOSchema): """Input schema for the ChoiceAgent.""" user_message: str = Field(..., description="The user's latest message or question") decision_type: str = Field(..., description="Explanation of the type of decision to make") class ChoiceAgentOutputSchema(BaseIOSchema): """Output schema for the ChoiceAgent.""" reasoning: str = Field(..., description="Detailed explanation of the decision-making process") decision: bool = Field(..., description="The final decision based on the analysis") choice_agent = AtomicAgent[ChoiceAgentInputSchema, ChoiceAgentOutputSchema]( AgentConfig( client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)), model=ChatConfig.model, model_api_parameters={"reasoning_effort": ChatConfig.reasoning_effort, "temperature": 0.1}, system_prompt_generator=SystemPromptGenerator( background=[ "You are a decision-making agent that determines whether a new web search is needed to answer the user's question.", "Your primary role is to analyze whether the existing context contains sufficient, up-to-date information to answer the question.", "You must output a clear TRUE/FALSE decision - TRUE if a new search is needed, FALSE if existing context is sufficient.", ], steps=[ "1. Analyze the user's question to determine whether or not an answer warrants a new search", "2. Review the available web search results", "3. Determine if existing information is sufficient and relevant", "4. Make a binary decision: TRUE for new search, FALSE for using existing context", ], output_instructions=[ "Your reasoning must clearly state WHY you need or don't need new information", "If the web search context is empty or irrelevant, always decide TRUE for new search", "If the question is time-sensitive, check the current date to ensure context is recent", "For ambiguous cases, prefer to gather fresh information", "Your decision must match your reasoning - don't contradict yourself", ], ), ) ) if __name__ == "__main__": # Example usage for search decision search_example = choice_agent.run( ChoiceAgentInputSchema(user_message="Who won the nobel prize in physics in 2024?", decision_type="needs_search") ) print(search_example) ``` ### File: atomic-examples/deep-research/deep_research/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 deep_research.config import ChatConfig class QuestionAnsweringAgentInputSchema(BaseIOSchema): """This is the input schema for the QuestionAnsweringAgent.""" question: str = Field(..., description="The question to answer.") class QuestionAnsweringAgentOutputSchema(BaseIOSchema): """This is the output schema for the QuestionAnsweringAgent.""" answer: str = Field(..., description="The answer to the question.") follow_up_questions: list[str] = Field( ..., description=( "Specific questions about the topic that would help the user learn more details about the subject matter. " "For example, if discussing a Nobel Prize winner, suggest questions about their research, impact, or " "related scientific concepts." ), ) question_answering_agent = AtomicAgent[QuestionAnsweringAgentInputSchema, QuestionAnsweringAgentOutputSchema]( 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 question answering agent focused on providing factual information and encouraging deeper topic exploration.", "For general greetings or non-research questions, provide relevant questions about the system's capabilities and research functions.", ], steps=[ "Analyze the question and identify the core topic", "Answer the question using available information", "For topic-specific questions, generate follow-up questions that explore deeper aspects of the same topic", "For general queries about the system, suggest questions about research capabilities and functionality", ], output_instructions=[ "Answer in a direct, informative manner", "NEVER generate generic conversational follow-ups like 'How are you?' or 'What would you like to know?'", "For topic questions, follow-up questions MUST be about specific aspects of that topic", "For system queries, follow-up questions should be about specific research capabilities", "Example good follow-ups for a Nobel Prize question:", "- What specific discoveries led to their Nobel Prize?", "- How has their research influenced their field?", "- What other scientists collaborated on this research?", "Example good follow-ups for system queries:", "- What types of sources do you use for research?", "- How do you verify information accuracy?", "- What are the limitations of your search capabilities?", ], ), model_api_parameters={"temperature": 0.1}, ) ) ``` ### File: atomic-examples/deep-research/deep_research/agents/query_agent.py ```python from deep_research.config import ChatConfig import instructor import openai from pydantic import Field from atomic_agents import BaseIOSchema, AtomicAgent, AgentConfig from atomic_agents.context import SystemPromptGenerator from deep_research.tools.searxng_search import SearXNGSearchToolInputSchema class QueryAgentInputSchema(BaseIOSchema): """This is the input schema for the QueryAgent.""" instruction: str = Field(..., description="A detailed instruction or request to generate search engine queries for.") num_queries: int = Field(..., description="The number of search queries to generate.") query_agent = AtomicAgent[QueryAgentInputSchema, SearXNGSearchToolInputSchema]( 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 search engine query generator with a deep understanding of which" "queries will maximize the number of relevant results." ) ], steps=[ "Analyze the given instruction to identify key concepts and aspects that need to be researched", "For each aspect, craft a search query using appropriate search operators and syntax", "Ensure queries cover different angles of the topic (technical, practical, comparative, etc.)", ], output_instructions=[ "Return exactly the requested number of queries", "Format each query like a search engine query, not a natural language question", "Each query should be a concise string of keywords and operators", ], ), ) ) ``` ### File: atomic-examples/deep-research/deep_research/config.py ```python import os from dataclasses import dataclass from typing import Set 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_searxng_base_url() -> str: """Retrieve SearXNG base URL from environment or use default""" base_url = os.getenv("SEARXNG_BASE_URL", "http://localhost:8080") return base_url def get_searxng_api_key() -> str: """Retrieve SearXNG API key from environment""" api_key = os.getenv("SEARXNG_API_KEY") return api_key @dataclass class ChatConfig: """Configuration for the chat application""" api_key: str = get_api_key() # This becomes a class variable model: str = "gpt-5-mini" reasoning_effort: str = "low" exit_commands: Set[str] = frozenset({"/exit", "/quit"}) searxng_base_url: str = get_searxng_base_url() searxng_api_key: str = get_searxng_api_key() def __init__(self): # Prevent instantiation raise TypeError("ChatConfig is not meant to be instantiated") ``` ### File: atomic-examples/deep-research/deep_research/context_providers.py ```python from dataclasses import dataclass from datetime import datetime from typing import List from atomic_agents.context import BaseDynamicContextProvider @dataclass class ContentItem: content: str url: str class ScrapedContentContextProvider(BaseDynamicContextProvider): def __init__(self, title: str): super().__init__(title=title) self.content_items: List[ContentItem] = [] def get_info(self) -> str: return "\n\n".join( [ f"Source {idx}:\nURL: {item.url}\nContent:\n{item.content}\n{'-' * 80}" for idx, item in enumerate(self.content_items, 1) ] ) class CurrentDateContextProvider(BaseDynamicContextProvider): def __init__(self, title: str, date_format: str = "%A %B %d, %Y"): super().__init__(title=title) self.date_format = date_format def get_info(self) -> str: return f"The current date in the format {self.date_format} is {datetime.now().strftime(self.date_format)}." ``` ### File: atomic-examples/deep-research/deep_research/main.py ```python from deep_research.agents.query_agent import QueryAgentInputSchema, query_agent from deep_research.agents.qa_agent import ( QuestionAnsweringAgentInputSchema, question_answering_agent, QuestionAnsweringAgentOutputSchema, ) from deep_research.agents.choice_agent import choice_agent, ChoiceAgentInputSchema from deep_research.tools.searxng_search import SearXNGSearchTool, SearXNGSearchToolConfig, SearXNGSearchToolInputSchema from deep_research.tools.webpage_scraper import WebpageScraperTool, WebpageScraperToolInputSchema from deep_research.context_providers import ContentItem, CurrentDateContextProvider, ScrapedContentContextProvider 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 console = Console() WELCOME_MESSAGE = ( "Welcome to Deep Research - your AI-powered research assistant! I can help you explore and " "understand any topic through detailed research and interactive discussion." ) STARTER_QUESTIONS = [ "Can you help me research the latest AI news?", "Who won the Nobel Prize in Physics this year?", "Where can I learn more about quantum computing?", ] def perform_search_and_update_context( user_message: str, scraped_content_context_provider: ScrapedContentContextProvider ) -> None: with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), console=console, ) as progress: # Generate search queries task = progress.add_task("[cyan]Generating search queries...", total=None) console.print("\n[bold yellow]🤔 Analyzing your question to generate relevant search queries...[/bold yellow]") query_agent_output = query_agent.run(QueryAgentInputSchema(instruction=user_message, num_queries=3)) progress.remove_task(task) console.print("\n[bold green]🔍 Generated search queries:[/bold green]") for i, query in enumerate(query_agent_output.queries, 1): console.print(f" {i}. [italic]{query}[/italic]") # Perform the search task = progress.add_task("[cyan]Searching the web...", total=None) console.print("\n[bold yellow]🌐 Searching across the web using SearXNG...[/bold yellow]") searxng_search_tool = SearXNGSearchTool(SearXNGSearchToolConfig(base_url="http://localhost:8080/")) search_results = searxng_search_tool.run(SearXNGSearchToolInputSchema(queries=query_agent_output.queries)) progress.remove_task(task) # Scrape content from search results console.print("\n[bold green]📑 Found relevant web pages:[/bold green]") for i, result in enumerate(search_results.results[:3], 1): console.print(f" {i}. [link={result.url}]{result.title}[/link]") task = progress.add_task("[cyan]Scraping webpage content...", total=None) console.print("\n[bold yellow]📥 Extracting content from web pages...[/bold yellow]") webpage_scraper_tool = WebpageScraperTool() results_for_context_provider = [] for result in search_results.results[:3]: scraped_content = webpage_scraper_tool.run(WebpageScraperToolInputSchema(url=result.url, include_links=True)) results_for_context_provider.append(ContentItem(content=scraped_content.content, url=result.url)) progress.remove_task(task) # Update the context provider with new content console.print("\n[bold green]🔄 Updating research context with new information...[/bold green]") scraped_content_context_provider.content_items = results_for_context_provider def initialize_conversation() -> None: initial_answer = QuestionAnsweringAgentOutputSchema( answer=WELCOME_MESSAGE, follow_up_questions=STARTER_QUESTIONS, ) question_answering_agent.history.add_message("assistant", initial_answer) def display_welcome() -> None: welcome_panel = Panel( WELCOME_MESSAGE, title="[bold blue]Deep Research Chat[/bold blue]", border_style="blue", padding=(1, 2) ) console.print("\n") console.print(welcome_panel) # Create a table for starter questions 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_search_status(is_new_search: bool, reasoning: str) -> None: if is_new_search: panel = Panel( f"[white]{reasoning}[/white]", title="[bold yellow]Performing New Search[/bold yellow]", border_style="yellow", padding=(1, 2), ) else: panel = Panel( f"[white]{reasoning}[/white]", title="[bold green]Using Existing Context[/bold green]", border_style="green", padding=(1, 2), ) console.print("\n") console.print(panel) def display_answer(answer: str, follow_up_questions: list[str]) -> None: # Display the main answer in a panel answer_panel = Panel(Markdown(answer), title="[bold blue]Answer[/bold blue]", border_style="blue", padding=(1, 2)) console.print("\n") console.print(answer_panel) # Display follow-up questions if available if follow_up_questions: questions_table = Table( show_header=True, header_style="bold cyan", box=box.ROUNDED, title="[bold]Follow-up Questions[/bold]" ) questions_table.add_column("№", style="dim", width=4) questions_table.add_column("Question", style="green") for i, question in enumerate(follow_up_questions, 1): questions_table.add_row(str(i), question) console.print("\n") console.print(questions_table) def chat_loop() -> None: console.print("\n[bold magenta]🚀 Initializing Deep Research System...[/bold magenta]") # Initialize context providers console.print("[dim]• Creating context providers...[/dim]") scraped_content_context_provider = ScrapedContentContextProvider("Scraped Content") current_date_context_provider = CurrentDateContextProvider("Current Date") # Register context providers console.print("[dim]• Registering context providers with agents...[/dim]") choice_agent.register_context_provider("current_date", current_date_context_provider) question_answering_agent.register_context_provider("current_date", current_date_context_provider) query_agent.register_context_provider("current_date", current_date_context_provider) choice_agent.register_context_provider("scraped_content", scraped_content_context_provider) question_answering_agent.register_context_provider("scraped_content", scraped_content_context_provider) query_agent.register_context_provider("scraped_content", scraped_content_context_provider) console.print("[dim]• Initializing conversation history...[/dim]") initialize_conversation() console.print("[bold green]✨ System initialized successfully![/bold green]\n") display_welcome() while True: 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 Deep Research.[/bold]") break console.print("\n[bold yellow]🤖 Processing your question...[/bold yellow]") # Determine if we need a new search console.print("[dim]• Evaluating if new research is needed...[/dim]") choice_agent_output = choice_agent.run( ChoiceAgentInputSchema( user_message=user_message, decision_type=( "Should we perform a new web search? TRUE if we need new or updated information, FALSE if existing " "context is sufficient. Consider: 1) Is the context empty? 2) Is the existing information relevant? " "3) Is the information recent enough?" ), ) ) # Display search status with new formatting display_search_status(choice_agent_output.decision, choice_agent_output.reasoning) if choice_agent_output.decision: perform_search_and_update_context(user_message, scraped_content_context_provider) # Get and display the answer with new formatting console.print("\n[bold yellow]🎯 Generating comprehensive answer...[/bold yellow]") question_answering_agent_output = question_answering_agent.run( QuestionAnsweringAgentInputSchema(question=user_message) ) display_answer(question_answering_agent_output.answer, question_answering_agent_output.follow_up_questions) if __name__ == "__main__": chat_loop() ``` ### 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 %% Decision Flow Diagram subgraph DecisionFlow["Research Decision Flow"] Start([User Question]) --> B{Need Search?} B -->|Yes| C[Generate Search Queries] C --> D[Perform Web Search] D --> E[Scrape Webpages] E --> F[Update Context] F --> G[Generate Answer] B -->|No| G G --> H[Show Answer & Follow-ups] H --> End([End]) end classDef default fill:#f9f9f9,stroke:#333,stroke-width:2px; classDef decision fill:#ff9800,stroke:#f57c00,color:#fff; classDef process fill:#4caf50,stroke:#388e3c,color:#fff; classDef terminator fill:#9c27b0,stroke:#7b1fa2,color:#fff; class B decision; class C,D,E,F,G process; class Start,End terminator; ``` ```mermaid graph TD %% System Architecture Diagram subgraph Agents["AI Agents"] CA[ChoiceAgent] QA[QueryAgent] AA[AnswerAgent] end subgraph Tools["External Tools"] ST[SearXNG Search] WS[Webpage Scraper] end subgraph Context["Context Providers"] SC[Scraped Content] CD[Current Date] end %% Connections User -->|Question| CA CA -->|Search Request| QA QA -->|Queries| ST ST -->|URLs| WS WS -->|Content| SC SC -.->|Context| CA & QA & AA CD -.->|Date Info| CA & QA & AA CA -->|Direct Answer| AA AA -->|Response| User %% Styling classDef agent fill:#4CAF50,stroke:#2E7D32,color:#fff; classDef tool fill:#FF9800,stroke:#EF6C00,color:#fff; classDef context fill:#F44336,stroke:#C62828,color:#fff; classDef user fill:#9C27B0,stroke:#6A1B9A,color:#fff; class CA,QA,AA agent; class ST,WS tool; class SC,CD context; class User user; ``` ### File: atomic-examples/deep-research/pyproject.toml ```toml [tool.poetry] name = "deep-research" version = "0.1.0" description = "" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} requests = "^2.32.3" beautifulsoup4 = "^4.12.3" markdownify = "^0.13.1" readability-lxml = "^0.8.1" lxml-html-clean = "^0.4.0" lxml = "^5.3.0" python-dotenv = ">=1.0.1,<2.0.0" openai = ">=1.35.12,<2.0.0" trafilatura = "^1.6.3" [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" ``` -------------------------------------------------------------------------------- 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 Poetry: ```bash poetry install ``` 4. Set up your OpenAI API key: ```bash export OPENAI_API_KEY="your-api-key-here" ``` 5. Run the example: ```bash poetry 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 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): chat_message: str = Field(..., description="User's question or message") class AgentResponse(BaseIOSchema): 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): 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)) config = AgentConfig( client=client, model="gpt-5-mini", model_api_parameters={"reasoning_effort": "low"}, history=ChatHistory(), system_prompt=system_prompt, ) 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. You must respond with: - A main answer - Alternative suggestions (list) - Confidence level (exactly 'low', 'medium', or 'high') - Whether follow-up is needed (boolean) Be very strict about the confidence_level field - it must be exactly one of the three allowed values.""", ) 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 [tool.poetry] name = "hooks-example" version = "1.0.0" description = "AtomicAgent hooks system example demonstrating monitoring, error handling, and retry mechanisms" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} instructor = "==1.9.2" openai = ">=1.35.12,<2.0.0" python-dotenv = ">=1.0.1,<2.0.0" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` -------------------------------------------------------------------------------- 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 poetry install ``` 3. Set up the client: ```bash cd ../example-client poetry install ``` 4. Run the example: **Using STDIO transport (default):** ```bash cd example-client poetry run python -m example_client.main --client stdio # or simply: poetry run python -m example_client.main ``` **Using async STDIO transport:** ```bash cd example-client poetry run python -m example_client.main --client stdio_async ``` **Using SSE transport (Deprecated):** ```bash # First terminal: Start the server cd example-mcp-server poetry run -m example_mcp_server.server --mode=sse # Second terminal: Run the client with SSE transport cd example-client poetry run python -m example_client.main --client sse ``` **Using HTTP Stream transport:** ```bash # First terminal: Start the server cd example-mcp-server poetry run python -m example_mcp_server.server --mode=http_stream # Second terminal: Run the client with HTTP Stream transport cd example-client poetry run python -m example_client.main --client http_stream ``` **Using FastAPI client:** ```bash # First terminal: Start the MCP server cd example-mcp-server poetry run python -m example_mcp_server.server --mode=http_stream # Second terminal: Run the FastAPI client cd example-client poetry 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, 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, 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] 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 = {} tool_schema_map: Dict[Type[BaseIOSchema], Type] = {} 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) 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())}") tool_schema_map.update( {ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")} ) available_schemas = tuple(tool_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(" poetry 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() 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]]: """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 = [] try: available_schemas = tuple(tool_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 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.", ], steps=[ "1. Use the reasoning field to determine if one or more successive tool calls could be used to handle the user's query.", "2. If so, choose the appropriate tool(s) one at a time and extract all necessary parameters from the query.", "3. If a single tool 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).", "4. If no sequence of tools could be used, or if you are finished processing the user's query, 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 input or FinalResponseSchema).", "3. Ensure all required parameters for the chosen tool are properly extracted and validated.", "4. Maintain a professional and helpful tone in all responses.", "5. Break down complex queries into sequential tool 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 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)}") tool_class = tool_schema_map.get(type(action_instance)) if not tool_class: print(f"Debug - Error: No tool found for schema {type(action_instance)}") print(f"Debug - Available schemas: {list(tool_schema_map.keys())}") return "I encountered an internal error. Could not find the appropriate tool.", tools_used tool_name = tool_class.mcp_tool_name tools_used.append(tool_name) print(f"Debug - Executing {tool_class.mcp_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 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 else: return "Error: Expected final response but got something else", tools_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 and following the schema structure.""" return { "message": "MCP FastAPI Client Example - Agent-based Architecture", "available_tools": list(mcp_tools.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", } }, "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 = await execute_with_orchestrator_async(request.query) return CalculationResponse(result=result_text, tools_used=tools_used, query=request.query) except Exception as e: raise HTTPException(status_code=500, detail=f"Agent calculation failed: {e}") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000) ``` ### 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, 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) if not tools: console.print(f"[bold red]No MCP tools 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) # 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") } available_schemas = tuple(tool_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 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.", ], steps=[ "1. Use the reasoning field to determine if one or more successive tool calls could be used to handle the user's query.", "2. If so, choose the appropriate tool(s) one at a time and extract all necessary parameters from the query.", "3. If a single tool 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).", "4. If no sequence of tools could be used, or if you are finished processing the user's query, 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 input or FinalResponseSchema).", "3. Ensure all required parameters for the chosen tool are properly extracted and validated.", "4. Maintain a professional and helpful tone in all responses.", "5. Break down complex queries into sequential tool 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): # Find the matching tool class tool_class = tool_schema_map.get(type(action_instance)) if not tool_class: console.print(f"[red]Error: No tool found for schema {type(action_instance)}[/red]") action_instance = FinalResponseSchema( response_text="I encountered an internal error. Could not find the appropriate tool." ) break # Execute the tool console.print(f"[blue]Executing {tool_class.mcp_tool_name}...[/blue]") console.print(f"[dim]Parameters: {action_instance.model_dump()}") tool_instance = tool_class() try: result = tool_instance.run(action_instance) console.print(f"[bold green]Result:[/bold green] {result.result}") # Ask orchestrator what to do next with the 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)) # Debug output for subsequent responses console.print( f"[dim]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"): 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, 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-mini and not gpt-4o-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" 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, ) if not tools: raise RuntimeError("No MCP tools 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()) # Available schemas include all tool input schemas and the final response schema available_schemas = tool_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'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]") # 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) # 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.", ], steps=[ "1. Use the reasoning field to determine if one or more successive tool calls could be used to handle the user's query.", "2. If so, choose the appropriate tool(s) one at a time and extract all necessary parameters from the query.", "3. If a single tool 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).", "4. If no sequence of tools could be used, or if you are finished processing the user's query, 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 input or FinalResponseSchema).", "3. Ensure all required parameters for the chosen tool are properly extracted and validated.", "4. Maintain a professional and helpful tone in all responses.", "5. Break down complex queries into sequential tool 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) ToolClass = tool_schema_to_class_map.get(schema_type) if not ToolClass: 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 type could not be recognized." ) break 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.history.add_message("system", result_message) # 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, 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 = "poetry 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 ) if not tools: raise RuntimeError("No MCP tools 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()) # Available schemas include all tool input schemas and the final response schema available_schemas = tool_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'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) # 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.", ], steps=[ "1. Use the reasoning field to determine if one or more successive tool calls could be used to handle the user's query.", "2. If so, choose the appropriate tool(s) one at a time and extract all necessary parameters from the query.", "3. If a single tool 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).", "4. If no sequence of tools could be used, or if you are finished processing the user's query, 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 input or FinalResponseSchema).", "3. Ensure all required parameters for the chosen tool are properly extracted and validated.", "4. Maintain a professional and helpful tone in all responses.", "5. Break down complex queries into sequential tool 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). 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) ToolClass = tool_schema_to_class_map.get(schema_type) if not ToolClass: raise ValueError(f"Unknown schema type '" f"{schema_type.__name__}" f"' returned by orchestrator") 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.history.add_message("system", result_message) # 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, 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 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-4o-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 = "poetry 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 - factory sees running loop tools = await fetch_mcp_tools_async( transport_type=MCPTransportType.STDIO, client_session=session, # factory sees running loop ) if not tools: raise RuntimeError("No MCP tools 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()) # Available schemas include all tool input schemas and the final response schema available_schemas = tool_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'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) # 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.", ], steps=[ "1. Use the reasoning field to determine if one or more successive tool calls could be used to handle the user's query.", "2. If so, choose the appropriate tool(s) one at a time and extract all necessary parameters from the query.", "3. If a single tool 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).", "4. If no sequence of tools could be used, or if you are finished processing the user's query, 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 input or FinalResponseSchema).", "3. Ensure all required parameters for the chosen tool are properly extracted and validated.", "4. Maintain a professional and helpful tone in all responses.", "5. Break down complex queries into sequential tool 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) ToolClass = tool_schema_to_class_map.get(schema_type) if not ToolClass: raise ValueError(f"Unknown schema type '" f"{schema_type.__name__}" f"' returned by orchestrator") 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.history.add_message("system", result_message) # 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 [tool.poetry] name = "example-client" version = "0.1.0" description = "Example: Choosing the right MCP tool for a user query using the MCP Tool Factory." authors = ["Your Name "] [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = { path = "../../../", develop = true } example-mcp-server = { path = "../example-mcp-server", develop = true } pydantic = ">=2.10.3,<3.0.0" rich = ">=13.0.0" openai = ">=1.0.0" mcp = {extras = ["cli"], version = "^1.9.4"} fastapi = "^0.115.14" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` ### 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 __all__ = ["Tool", "BaseToolInput", "ToolResponse", "ToolContent", "Resource"] ``` ### 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 List, Optional, ClassVar from pydantic import BaseModel, Field class ResourceContent(BaseModel): """Model for content in resource responses.""" type: str = Field(default="text") text: str uri: str mime_type: Optional[str] = None class ResourceResponse(BaseModel): """Model for resource responses.""" contents: List[ResourceContent] class Resource(ABC): """Abstract base class for all resources.""" name: ClassVar[str] description: ClassVar[str] uri: ClassVar[str] mime_type: ClassVar[Optional[str]] = None @abstractmethod async def read(self) -> ResourceResponse: """Read the resource content.""" pass ``` ### 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/resources/__init__.py ```python """Resource exports.""" __all__ = [] ``` ### 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.interfaces.tool import Tool from example_mcp_server.interfaces.resource import Resource from example_mcp_server.tools import ( AddNumbersTool, SubtractNumbersTool, MultiplyNumbersTool, DivideNumbersTool, BatchCalculatorTool, ) 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 [] def create_mcp_server() -> FastMCP: """Create and configure the MCP server.""" mcp = FastMCP("example-mcp-server") tool_service = ToolService() resource_service = ResourceService() # 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) 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.interfaces.tool import Tool from example_mcp_server.interfaces.resource import Resource from example_mcp_server.tools import AddNumbersTool, SubtractNumbersTool, MultiplyNumbersTool, DivideNumbersTool 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 [] 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() # 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) # 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.interfaces.tool import Tool from example_mcp_server.interfaces.resource import Resource # from example_mcp_server.tools import HelloWorldTool # Removed from example_mcp_server.tools import ( # Added imports for new tools AddNumbersTool, SubtractNumbersTool, MultiplyNumbersTool, DivideNumbersTool, ) 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 [ # Add more resources here as you create them ] def main(): """Entry point for the server.""" mcp = FastMCP("example-mcp-server") tool_service = ToolService() resource_service = ResourceService() # 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) 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/resource_service.py ```python """Service layer for managing resources.""" from typing import Dict, List import re 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[^/]+) 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[^/]+) 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.""" return await resource.read() # Set metadata for the handler static_handler.__name__ = resource.name static_handler.__doc__ = resource.description return static_handler else: # For resources with parameters # Define a dynamic function with named parameters matching URI placeholders params_str = ", ".join(uri_params) func_def = f"async def param_handler({params_str}) -> ResourceResponse:\n" func_def += f' """{resource.description}"""\n' func_def += f" return await resource.read({params_str})" # Create namespace for execution namespace = {"resource": resource, "ResourceResponse": ResourceResponse} exec(func_def, namespace) # Get the handler and set its name handler = namespace["param_handler"] handler.__name__ = resource.name return 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 [project] name = "example-mcp-server" version = "0.1.0" description = "example-mcp-server MCP server" authors = [] requires-python = ">=3.12.0" dependencies = [ "mcp[cli]>=1.9.4", "rich>=13.0.0", "pydantic>=2.0.0", "uvicorn>=0.15.0", ] [build-system] requires = ["hatchling"] build-backend = "hatchling.build" [project.scripts] example-mcp-server = "example_mcp_server.server:main" ``` -------------------------------------------------------------------------------- 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 Poetry: ```bash poetry install ``` 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 poetry 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.history.add_message("system", 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 [tool.poetry] name = "orchestration-agent" version = "0.1.0" description = "" authors = ["KennyVaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} instructor = "==1.9.2" pydantic = ">=2.10.3,<3.0.0" sympy = "^1.13.3" python-dotenv = ">=1.0.1,<2.0.0" openai = ">=1.35.12,<2.0.0" [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" ``` -------------------------------------------------------------------------------- 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 Poetry: ```bash poetry install ``` 4. Run the examples using Poetry: ```bash poetry run python quickstart/1_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 poetry run python quickstart/.py ``` Replace `` 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 [tool.poetry] name = "quickstart" version = "1.0.0" description = "Quickstart example for Atomic Agents" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} instructor = "==1.9.2" openai = ">=1.35.12,<2.0.0" groq = ">=0.11.0,<1.0.0" mistralai = ">=1.1.0,<2.0.0" anthropic = ">=0.39.0,<1.0.0" python-dotenv = ">=1.0.1,<2.0.0" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` ### 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 # 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() # 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) # 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"} 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 = {} 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 = {} 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 = {} elif provider == "5" or provider == "gemini": from openai import OpenAI api_key = os.getenv("GEMINI_API_KEY") client = instructor.from_openai( OpenAI(api_key=api_key, base_url="https://generativelanguage.googleapis.com/v1beta/openai/"), mode=instructor.Mode.JSON, ) model = "gpt-5-mini" model_api_parameters = {"reasoning_effort": "low"} 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 = {} else: raise ValueError(f"Unsupported provider: {provider}") return client, model, model_api_parameters # Prompt the user to choose a provider from one in the list below. providers_list = ["openai", "anthropic", "groq", "ollama", "gemini", "openrouter"] 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 = setup_client(provider) # Agent setup with specified configuration agent = AtomicAgent[BasicChatInputSchema, BasicChatOutputSchema]( config=AgentConfig( client=client, model=model, history=history, model_api_parameters={**model_api_parameters, "max_tokens": 2048} ) ) # 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 # 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 Poetry:** ```bash poetry install ``` 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 poetry 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 poetry run python rag_chatbot/main.py ``` ### Using Qdrant (Local) ```bash export VECTOR_DB_TYPE=qdrant poetry 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 [tool.poetry] name = "rag-chatbot" version = "0.1.0" description = "A RAG chatbot example using Atomic Agents and ChromaDB/Qdrant" authors = ["Your Name "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} chromadb = "^1.0.20" qdrant-client = "^1.15.1" numpy = "^2.3.2" python-dotenv = "^1.0.1" openai = "^1.100.2" pulsar-client = "^3.8.0" rich = "^13.7.0" wget = "^3.2" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` ### 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 Poetry: ```bash poetry install ``` 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 poetry 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 [tool.poetry] name = "web-search-agent" version = "1.0.0" description = "Web search agent example for Atomic Agents" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<4.0" atomic-agents = {path = "../..", develop = true} openai = ">=1.35.12,<2.0.0" pydantic = ">=2.9.2,<3.0.0" instructor = "==1.9.2" python-dotenv = ">=1.0.1,<2.0.0" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` ### 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 Poetry:** ```bash poetry install ``` 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 poetry run python youtube_summarizer/main.py ``` or ```bash poetry 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 [tool.poetry] name = "youtube-summarizer" version = "1.0.0" description = "Youtube Summarizer example for Atomic Agents" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<3.14" atomic-agents = {path = "../..", develop = true} openai = ">=1.35.12,<2.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" instructor = "==1.9.2" python-dotenv = ">=1.0.1,<2.0.0" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` ### 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 Poetry:** ```bash poetry install ``` 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 poetry 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 [tool.poetry] name = "youtube-to-recipe" version = "1.0.0" description = "Youtube Recipe Extractor example for Atomic Agents" authors = ["Kenny Vaneetvelde "] readme = "README.md" [tool.poetry.dependencies] python = ">=3.12,<3.14" atomic-agents = {path = "../..", develop = true} openai = ">=1.35.12,<2.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" instructor = "==1.9.2" python-dotenv = ">=1.0.1,<2.0.0" [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" ``` ### 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 ================================================================================