pip install ag2[anthrpic]

import os
import textwrap

from dotenv import load_dotenv
from pydantic import BaseModel

from autogen import AssistantAgent, UserProxyAgent
from autogen.io.run_response import Cost

load_dotenv()

# Helper function to extract total cost from ChatResult.cost dict
def get_total_cost(cost_dict):
    """Extract total cost from ChatResult.cost dict structure."""
    total = 0.0
    for usage_type in cost_dict.values():
        if isinstance(usage_type, dict):
            for model_usage in usage_type.values():
                if isinstance(model_usage, dict) and "cost" in model_usage:
                    total += model_usage["cost"]
    return total

# Helper function to extract cost from run response
def get_total_cost_from_run(run_response_cost):
    """Extract total cost from run response object."""
    if isinstance(run_response_cost, Cost):
        return run_response_cost.usage_including_cached_inference.total_cost
    return 0.0

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 using V2 client
llm_config = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.getenv("ANTHROPIC_API_KEY"),
            "api_type": "anthropic_v2",  # <-- Use V2 client
            "response_format": MathReasoning,  # Enable structured outputs
        }
    ],
}

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

math_assistant = 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.run(
    math_assistant,
    message="Solve the equation: 3x + 7 = 22",
    max_turns=1,
)

chat_result.process()

# 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 using V2 client
llm_config_strict = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.getenv("ANTHROPIC_API_KEY"),
            "api_type": "anthropic_v2",  # <-- Use V2 client
        }
    ],
    "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 = AssistantAgent(
    name="WeatherAssistant",
    system_message="You help users get weather information. Use the get_weather function.",
    llm_config=llm_config_strict,
)

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

# Register function on both agents
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 using V2 client
llm_config_combined = {
    "config_list": [
        {
            "model": "claude-sonnet-4-5",
            "api_key": os.getenv("ANTHROPIC_API_KEY"),
            "api_type": "anthropic_v2",  # <-- Use V2 client
            "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 = AssistantAgent(
    name="MathAssistant",
    system_message="You solve math problems using tools and provide structured results.",
    llm_config=llm_config_combined,
)

user_proxy_3 = 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")

chat_result = user_proxy_3.run(
    calc_assistant,
    message="add 3 and 555",
    max_turns=2,
)
chat_result.process()

# Configure LLM to use V2 client for vision
llm_config_vision = {
    "config_list": [
        {
            "api_type": "anthropic_v2",  # <-- Key: use V2 client architecture
            "model": "claude-3-5-haiku-20241022",  # Vision-capable model
            "api_key": os.getenv("ANTHROPIC_API_KEY"),
        }
    ],
    "temperature": 0.3,
}

# Create vision assistant
vision_assistant = AssistantAgent(
    name="VisionBot",
    llm_config=llm_config_vision,
    system_message=textwrap.dedent("""
        You are an AI assistant with vision capabilities.
        You can analyze images and provide detailed, accurate descriptions.
    """).strip(),
)

# Create user proxy
user_proxy_vision = UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=0,
    code_execution_config=False,
)

# Test image URL
IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3b/BlkStdSchnauzer2.jpg"

print("✅ Vision assistant with V2 client created")

# Formal image input format (recommended)
message_with_image = {
    "role": "user",
    "content": [
        {"type": "text", "text": "Describe this image in one sentence."},
        {"type": "image_url", "image_url": {"url": IMAGE_URL{{ "}}" }},
    ],
}

# Initiate chat with image
chat_result = user_proxy_vision.initiate_chat(
    vision_assistant, message=message_with_image, max_turns=1, summary_method="last_msg"
)

print("\n=== Response ===")
print(chat_result.summary)
print(f"\nCost: ${get_total_cost(chat_result.cost):.4f}")

detailed_message = {
    "role": "user",
    "content": [
        {"type": "text", "text": "Analyze this image in detail. What breed is this dog? What are its characteristics?"},
        {"type": "image_url", "image_url": {"url": IMAGE_URL{{ "}}" }},
    ],
}

chat_result = user_proxy_vision.initiate_chat(
    vision_assistant,
    message=detailed_message,
    max_turns=1,
    clear_history=True,  # Start fresh conversation
)

print(chat_result.summary)
print(f"\nCost: ${get_total_cost(chat_result.cost):.4f}")

llm_config = {
    "config_list": [{
        "api_type": "anthropic_v2",  # <-- That's it!
        "model": "claude-sonnet-4-5",
        "api_key": "...",
        "response_format": YourPydanticModel,  # Optional: for structured outputs
    }]
}

assistant = AssistantAgent(llm_config=llm_config)
# Everything works as before, but with rich UnifiedResponse internally