Anthropic Structured Outputs with AG2#

import os

from pydantic import BaseModel

import autogen

# Ensure you have your Anthropic API key set
# os.environ["ANTHROPIC_API_KEY"] = "your-api-key-here"

# Verify the API key is set
assert os.getenv("ANTHROPIC_API_KEY"), "Please set ANTHROPIC_API_KEY environment variable"
print("✅ Environment configured")

# Define the structured output schema using Pydantic
class Step(BaseModel):
    """A single step in mathematical reasoning."""

    explanation: str
    output: str

class MathReasoning(BaseModel):
    """Structured output for mathematical problem solving."""

    steps: list[Step]
    final_answer: str

    def format(self) -> str:
        """Format the response for display."""
        steps_output = "\n".join(
            f"Step {i + 1}: {step.explanation}\n  Output: {step.output}" for i, step in enumerate(self.steps)
        )
        return f"{steps_output}\n\nFinal Answer: {self.final_answer}"

# Configure LLM with structured output
llm_config = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.environ["ANTHROPIC_API_KEY"],
            "api_type": "anthropic",
            "response_format": MathReasoning,  # Enable structured outputs
        }
    ],
}

# Create agents
user_proxy = autogen.UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=0,
    code_execution_config=False,
)

math_assistant = autogen.AssistantAgent(
    name="MathAssistant",
    system_message="You are a math tutor. Solve problems step by step.",
    llm_config=llm_config,
)

print("✅ Example 1 configured: Math reasoning with structured outputs")

# Ask the assistant to solve a math problem
chat_result = user_proxy.initiate_chat(
    math_assistant,
    message="Solve the equation: 3x + 7 = 22",
    max_turns=1,
)

# The response is automatically formatted using the format() method
print("\n" + "=" * 60)
print("STRUCTURED OUTPUT RESULT:")
print("=" * 60)
print(chat_result.chat_history[-1]["content"])

# Define a tool function
def get_weather(location: str, unit: str = "celsius") -> str:
    """Get the weather for a location.

    Args:
        location: The city and state, e.g. San Francisco, CA
        unit: Temperature unit (celsius or fahrenheit)
    """
    # In a real application, this would call a weather API
    return f"Weather in {location}: 22°{unit.upper()[0]}, partly cloudy"

# Configure LLM with strict tool
llm_config_strict = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.environ["ANTHROPIC_API_KEY"],
            "api_type": "anthropic",
        }
    ],
    "functions": [
        {
            "name": "get_weather",
            "description": "Get the weather for a location",
            "strict": True,  # Enable strict schema validation ✨
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"},
                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature unit"},
                },
                "required": ["location"],
            },
        }
    ],
}

# Create agents
weather_assistant = autogen.AssistantAgent(
    name="WeatherAssistant",
    system_message="You help users get weather information. Use the get_weather function.",
    llm_config=llm_config_strict,
)

user_proxy_2 = autogen.UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=1,
    code_execution_config=False,
)

# Register function on both agents
# Assistant needs it for LLM awareness, UserProxy executes it
weather_assistant.register_function({"get_weather": get_weather})
user_proxy_2.register_function({"get_weather": get_weather})

print("✅ Example 2 configured: Strict tool use for weather queries")

# Query the weather
chat_result = user_proxy_2.initiate_chat(
    weather_assistant,
    message="What's the weather in Boston, MA?",
    max_turns=2,
)

# Verify tool call had strict typing
print("\n" + "=" * 60)
print("TOOL CALL VERIFICATION:")
print("=" * 60)

import json

for message in chat_result.chat_history:
    if message.get("tool_calls"):
        tool_call = message["tool_calls"][0]
        args = json.loads(tool_call["function"]["arguments"])
        print(f"Function: {tool_call['function']['name']}")
        print(f"Arguments: {args}")
        print(f"✅ location type: {type(args['location']).__name__}")
        if "unit" in args:
            print(f"✅ unit value: {args['unit']} (valid enum)")
        break

