11 | 高级 RAG 与记忆系统——从能用到好用

用户问题 → Embedding → 向量数据库检索 Top-K → 注入 Prompt → LLM 生成回答

问题："Transformer 论文的第一作者现在在哪家公司？"

用户问题："李飞飞创立的 AI 实验室获得了哪家公司的投资？"

第 1 步：用原始问题检索
→ 找到文档：李飞飞是斯坦福 HAI 的联合创始人。
→ 提取关键实体：HAI（Stanford Institute for Human-Centered AI）

第 2 步：用"HAI 投资 融资"作为新查询检索
→ 找到文档：HAI 获得了 Google、Microsoft 等多家公司的资助。
→ 综合两步结果，生成最终回答。

def multi_hop_rag(question, llm, retriever, max_hops=3):
    """多步检索：每一步基于上一步的结果生成新的查询"""

    context = []
    current_query = question

    for hop in range(max_hops):
        # 1. 检索
        docs = retriever.search(current_query, top_k=3)
        context.extend(docs)

        # 2. 让 LLM 判断：信息够了吗？还缺什么？
        reflection = llm.invoke(f"""
基于以下已检索的信息：
{format_docs(context)}

原始问题：{question}

请判断：
1. 已有信息是否足够回答这个问题？（是/否）
2. 如果不够，还需要检索什么信息？请生成一个新的搜索查询。
""")

        # 3. 如果够了，直接生成回答
        if "是" in reflection.content:
            break

        # 4. 不够的话，用新生成的查询继续检索
        current_query = extract_new_query(reflection.content)

    # 最终回答
    answer = llm.invoke(f"""
基于以下检索到的信息：
{format_docs(context)}

请回答：{question}
""")
    return answer

用户问题
  │
  ├─ 模型判断：需要检索吗？
  │   ├─ 不需要 → 直接生成回答
  │   └─ 需要 ↓
  │
  ├─ 检索文档
  │
  ├─ 模型判断：文档相关吗？
  │   ├─ 不相关 → 换策略 / 告知用户
  │   └─ 相关 ↓
  │
  ├─ 基于文档生成回答
  │
  ├─ 模型判断：回答被文档支持吗？
  │   ├─ 不支持 → 重新生成
  │   └─ 支持 ↓
  │
  └─ 输出最终回答

def self_rag(question, llm, retriever):
    """Self-RAG：LLM 自主决定是否需要检索"""

    # Step 1: 判断是否需要检索
    need_retrieval = llm.invoke(f"""
问题：{question}

判断这个问题是否需要外部知识来回答。
- 如果是常识/通用知识，回答"不需要"
- 如果是特定领域/时效性/私有知识，回答"需要"
""")

    if "不需要" in need_retrieval.content:
        # 直接回答
        return llm.invoke(f"请回答：{question}")

    # Step 2: 检索
    docs = retriever.search(question, top_k=5)

    # Step 3: 过滤不相关文档
    relevant_docs = []
    for doc in docs:
        relevance = llm.invoke(f"""
问题：{question}
文档：{doc.text}
这个文档与问题直接相关吗？回答"相关"或"不相关"。
""")
        if "相关" in relevance.content:
            relevant_docs.append(doc)

    if not relevant_docs:
        return "抱歉，我没有找到与您问题相关的信息。"

    # Step 4: 基于相关文档生成回答
    answer = llm.invoke(f"""
基于以下参考文档回答问题。如果文档中没有足够信息，请明确说明。

{format_docs(relevant_docs)}

问题：{question}
""")

    # Step 5: 验证回答是否被文档支持
    verification = llm.invoke(f"""
回答：{answer.content}
参考文档：{format_docs(relevant_docs)}
这个回答是否完全被参考文档支持？回答"支持"或"不支持"。
""")

    if "不支持" in verification.content:
        # 重新生成，更严格地约束
        answer = llm.invoke(f"""
严格基于以下文档回答问题，不要添加文档中没有的信息。
{format_docs(relevant_docs)}
问题：{question}
""")

    return answer

