Autogen_core: Tool Use

完整代码

import random

from autogen_core import CancellationToken
from autogen_core.tools import FunctionTool
from typing_extensions import Annotated


async def get_stock_price(ticker: str, date: Annotated[str, "Date in YYYY/MM/DD"]) -> float:
    # Returns a random stock price for demonstration purposes.
    return random.uniform(10, 200)


# Create a function tool.
stock_price_tool = FunctionTool(get_stock_price, description="Get the stock price.")

# Run the tool.
cancellation_token = CancellationToken()
result = await stock_price_tool.run_json({"ticker": "AAPL", "date": "2021/01/01"}, cancellation_token)

# Print the result.
print(stock_price_tool.return_value_as_string(result))

23.67341202433553

from dataclasses import dataclass
from typing import List

from autogen_core import (
    AgentId,
    MessageContext,
    RoutedAgent,
    SingleThreadedAgentRuntime,
    message_handler,
)
from autogen_core.models import (
    ChatCompletionClient,
    LLMMessage,
    SystemMessage,
    UserMessage,
)
from autogen_core.tool_agent import ToolAgent, tool_agent_caller_loop
from autogen_core.tools import FunctionTool, Tool, ToolSchema
from autogen_ext.models.openai import OpenAIChatCompletionClient


@dataclass
class Message:
    content: str


class ToolUseAgent(RoutedAgent):
    def __init__(self, model_client: ChatCompletionClient, tool_schema: List[ToolSchema], tool_agent_type: str) -> None:
        super().__init__("An agent with tools")
        self._system_messages: List[LLMMessage] = [SystemMessage(content="You are a helpful AI assistant.")]
        self._model_client = model_client
        self._tool_schema = tool_schema
        self._tool_agent_id = AgentId(tool_agent_type, self.id.key)

    @message_handler
    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:
        # Create a session of messages.
        session: List[LLMMessage] = self._system_messages + [UserMessage(content=message.content, source="user")]
        # Run the caller loop to handle tool calls.
        messages = await tool_agent_caller_loop(
            self,
            tool_agent_id=self._tool_agent_id,
            model_client=self._model_client,
            input_messages=session,
            tool_schema=self._tool_schema,
            cancellation_token=ctx.cancellation_token,
        )
        # Return the final response.
        assert isinstance(messages[-1].content, str)
        return Message(content=messages[-1].content)

# Create a runtime.
runtime = SingleThreadedAgentRuntime()
# Create the tools.
tools: List[Tool] = [FunctionTool(get_stock_price, description="Get the stock price.")]

model_client = OpenAIChatCompletionClient(
                    model="GLM-4-Air-0111",
                    api_key = "your api key",
                    base_url="https://open.bigmodel.cn/api/paas/v4/",
                    model_capabilities={
                "vision": False,
                "function_calling": True,
                "json_output": True,
            }
                    )
# Register the agents.
await ToolAgent.register(runtime, "tool_executor_agent", lambda: ToolAgent("tool executor agent", tools))
await ToolUseAgent.register(
    runtime,
    "tool_use_agent",
    lambda: ToolUseAgent(
        model_client , [tool.schema for tool in tools], "tool_executor_agent"
    ),
)

AgentType(type='tool_use_agent')

# Start processing messages.
runtime.start()
# Send a direct message to the tool agent.
tool_use_agent = AgentId("tool_use_agent", "default")
response = await runtime.send_message(Message("What is the stock price of NVDA on 2024/06/01?"), tool_use_agent)
print(response.content)
# Stop processing messages.
await runtime.stop()

The stock price of NVDA on 2024/06/01 is $182.6841806011289.

代码解释

这段代码演示了一个基于 autogen 的框架，使用 FunctionTool 工具和一个 ToolUseAgent 智能体来获取和返回特定日期的股票价格。代码分为几个部分，我们将逐一解释每个部分的功能。

第一部分：获取股票价格

import random
from autogen_core import CancellationToken
from autogen_core.tools import FunctionTool
from typing_extensions import Annotated

async def get_stock_price(ticker: str, date: Annotated[str, "Date in YYYY/MM/DD"]) -> float:
    # Returns a random stock price for demonstration purposes.
    return random.uniform(10, 200)