# Define calculator tool
def calculate(operation: str, a: float, b: float) -> float:
    """Perform a calculation.

    Args:
        operation: The operation to perform (add, subtract, multiply, divide)
        a: First number
        b: Second number
    """
    if operation == "add":
        return a + b
    elif operation == "subtract":
        return a - b
    elif operation == "multiply":
        return a * b
    elif operation == "divide":
        return a / b if b != 0 else 0
    return 0

# Result model for structured output
class CalculationResult(BaseModel):
    """Structured output for calculation results."""

    problem: str
    steps: list[str]
    result: float
    verification: str

# Configure with BOTH features
llm_config_combined = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.environ["ANTHROPIC_API_KEY"],
            "api_type": "anthropic",
            "response_format": CalculationResult,  # 1. Structured JSON output
        }
    ],
    "functions": [
        {  # 2. Strict tool validation
            "name": "calculate",
            "description": "Perform arithmetic calculation",
            "strict": True,  # Enable strict mode
            "parameters": {
                "type": "object",
                "properties": {
                    "operation": {"type": "string", "enum": ["add", "subtract", "multiply", "divide"]},
                    "a": {"type": "number"},
                    "b": {"type": "number"},
                },
                "required": ["operation", "a", "b"],
            },
        }
    ],
}

# Create agents
calc_assistant = autogen.AssistantAgent(
    name="MathAssistant",
    system_message="You solve math problems using tools and provide structured results.",
    llm_config=llm_config_combined,
)

user_proxy_3 = autogen.UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=5,
    code_execution_config=False,
)

# Register function on both agents
calc_assistant.register_function({"calculate": calculate})
user_proxy_3.register_function({"calculate": calculate})

print("✅ Example 3 configured: Combined strict tools + structured output")

# Ask for a calculation with explanation
chat_result = user_proxy_3.initiate_chat(
    calc_assistant,
    message="Calculate (15 + 7) * 3 and explain your steps",
    max_turns=6,
)

print("\n" + "=" * 60)
print("COMBINED FEATURES RESULT:")
print("=" * 60)

# Check for tool calls (strict validation)
found_tool_call = False
found_structured_output = False

for message in chat_result.chat_history:
    if message.get("tool_calls"):
        found_tool_call = True
        tool_call = message["tool_calls"][0]
        args = json.loads(tool_call["function"]["arguments"])
        print(f"\n✅ Tool Call: {tool_call['function']['name']}")
        print(f"   Arguments: {args}")
        print(f"   Types verified: a={type(args['a']).__name__}, b={type(args['b']).__name__}")

    # Check for structured output
    if message.get("role") == "assistant" and message.get("content"):
        result = CalculationResult.model_validate_json(message["content"])
        found_structured_output = True
        print("\n✅ Structured Output:")
        print(f"   Problem: {result.problem}")
        print(f"   Steps: {len(result.steps)} steps")
        print(f"   Result: {result.result}")
        print(f"   Verification: {result.verification}")

print(f"\n{'=' * 60}")
print(
    f"Features used: Tool Calls={'✅' if found_tool_call else '❌'} | Structured Output={'✅' if found_structured_output else '❌'}"
)

# Define structured output for research report
class ResearchFinding(BaseModel):
    """A single research finding."""

    category: str
    finding: str
    confidence: str  # high, medium, low

class ResearchReport(BaseModel):
    """Structured output for collaborative research."""

    topic: str
    summary: str
    findings: list[ResearchFinding]
    recommendations: list[str]
    contributors: list[str]

    def format(self) -> str:
        """Format the research report for display."""
        output = f"# Research Report: {self.topic}\n\n"
        output += f"## Summary\n{self.summary}\n\n"
        output += f"## Findings ({len(self.findings)} total)\n"
        for i, finding in enumerate(self.findings):
            output += f"{i + 1}. [{finding.category}] {finding.finding} (Confidence: {finding.confidence})\n"
        output += "\n## Recommendations\n"
        for i, rec in enumerate(self.recommendations):
            output += f"{i + 1}. {rec}\n"
        output += f"\n## Contributors: {', '.join(self.contributors)}"
        return output

