Atomic Agents - Even More Dynamic Orchestration Example

dynamic_orchestration_v2.py

import os
from typing import Union, List, Type
import instructor
import openai
from pydantic import Field
from rich.console import Console
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig
from atomic_agents.lib.base.base_io_schema import BaseIOSchema
from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator
from atomic_agents.lib.base.base_tool import BaseTool

# Import tools from local directory
from tools import (
    SearxNGSearchTool,
    SearxNGSearchToolConfig,
    CalculatorTool,
    CalculatorToolConfig,
)

console = Console()


# Input schema for the agent
class AgentInputSchema(BaseIOSchema):
    """Schema for the input to the Agent."""
    user_input: str = Field(..., description="The user's input or question")


# Final answer schema
class FinalAnswerSchema(BaseIOSchema):
    """Schema for providing the final answer to the user."""
    response: str = Field(..., description="The final response to the user")


def create_dynamic_output_schema(tools: List[BaseTool]) -> Type[BaseIOSchema]:
    """Create a dynamic output schema that includes all tools plus final answer."""
    # Get all tool input schemas
    tool_schemas = [tool.input_schema for tool in tools]
    
    # Add FinalAnswerSchema to the list
    all_schemas = tool_schemas + [FinalAnswerSchema]
    
    # Create Union type
    if len(all_schemas) == 1:
        union_type = all_schemas[0]
    else:
        union_type = Union[tuple(all_schemas)]
    
    # Create new schema class dynamically
    class DynamicAgentOutputSchema(BaseIOSchema):
        """Schema for agent output - either tool parameters or final answer."""
        action: union_type = Field(
            ..., 
            description="Either tool parameters to execute or final answer to provide"
        )
    
    return DynamicAgentOutputSchema


def execute_tool_by_type(tools: List[BaseTool], action: BaseIOSchema) -> BaseIOSchema:
    """Execute the appropriate tool based on the action type."""
    for tool in tools:
        if isinstance(action, tool.input_schema):
            console.print(f"[cyan]Using {tool.__class__.__name__} (detected by type)[/cyan]")
            return tool.run(action)
    
    raise ValueError(f"No tool found for action type: {type(action)}")


def main():
    # Initialize the client
    client = instructor.from_openai(openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")))
    
    # Create tools list
    tools = [
        SearxNGSearchTool(
            config=SearxNGSearchToolConfig(base_url="http://localhost:8888")
        ),
        CalculatorTool(config=CalculatorToolConfig()),
    ]
    
    # Create dynamic output schema based on tools
    DynamicOutputSchema = create_dynamic_output_schema(tools)
    
    # Create the single agent with type-based selection
    system_prompt = SystemPromptGenerator(
        background=[
            "You are an intelligent AI assistant that can use tools to help answer questions.",
            "You have access to various tools through the schema.",
            "You can either use a tool OR provide a final answer directly."
        ],
        steps=[
            "Analyze the user's input to determine if you need to use a tool",
            "If you need more information, return the appropriate tool schema",
            "If you have enough information (including from previous tool results), return FinalAnswerSchema with your response",
            "The system will automatically detect your choice based on the schema type"
        ],
        output_instructions=[
            "Return ONLY ONE action - either tool parameters OR final answer",
            "Do NOT include any wrapper fields - the action type determines what happens",
            "Make sure to fill in all required fields for your chosen action"
        ]
    )
    
    agent = BaseAgent(
        config=BaseAgentConfig(
            client=client,
            model="gpt-4o-mini",
            system_prompt_generator=system_prompt,
            input_schema=AgentInputSchema,
            output_schema=DynamicOutputSchema
        )
    )
    
    # Main interaction loop
    console.print("[bold green]Single Agent with Type-Based Tool Selection[/bold green]")
    console.print("Ask me anything! I'll use tools as needed and provide answers.")
    console.print("Type 'exit' to quit.\n")
    
    while True:
        user_input = console.input("[bold blue]You:[/bold blue] ")
        
        if user_input.lower() in ['exit', 'quit']:
            console.print("Goodbye!")
            break
        
        try:
            # Initial input from user
            agent_input = AgentInputSchema(user_input=user_input)
            agent_output = agent.run(agent_input)
            
            # Keep processing until we get a final answer
            while not isinstance(agent_output.action, FinalAnswerSchema):
                # Execute the tool
                tool_result = execute_tool_by_type(tools, agent_output.action)
                
                # Add tool result to memory and let agent continue
                # The agent will see the tool result in its memory
                agent.memory.add_message("assistant", agent_output)
                agent.memory.add_message("user", tool_result)
                
                # Run agent again without new input to process the tool result
                agent_output = agent.run()
            
            # We have a final answer
            console.print(f"\n[bold green]Assistant:[/bold green] {agent_output.action.response}\n")
            
        except Exception as e:
            console.print(f"[bold red]Error:[/bold red] {str(e)}\n")


# To add more tools, just add them to the tools list:
# from tools.another_tool import AnotherTool, AnotherToolConfig
# tools.append(AnotherTool(config=AnotherToolConfig()))


if __name__ == "__main__":
    main()