LangChain与AutoGPT双驱动：手写Agent循环控制与记忆管理

Python: 3.11.7
langchain: 0.1.0 (关键：0.1.x的BaseTool必须实现`_run`而不是`run`)
langchain-openai: 0.0.5
openai: 1.12.0
tiktoken: 0.5.2 (用于记忆Token裁剪)

# tools.py - 基于LangChain的Tool抽象
from langchain.tools import BaseTool
from pydantic import BaseModel, Field
import random

class DBQueryInput(BaseModel):
    sql: str = Field(description="SQL query string")
    limit: int = Field(default=10, description="max rows")

class DBQueryTool(BaseTool):
    name = "database_query"
    description = "Execute SQL-like query on sales DB. Returns text result."
    args_schema = DBQueryInput

    def _run(self, sql: str, limit: int = 10) -> str:
        # 模拟真实数据库延迟
        import time; time.sleep(0.3)
        # 危险操作拦截
        if "delete" in sql.lower() or "drop" in sql.lower():
            raise ValueError("Destructive SQL blocked")
        fake_result = [{"product": f"item_{i}", "sales": random.randint(100,999)} for i in range(limit)]
        return str(fake_result[:limit])

    async def _arun(self, *args, **kwargs):
        return self._run(*args, **kwargs)

# 注册时用handle_tool_error包裹
from langchain.tools import tool
@tool("math_calc", handle_tool_error=True)
def math_calc(expression: str) -> str:
    """Evaluate math expression safely. Use python eval with restrictions."""
    allowed = set("0123456789+-*/(). ")
    if not all(c in allowed for c in expression):
        return "ERROR: Invalid characters"
    try:
        result = eval(expression, {"__builtins__": {}}, {})
        return f"Result: {result}"
    except Exception as e:
        return f"ERROR: {str(e)}"

# agent_core.py - 手写Agent循环
from langchain_openai import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage, AIMessage
import tiktoken, json, re

class SimpleAgent:
    def __init__(self, tools, llm, memory_budget=2000):
        self.tools = {t.name: t for t in tools}
        self.llm = llm
        self.memory = []
        self.encoder = tiktoken.encoding_for_model("gpt-4")
        self.budget = memory_budget

    def _trim_memory(self):
        """裁剪最旧的非System非Tool记录"""
        # 简单实现：如果总Token超预算，移除第一条Human/AI消息
        total_tokens = sum(len(self.encoder.encode(m.content)) for m in self.memory)
        while total_tokens > self.budget and len(self.memory) > 2:
            removed = self.memory.pop(1)  # 移除第一条非System
            total_tokens -= len(self.encoder.encode(removed.content))

    def _parse_action(self, text):
        """从LLM输出中解析Action。格式为: 
        Action: tool_name\nAction Input: {"arg": "value"}
        """
        match = re.search(r"Action: (\w+)\nAction Input: (\{.*\})", text, re.DOTALL)
        if match:
            tool_name, input_str = match.group(1), match.group(2)
            try:
                input_dict = json.loads(input_str)
                return tool_name, input_dict
            except json.JSONDecodeError:
                return None, None
        return None, None

    def run(self, task: str):
        # 初始化系统提示词
        sys_prompt = SystemMessage(content="""You are an agent with tools.
Available tools: {tool_desc}
You must respond with:
Thought: your reasoning
Action: tool_name
Action Input: {{"arg": "value"}}

If you have the final answer, respond with:
Final Answer: your answer""".format(tool_desc=", ".join(self.tools.keys())))

        self.memory = [sys_prompt, HumanMessage(content=f"Task: {task}")]

        for iteration in range(8):  # 循环守护
            # 调用LLM
            response = self.llm.invoke(self.memory)
            ai_text = response.content
            self.memory.append(AIMessage(content=ai_text))

            # 判断是否是最终回答
            if "Final Answer:" in ai_text:
                final = ai_text.split("Final Answer:")[1].strip()
                return {"success": True, "answer": final, "iterations": iteration+1, "tokens_used": self._count_tokens()}

            # 解析工具调用 - 带错误重试逻辑
            tool_name, input_dict = self._parse_action(ai_text)
            if tool_name is None or tool_name not in self.tools:
                error_msg = "Invalid action format. Please use correct format."
                self.memory.append(HumanMessage(content=f"Observation: {error_msg}"))
                continue  # 直接进入下一轮，让模型修正

            # 执行工具，捕获异常作为观察值
            try:
                tool_result = self.tools[tool_name].run(input_dict)
                observation = f"Observation: {tool_result}"
            except Exception as e:
                observation = f"Observation: Tool ERROR - {str(e)}"

            self.memory.append(HumanMessage(content=observation))
            self._trim_memory()  # 记忆管理

        return {"success": False, "answer": "Max iterations reached", "iterations": 8}

    def _count_tokens(self):
        return sum(len(self.encoder.encode(m.content)) for m in self.memory)