# Define a research tool
def search_literature(query: str, field: str) -> str:
    """Search academic literature for a query in a specific field.

    Args:
        query: The search query
        field: The field to search (computer_science, biology, physics)
    """
    # Simulated literature search results
    results = {
        "computer_science": "Recent advances in LLMs show 40% improvement in reasoning tasks.",
        "biology": "Studies indicate protein folding accuracy increased by 35% with AI models.",
        "physics": "Quantum computing simulations demonstrate 50x speedup on specific problems.",
    }
    return results.get(field, "No results found for this field.")

# Configure LLM for group agents with strict tools
llm_config_group = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.environ["ANTHROPIC_API_KEY"],
            "api_type": "anthropic",
        }
    ],
    "functions": [
        {
            "name": "search_literature",
            "description": "Search academic literature in a specific field",
            "strict": True,
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "The search query"},
                    "field": {
                        "type": "string",
                        "enum": ["computer_science", "biology", "physics"],
                        "description": "The academic field",
                    },
                },
                "required": ["query", "field"],
            },
        }
    ],
}

# Configure LLM for report writer with structured output
llm_config_report = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.environ["ANTHROPIC_API_KEY"],
            "api_type": "anthropic",
            "response_format": ResearchReport,  # Structured output for final report
        }
    ],
}

# Create specialized research agents
cs_researcher = autogen.AssistantAgent(
    name="CS_Researcher",
    system_message="You are a computer science researcher. Analyze AI and ML topics. Use search_literature when needed.",
    llm_config=llm_config_group,
)

bio_researcher = autogen.AssistantAgent(
    name="Bio_Researcher",
    system_message="You are a biology researcher. Analyze biological and medical topics. Use search_literature when needed.",
    llm_config=llm_config_group,
)

report_writer = autogen.AssistantAgent(
    name="Report_Writer",
    system_message="You synthesize research from other agents into comprehensive structured reports. Wait for all researchers to contribute before writing the final report.",
    llm_config=llm_config_report,
)

# Register search function on all researchers
for agent in [cs_researcher, bio_researcher]:
    agent.register_function({"search_literature": search_literature})

print("✅ Example 4 configured: GroupChat with AutoPattern and structured outputs")

# Import AutoPattern for intelligent speaker selection
from autogen.agentchat.group.multi_agent_chat import initiate_group_chat
from autogen.agentchat.group.patterns import AutoPattern

llm_config_manager = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.environ["ANTHROPIC_API_KEY"],
            "api_type": "anthropic",
        }
    ],
}

# Initialize pattern with agents - AutoPattern uses agents' existing llm_config
pattern = AutoPattern(
    initial_agent=cs_researcher,
    agents=[cs_researcher, bio_researcher, report_writer],
    group_manager_args={"llm_config": llm_config_manager},
)

# Create initial research task
research_task = """
Analyze the impact of AI on scientific research across different fields.
Each researcher should contribute findings from their domain, then the report writer
should create a comprehensive structured report with all findings and recommendations.
"""

# Initiate group chat
chat_result, context_variables, last_agent = initiate_group_chat(
    pattern=pattern,
    messages=research_task,
    max_rounds=8,
)

print("\n" + "=" * 60)
print("GROUPCHAT WITH STRUCTURED OUTPUT:")
print("=" * 60)
print(f"\nTotal messages: {len(chat_result.chat_history)}")
print(f"Last agent: {last_agent.name}")

# Display the structured research report
for message in chat_result.chat_history:
    if message.get("name") == "Report_Writer" and message.get("content"):
        # Try to parse as ResearchReport
        report = ResearchReport.model_validate_json(message["content"])
        print(f"\n{report.format()}")
        print(f"\n✅ Structured report generated with {len(report.findings)} findings")
        break