if (
    finish_reason == "tool_calls" or finish_reason == "stop"
) and function_calls:
    # ... builds and yields aggregated_llm_response_with_tool_call
elif finish_reason == "stop" and (text or reasoning_parts):
    # ... builds and yields aggregated_llm_response

import asyncio
import os

from dotenv import load_dotenv
from google.adk.agents import Agent
from google.adk.agents.run_config import RunConfig, StreamingMode
from google.adk.models.lite_llm import LiteLlm
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types
from google.genai.types import GenerateContentConfig

load_dotenv()


def get_litellm_model() -> LiteLlm:
    deployment_id = os.getenv("AZURE_OPENAI_DEPLOYMENT_ID")
    return LiteLlm(
        model=f"azure/{deployment_id}",
        stream=True,
    )


async def add(a: float, b: float) -> float:
    return a + b


def create_simple_agent() -> Agent:
    model = get_litellm_model()

    instructions = """
    You are a helpful mathematical assistant with access to a calculator tool.
    """

    agent = Agent(
        name="Claudia",
        model=model,
        generate_content_config=GenerateContentConfig(
            temperature=0.0, max_output_tokens=10
        ),
        instruction=instructions,
        description="A simple agent that can perform basic mathematical calculations",
        tools=[add],
    )

    return agent


async def run_agent_single_query(query: str):
    agent = create_simple_agent()
    session_service = InMemorySessionService()

    runner = Runner(
        agent=agent,
        app_name="SimpleADKApp",
        session_service=session_service,
        auto_create_session=True,
    )

    user_id = "user_123"
    session_id = "session_002"

    content = types.Content(role="user", parts=[types.Part(text=query)])

    response_text = ""
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=content,
        run_config=(
            RunConfig(
                streaming_mode=StreamingMode.SSE,
                response_modalities=["TEXT"],
            )
        ),
    ):
        if event.content and event.content.parts:
            response_text = event.content.parts[0].text or ""

    print(f"Agent Response:\n{response_text}\n")
    return response_text


async def main():
    await run_agent_single_query("Ciao! 3 + 9")


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

          elif finish_reason == "stop" and (text or reasoning_parts):
            message_content = text if text else None
            aggregated_llm_response = _message_to_generate_content_response(
                # ... existing code ...
            )
            aggregated_llm_response.finish_reason = _map_finish_reason(
                finish_reason
            )
            text = ""
            reasoning_parts = []

          elif finish_reason == "length" and function_calls:
            raise ValueError(
                "LLM response was truncated due to max output token limit "
                "while generating a tool call. The partial tool call data "
                "cannot be executed. Consider increasing `max_output_tokens` "
                "in your GenerateContentConfig, or reducing the number/complexity "
                "of available tools to allow the model to complete its response."
            )