from orxhestra.events.event import Event, EventType

async for event in agent.astream("Write about distributed systems"):
    if event.type == EventType.AGENT_MESSAGE and event.partial:
        print(event.text, end="", flush=True)
    elif event.is_final_response():
        print(f"\n[DONE] {event.text}")

from orxhestra import LlmAgent
from orxhestra.tools.agent_tool import AgentTool

weather_agent = LlmAgent(name="WeatherAgent", model=model, tools=[get_weather])
travel_agent = LlmAgent(name="TravelAgent", model=model, tools=[get_attractions])

planner = LlmAgent(
    name="TripPlanner",
    model=model,
    tools=[AgentTool(weather_agent), AgentTool(travel_agent)],
    instructions="Use the sub-agents to plan a trip.",
)

async for event in planner.astream("Plan a trip to Copenhagen"):
    if event.branch:
        print(f"  [{event.agent_name}] {event.text}", end="")
    elif event.is_final_response():
        print(f"\nFinal: {event.text}")

async for event in runner.astream(
    user_id="user-1",
    session_id="session-1",
    new_message="Write me a long essay.",
):
    if event.is_final_response():
        print(f"\n[DONE] {event.text}")
    elif event.type == EventType.AGENT_MESSAGE and event.partial:
        print(event.text, end="", flush=True)