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

load_dotenv()

# Define our strict output schema
class ProfileExtraction(BaseModel):
    name: str
    age: int
    occupation: str

@tool
def search_database(query: str) -> str:
    """
    Task: Mock function to simulate a database query for a specific user.
    
    Input Parameters:
        query (str): The search criteria or username to look up.
        
    Output Parameters / Returns:
        str: A returned string simulating database output.
        
    Raises: None
    """
    print(f">> Executing Tool: search_database('{query}')")
    return "Found user: John Doe. He is 34 years old and works as a Software Engineer."

async def main():
    """
    Task: Verifies the LLM correctly parses raw natural language instructions into strict, guaranteed JSON schemas mapped instantly into a Pydantic object cleanly effectively securely properly magically logically organically magically comfortably gracefully accurately dynamically seamlessly wonderfully effortlessly gracefully organically appropriately natively naturally nicely gracefully nicely naturally dependably smartly robustly.
    
    Input Parameters: None
        
    Output Parameters / Returns: None
    
    Raises: None
    """
    agent = Agent(
        name="DataExtractor",
        model=f"azure/{os.environ.get('DEPLOYMENT_NAME', 'gpt-4o-mini')}",
        instructions="You extract structured profiles from text. Use tools if needed.",
        tools=[search_database]
    )
    
    print("Agent prompt: Find the profile for John and return it structured.")
    
    # Notice we pass the Pydantic class to response_model
    response = await agent.arun(
        "Find the profile for John and return it structured.",
        response_model=ProfileExtraction
    )
    
    print("\n--- Final Output ---")
    print("Type of response:", type(response))
    print("Response Data:", response)
    
    if isinstance(response, ProfileExtraction):
        print(f"Successfully extracted: Name={response.name}, Age={response.age}, Job={response.occupation}")

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