Multi-Agent Systems

from pydantic_ai import Agent, RunContext

# Specialist agent - does ONE thing well
calculator_agent = Agent(
    'openai:gpt-4o-mini',  # Cheaper model for simple task
    instructions='You are a calculator. Return only the numeric result.',
)

# Main agent - coordinates everything
main_agent = Agent(
    'openai:gpt-4o',
    instructions='You help with various tasks. Use the calculator for math.'
)

@main_agent.tool
async def calculate(ctx: RunContext[None], expression: str) -> str:
    """Calculate a mathematical expression."""
    result = await calculator_agent.run(
        f'Calculate: {expression}',
        usage=ctx.usage  # Track combined token usage
    )
    return result.output

# Use the system
result = main_agent.run_sync('What is 15% of 250?')
print(result.output)
# Output: 15% of 250 is 37.5

from pydantic_ai import Agent, RunContext

# Agent 1: Researcher - finds information
research_agent = Agent(
    'openai:gpt-4o',
    instructions='''You are a research specialist.
    Find 2-3 key facts about the given topic.
    Be thorough but concise.'''
)

@research_agent.tool_plain
def web_search(query: str) -> str:
    """Search the web for information."""
    return f"Search results for '{query}': [relevant information here]"

# Agent 2: Summarizer - creates summaries
summarizer_agent = Agent(
    'openai:gpt-4o-mini',
    instructions='''You create concise summaries.
    Format as 3-5 bullet points.
    Keep it simple and clear.'''
)

# Coordinator agent
coordinator = Agent(
    'openai:gpt-4o',
    instructions='Coordinate research and summarization tasks.'
)

@coordinator.tool
async def research_topic(ctx: RunContext[None], topic: str) -> str:
    """Research a topic thoroughly."""
    result = await research_agent.run(f'Research: {topic}', usage=ctx.usage)
    return result.output

@coordinator.tool
async def summarize_text(ctx: RunContext[None], text: str) -> str:
    """Summarize the given text."""
    result = await summarizer_agent.run(f'Summarize:\n{text}', usage=ctx.usage)
    return result.output

# Use it
result = coordinator.run_sync('Research AI agents and give me a summary')
print(result.output)

from pydantic_ai import Agent

# Three specialist agents
researcher = Agent('openai:gpt-4o',
    instructions='Research the topic. Provide key facts.')

writer = Agent('openai:gpt-4o',
    instructions='Write a clear article based on research.')

editor = Agent('openai:gpt-4o',
    instructions='Edit for clarity and grammar. Return improved version.')

async def content_pipeline(topic: str) -> str:
    # Step 1: Research
    research_result = await researcher.run(f'Research: {topic}')
    research = research_result.output

    # Step 2: Write (needs research)
    article_result = await writer.run(
        f'Write an article based on this research:\n{research}'
    )
    article = article_result.output

    # Step 3: Edit (needs draft)
    final_result = await editor.run(f'Edit this article:\n{article}')

    return final_result.output

# Use it
import asyncio
article = asyncio.run(content_pipeline('The future of AI'))
print(article)

from pydantic_ai import Agent
import asyncio

# Three analysts
market_analyst = Agent('openai:gpt-4o',
    instructions='Analyze market trends. Be specific with numbers.')

tech_analyst = Agent('openai:gpt-4o',
    instructions='Analyze technical aspects. Focus on innovation.')

risk_analyst = Agent('openai:gpt-4o',
    instructions='Identify potential risks. Be thorough.')

async def parallel_analysis(company: str) -> dict:
    # Run all three IN PARALLEL using asyncio.gather
    results = await asyncio.gather(
        market_analyst.run(f'Market analysis for {company}'),
        tech_analyst.run(f'Technical analysis for {company}'),
        risk_analyst.run(f'Risk analysis for {company}')
    )

    return {
        'market': results[0].output,
        'technology': results[1].output,
        'risks': results[2].output
    }

# Much faster than running one by one
analysis = asyncio.run(parallel_analysis('Tesla'))
print(analysis)

from pydantic import BaseModel
from pydantic_ai import Agent

class Review(BaseModel):
    approved: bool
    feedback: str

writer = Agent('openai:gpt-4o',
    instructions='Write or improve content based on feedback.')

critic = Agent('openai:gpt-4o',
    output_type=Review,
    instructions='''Review the content critically.
    Set approved=True ONLY if it is excellent.
    Otherwise, give specific feedback for improvement.'''
)

async def iterative_writing(topic: str, max_rounds: int = 3) -> str:
    # Initial draft
    result = await writer.run(f'Write about: {topic}')
    content = result.output

    for round_num in range(max_rounds):
        # Get critique
        review_result = await critic.run(f'Review this:\n{content}')
        review = review_result.output

        if review.approved:
            print(f'Approved after {round_num + 1} round(s)')
            return content

        print(f'Round {round_num + 1}: {review.feedback}')

        # Improve based on feedback
        result = await writer.run(
            f'Improve this based on feedback.\n\n'
            f'Content: {content}\n\n'
            f'Feedback: {review.feedback}'
        )
        content = result.output

    print('Max rounds reached')
    return content

# Use it
import asyncio
final_content = asyncio.run(iterative_writing('Benefits of meditation'))