释放自定义 Ai 代理的力量：在 Ollama 中掌握工具集成

Rifx.Online
Large Language Models , Software Development , AI Applications
16 Feb, 2025

一只配备工具的羊驼

介绍

今天，我将分享一个我开发的解决方案，它提供了一种简单的方法来集成自定义工具 ⚙️，而无需依赖流行的第三方库。

这是必要的，因为 agents，即配备工具的模型，无法自行检索新数据。因此，它们的知识范围是有限的，除非从互联网提供新的信息。通过实现自定义函数，我们可以有效地扩展这些 agents 的能力，使它们能够访问和利用实时数据，以增强其性能和决策过程。

Ollama 最近通过引入一个新功能来增强其能力，该功能允许直接从 Python 方法调用函数。这使得更好地理解如何将自定义函数集成到模型中成为一个很好的时机。

通过跟随本教程，您将看到在项目中实现这些工具是多么简单，使您能够专注于构建应用程序，而不会被复杂的设置或依赖关系所困扰。

代码

如果您想运行代码，可以通过 GitHub 在这里访问它：带工具的代理

要求

第一步是安装必要的库。使用的Python版本是3.12。将以下数据保存为requirements.txt：

httpx==0.27.2
ollama==0.4.1
openai==1.55.1
pydantic==2.9.2
python-dotenv==1.0.1

然后运行以下命令：

pip install -r requirements.txt

接下来，访问 https://ollama.com/download 并在本地安装Ollama。完成此步骤后，打开CMD并输入以下命令以安装llama3.2模型：

如果您需要不同的配置或希望将本地保存的模型重定向到另一个驱动器，请查看本文：逐步运行Ollama模型。

要检查模型是否正确安装，请在CMD中输入：

核心组件

该解决方案由三个主要类组成：

BaseChatModel: 处理通用功能的基础类
OllamaChatModel: Ollama模型的实现
OpenAIChatModel: OpenAI模型的实现

让我们来分析一下它是如何工作的。

1. 打造坚实基础：基础聊天模型类

基础聊天模型 类作为这个基本支柱，提供了核心功能，使用户与AI之间的无缝通信成为可能。让我们深入了解它的结构和功能，如记忆管理和消息格式化：

class BaseChatModel:
    def __init__(
        self,
        client: Any,
        model: str,
        temperature: float = 0.0,
        keep_alive: int = -1,
        format: Optional[str] = None,
        max_stored_memory_messages: int = 50,
    ) -> None:
        self.client = client
        self.model = model
        self.format = format
        self.temperature = temperature
        self.memory = 基础记忆(max_size=max_stored_memory_messages)

client: 该参数表示与底层语言模型API的连接，允许类发送和接收消息。
model: 正在使用的特定模型（例如，Ollama或OpenAI），决定了聊天的能力和响应。
temperature: 一个浮点值，影响模型响应的随机性。较低的值会产生更确定的输出，而较高的值则引入变异性。
keep_alive: 一个整数，指定与客户端保持连接的时间，确保高效通信而不产生不必要的开销。
format: 一个可选字符串，可以定义在发送或显示消息之前如何格式化消息。如果您决定将输出格式更改为 json，请记得调整其余代码以根据请求的键处理代理的响应。
max_stored_memory_messages: 这个整数设置了可以在记忆中存储的过去消息的数量限制，从而在对话中促进上下文保留。

基础聊天模型 的一个特点是其内置的记忆管理系统。通过初始化一个指定最大大小的 基础记忆 实例，模型可以有效地记住先前的交互，增强其根据上下文提供相关响应的能力。

为了引导AI在交互过程中的行为，包含了一个预定义的系统消息：

self.system = """You are a helpful AI assistant. When using tools, always provide a natural, complete response using the information gathered.
Format your response as a short, coherent sentence."""

这条指令有助于塑造AI如何构造其回复，确保它们不仅信息丰富，而且以用户友好的方式呈现。