stock_price_tool = FunctionTool(get_stock_price, description="Get the stock price.")
cancellation_token = CancellationToken()
result = await stock_price_tool.run_json({"ticker": "AAPL", "date": "2021/01/01"}, cancellation_token)
print(stock_price_tool.return_value_as_string(result))

这段代码定义了一个异步函数 get_stock_price，它接受股票代码和日期，返回一个随机生成的股票价格（在10到200之间）。FunctionTool 类将这个函数封装成一个工具，使其可以在 autogen 框架中调用。然后，代码执行这个工具，传入参数 {"ticker": "AAPL", "date": "2021/01/01"}，并打印结果。

第二部分：定义智能体

from dataclasses import dataclass
from typing import List

from autogen_core import (
    AgentId,
    MessageContext,
    RoutedAgent,
    SingleThreadedAgentRuntime,
    message_handler,
)
from autogen_core.models import (
    ChatCompletionClient,
    LLMMessage,
    SystemMessage,
    UserMessage,
)
from autogen_core.tool_agent import ToolAgent, tool_agent_caller_loop
from autogen_core.tools import FunctionTool, Tool, ToolSchema
from autogen_ext.models.openai import OpenAIChatCompletionClient

@dataclass
class Message:
    content: str

class ToolUseAgent(RoutedAgent):
    def __init__(self, model_client: ChatCompletionClient, tool_schema: List[ToolSchema], tool_agent_type: str) -> None:
        super().__init__("An agent with tools")
        self._system_messages: List[LLMMessage] = [SystemMessage(content="You are a helpful AI assistant.")]
        self._model_client = model_client
        self._tool_schema = tool_schema
        self._tool_agent_id = AgentId(tool_agent_type, self.id.key)

    @message_handler
    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:
        session: List[LLMMessage] = self._system_messages + [UserMessage(content=message.content, source="user")]
        messages = await tool_agent_caller_loop(
            self,
            tool_agent_id=self._tool_agent_id,
            model_client=self._model_client,
            input_messages=session,
            tool_schema=self._tool_schema,
            cancellation_token=ctx.cancellation_token,
        )
        assert isinstance(messages[-1].content, str)
        return Message(content=messages[-1].content)

这段代码定义了一个 ToolUseAgent 类，它继承自 RoutedAgent。这个智能体可以接收用户消息，使用工具，并返回响应。它使用一个 model_client（一个聊天完成客户端）来处理消息，并调用 tool_agent_caller_loop 来处理工具调用。

第三部分：创建运行时和工具

runtime = SingleThreadedAgentRuntime()
tools: List[Tool] = [FunctionTool(get_stock_price, description="Get the stock price.")]
model_client = OpenAIChatCompletionClient(
                    model="GLM-4-Air-0111",
                    api_key = "your api key",
                    base_url="https://open.bigmodel.cn/api/paas/v4/",
                    model_capabilities={
                "vision": False,
                "function_calling": True,
                "json_output": True,
            }
                    )
await ToolAgent.register(runtime, "tool_executor_agent", lambda: ToolAgent("tool executor agent", tools))
await ToolUseAgent.register(
    runtime,
    "tool_use_agent",
    lambda: ToolUseAgent(
        model_client , [tool.schema for tool in tools], "tool_executor_agent"
    ),
)

这段代码创建了一个单线程的智能体运行时 runtime，并注册了一个 ToolAgent（工具执行智能体）和一个 ToolUseAgent（工具使用智能体）。ToolAgent 负责执行工具，而 ToolUseAgent 负责接收用户消息并调用工具。

第四部分：处理消息

# Start processing messages.
runtime.start()
# Send a direct message to the tool agent.
tool_use_agent = AgentId("tool_use_agent", "default")
response = await runtime.send_message(Message("What is the stock price of NVDA on 2024/06/01?"), tool_use_agent)
print(response.content)
# Stop processing messages.
await runtime.stop()

The stock price of NVDA on 2024/06/01 is $182.6841806011289.

