第4章：Memory — 让 Agent 拥有"记忆"

📌 本章目标

• 理解短期记忆和长期记忆的区别
• 实现带上下文窗口管理的短期记忆
• 实现简单的长期记忆（关键事实提取）
• 处理上下文超长的问题

---

4.1 为什么 Agent 需要记忆？

想象一个没有记忆的对话：

用户: 我叫小明
Agent: 你好，小明！
用户: 我喜欢 Python
Agent: Python 很棒！
用户: 我之前说我叫什么来着？
Agent: 抱歉，我不知道你的名字。（忘了！）

没有记忆的 Agent，每一次对话都是"重新认识你"。这在真实场景中是完全不可接受的。

---

4.2 短期记忆：对话历史

最简单的方式：存 messages 数组

其实我们在前两章已经在用短期记忆了——就是那个 messages 列表！

# 短期记忆 = 不断追加的 messages 列表
messages = [
    {"role": "system", "content": "你是一个助手"},
    {"role": "user",   "content": "我叫小明"},       # ← 第1轮
    {"role": "assistant", "content": "你好小明！"},   # ← 第1轮
    {"role": "user",   "content": "我喜欢Python"},    # ← 第2轮
    {"role": "assistant", "content": "Python很棒！"},  # ← 第2轮
    {"role": "user",   "content": "我叫什么？"},       # ← 第3轮（能回答！）
]

问题来了：messages 越来越长怎么办？

DeepSeek 的上下文窗口是 128K token（约 10 万字），看起来很够用。但问题是：

1. 成本：每次把全部历史发过去，token 消耗越来越多
2. 速度：消息越长，响应越慢
3. 注意力稀释：模型对越靠前的信息关注越少

解决方案：滑动窗口

"""
带滑动窗口的对话记忆
"""
from collections import deque

class SlidingWindowMemory:
    """短期记忆：保留最近 N 条消息"""
    
    def __init__(self, max_messages: int = 20, system_prompt: str = ""):
        """
        参数:
            max_messages: 最多保留多少条消息（不含 system prompt）
            system_prompt: 系统提示词，始终保留在最前面
        """
        self.max_messages = max_messages
        self.system_prompt = system_prompt
        self.messages = deque(maxlen=max_messages)  # 自动丢弃旧消息
    
    def add_user(self, content: str):
        """添加用户消息"""
        self.messages.append({"role": "user", "content": content})
    
    def add_assistant(self, content: str):
        """添加助手消息"""
        self.messages.append({"role": "assistant", "content": content})
    
    def add_tool_result(self, tool_call_id: str, content: str):
        """添加工具执行结果"""
        self.messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "content": content
        })
    
    def build_context(self) -> list[dict]:
        """构建发送给 LLM 的完整上下文"""
        context = []
        if self.system_prompt:
            context.append({"role": "system", "content": self.system_prompt})
        context.extend(list(self.messages))
        return context
    
    def __len__(self):
        return len(self.messages)


# ── 测试滑动窗口 ──
if __name__ == "__main__":
    memory = SlidingWindowMemory(max_messages=6, system_prompt="你是一个有记忆的助手")
    
    # 模拟 5 轮对话
    for i in range(5):
        memory.add_user(f"第{i+1}轮的问题")
        memory.add_assistant(f"第{i+1}轮的回答")
    
    context = memory.build_context()
    print(f"总共 {len(context)} 条消息（含 system prompt）")
    for msg in context:
        print(f"  [{msg['role']}] {msg['content'][:30]}...")
    
    # 输出：
    # [system] 你是一个有记忆的助手
    # [user] 第3轮的问题...          ← 前两轮被自动丢弃了！
    # [assistant] 第3轮的回答...
    # [user] 第4轮的问题...
    # [assistant] 第4轮的回答...
    # [user] 第5轮的问题...
    # [assistant] 第5轮的回答...

---

4.3 长期记忆：记住关键信息

滑动窗口能解决消息数量问题，但不能解决"记住重要信息"的问题。

我们需要一个长期记忆系统：自动从对话中提取关键信息，存下来，下次对话时注入。

简易长期记忆实现

"""
长期记忆 + 带记忆的 Agent —— 可直接运行
使用前请先: pip install openai python-dotenv
并在 .env 文件中设置: DEEPSEEK_API_KEY=sk-你的密钥
"""
import json
import os
from collections import deque
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com/v1"
)