文档 → 实体提取 → 构建知识图谱 → 图检索 + 文本检索 → LLM 回答

# 伪代码：GraphRAG 的实体和关系提取
def extract_entities_and_relations(text, llm):
    """用 LLM 从文本中提取知识图谱的节点和边"""

    result = llm.invoke(f"""
从以下文本中提取所有实体和它们之间的关系。

文本：{text}

请按以下格式输出：
实体：[名称, 类型]
关系：[实体A, 关系类型, 实体B]

示例：
实体：[Satya Nadella, Person]
实体：[Microsoft, Company]
关系：[Satya Nadella, CEO_of, Microsoft]
""")
    return parse_graph_elements(result.content)


# 查询时的图遍历
def graph_rag_query(question, knowledge_graph, text_retriever, llm):
    """Graph RAG：结合图检索和文本检索"""

    # 1. 从问题中提取关键实体
    entities = extract_entities(question, llm)

    # 2. 在知识图谱中找到相关实体和它们的关系
    subgraph = knowledge_graph.get_subgraph(
        entities,
        depth=2  # 探索 2 跳以内的关系
    )

    # 3. 同时做文本检索
    text_docs = text_retriever.search(question, top_k=3)

    # 4. 把图结构和文本一起给 LLM
    answer = llm.invoke(f"""
知识图谱信息：
{format_subgraph(subgraph)}

相关文档：
{format_docs(text_docs)}

请回答：{question}
""")
    return answer

from rank_bm25 import BM25Okapi
import numpy as np


class HybridRetriever:
    """混合检索：向量 + BM25"""

    def __init__(self, documents, embeddings, alpha=0.5):
        self.documents = documents
        self.embeddings = embeddings  # 预计算的文档 embedding
        self.bm25 = BM25Okapi([doc.split() for doc in documents])
        self.alpha = alpha  # 向量 vs BM25 的权重

    def search(self, query, query_embedding, top_k=5):
        # 1. 向量检索分数
        vector_scores = cosine_similarity(query_embedding, self.embeddings)

        # 2. BM25 关键词检索分数
        bm25_scores = self.bm25.get_scores(query.split())

        # 3. 归一化到 [0, 1]
        vector_scores = normalize(vector_scores)
        bm25_scores = normalize(bm25_scores)

        # 4. 加权融合
        combined_scores = (
            self.alpha * vector_scores +
            (1 - self.alpha) * bm25_scores
        )

        # 5. 取 Top-K
        top_indices = np.argsort(combined_scores)[-top_k:][::-1]
        return [
            {"text": self.documents[i], "score": combined_scores[i]}
            for i in top_indices
        ]


# 使用
retriever = HybridRetriever(documents, embeddings, alpha=0.6)
results = retriever.search(
    query="Python asyncio TaskGroup",
    query_embedding=embed_model("Python asyncio TaskGroup"),
    top_k=5
)

第 1 次对话：
你："帮我写一个 Python 脚本处理 CSV 数据。"
助手："好的，我来用 pandas 写..."

第 2 次对话（第二天）：
你："帮我写一个脚本处理 JSON 数据。"
助手："好的，请问你想用什么语言？用什么库？"
你："……昨天不是刚说过用 Python + pandas 吗？"

# 短期记忆 = 当前对话的消息列表
working_memory = [
    {"role": "user", "content": "帮我优化这个 SQL 查询"},
    {"role": "assistant", "content": "好的，请给我看看查询语句..."},
    {"role": "user", "content": "SELECT * FROM orders WHERE ..."},
    # ...当前会话的所有消息
]

