2. 实现Agent的记忆模块

任务目标

为Agent实现双层记忆系统：短期记忆（当前会话上下文）与长期记忆（跨会话持久化），使Agent在生成回复时能够检索并利用历史对话信息。

核心实现

import json
import os
from datetime import datetime
from typing import List, Dict, Optional

class ShortTermMemory:
    """短期记忆：维护当前会话的对话历史，有最大长度限制"""
    def __init__(self, max_turns: int = 10):
        self.max_turns = max_turns
        self.history: List[Dict] = []
    
    def add(self, role: str, content: str) -> None:
        self.history.append({"role": role, "content": content, "time": datetime.now().isoformat()})
        if len(self.history) > self.max_turns * 2:
            self.history = self.history[-(self.max_turns * 2):]
    
    def get_recent(self, turns: Optional[int] = None) -> List[Dict]:
        if turns is None:
            return self.history
        return self.history[-(turns * 2):]
    
    def clear(self) -> None:
        self.history = []

class LongTermMemory:
    """长期记忆：基于关键词的持久化记忆存储，跨会话保留"""
    def __init__(self, storage_path: str = "agent_memory.json"):
        self.storage_path = storage_path
        self.memories: Dict[str, List[Dict]] = self._load()
    
    def _load(self) -> Dict[str, List[Dict]]:
        if os.path.exists(self.storage_path):
            with open(self.storage_path, "r", encoding="utf-8") as f:
                return json.load(f)
        return {}
    
    def _save(self) -> None:
        with open(self.storage_path, "w", encoding="utf-8") as f:
            json.dump(self.memories, f, ensure_ascii=False, indent=2)
    
    def _extract_keywords(self, text: str) -> List[str]:
        """简单关键词提取：按长度过滤的词语（实际可替换为更智能的抽取）"""
        stopwords = {"的", "了", "是", "我", "你", "他", "她", "它", "们", "在", "有", "和", "就", "不", "人", "都", "一", "一个", "上", "也", "很", "到", "说", "要", "去", "会", "着", "没有", "看", "好", "自己"}
        words = [w for w in text.replace("，", " ").replace("。", " ").replace("？", " ").replace("！", " ").split() if len(w) > 1 and w not in stopwords]
        return list(set(words))[:5]
    
    def add(self, content: str, metadata: Dict = None) -> None:
        """将重要信息存入长期记忆，按关键词索引"""
        keywords = self._extract_keywords(content)
        entry = {
            "content": content,
            "time": datetime.now().isoformat(),
            "metadata": metadata or {}
        }
        for kw in keywords:
            if kw not in self.memories:
                self.memories[kw] = []
            self.memories[kw].append(entry)
        self._save()
    
    def recall(self, query: str, top_k: int = 3) -> List[Dict]:
        """根据查询召回相关长期记忆"""
        keywords = self._extract_keywords(query)
        results = []
        seen = set()
        for kw in keywords:
            for mem in self.memories.get(kw, []):
                mem_id = id(mem)
                if mem_id not in seen:
                    seen.add(mem_id)
                    results.append(mem)
        # 按时间排序，返回最近top_k条
        results.sort(key=lambda x: x["time"], reverse=True)
        return results[:top_k]

class AgentWithMemory:
    """带记忆能力的Agent：结合短期与长期记忆生成回复"""
    def __init__(self, name: str = "Agent"):
        self.name = name
        self.short_term = ShortTermMemory(max_turns=10)
        self.long_term = LongTermMemory()
    
    def _build_prompt(self, user_input: str) -> str:
        """构建包含记忆的提示词"""
        # 1. 召回相关长期记忆
        relevant_memories = self.long_term.recall(user_input)
        memory_context = ""
        if relevant_memories:
            memory_context = "【相关历史记忆】\n" + "\n".join([f"- {m['content']}" for m in relevant_memories]) + "\n\n"
        
        # 2. 获取短期对话历史
        history = self.short_term.get_recent()
        history_text = ""
        if history:
            history_text = "【当前对话历史】\n" + "\n".join([f"{h['role']}: {h['content']}" for h in history]) + "\n\n"
        
        # 3. 组装完整提示
        prompt = f"{memory_context}{history_text}【用户当前输入】\n{user_input}\n\n请基于以上记忆和对话历史，以{self.name}的身份回复用户。"
        return prompt
    
    def respond(self, user_input: str) -> str:
        """处理用户输入并生成回复"""
        # 将用户输入加入短期记忆
        self.short_term.add("user", user_input)
        
        # 构建带记忆的提示词
        prompt = self._build_prompt(user_input)
        
        # 模拟LLM调用（实际项目中替换为真实模型API）
        # 此处演示记忆如何影响回复——根据记忆内容构造回复
        relevant = self.long_term.recall(user_input)
        response = f"[{self.name}回复] 我记住了您刚才说的内容。"
        if relevant:
            response += f" 根据我的长期记忆，您之前提到过：{relevant[0]['content']}"
        if len(self.short_term.history) > 2:
            response += f" 这是我们对话的第{len(self.short_term.history)//2}轮。"
        
        # 将回复加入短期记忆
        self.short_term.add("assistant", response)
        return response
    
    def remember_important(self, content: str) -> None:
        """主动将重要信息存入长期记忆"""
        self.long_term.add(content, {"source": "user_stated"})

# ===== 使用示例 =====
if __name__ == "__main__":
    agent = AgentWithMemory("记忆助手")
    
    # 第一轮对话：用户告知偏好
    print("用户：我喜欢喝美式咖啡，不加糖")
    agent.remember_important("用户喜欢喝美式咖啡，不加糖")
    r1 = agent.respond("你好，请记住我的咖啡偏好")
    print(r1)
    
    # 第二轮对话：测试短期记忆
    print("\n用户：我刚才说了什么？")
    r2 = agent.respond("我刚才说了什么？")
    print(r2)
    
    # 第三轮对话：测试长期记忆召回
    print("\n用户：帮我推荐一杯咖啡")
    r3 = agent.respond("帮我推荐一杯咖啡")
    print(r3)
    
    # 验证长期记忆持久化
    print("\n=== 验证 ===")
    print(f"长期记忆文件已保存: {os.path.exists('agent_memory.json')}")
    with open("agent_memory.json", "r", encoding="utf-8") as f:
        data = json.load(f)
    print(f"记忆关键词: {list(data.keys())}")
    print(f"记忆条数: {sum(len(v) for v in data.values())}")

运行与验证

运行方式：

python agent_memory.py

验证路径：

程序执行后，观察三轮对话输出：第一轮建立长期记忆；第二轮通过短期记忆回答“我刚才说了什么”；第三轮通过长期记忆召回咖啡偏好。
检查当前目录生成 agent_memory.json 文件，内容包含关键词“咖啡”“美式”等及对应记忆条目。
重新运行程序（新会话），输入“帮我推荐一杯咖啡”，若仍能召回之前的偏好，说明长期记忆跨会话生效。

通过条件：短期记忆能正确引用当前会话前文；长期记忆在程序重启后仍可召回；记忆文件结构完整且可读。

关键说明

短期记忆采用滑动窗口（默认保留最近10轮），超出后自动丢弃最早记录，模拟有限上下文窗口。
长期记忆使用关键词索引的JSON文件持久化，关键词提取为简化实现，实际生产可替换为嵌入向量检索。
示例中respond()方法用规则模拟LLM回复以演示记忆机制，接入真实模型时只需将_build_prompt()的输出传给模型API即可。

📚 课程目录

2. 实现Agent的记忆模块

任务目标

核心实现

运行与验证

关键说明