def chat(messages: list, model: str = "deepseek-chat", temperature: float = 0.7) -> str:
    """封装 LLM 调用"""
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
    )
    return response.choices[0].message.content

# ============================================
# 短期记忆：滑动窗口
# ============================================
class SlidingWindowMemory:
    """短期记忆：保留最近 N 条消息"""
    
    def __init__(self, max_messages: int = 20, system_prompt: str = ""):
        self.max_messages = max_messages
        self.system_prompt = system_prompt
        self.messages = deque(maxlen=max_messages)
    
    def add_user(self, content: str):
        self.messages.append({"role": "user", "content": content})
    
    def add_assistant(self, content: str):
        self.messages.append({"role": "assistant", "content": content})
    
    def add_tool_result(self, tool_call_id: str, content: str):
        self.messages.append({"role": "tool", "tool_call_id": tool_call_id, "content": content})
    
    def build_context(self) -> list:
        context = []
        if self.system_prompt:
            context.append({"role": "system", "content": self.system_prompt})
        context.extend(list(self.messages))
        return context

# ============================================
# 长期记忆：自动提取关键信息
# ============================================
class LongTermMemory:
    """用 LLM 自动从对话中提取关键信息"""
    
    def __init__(self):
        self.facts: dict[str, str] = {}  # {事实描述: 详情}
    
    def extract_facts(self, user_message: str, assistant_response: str):
        """
        从一轮对话中自动提取用户的关键信息

        思路：让 LLM 判断这段对话里有没有值得记住的信息
        """
        prompt = f"""从以下对话中提取关于用户的关键信息。如果没有值得长期记住的信息，回复"无"。

用户: {user_message}
助手: {assistant_response}

请以 JSON 格式输出，每个信息包含 key 和 value：
{{"facts": [{{"key": "用户名", "value": "小明"}}, {{"key": "职业", "value": "程序员"}}]}}

注意：
- 只提取持久性的信息（偏好、身份、重要背景等）
- 不要提取一次性的请求（如"帮我查天气"）
- 如果没有值得记住的信息，返回 {{"facts": []}}
"""
        result = chat([
            {"role": "system", "content": "你是信息提取器。只输出 JSON。"},
            {"role": "user", "content": prompt}
        ])
        
        try:
            data = json.loads(result)
            for fact in data.get("facts", []):
                self.remember(fact["key"], fact["value"])
        except json.JSONDecodeError:
            pass  # 解析失败就跳过
    
    def remember(self, key: str, value: str):
        """手动记住一个事实"""
        self.facts[key] = value
        print(f"🧠 记住: {key} = {value}")
    
    def forget(self, key: str):
        """忘记一个事实"""
        if key in self.facts:
            del self.facts[key]
            print(f"🗑️ 忘记: {key}")
    
    def recall_all(self) -> str:
        """把所有记住的信息格式化为一段文本"""
        if not self.facts:
            return ""
        
        lines = ["\n【你已知的关于用户的长期信息】"]
        for key, value in self.facts.items():
            lines.append(f"- {key}: {value}")
        return "\n".join(lines)
    
    def get(self, key: str) -> str | None:
        """查询特定信息"""
        return self.facts.get(key)


# ── 集成到 Agent 中 ──
class AgentWithMemory:
    """带记忆的 Agent"""
    
    def __init__(self, system_prompt: str = "你是一个贴心的 AI 助手"):
        self.short_term = SlidingWindowMemory(
            max_messages=30,
            system_prompt=system_prompt
        )
        self.long_term = LongTermMemory()
    
    def chat(self, user_input: str) -> str:
        """处理一轮对话"""
        # 构建上下文（短期记忆 + 长期记忆）
        context = self.short_term.build_context()
        
        # 如果有长期记忆，注入到 system prompt 中
        facts_text = self.long_term.recall_all()
        if facts_text and context[0]["role"] == "system":
            context[0]["content"] += facts_text
        
        # 添加用户输入
        self.short_term.add_user(user_input)
        context = self.short_term.build_context()
        # 重新注入长期记忆
        if facts_text and context[0]["role"] == "system":
            context[0]["content"] += facts_text
        
        # 调用 LLM（使用前面定义的 chat 函数）
        response = chat(context)
        
        # 保存回复
        self.short_term.add_assistant(response)
        
        # 自动提取长期记忆
        self.long_term.extract_facts(user_input, response)
        
        return response