# 长期记忆 = 持久化存储的用户画像和项目知识
long_term_memory = {
    "user_preferences": {
        "language": "Python",
        "framework": "FastAPI",
        "database": "PostgreSQL",
        "style": "喜欢简洁的代码，不喜欢过度注释"
    },
    "project_context": {
        "name": "DataPilot",
        "description": "数据分析平台",
        "tech_stack": ["Python", "FastAPI", "React", "PostgreSQL"],
        "stage": "MVP 阶段"
    },
    "past_decisions": [
        {"date": "2026-06-15", "decision": "选择 PostgreSQL 而非 MongoDB，因为需要复杂 JOIN"},
        {"date": "2026-06-20", "decision": "用 Celery 做异步任务，而非自建队列"},
    ]
}

# 程序性记忆 = Agent 学到的操作模式和最佳实践
procedural_memory = [
    {
        "trigger": "用户要求处理 CSV 文件",
        "pattern": "先检测文件编码（chardet），再用 pandas 读取",
        "reason": "中文 CSV 经常是 GBK 编码，直接读会报错"
    },
    {
        "trigger": "用户要求部署到 AWS",
        "pattern": "先检查 IAM 权限，再检查 VPC 配置，最后部署",
        "reason": "权限问题是最常见的部署失败原因"
    },
    {
        "trigger": "用户要求写 SQL 查询",
        "pattern": "先确认表结构和索引情况，再写查询",
        "reason": "不知道索引情况就写查询可能导致性能问题"
    }
]

import sqlite3
import chromadb


class AgentMemory:
    """Agent 记忆系统：SQLite + 向量数据库"""

    def __init__(self, user_id):
        self.user_id = user_id

        # 结构化记忆用 SQLite
        self.db = sqlite3.connect(f"memory_{user_id}.db")
        self.db.execute("""
            CREATE TABLE IF NOT EXISTS user_profile (
                key TEXT PRIMARY KEY,
                value TEXT,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        self.db.execute("""
            CREATE TABLE IF NOT EXISTS decisions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                date TEXT,
                decision TEXT,
                context TEXT
            )
        """)

        # 语义记忆用向量数据库
        self.vector_db = chromadb.Client()
        self.conversation_store = self.vector_db.get_or_create_collection(
            name=f"conversations_{user_id}",
            metadata={"hnsw:space": "cosine"}
        )

    def save_preference(self, key, value):
        """保存用户偏好"""
        self.db.execute(
            "INSERT OR REPLACE INTO user_profile (key, value) VALUES (?, ?)",
            (key, value)
        )
        self.db.commit()

    def get_preferences(self):
        """获取所有用户偏好"""
        cursor = self.db.execute("SELECT key, value FROM user_profile")
        return dict(cursor.fetchall())

    def save_conversation(self, conversation_text, metadata=None):
        """保存对话到向量数据库（用于语义检索）"""
        self.conversation_store.add(
            documents=[conversation_text],
            ids=[f"conv_{len(self.conversation_store.get()['ids'])}"],
            metadatas=[metadata or {}]
        )

    def recall_conversation(self, query, top_k=3):
        """语义检索历史对话"""
        results = self.conversation_store.query(
            query_texts=[query],
            n_results=top_k
        )
        return results["documents"][0]

class MemoryRetriever:
    """记忆检索器：根据当前对话自动检索相关记忆"""

    def __init__(self, memory: AgentMemory):
        self.memory = memory

    def retrieve_relevant_memories(self, current_query, llm):
        """综合检索：结构化偏好 + 语义历史"""

        memories = []

        # 1. 始终注入用户偏好（数据量小，直接全量注入）
        prefs = self.memory.get_preferences()
        if prefs:
            pref_text = ", ".join(f"{k}: {v}" for k, v in prefs.items())
            memories.append(f"[用户偏好] {pref_text}")

        # 2. 语义检索历史对话（数据量大，按相关度检索）
        relevant_convos = self.memory.recall_conversation(current_query, top_k=3)
        for i, convo in enumerate(relevant_convos):
            memories.append(f"[历史对话 {i+1}] {convo}")

        # 3. 检索相关决策记录
        decisions = self._search_decisions(current_query, llm)
        for d in decisions:
            memories.append(f"[历史决策] {d['date']}: {d['decision']}")

        return memories

    def _search_decisions(self, query, llm):
        """用 LLM 辅助判断哪些历史决策与当前问题相关"""
        cursor = self.memory.db.execute(
            "SELECT date, decision, context FROM decisions ORDER BY date DESC LIMIT 20"
        )
        all_decisions = [
            {"date": r[0], "decision": r[1], "context": r[2]}
            for r in cursor.fetchall()
        ]

        if not all_decisions:
            return []

        # 让 LLM 筛选相关的决策
        relevant = llm.invoke(f"""
