LLM Agent 记忆系统设计指南：让AI拥有"记忆"的技术原理与实战

messages = [
    {"role": "system", "content": "你是助手"},
    {"role": "user", "content": "第一轮对话"},
    {"role": "assistant", "content": "第一轮回复"},
    {"role": "user", "content": "第二轮对话"},
    # ... 所有历史
]
response = chat(messages)

def summarize_if_needed(conversation_history: list, threshold: int = 10):
    if len(conversation_history) >= threshold:
        summary_prompt = f"""请摘要以下对话的核心要点，保留关键信息：
        {conversation_history}"""
        summary = llm.chat([{"role": "user", "content": summary_prompt}])
        return {"summary": summary, "original_count": len(conversation_history)}
    return None

from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings

# 写入记忆
def write_memory(text: str, vectorstore: Chroma):
    vectorstore.add_texts([text])

# 检索记忆
def retrieve_memory(query: str, vectorstore: Chroma, k: int = 5) -> list[str]:
    docs = vectorstore.similarity_search(query, k=k)
    return [doc.page_content for doc in docs]

class AttentionMemory:
    def __init__(self, max_turns: int = 10):
        self.messages: list[dict] = []
        self.max_turns = max_turns
    
    def add(self, role: str, content: str):
        self.messages.append({"role": role, "content": content})
        # 超过阈值，删除最早的回合
        if len(self.messages) > self.max_turns * 2:
            self.messages = self.messages[2:]
    
    def get_context(self) -> list[dict]:
        return self.messages

class SemanticMemory:
    def __init__(self, vectorstore: Chroma):
        self.store = vectorstore
        self.metadata_filter = {}  # 可按 type 过滤
    
    def write(self, text: str, memory_type: str, session_id: str):
        self.store.add_texts(
            texts=[text],
            metadatas=[{"type": memory_type, "session_id": session_id}]
        )
    
    def retrieve(self, query: str, memory_types: list[str] | None = None, k: int = 5) -> list[dict]:
        return self.store.similarity_search_with_score(
            query, k=k, filter={"type": {"$in": memory_types}} if memory_types else None
        )

SYSTEM_PROMPT = """你是一个专业的 AI 编程助手。

## 系统能力
- 可以使用 Read/Edit/Write 工具操作本地文件
- 可以使用 Exec 工具执行 Shell 命令
- 可以使用 Browser 工具控制浏览器

## 工作流程
当用户提出需求时，你应该：
1. 理解需求并拆解任务步骤
2. 按步骤执行，及时反馈进度
3. 完成后总结做了什么、结果是什么

## 约束
- 不要执行删除、格式化等危险操作
- 遇到错误先分析原因，再尝试修复
"""

from langchain.retrievers import EnsembleRetriever

# 混合检索
retriever = EnsembleRetriever(
    retrievers=[
        bm25_retriever,  # 关键词
        vectorstore.as_retriever(search_kwargs={"k": 10})  # 向量
    ],
    weights=[0.3, 0.7]  # 权重分配
)
results = retriever.get_relevant_documents(query)

from langchain.compressors import LLMChainExtractor

compressor = LLMChainExtractor.from_llm(llm)

from langchain.retrievers import ContextualCompressionRetriever
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=retriever
)
# 返回的是压缩后的精简片段
compressed_docs = compression_retriever.get_relevant_documents(query)

class MemoryWritePolicy:
    def should_write(self, message: dict, context: dict) -> bool:
        content = message["content"]
        
        # 规则1：用户明确声明
        if re.match(r"记住，.*", content):
            return True
        
        # 规则2：执行了外部操作
        if message.get("role") == "system" and "tool_output" in content:
            return True
        
        # 规则3：识别到关键实体（项目名、人名、技术栈）
        entities = extract_entities(content)  # 用 NER 或规则
        if len(entities) > 2:
            return True
        
        # 规则4：对话涉及长期项目上下文
        if context.get("task_type") == "project_work":
            return True
        
        return False

from langchain.llms import OpenAI
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain.agents import AgentExecutor, ZeroShotAgent
from langchain.prompts import PromptTemplate

class SimpleAgentMemory:
    def __init__(self, llm, vectorstore_path: str = "./memory_db"):
        self.llm = llm
        self.embeddings = OpenAIEmbeddings()
        self.vectorstore = Chroma(
            persist_directory=vectorstore_path,
            embedding_function=self.embeddings
        )
        # 注意力记忆：最近 10 轮
        self.attention = ConversationBufferMemory(memory_key="chat_history", k=10)
    
    def retrieve(self, query: str, k: int = 5) -> str:
        """从向量数据库检索相关记忆"""
        docs = self.vectorstore.similarity_search(query, k=k)
        return "\n".join([f"- {d.page_content}" for d in docs])
    
    def write(self, text: str, memory_type: str = "general"):
        """写入记忆到向量数据库"""
        self.vectorstore.add_texts(
            texts=[text],
            metadatas=[{"type": memory_type}]
        )
        self.vectorstore.persist()
    
    def chat(self, user_input: str) -> str:
        # Step 1: 检索相关记忆
        relevant_memory = self.retrieve(user_input)
        
        # Step 2: 构建 Prompt（注意力记忆 + 语义记忆）
        attention_context = self.attention.load_memory_variables({})["chat_history"]
        
        prompt = f"""你是 AI 助手。以下是相关记忆：
{relevant_memory}

当前对话历史：
{attention_context}

用户：{user_input}
助手："""
        
        # Step 3: LLM 推理
        response = self.llm(prompt)
        
        # Step 4: 更新注意力记忆
        self.attention.save_context(
            {"input": user_input},
            {"output": response}
        )
        
        # Step 5: 关键内容写入语义记忆
        if self._is_important(user_input) or self._is_important(response):
            self.write(f"用户问：{user_input}\n助手答：{response}")
        
        return response
    
    def _is_important(self, text: str) -> bool:
        # 简单规则：包含项目名、技术决策、用户偏好等
        keywords = ["项目", "记住", "偏好", "配置", "名字", "我的"]
        return any(kw in text for kw in keywords)

# 使用
memory = SimpleAgentMemory(llm=OpenAI(temperature=0))
response = memory.chat("我叫张三，在做一个叫code-minions的项目")
print(response)

# 索引结构设计
collection_name = f"{user_id}_{topic}_{session_id}"
# 例如：user123_code-minions_session-001

# 乐观锁写入
def write_with_version(vectorstore, text, expected_version):
    current = vectorstore.get_version()
    if current != expected_version:
        raise MemoryConflictError("记忆已被其他 Agent 修改")
    vectorstore.add_texts([text])
    vectorstore.set_version(current + 1)