import asyncio
 
from braintrust_adk import setup_braintrust
from google.adk import Runner
from google.adk.agents import LlmAgent
from google.adk.sessions import InMemorySessionService
from google.genai import types
 
setup_braintrust(
    project_name="my-adk-project",
)
 
 
# Create your ADK agent as normal
def get_weather(city: str) -> dict:
    """Get weather for a city."""
    return {"temperature": 72, "condition": "sunny", "city": city}
 
 
def get_current_time() -> str:
    """Get the current time."""
    from datetime import datetime
 
    return datetime.now().strftime("%I:%M %p")
 
 
async def main():
    # Create the agent
    agent = LlmAgent(
        name="weather_time_assistant",
        tools=[get_weather, get_current_time],
        model="gemini-2.0-flash-exp",
        instruction="You are a helpful assistant that can check weather and time.",
    )
    # Create a session service and a runner
    session_service = InMemorySessionService()
    runner = Runner(app_name="weather_app", agent=agent, session_service=session_service)
    # Create a fake session
    user_id = "user123"
    session_id = "session123"
    await session_service.create_session(app_name="weather_app", user_id=user_id, session_id=session_id)
    # Create the message to send
    new_message = types.Content(
        parts=[types.Part(text="What's the weather like in New York?")],
        role="user",
    )
    # Run the agent with the query
    events = runner.run(
        user_id=user_id,
        session_id=session_id,
        new_message=new_message,
    )
    # Process the events and print the agent's response
    for event in events:
        print(event)
 
 
if __name__ == "__main__":
    asyncio.run(main())