augmented_llm.py

import asyncio
import os

from coagent.agents import ChatAgent, ModelClient, tool
from coagent.agents.messages import ChatMessage
from coagent.core import AgentSpec, new, set_stderr_logger
from coagent.runtimes import LocalRuntime


class Assistant(ChatAgent):
    system = """You are an agent who can use tools."""
    client = ModelClient(
        model=os.getenv("MODEL_NAME"),
        api_base=os.getenv("MODEL_API_BASE"),
        api_version=os.getenv("MODEL_API_VERSION"),
        api_key=os.getenv("MODEL_API_KEY"),
    )

    @tool
    async def query_weather(self, city: str) -> str:
        """Query the weather in the given city."""
        return f"The weather in {city} is sunny."


assistant = AgentSpec("assistant", new(Assistant))


async def main():
    async with LocalRuntime() as runtime:
        await runtime.register(assistant)

        result = await assistant.run(
            ChatMessage(
                role="user",
                content="What's the weather like in Beijing?",
            ).encode(),
            stream=True,
        )
        async for chunk in result:
            msg = ChatMessage.decode(chunk)
            print(msg.content, end="", flush=True)


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