LangChain与AutoGPT启发下从零构建AI Agent循环控制与记忆机制

# requirements.txt
langchain==0.1.0
langchain-openai==0.0.2.post1
openai==1.10.0
pydantic==2.5.3
redis==5.0.1
python-dotenv==1.0.0

from pydantic import BaseModel, Field, validator
from typing import Optional, List, Dict
import json

class CodeExecutorTool(BaseModel):
    """代码执行工具，带沙箱限制"""
    code: str = Field(description="Python代码，必须是纯函数")
    timeout: int = Field(default=10, ge=1, le=30, description="超时秒数")
    memory_limit: int = Field(default=128, ge=64, le=1024, description="内存MB")

    @validator('code')
    def validate_code(cls, v):
        if 'os.system' in v or 'subprocess' in v:
            raise ValueError("禁止执行系统命令")
        if len(v) > 2000:
            raise ValueError("代码长度超限，请精简逻辑")
        return v

# 工具注册表
TOOL_REGISTRY = {
    "code_executor": {
        "name": "code_executor",
        "description": "执行Python代码并返回结果，适合数学计算和数据处理",
        "schema": CodeExecutorTool,
        "handler": lambda ctx: _safe_execute(ctx.code, ctx.timeout, ctx.memory_limit)
    }
}

def _safe_execute(code: str, timeout: int, memory_limit: int):
    """使用resource模块限制资源，避免Agent写死循环"""
    import resource, signal, subprocess

    def handler(signum, frame):
        raise TimeoutError("代码执行超时")

    signal.signal(signal.SIGALRM, handler)
    signal.alarm(timeout)

    try:
        # 用exec隔离执行，注入受限全局变量
        local_ns = {}
        exec(code, {"__builtins__": __import__('builtins')}, local_ns)
        return local_ns.get('result', "执行成功，无返回值")
    except TimeoutError:
        return "错误：代码执行超时，请检查是否有死循环"
    finally:
        signal.alarm(0)

import redis, sqlite3, json
from collections import deque

class MemoryManager:
    def __init__(self, short_mem_size=20, redis_url="redis://localhost:6379/0"):
        self.short_mem = deque(maxlen=short_mem_size)
        self.redis = redis.from_url(redis_url)
        self.sqlite = sqlite3.connect('long_term_memory.db')
        self._init_db()

    def _init_db(self):
        self.sqlite.execute('''CREATE TABLE IF NOT EXISTS long_memory 
            (id INTEGER PRIMARY KEY AUTOINCREMENT,
             summary TEXT,
             created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
        self.sqlite.commit()

    def add_short(self, role: str, content: str):
        """添加短期记忆，自动淘汰最旧的"""
        entry = {"role": role, "content": content}
        self.short_mem.append(entry)
        # 当短期记忆满时，生成摘要存入长期
        if len(self.short_mem) == self.short_mem.maxlen:
            self._summarize()

    def _summarize(self):
        """用LLM生成摘要存入长期记忆"""
        conversation = "\n".join([f"{e['role']}: {e['content']}" 
                                  for e in self.short_mem])
        try:
            summary = self.llm.invoke(f"请压缩以下对话为100字摘要：\n{conversation}")
            self.sqlite.execute(
                "INSERT INTO long_memory (summary) VALUES (?)",
                (summary.content,))
            self.sqlite.commit()
            # 关键：清空短期记忆，保留最近2条作为上下文衔接
            recent = list(self.short_mem)[-2:]
            self.short_mem.clear()
            self.short_mem.extend(recent)
        except Exception as e:
            print(f"摘要生成失败: {e}，保留短期记忆")

    def get_context(self) -> str:
        """组装提示词上下文"""
        short_text = "\n".join([f"{e['role']}: {e['content']}" 
                                for e in self.short_mem])
        # 从长期记忆取最近3条摘要
        cursor = self.sqlite.execute(
            "SELECT summary FROM long_memory ORDER BY id DESC LIMIT 3")
        long_text = "\n".join([row[0] for row in cursor.fetchall()])
        return f"[长期记忆]\n{long_text}\n[短期记忆]\n{short_text}"

class AgentLoopController:
    def __init__(self, max_iters=10, max_tokens_budget=5000):
        self.max_iters = max_iters
        self.iter_count = 0
        self.tokens_used = 0
        self.max_tokens = max_tokens_budget
        self.history = []  # 记录每轮决策

    def should_continue(self, response_tokens: int) -> str:
        """
        返回值：'continue' / 'summarize' / 'terminate'
        """
        self.iter_count += 1
        self.tokens_used += response_tokens

        # 阈值1：迭代次数超限
        if self.iter_count >= self.max_iters:
            return 'terminate'

        # 阈值2：Token预算消耗达80%
        if self.tokens_used >= self.max_tokens * 0.8:
            return 'summarize'  # 强制总结模式

        # 阈值3：连续3轮没有新工具调用
        recent_tools = [h['tool'] for h in self.history[-3:]]
        if len(recent_tools) == 3 and len(set(recent_tools)) == 1:
            return 'terminate'  # 陷入同一个工具死循环

        return 'continue'

    def log_decision(self, tool: str, reasoning: str):
        self.history.append({
            'iter': self.iter_count,
            'tool': tool,
            'reasoning': reasoning,
            'tokens': self.tokens_used
        })