基础聊天模型 还包括一个用于格式化用户与AI之间交换消息的方法：

def message(self, human: str, ai: str) -> ChatMessage:
    return ChatMessage(message={"human": human, "ai": ai})

此方法将人类和AI消息封装成结构化格式，使管理和显示对话历史变得更加容易。

2. Ollama 实现

类初始化

OllamaChatModel 类是 基础聊天模型 的扩展，旨在促进与 Ollama 语言模型的交互，提供一种结构化的方式来管理工具和响应。让我们详细探讨它的特性和功能。

class OllamaChatModel(BaseChatModel):
    """http://localhost:11434 是默认的 Ollama 端口，用于提供 API。"""
    
    def __init__(self, tools: list[dict], model: str = "llama3.2") -> None:
        self.model = model
        self.tools = tools
        self.client = ollama.Client(host="http://localhost:11434")
        super().__init__(client=self.client, model=self.model)

model: model 参数允许用户指定要使用的 Ollama 模型版本，默认为 “llama3.2”。
tools: tools 参数接受一个自定义工具的列表，可以在聊天会话中调用，从而增强模型的功能。
client: client 被初始化以连接到 Ollama API，该 API 在指定的主机和端口上运行，确保所有请求正确导向。如果您在 Docker 容器中运行 Ollama，您可以轻松配置 Ollama 主机以指向适当的网络设置。

OllamaChatModel 的一个特性是能够动态提取和执行工具调用。提取多个工具调用使我们能够给代理一个机会，根据用户请求进行多次工具调用。

def extract(self, tool_call) -> list:
    """提取并执行工具调用"""
    data = []

    if not isinstance(tool_call, list):
        tool_call = [tool_call]

    for tool in tool_call:
        func_name = tool.function.name
        if isinstance(tool.function.arguments, str):
            func_arguments = json.loads(tool.function.arguments)
        else:
            func_arguments = tool.function.arguments
        result = run_callable(func_name, func_arguments)
        data.append(result)
    
    return data

这种灵活性允许开发人员创建丰富的交互，模型可以根据用户输入执行操作。

生成响应

response 方法对于根据用户提示生成模型回复至关重要：

def response(self, user_prompt: str, system_message: str = None) -> ollama.ChatResponse:
    messages = [
        {
            "role": "system",
            "content": system_message if system_message else self.system,
        }
    ]

    for msg in self.memory.get():
        if isinstance(msg, ChatMessage):
            messages.extend(
                [
                    {"role": "user", "content": msg.human},
                    {"role": "assistant", "content": msg.ai},
                ]
            )

    messages.append({"role": "user", "content": user_prompt})

    return self.client.chat(
        model=self.model,
        messages=messages,
        format=self.format,
        keep_alive=self.keep_alive,
        tools=self.tools,
    )

该方法构建了一个消息历史记录，包括系统指令和先前的交互。它将此上下文发送到Ollama客户端，以根据当前用户提示生成连贯的响应。

互动聊天功能

chat 方法为用户与模型的互动提供了一个循环。它实时处理用户输入，并在生成和显示响应之前处理任何必要的工具调用。将聊天保存为 .json 文件的历史记录为希望重温先前对话的用户增加了额外的功能。

def chat(self, system_message: str = None, save_chat: bool = False) -> None:
    system_message = system_message if system_message else self.system

    while True:
        print("current memory: ", self.memory.get(), "\n\n")

        user_prompt = input("User: ")
        if user_prompt == "bye":
            self.memory.add(self.message(human=user_prompt, ai="Bye"))
            if save_chat:
                self.memory.save(model_name=str(self.model))
            self.memory.clear()
            print("AI: Bye")
            break

        response = self.response(user_prompt, system_message)

        if hasattr(response.message, "tool_calls") and response.message.tool_calls:
            collected_data = {}

            for tool_call in response.message.tool_calls:
                result = self.extract(tool_call)
                collected_data[tool_call.function.name] = result

            final_prompt = (
                f"Based on the following information:\n"
                f"{collected_data}"
                f"Please provide a natural response to the original question: '{user_prompt}'"
            )

            final_response = self.client.chat(
                model=self.model,
                messages=[
                    {
                        "role": "system",
                        "content": "Sum up your response before sending back. Make it short and concise"
                        + " "
                        + system_message,
                    },
                    {"role": "user", "content": final_prompt},
                ],
            )
            response_content = final_response.message.content
        else:
            response_content = response.message.content

        if response_content:
            print(f"AI: {response_content}", end="\n\n")
            self.memory.add(self.message(human=user_prompt, ai=response_content))