当前问题：{query}

历史决策列表：
{format_list(all_decisions)}

请选出与当前问题相关的决策，返回编号（逗号分隔）。如果没有相关的，返回"无"。
""")

        if "无" in relevant.content:
            return []

        indices = parse_indices(relevant.content)
        return [all_decisions[i] for i in indices if i < len(all_decisions)]

def decay_memories(conversation_store, days_threshold=30):
    """时间衰减：超过 N 天的记忆降低优先级"""
    all_docs = conversation_store.get(include=["metadatas"])

    for i, doc_id in enumerate(all_docs["ids"]):
        metadata = all_docs["metadatas"][i]
        days_old = (datetime.now() - parse_date(metadata["date"])).days

        if days_old > days_threshold:
            # 标记为"低优先级"而非删除
            metadata["priority"] = "low"
            conversation_store.update(
                ids=[doc_id],
                metadatas=[metadata]
            )

def compress_topic_conversations(conversation_store, topic, llm):
    """把同一主题的多段对话压缩为一段摘要"""

    # 找到该主题的所有对话
    results = conversation_store.query(
        query_texts=[topic],
        n_results=20
    )

    if len(results["documents"][0]) < 3:
        return  # 太少，不需要压缩

    # 让 LLM 压缩
    summary = llm.invoke(f"""
将以下关于"{topic}"的多段对话压缩为一段简洁的摘要。
保留关键结论和决策，省略具体的来回讨论过程。

对话内容：
{chr(10).join(results['documents'][0])}
""")

    # 删除原始对话，存入压缩后的摘要
    old_ids = results["ids"][0]
    conversation_store.delete(ids=old_ids)
    conversation_store.add(
        documents=[summary.content],
        ids=[f"summary_{topic}_{datetime.now().isoformat()}"],
        metadatas=[{"type": "compressed_summary", "topic": topic}]
    )

    return summary.content

def resolve_conflict(new_info, old_memory, llm):
    """检测并解决记忆冲突"""

    check = llm.invoke(f"""
新信息：{new_info}
旧记忆：{old_memory}

这两条信息是否矛盾？如果矛盾，请指出具体矛盾点。
回答格式：矛盾/不矛盾 + 原因
""")

    if "矛盾" in check.content and "不矛盾" not in check.content:
        # 以新信息为准，更新旧记忆
        return new_info
    return old_memory

from langchain.memory import ConversationBufferMemory
from langchain.memory import ConversationSummaryMemory
from langchain.memory import VectorStoreRetrieverMemory

# 1. 简单缓冲记忆（保留最近 N 轮）
buffer_memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True
)

# 2. 摘要记忆（自动压缩历史对话）
summary_memory = ConversationSummaryMemory(
    llm=llm,
    memory_key="chat_history",
    return_messages=True
)

# 3. 向量存储记忆（语义检索历史）
vector_memory = VectorStoreRetrieverMemory(
    retriever=vectorstore.as_retriever(search_kwargs=dict(k=3)),
    memory_key="chat_history"
)

from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.storage.chat_store import SimpleChatStore

# 持久化聊天存储
chat_store = SimpleChatStore()

memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,        # 限制记忆 token 数
    chat_store=chat_store,
    chat_store_key="user_1"  # 每个用户一个 key
)

# 保存到磁盘
chat_store.persist(persist_path="./chat_store.json")

# 加载
chat_store = SimpleChatStore.from_persist_path("./chat_store.json")