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

from pydantic import BaseModel

from liteswarm.chat import SwarmChat
from liteswarm.types import LLM, Agent


class MathResult(BaseModel):
    result: int
    thoughts: str


async def run() -> None:
    agent = Agent(
        id="math_expert",
        instructions="You are a math expert.",
        llm=LLM(
            model="gpt-4o",
            response_format=MathResult,
        ),
        output_type=MathResult,
    )

    chat = SwarmChat()
    stream = chat.send_message(
        "What is 2 + 2 * 2?",
        agent=agent,
        final_output_type=MathResult,
    )

    # Get streaming partial parsed responses
    async for event in stream:
        if event.type == "agent_response_chunk":
            if event.chunk.parsed:
                print(event.chunk.parsed)

    # Get final response output
    result = await stream.get_return_value()
    if result.final_response.output:
        print(result.final_response.output.model_dump_json(indent=2))


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