3. OpenAI 实现

类 OpenAIChatModel 与 OllamaChatModel 非常相似，并包含相同的函数。我们将不讨论它们，而只关注导致单独实现的差异。

API 密钥与客户端

OpenAIChatModel 在初始化时需要一个 api_key，这是对 OpenAI API 进行请求身份验证所必需的。相比之下，OllamaChatModel 连接到本地的 Ollama 服务器实例，而不需要 API 密钥。

self.client = openai.OpenAI(api_key=api_key, timeout=timeout, max_retries=max_retries)

self.client = ollama.Client(host="http://localhost:11434")

工具调用提取

两个模型中的 extract 方法处理工具调用，但 OpenAIChatModel 直接从 JSON 字符串格式解析工具调用参数，而 OllamaChatModel 则检查输入是否为列表并相应处理。

func_arguments = json.loads(tool_call.function.arguments)

if not isinstance(tool_calls, list):
    tool_call = [tool_calls]

响应方法

生成响应的方式在两个模型之间有所不同。OpenAIChatModel 使用 self.client.chat.completions.create() 来生成响应，而 OllamaChatModel 则调用 self.client.chat()。

return self.client.chat.completions.create(model=self.model, messages=messages, tools=self.tools)

return self.client.chat(model=self.model, messages=messages)

最终提示处理

两个模型都构建最终提示，以根据从工具调用收集的数据生成简洁的响应。然而，它们在格式化和将此提示发送回各自客户端的方式上有所不同。

4. 扩展自定义工具

此外，我们使用两个模块：memory.py，在这里我们找到传递给代理的记忆的简化实现，以及 tools.py，在这里我们保留我们的工具。

在文件的开头，我们有两个函数负责所有使用工具的全局注册。如果您希望从不同于 tools.py 的文件提供工具，请记得在以下位置调整更改：

module = importlib.import_module("tools")

registered_functions = {}

def register_function(func) -> Callable:
    registered_functions[func.__name__] = func
    return func

def run_callable(name: str, arguments: dict) -> Callable | dict:
    try:
        module = importlib.import_module("tools")
        func = getattr(module, name, None)
        func = registered_functions.get(name)

        if func and callable(func):
            return func(**arguments)
        else:
            return {"error": f"Function '{name}' is not callable or not found"}
    except Exception as e:
        traceback.print_exc()
        return {"error": f"Error executing {name}: {str(e)}"}

如何添加新工具

转到 tools.py，创建并注册一个新函数。我提供了简单的方法，但您可以结合更强大的解决方案，例如像 SerpApi 这样的付费 API，或者尝试像 CatFacts Api 这样的免费 API。

@register_function
def say_hello(name: str) -> str:
    """
    根据用户的名字打招呼。

    返回：
        str: 用户问候
    """
    return f"Hello {name}!"

每个函数需要一个第二部分，描述性部分。您可以为字典变量选择任何名称，但与键 “name” 关联的值必须与被调用函数的名称完全匹配。此修订明确指出，尽管字典名称是灵活的，但 “name” 的值必须精确对应它所代表的函数。

say_hello_tool = {
    "type": "function",
    "function": {
        "name": "say_hello",
        "description": "根据用户的名字打招呼。",
        "parameters": {
            "type": "object",
            "properties": {"name": {"type": "string", "description": "用户的名字"}},
            "required": ["name"],
        },
    },
}