这段代码启动了运行时，并发送了一条消息给 tool_use_agent。tool_use_agent 接收到消息后，调用 tool_agent 来执行工具，并返回结果。最后，运行时停止处理消息。

类似例子

import random

from autogen_core import CancellationToken
from autogen_core.tools import FunctionTool
from typing_extensions import Annotated


async def get_weather_data(location: str, date: Annotated[str, "Date in YYYY/MM/DD"]) -> str:
    # Simulating fetching weather data.
    weather_conditions = ["Sunny", "Cloudy", "Rainy", "Windy", "Snowy"]
    temperature = random.uniform(-10, 40)  # Temperature in Celsius
    condition = random.choice(weather_conditions)
    return f"{condition}, {temperature:.1f}°C"


# Create a function tool.
weather_tool = FunctionTool(get_weather_data, description="Get the weather data.")

# Run the tool.
cancellation_token = CancellationToken()
result = await weather_tool.run_json({"location": "New York", "date": "2024/06/01"}, cancellation_token)

# Print the result.
print(weather_tool.return_value_as_string(result))

Windy, 2.5°C

from dataclasses import dataclass
from typing import List

from autogen_core import (
    AgentId,
    MessageContext,
    RoutedAgent,
    SingleThreadedAgentRuntime,
    message_handler,
)
from autogen_core.models import (
    ChatCompletionClient,
    LLMMessage,
    SystemMessage,
    UserMessage,
)
from autogen_core.tool_agent import ToolAgent, tool_agent_caller_loop
from autogen_core.tools import FunctionTool, Tool, ToolSchema
from autogen_ext.models.openai import OpenAIChatCompletionClient


@dataclass
class Message:
    content: str


class ToolUseAgent(RoutedAgent):
    def __init__(self, model_client: ChatCompletionClient, tool_schema: List[ToolSchema], tool_agent_type: str) -> None:
        super().__init__("An agent with tools")
        self._system_messages: List[LLMMessage] = [SystemMessage(content="You are a helpful AI assistant.")]
        self._model_client = model_client
        self._tool_schema = tool_schema
        self._tool_agent_id = AgentId(tool_agent_type, self.id.key)

    @message_handler
    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:
        # Create a session of messages.
        session: List[LLMMessage] = self._system_messages + [UserMessage(content=message.content, source="user")]
        # Run the caller loop to handle tool calls.
        messages = await tool_agent_caller_loop(
            self,
            tool_agent_id=self._tool_agent_id,
            model_client=self._model_client,
            input_messages=session,
            tool_schema=self._tool_schema,
            cancellation_token=ctx.cancellation_token,
        )
        # Return the final response.
        assert isinstance(messages[-1].content, str)
        return Message(content=messages[-1].content)


# Create a runtime.
runtime = SingleThreadedAgentRuntime()
# Create the tools.
tools: List[Tool] = [FunctionTool(get_weather_data, description="Get the weather data.")]

model_client = OpenAIChatCompletionClient(
                    model="GLM-4-Air-0111",
                    api_key="your api key",
                    base_url="https://open.bigmodel.cn/api/paas/v4/",
                    model_capabilities={
                        "vision": False,
                        "function_calling": True,
                        "json_output": True,
                    }
                    )
# Register the agents.
await ToolAgent.register(runtime, "tool_executor_agent", lambda: ToolAgent("tool executor agent", tools))
await ToolUseAgent.register(
    runtime,
    "tool_use_agent",
    lambda: ToolUseAgent(
        model_client , [tool.schema for tool in tools], "tool_executor_agent"
    ),
)

AgentType(type='tool_use_agent')

# Start processing messages.
runtime.start()
# Send a direct message to the tool agent.
tool_use_agent = AgentId("tool_use_agent", "default")
response = await runtime.send_message(Message("What is the weather in New York on 2024/06/01?"), tool_use_agent)
print(response.content)
# Stop processing messages.
await runtime.stop()

According to the API call result, the weather in New York on June 1, 2024 is sunny with a temperature of 9.2°C. I hope this information is helpful to you.

参考链接：https://microsoft.github.io/autogen/stable/user-guide/core-user-guide/framework/tools.html