LangChain与AutoGPT之外：从零手写一个带记忆与循环控制的AI Agent

用户输入 → 构建prompt(系统提示+记忆+工具描述) → 调用LLM
    → 解析输出 → 如果是工具调用 → 执行工具 → 结果写入记忆 → 回到调用LLM
    → 如果是最终答案 → 返回用户

from pydantic import BaseModel, Field
from typing import Callable, Any
import json

class Tool:
    def __init__(self, name: str, description: str, args_model: type[BaseModel], func: Callable):
        self.name = name
        self.description = description
        self.args_model = args_model
        self.func = func

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

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


# 示例工具：计算器
class CalcArgs(BaseModel):
    expression: str = Field(..., description="要计算的数学表达式，例如 '2+3*4'")

def calc(expression: str) -> str:
    # 生产环境请用 ast.literal_eval 或 sympy，这里为演示简化
    allowed = set("0123456789+-*/(). ")
    if not set(expression)  int:
        total = 0
        for m in msgs:
            content = m.get("content") or ""
            total += len(ENC.encode(str(content)))
            if m.get("tool_calls"):
                total += len(ENC.encode(str(m["tool_calls"])))
        return total

    def build(self) -> list[dict]:
        # 从后往前保留，直到超过预算
        kept = []
        budget = self.max_tokens
        for m in reversed(self.messages):
            cost = self._count([m])
            if cost > budget and kept:
                break
            kept.append(m)
            budget -= cost
        kept.reverse()
        # 关键：保证 assistant(tool_calls) 和紧随其后的 tool 消息成对出现
        return [self.system] + self._fix_tool_pairs(kept)

    def _fix_tool_pairs(self, msgs: list[dict]) -> list[dict]:
        result = []
        i = 0
        while i < len(msgs):
            m = msgs[i]
            if m.get("role") == "assistant" and m.get("tool_calls"):
                # 找到后续所有 tool 消息
                j = i + 1
                tools = []
                while j < len(msgs) and msgs[j].get("role") == "tool":
                    tools.append(msgs[j])
                    j += 1
                if len(tools) == len(m["tool_calls"]):
                    result.append(m)
                    result.extend(tools)
                i = j
            else:
                result.append(m)
                i += 1
        return result

import time
from openai import OpenAI, APIError, RateLimitError

client = OpenAI(api_key="sk-xxx", timeout=30.0)

SYSTEM_PROMPT = """你是一个会使用工具的助手。
规则：
1. 需要计算时调用 calculator 工具。
2. 工具返回以 [TOOL_ERROR] 或 [ARG_ERROR] 开头时，说明调用失败，请修正参数后重试，最多2次。
3. 拿到足够信息后，直接给出最终答案，不要再调用工具。"""

def call_llm_with_retry(messages, tools, max_retry=3):
    for attempt in range(max_retry):
        try:
            return client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                tools=[t.to_schema() for t in tools],
                tool_choice="auto",
                temperature=0.2,
            )
        except RateLimitError:
            wait = 2 ** attempt
            print(f"[WARN] 限流，{wait}s 后重试")
            time.sleep(wait)
        except APIError as e:
            if attempt == max_retry - 1:
                raise
            time.sleep(1)
    raise RuntimeError("LLM调用失败")


def run_agent(user_input: str, tools: list[Tool], max_steps: int = 8):
    tool_map = {t.name: t for t in tools}
    mem = Memory(SYSTEM_PROMPT, max_tokens=3000)
    mem.add({"role": "user", "content": user_input})

    for step in range(max_steps):
        print(f"\n===== Step {step + 1}/{max_steps} =====")
        resp = call_llm_with_retry(mem.build(), tools)
        msg = resp.choices[0].message

        # 没有工具调用 → 最终答案
        if not msg.tool_calls:
            mem.add({"role": "assistant", "content": msg.content})
            return msg.content

        # 有工具调用 → 逐个执行
        mem.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 = tc.function.arguments
            print(f"[TOOL] {name}({args})")
            if name not in tool_map:
                result = f"[TOOL_ERROR] 未知工具: {name}"
            else:
                result = tool_map[name].run(args)
            print(f"[RESULT] {result[:200]}")
            mem.add({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": result,
            })

    return "[FAILED] 达到最大步数仍未完成，请简化任务或检查工具定义。"