因为我们向函数传递了一个参数，所以必须指定哪些字段是 required。

现在，让我们在 agents_with_tools.py 文件中扩展工具集：

tools = [today_is_tool, weather_tool, day_of_week_tool, say_hello_tool]

替换记忆

在 memory.py 中，我们使用自定义类处理用户与 AI 之间的通信，以及存储历史数据。通过利用 LangChain 记忆类型，这一功能可以简化；然而，我选择不这样做，以帮助您更好地理解其从零开始的工作原理。欢迎您为您的代理实现自己的记忆。

ChatMessage 类提供了一个简单的通信接口，确保传递的数据始终包含预期的键：human 和 ai。

class ChatMessage(BaseModel):
    message: Dict[str, str] = Field(
        ..., description="A dictionary containing 'human' and 'ai' messages."
    )

    @property
    def human(self) -> str:
        """Return the human message."""
        return self.message.get("human", "")

    @property
    def ai(self) -> str:
        """Return the AI message."""
        return self.message.get("ai", "")

    def cleanup(memory: Any):
        memory.clear()

    def __str__(self) -> str:
        """Custom string representation for ChatMessage."""
        return json.dumps({"human": self.human, "ai": self.ai}, ensure_ascii=False)

BaseMemory 是一个双端队列结构，确保您最多存储 N 条消息。这一点很重要，原因有二：

向付费 LLM 发送过多的 tokens 是昂贵的。
大多数 LLM 具有中到大的 token 窗口，但很容易忘记这一点并超出限制。

此外，如果我们发送过多的历史信息，LLM 的效率会降低。

class BaseMemory:
    def __init__(self, max_size=100):
        """Initialize the memory with a maximum size."""
        self.memory = deque(maxlen=max_size)

    def add(self, message: ChatMessage):
        """Add an item to the memory."""
        self.memory.append(message)

    def peek(self) -> str:
        if self.memory:
            return str(self.memory[-1])
        return None

    def get(self):
        """Retrieve all items in memory."""
        return list(self.memory)

    def get_as_str(self):
        """Retrieve all items in memory."""
        return ";".join(f"{x}" for x in list(self.memory))

    def get_max(self, limit: int):
        """Retrieve all items in memory."""
        if self.memory:
            return list(self.memory)[:limit]
        return None

    def clear(self):
        """Clear all items from memory."""
        self.memory.clear()

    def save(self, model_name: str = None, file_path: str = None):
        """Dump all memory contents to a local JSON file."""
        if not file_path or not os.path.exists(file_path):
            file_path = Path(
                os.path.abspath(os.path.join(os.path.dirname(__file__)))
            ).as_posix()

        if self.memory:
            timestamp = (
                str(datetime.now().replace(microsecond=0))
                .replace(":", "")
                .replace(" ", "")
                .replace("-", "")
                .replace("_", "")
            )
            name = (
                f"chat_{model_name.replace('-', '')}_{timestamp}.json"
                if model_name
                else f"chat_{timestamp}.json"
            )
            with open(
                os.path.join(file_path, name), "w", encoding="utf-8"
            ) as json_file:
                json.dump([msg.message for msg in self.memory], json_file, indent=4)

    def __str__(self):
        """String representation of the memory."""
        if self.memory:
            return f"Memory({list(self.memory)})"
        return None

运行脚本

选择您喜欢的模型，然后简单地运行脚本以开始与代理聊天！

def load_model(model: str = None) -> OllamaChatModel | OpenAIChatModel:
    tools = [today_is_tool, weather_tool, day_of_week_tool]
    match model:
        case "ollama":
            return OllamaChatModel(tools=tools)
        case "openai":
            return OpenAIChatModel(api_key=os.environ.get("openai_apikey"), tools=tools)
        case _:
            return OllamaChatModel(tools=tools)

def main():
    model = load_model("openai") 
    model.chat(save_chat=True)

if __name__ == "__main__":
    main()