store.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. import os
  2. import logging
  3. from datetime import datetime
  4. logger = logging.getLogger(__name__)
  5. class MemoryManager:
  6. """轻量级记忆管理器 - 仅保存对话历史到向量数据库
  7. 不再使用 langchain.memory(已废弃),直接实现核心功能:
  8. - Layer 3 - 语义记忆:每次对话后向量化存储,供后续会话召回
  9. """
  10. def __init__(self, embeddings):
  11. """初始化记忆管理器
  12. Args:
  13. embeddings: 嵌入模型实例 (OpenAIEmbeddings 或其他兼容接口)
  14. """
  15. self.embeddings = embeddings
  16. # Layer 3 向量库:历史对话记忆
  17. from langchain_chroma import Chroma
  18. self.memory_vectorstore = Chroma(
  19. collection_name="user_memory",
  20. embedding_function=self.embeddings,
  21. persist_directory=os.path.expanduser("~/.cfc/langgraph/memory"),
  22. )
  23. async def save_conversation(self, user_id: int, conversation_id: str,
  24. messages: list[dict]):
  25. """会话结束后保存到向量记忆库"""
  26. texts = []
  27. for msg in messages:
  28. role = msg.get("role", "unknown")
  29. content = msg.get("content", "")
  30. texts.append(f"[{role}] {content}")
  31. full_text = "\n".join(texts)
  32. metadata = {
  33. "user_id": user_id,
  34. "conversation_id": conversation_id,
  35. "timestamp": str(datetime.now()),
  36. }
  37. await self.memory_vectorstore.aadd_texts(
  38. texts=[full_text],
  39. metadatas=[metadata],
  40. )
  41. self.memory_vectorstore.persist()
  42. logger.info("已保存对话到向量记忆:conv=%s, user=%s", conversation_id, user_id)
  43. async def recall(self, user_id: int, query: str, k: int = 3) -> list[str]:
  44. """语义召回:查询与 query 最相似的历史对话片段"""
  45. results = self.memory_vectorstore.similarity_search(
  46. query,
  47. k=k,
  48. filter={"user_id": user_id},
  49. )
  50. return [doc.page_content for doc in results]