import asyncio

from liteswarm import LLM, Agent, Message, Swarm


async def main() -> None:
    # Create a simple agent
    agent = Agent(
        id="assistant",
        instructions="You are a helpful assistant.",
        llm=LLM(model="gpt-4o"),
    )

    # Create swarm and run
    swarm = Swarm()
    result = await swarm.run(
        agent=agent,
        messages=[Message(role="user", content="Hello!")],
    )
    print(result.final_response.content)


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

import asyncio

from liteswarm import LLM, Agent, Message, Swarm, ToolResult, tool_plain


async def main() -> None:
    # Define a tool that can switch to another agent
    @tool_plain
    def switch_to_expert(domain: str) -> ToolResult:
        expert_agent = Agent(
            id=f"{domain}-expert",
            instructions=f"You are a {domain} expert.",
            llm=LLM(
                model="gpt-4o",
                temperature=0.0,
            ),
        )

        return ToolResult.switch_to(expert_agent)

    # Create a router agent that can switch to experts
    router = Agent(
        id="router",
        instructions="Route questions to appropriate experts.",
        llm=LLM(model="gpt-4o"),
        tools=[switch_to_expert],
    )

    # Stream responses in real-time
    swarm = Swarm()
    stream = swarm.stream(
        agent=router,
        messages=[Message(role="user", content="Explain quantum physics like I'm 5")],
    )

    async for event in stream:
        if event.type == "agent_response_chunk":
            completion = event.chunk.completion
            if completion.delta.content:
                print(completion.delta.content, end="", flush=True)
            if completion.finish_reason == "stop":
                print()

    # Optionally, get complete run result from stream
    result = await stream.get_return_value()
    print(result.final_response.content)


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

import asyncio

from liteswarm import LLM, Agent, SwarmChat, SwarmEvent


def handle_event(event: SwarmEvent) -> None:
    if event.type == "agent_response_chunk":
        completion = event.chunk.completion
        if completion.delta.content:
            print(completion.delta.content, end="", flush=True)
        if completion.finish_reason == "stop":
            print()


async def main() -> None:
    # Create an agent
    agent = Agent(
        id="assistant",
        instructions="You are a helpful assistant. Provide short answers.",
        llm=LLM(model="gpt-4o"),
    )

    # Create stateful chat
    chat = SwarmChat()

    # First message
    print("First message:")
    async for event in chat.send_message("Tell me about Python", agent=agent):
        handle_event(event)

    # Second message - chat remembers the context
    print("\nSecond message:")
    async for event in chat.send_message("What are its key features?", agent=agent):
        handle_event(event)

    # Access conversation history
    messages = await chat.get_messages()
    print(f"\nMessages in history: {len(messages)}")


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