from google.adk.agents import LlmAgent
from google.adk.models import Gemini
from google.adk.runners import InMemoryRunner
from google.adk.agents.run_config import RunConfig, StreamingMode
from google.genai import types
import asyncio

def get_weather(city: str) -> dict:
    """Get weather for a city."""
    return {"city": city, "temp": 72, "condition": "Sunny"}

agent = LlmAgent(
    model=Gemini(model="gemini-2.5-flash", use_interactions_api=True),
    name="test_agent",
    instruction="Use get_weather when asked about weather. Be concise.",
    tools=[get_weather],
)

async def test():
    runner = InMemoryRunner(agent=agent, app_name="test")
    session = await runner.session_service.create_session(
        user_id="user1", app_name="test"
    )

    # WORKS
    # run_config = RunConfig(streaming_mode=StreamingMode.NONE)

    # FAILS
    run_config = RunConfig(streaming_mode=StreamingMode.SSE)

    async for event in runner.run_async(
        user_id="user1",
        session_id=session.id,
        new_message=types.Content(
            role="user", parts=[types.Part(text="Weather in Tokyo?")]
        ),
        run_config=run_config,
    ):
        if event.get_function_calls():
            for fc in event.get_function_calls():
                print(f"Tool: {fc.name}({fc.args})")
        if (event.content and event.content.parts
            and not event.partial and event.author != "user"):
            for p in event.content.parts:
                if p.text:
                    print(f"Response: {p.text[:100]}")

asyncio.run(test())