import asyncio
import os
from dotenv import load_dotenv
from apexagents import Agent, tool

load_dotenv()

@tool
def calculate_sum(a: int, b: int) -> int:
    """
    Task: Calculates the sum of two integers.
    
    Input Parameters:
        a (int): The first integer.
        b (int): The second integer.
        
    Output Parameters / Returns:
        int: The resulting sum of the two integers.
        
    Raises: None
    """
    print(f">> Executing Tool: calculate_sum({a}, {b})")
    return a + b

@tool
def get_weather(location: str) -> str:
    """
    Task: Gets the simulated weather for a given location.
    
    Input Parameters:
        location (str): The name of the city or location.
        
    Output Parameters / Returns:
        str: A string describing the current weather condition.
        
    Raises: None
    """
    print(f">> Executing Tool: get_weather('{location}')")
    return f"The weather in {location} is sunny and 75 degrees."

async def main():
    """
    Task: Initializes a basic Agent with two tools, running an asynchronous loop to verify conversational tool usage.
    
    Input Parameters: None
        
    Output Parameters / Returns: None
    
    Raises: None
    """
    agent = Agent(
        name="TestAgent",
        model=f"azure/{os.environ.get('DEPLOYMENT_NAME', 'gpt-4o-mini')}",
        instructions="You are a helpful assistant. Use tools when necessary to answer the user's questions.",
        tools=[calculate_sum, get_weather]
    )
    
    print("Agent prompt 1: What is the sum of 1500 and 3200?")
    response = await agent.arun("What is the sum of 1500 and 3200?")
    print("Agent Response:", response)
    
    print("\n-------------------\n")
    
    print("Agent prompt 2: What is the weather in Seattle?")
    response = await agent.arun("What is the weather in Seattle?")
    print("Agent Response:", response)

if __name__ == "__main__":
    # Ensure OPENAI_API_KEY is set in environment, or use another litellm compatible key/model
    asyncio.run(main())
