run.py

# Copyright 2025 GlyphyAI
#
# Use of this source code is governed by an MIT-style
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT.

import asyncio
import json

from pydantic import TypeAdapter

from liteswarm import LLM, Agent, AgentContext, ChatMessage, SwarmChat, ToolResult, tool

ChatMessageList = TypeAdapter(list[ChatMessage])


@tool
def fetch_weather(context: AgentContext, location: str) -> ToolResult:
    return ToolResult(content="The weather in New York is sunny.")


async def run() -> None:
    # Create a simple agent
    agent = Agent(
        id="assistant",
        instructions="You are a helpful assistant that provides clear and concise responses.",
        llm=LLM(model="gpt-4o"),
    )

    weather_agent = Agent(
        id="weather",
        instructions="You are a weather agent that provides weather information.",
        llm=LLM(model="gpt-4o"),
        tools=[fetch_weather],
    )

    @tool(agent=agent)
    def switch_to_weather_agent(context: AgentContext) -> ToolResult:
        return ToolResult.switch_to(
            content="Switched to weather agent",
            agent=weather_agent,
            params=context.params,
        )

    # Initialize chat
    chat = SwarmChat()

    # Send messages and stream responses
    user_messages = [
        "Hello! My name is John. What can you help me with?",
        "What is my name?",
        "What is the weather in New York?",
    ]

    for message in user_messages:
        print(f"\nUser: {message}")
        print("Assistant: ", end="", flush=True)

        async for event in chat.send_message(message, agent=agent):
            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()

    # Get conversation history
    chat_messages = await chat.get_messages()
    history = ChatMessageList.dump_python(chat_messages, exclude_none=True)
    print(f"\nConversation history:\n{json.dumps(history, indent=2)}")


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