从零手写AI Agent：工具定义、记忆管理与循环控制剖析

Python 3.11.9
openai==1.35.0          # 用新版client，支持tool_calls
pydantic==2.7.4         # 工具参数Schema校验
tiktoken==0.7.0         # Token计数，用于记忆裁剪

用户输入 → [LLM推理] → 有tool_call? → 执行工具 → 结果回填 → 继续循环
                ↓ 无tool_call
             返回最终答案

import json
from dataclasses import dataclass, field
from typing import Callable, Any
from pydantic import BaseModel, ValidationError

@dataclass
class Tool:
    name: str
    description: str
    args_schema: type[BaseModel]
    func: Callable[..., Any]

    def to_openai_schema(self) -> dict:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.args_schema.model_json_schema(),
            },
        }

    def run(self, raw_args: str) -> str:
        try:
            args = self.args_schema.model_validate_json(raw_args)
        except ValidationError as e:
            return f"[ARG_ERROR] 参数校验失败: {e.errors()[:2]}"
        try:
            return str(self.func(**args.model_dump()))
        except Exception as e:
            return f"[TOOL_ERROR] {type(e).__name__}: {e}"

# 示例工具
class SearchArgs(BaseModel):
    query: str
    top_k: int = 3

def fake_search(query: str, top_k: int = 3) -> str:
    # 实际项目里换成向量检索
    return f"关于'{query}'的{top_k}条结果：..."

class CalcArgs(BaseModel):
    expression: str

def calculator(expression: str) -> str:
    # 生产环境别用eval，这里仅演示
    allowed = set("0123456789+-*/(). ")
    if not set(expression)  self.max_msgs:
            self._compress()

    def _compress(self):
        old = self.messages[:-self.keep_recent]
        # 按tool_call pair成对裁剪，避免孤立tool消息
        cut = 0
        for i, m in enumerate(old):
            cut = i + 1
            if m.get("role") == "tool":
                break
        to_summarize = self.messages[:cut]
        self.messages = self.messages[cut:]
        text = "\n".join(
            f"{m['role']}: {m.get('content') or m.get('tool_calls')}"
            for m in to_summarize
        )
        new_summary = self.summarize_fn(self.summary, text)
        self.summary = new_summary

    def build(self) -> list[dict]:
        sys_content = self.system
        if self.summary:
            sys_content += f"\n\n[历史摘要]\n{self.summary}"
        return [{"role": "system", "content": sys_content}] + self.messages

import time
from openai import OpenAI

client = OpenAI()
MAX_STEPS = 8
MAX_RETRY_PER_TOOL = 2

def run_agent(user_input: str, memory: Memory, max_steps: int = MAX_STEPS):
    memory.add({"role": "user", "content": user_input})
    call_history = []  # (tool_name, args_str) 用于死循环检测

    for step in range(max_steps):
        resp = call_llm_with_retry(memory.build())
        msg = resp.choices[0].message

        if not msg.tool_calls:
            memory.add({"role": "assistant", "content": msg.content})
            return msg.content

        # 记录assistant的tool_call消息（必须原样回填）
        memory.add({
            "role": "assistant",
            "content": msg.content,
            "tool_calls": [tc.model_dump() for tc in msg.tool_calls],
        })

        for tc in msg.tool_calls:
            name = tc.function.name
            args_str = tc.function.arguments
            sig = (name, args_str)

            # 死循环检测：同签名连续出现2次，注入警告
            if call_history.count(sig) >= MAX_RETRY_PER_TOOL:
                result = f"[LOOP_WARN] 工具{name}已用相同参数调用{MAX_RETRY_PER_TOOL}次，请换策略或直接回答。"
            else:
                tool = TOOLS.get(name)
                result = tool.run(args_str) if tool else f"[UNKNOWN_TOOL] {name}不存在"
                call_history.append(sig)

            memory.add({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": result,
            })

    # 达到最大步数，强制收尾
    memory.add({"role": "user", "content": "已达到最大推理步数，请基于已有信息给出最终答案。"})
    final = call_llm_with_retry(memory.build())
    return final.choices[0].message.content

def call_llm_with_retry(messages: list[dict], retries: int = 3):
    for i in range(retries):
        try:
            return client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                tools=[t.to_openai_schema() for t in TOOLS.values()],
                tool_choice="auto",
                temperature=0,
                timeout=20,
            )
        except Exception as e:
            if i == retries - 1:
                raise
            time.sleep(2 ** i)  # 指数退避 1s, 2s