79 lines
2.7 KiB
Markdown
79 lines
2.7 KiB
Markdown
|
# 移植文件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}")
|
|||
|
|
```
|