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.core import Swarm
from liteswarm.types import LLM, Agent, Message


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,
    )

    swarm = Swarm()

    # Method 1: Use streaming API to get structured outputs
    stream = swarm.stream(
        agent=agent,
        messages=[Message(role="user", content="What is 2 + 2 * 2?")],
        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 parsed response
    result = await stream.get_return_value()
    if result.final_response.output:
        print(result.final_response.output.model_dump_json(indent=2, exclude_none=True))

    # # Method 2: Use convenience method to get final parsed response
    # result = await swarm.run(
    #     agent,
    #     messages=[Message(role="user", content="What is 2 + 2 * 2?")],
    #     final_output_type=MathResult,
    # )

    # if result.final_response.output:
    #     print(result.final_response.output.model_dump_json(indent=2, exclude_none=True))


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