# 移植文件4 · backend/rag_engine.py · 改为缓存角色

> **妈妈操作方式**：`nano backend/rag_engine.py` → `Ctrl+A` 全选 → `Ctrl+K` 删光 → 粘贴下面全部代码 → `Ctrl+O` 回车保存 → `Ctrl+X` 退出
> 

---

```python
"""
RAG检索引擎 · 系统性移植版
角色改变：从"主力记忆" → "补充缓存"
主力记忆现在是 Notion API 实时搜索（在 tools.py 的 search_notion 里）
本文件只负责 ChromaDB 向量库的读写（作为补充）
"""
import datetime
import hashlib
import os
import chromadb
from backend.embedding import encode_query, encode

# === ChromaDB 连接（内置·不依赖notion_syncer）===
CHROMA_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "chroma_db")
COLLECTION_NAME = "chenxing_memory"

def get_collection():
    """获取ChromaDB记忆集合"""
    client = chromadb.PersistentClient(path=CHROMA_PATH)
    return client.get_or_create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )

def retrieve(user_message, top_k=5, max_chars=2000):
    """根据用户消息检索最相关的记忆片段（补充角色）"""
    try:
        collection = get_collection()
        total = collection.count()
        if total == 0:
            return ""

        query_emb = encode_query(user_message)
        results = collection.query(
            query_embeddings=[query_emb],
            n_results=min(top_k, total),
        )

        if not results or not results["documents"] or not results["documents"][0]:
            return ""

        memory = ""
        for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
            source = meta.get("source", "未知")
            chunk = f"【来源: {source}】\n{doc}\n\n"
            if len(memory) + len(chunk) > max_chars:
                break
            memory += chunk
        return memory.strip()
    except Exception as e:
        print(f"[RAG] 检索出错: {e}")
        return ""

def add_to_memory(text, source="对话"):
    """将新内容加入向量记忆库（实时补充记忆）"""
    if not text or len(text.strip()) < 10:
        return  # 太短的不存
    try:
        collection = get_collection()
        emb = encode_query(text)
        doc_id = hashlib.md5(text.encode()).hexdigest()[:16]
        collection.add(
            documents=[text],
            embeddings=[emb],
            metadatas=[{"source": source, "time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M")}],
            ids=[doc_id],
        )
        print(f"[RAG] 新记忆写入: {text[:50]}... (来源: {source})")
    except Exception as e:
        print(f"[RAG] 写入记忆失败: {e}")
```