import os
from dotenv import load_dotenv
from openai import AsyncOpenAI
import asyncio
import json

load_dotenv()

async def main():
    api_key = os.environ.get("GROQ_API_KEY")
    client = AsyncOpenAI(
        api_key=api_key,
        base_url="https://api.groq.com/openai/v1"
    )

    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get the current weather in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA",
                        },
                    },
                    "required": ["location"],
                },
            },
        }
    ]

    print("Testing Tool Call with openai/gpt-oss-20b...")
    
    try:
        response = await client.chat.completions.create(
            model="openai/gpt-oss-20b",
            messages=[
                {"role": "user", "content": "What's the weather in San Francisco?"}
            ],
            tools=tools,
            tool_choice="auto"
        )
        
        message = response.choices[0].message
        print(f"Initial Response Role: {message.role}")
        
        if message.tool_calls:
            print(f"Tool Calls: {len(message.tool_calls)}")
            for tc in message.tool_calls:
                print(f"  - Function: {tc.function.name}")
                print(f"  - Args: {tc.function.arguments}")
        else:
            print("No tool calls triggered.")

    except Exception as e:
        print(f"Failed during tool call test: {e}")

if __name__ == "__main__":
    asyncio.run(main())