#!/usr/bin/env python3
"""Test custom tools with LLM."""

from pygeai import Client

def test_tool_calling():
    """Test tool calling workflow."""
    
    client = Client()
    
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get current weather for a location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "City name"},
                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
                    },
                    "required": ["location"]
                }
            }
        }
    ]
    
    messages = [
        {"role": "user", "content": "What's the weather in San Francisco?"}
    ]
    
    response = client.llm.chat(
        model="gpt-4o",
        messages=messages,
        tools=tools,
        tool_choice="auto"
    )
    
    if response.tool_calls:
        for tool_call in response.tool_calls:
            print(f"Function: {tool_call.function.name}")
            print(f"Arguments: {tool_call.function.arguments}")
            
            # Execute function
            result = execute_function(tool_call.function.name, tool_call.function.arguments)
            
            # Send result back
            messages.append(response)
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": result
            })
            
            final_response = client.llm.chat(model="gpt-4o", messages=messages)
            print(f"Final answer: {final_response.content}")

def execute_function(name, arguments):
    """Execute function and return result."""
    import json
    args = json.loads(arguments)
    
    if name == "get_weather":
        # Mock implementation
        return json.dumps({
            "location": args["location"],
            "temperature": 72,
            "condition": "sunny",
            "unit": args.get("unit", "fahrenheit")
        })

if __name__ == "__main__":
    test_tool_calling()
