从零手写AI Agent：工具定义、记忆管理与循环控制实现细节解析

while 未完成 and 步数  dict:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            }
        }

class ToolRegistry:
    def __init__(self):
        self._tools: dict[str, Tool] = {}

    def register(self, tool: Tool):
        if tool.name in self._tools:
            raise ValueError(f"工具 {tool.name} 已注册")
        self._tools[tool.name] = tool

    def get_schemas(self) -> list[dict]:
        return [t.to_openai_schema() for t in self._tools.values()]

    def execute(self, name: str, args: dict) -> str:
        tool = self._tools.get(name)
        if not tool:
            return f"[ERROR] 未知工具: {name}，可用工具: {list(self._tools)}"
        try:
            result = tool.func(**args)
            return str(result)
        except TypeError as e:
            return f"[ERROR] 参数不匹配: {e}，收到的参数: {args}"
        except Exception as e:
            return f"[ERROR] 工具执行异常: {type(e).__name__}: {e}"

def get_weather(city: str, unit: str = "celsius") -> str:
    # 假装调API
    data = {"北京": 22, "上海": 26, "深圳": 30}
    temp = data.get(city, 20)
    if unit == "fahrenheit":
        temp = temp * 9 / 5 + 32
    return f"{city}当前温度{temp}{'°C' if unit=='celsius' else '°F'}"

registry = ToolRegistry()
registry.register(Tool(
    name="get_weather",
    description="查询指定城市的当前天气温度",
    parameters={
        "type": "object",
        "properties": {
            "city": {"type": "string", "description": "城市名，如'北京'"},
            "unit": {"type": "string", "enum": ["celsius", "fahrenheit"],
                     "description": "温度单位，默认celsius"}
        },
        "required": ["city"]
    },
    func=get_weather
))

import tiktoken

class Memory:
    def __init__(self, max_messages=20, max_tokens=4000):
        self.messages: list[dict] = []
        self.max_messages = max_messages
        self.max_tokens = max_tokens
        self.enc = tiktoken.encoding_for_model("gpt-4o-mini")

    def add(self, msg: dict):
        self.messages.append(msg)
        self._trim()

    def _count_tokens(self) -> int:
        return sum(len(self.enc.encode(json.dumps(m, ensure_ascii=False)))
                   for m in self.messages)

    def _trim(self):
        # 1. 按条数裁剪，但保护 tool 配对
        while len(self.messages) > self.max_messages:
            # 如果第二条是tool结果，说明第一条是tool_call，成对删
            if len(self.messages) > 1 and self.messages[1].get("role") == "tool":
                self.messages = self.messages[2:]
            else:
                self.messages = self.messages[1:]
        # 2. 按token裁剪（简化版：直接丢最老的，实际可接摘要）
        while self._count_tokens() > self.max_tokens and len(self.messages) > 4:
            self.messages = self.messages[1:]

    def get(self) -> list[dict]:
        return self.messages

from openai import OpenAI

client = OpenAI(api_key="sk-xxx")

class Agent:
    def __init__(self, registry: ToolRegistry, system_prompt: str,
                 max_steps=8, max_retries=2):
        self.registry = registry
        self.system_prompt = system_prompt
        self.memory = Memory()
        self.max_steps = max_steps
        self.max_retries = max_retries

    def run(self, user_input: str) -> str:
        self.memory.add({"role": "user", "content": user_input})
        last_tool_sig = None
        repeat_count = 0

        for step in range(self.max_steps):
            messages = [{"role": "system", "content": self.system_prompt}] \
                       + self.memory.get()

            # --- LLM调用带重试 ---
            resp = None
            for attempt in range(self.max_retries + 1):
                try:
                    resp = client.chat.completions.create(
                        model="gpt-4o-mini",
                        messages=messages,
                        tools=self.registry.get_schemas(),
                        tool_choice="auto",
                        temperature=0.1,
                    )
                    break
                except Exception as e:
                    if attempt == self.max_retries:
                        return f"[Agent失败] LLM调用重试{self.max_retries}次仍失败: {e}"
                    continue

            msg = resp.choices[0].message

            # --- 情况1：最终答案 ---
            if not msg.tool_calls:
                self.memory.add({"role": "assistant", "content": msg.content})
                return msg.content

            # --- 情况2：工具调用 ---
            self.memory.add(msg.model_dump(exclude_none=True))

            for tc in msg.tool_calls:
                name = tc.function.name
                try:
                    args = json.loads(tc.function.arguments)
                except json.JSONDecodeError:
                    # LLM偶尔返回非法JSON，把错误喂回去
                    result = f"[ERROR] 参数JSON解析失败: {tc.function.arguments}"
                    self.memory.add({"role": "tool", "tool_call_id": tc.id,
                                     "content": result})
                    continue

                # --- 死循环检测：连续两次相同工具+相同参数 ---
                sig = f"{name}:{json.dumps(args, sort_keys=True)}"
                if sig == last_tool_sig:
                    repeat_count += 1
                    if repeat_count >= 2:
                        return f"[Agent终止] 检测到重复调用 {name}，参数未变化，主动退出"
                else:
                    repeat_count = 0
                last_tool_sig = sig

                result = self.registry.execute(name, args)
                self.memory.add({"role": "tool", "tool_call_id": tc.id,
                                 "content": result})

        return "[Agent终止] 达到最大步数限制，未能完成"

    def reset(self):
        self.memory = Memory()

agent = Agent(
    registry=registry,
    system_prompt="你是一个助手，可以查询天气。回答要简洁。",
    max_steps=6,
)

# 测试1：正常调用
print(agent.run("北京和上海哪个更热？"))
# 预期：Agent调用两次get_weather，比较后回答

agent.reset()

# 测试2：触发死循环检测（构造一个会反复调用的场景）
print(agent.run("查一下火星的天气"))
# 预期：工具返回默认20度，Agent可能反复查，被重复检测或步数限制拦下