import asyncio

from llama_index.core.agent.workflow import (
    AgentStream,
    FunctionAgent,
    ToolCall,
    ToolCallResult,
)
from llama_index.core.llms.mock import MockFunctionCallingLLM


def add_numbers(a: int = 6, b: int = 7) -> int:
    """Add two integers."""
    return a + b


async def main() -> None:
    agent = FunctionAgent(
        tools=[add_numbers],
        llm=MockFunctionCallingLLM(is_chat_model=True),
        system_prompt="Use tools for arithmetic.",
    )
    handler = agent.run(user_msg="Add 6 and 7.")

    async for event in handler.stream_events():
        if isinstance(event, ToolCall):
            print(f"Tool call: {event.tool_name} {event.tool_kwargs}")
        elif isinstance(event, ToolCallResult):
            print(f"Tool result: {event.tool_name} -> {event.tool_output.content}")
        elif isinstance(event, AgentStream) and event.delta:
            print(f"Agent delta: {event.delta}")

    response = await handler
    print(f"Final response: {response}")


asyncio.run(main())