# ── 测试 ──
if __name__ == "__main__":
    agent = AgentWithMemory()
    
    print("=== 第1轮 ===")
    print(agent.chat("你好！我叫小明，我是一名 Python 后端开发者"))
    
    print("\n=== 第2轮 ===")
    print(agent.chat("我喜欢用 FastAPI 框架"))
    
    print("\n=== 第3轮 ===")
    print(agent.chat("帮我推荐一个适合我的 Python 学习方向"))
    # Agent 应该能结合"Python后端"和"FastAPI"给出个性化推荐
    
    print("\n=== 第4轮 ===")
    print(agent.chat("你还记得我的名字和职业吗？"))
    # Agent 应该能回答出"小明"和"Python后端开发者"
    
    print(f"\n🧠 长期记忆内容: {agent.long_term.facts}")

---

4.4 记忆总结：当对话太长时

有时候对话虽然没超过上下文窗口，但已经很长了，每次发送完整历史太浪费。这时候我们可以对历史做摘要。

"""
记忆压缩 —— 将本节代码追加到上一节的完整代码中即可运行
"""
import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com/v1"
)

def chat(messages: list, model: str = "deepseek-chat", temperature: float = 0.7) -> str:
    response = client.chat.completions.create(
        model=model, messages=messages, temperature=temperature,
    )
    return response.choices[0].message.content

def summarize_history(messages: list[dict]) -> str:
    """
    对对话历史做摘要，压缩成一段文字
    
    这样旧的对话不用逐条发送，只发摘要就行
    """
    # 把历史格式化成文本
    history_text = ""
    for msg in messages:
        role = msg["role"]
        content = msg["content"][:200]  # 每条截取前200字符
        history_text += f"[{role}]: {content}\n"
    
    summary_prompt = f"""请用3-5句话总结以下对话的关键内容，只保留重要的信息：

{history_text}

总结："""
    
    return chat([
        {"role": "system", "content": "你是一个对话摘要专家，简洁准确。"},
        {"role": "user", "content": summary_prompt}
    ])


# 使用方式：
# 当 messages 超过一定长度时，对前半部分做摘要
def compress_memory(memory: SlidingWindowMemory, keep_last: int = 6):
    """
    压缩记忆：保留最近 N 条原文，更早的做摘要
    """
    messages = list(memory.messages)
    if len(messages) <= keep_last:
        return  # 不需要压缩
    
    old = messages[:-keep_last]
    recent = messages[-keep_last:]
    
    summary = summarize_history(old)
    
    # 用摘要替代旧消息
    memory.messages.clear()
    memory.messages.append({
        "role": "system",
        "content": f"[对话历史摘要] {summary}"
    })
    for msg in recent:
        memory.messages.append(msg)
    
    print(f"📦 压缩完成：{len(old)} 条消息 → 1 条摘要")

---

4.5 记忆系统架构总览

记忆类型	存储什么	保留时间	实现方式
短期记忆	最近 N 条对话	当前对话	deque(maxlen=N)
长期记忆	用户身份、偏好、关键事实	跨对话	LLM 提取 → dict 存储
摘要记忆	超过窗口的旧对话压缩版	当前对话	LLM 总结 → 一段文本

---

📝 本章小结

1. 短期记忆 = messages 列表，用滑动窗口控制长度
2. 长期记忆 = 用 LLM 自动提取关键事实，存下来反复使用
3. 记忆压缩 = 当对话太长时，用 LLM 对旧内容做摘要
4. 核心原则：不让 LLM 漏掉重要信息，也不让它被无关信息淹没

---

✏️ 练习题

1. 基础题：把 SlidingWindowMemory 集成到第3章的 agent_loop 中，让天气 Agent 也能记住之前的对话。

2. 进阶题：修改 LongTermMemory，让它能处理"更新"——比如用户先说"我喜欢 Python"，后来说"我转学 Go 了"，它能更新事实而不是追加矛盾的信息。

3. 挑战题：实现一个简单的持久化——把 LongTermMemory 的 facts 保存到 JSON 文件，下次启动时加载回来。

---

下一章预告：第5章：Planning — 让 Agent 学会"思考" —— 工具和记忆都有了，但 Agent 还需要学会"三思而后行"！