""" Mem0 配置数据类 用于管理 Mem0 长期记忆系统的配置参数 """ from datetime import datetime from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Optional # 避免循环导入 if TYPE_CHECKING: from langchain_core.language_models import BaseChatModel # 缓存已加载的提示词模板 _fact_extraction_prompt_template: Optional[str] = None def _load_fact_extraction_prompt() -> str: """从 prompt/FACT_RETRIEVAL_PROMPT.md 加载提示词模板 Returns: str: 提示词模板内容 """ global _fact_extraction_prompt_template if _fact_extraction_prompt_template is not None: return _fact_extraction_prompt_template # 获取项目根目录(假设 prompt/ 在项目根目录下) project_root = Path(__file__).parent.parent prompt_file = project_root / "prompt" / "FACT_RETRIEVAL_PROMPT.md" try: _fact_extraction_prompt_template = prompt_file.read_text(encoding="utf-8") except FileNotFoundError: # 如果文件不存在,返回默认提示词 _fact_extraction_prompt_template = ( "Extract relevant facts about the user from the conversation.\n" "Today's date is {current_time}.\n" "Return response in JSON format: {\"facts\": [\"fact1\", \"fact2\"]}" ) return _fact_extraction_prompt_template @dataclass class Mem0Config: """Mem0 长期记忆配置类""" # 功能开关 enabled: bool = False # 语义搜索配置 semantic_search_top_k: int = 20 memory_prompt_template: str = ( "\n\n=== 相关记忆 ===\n" "以下是从历史对话中检索到的相关信息,可以帮助你更好地回答用户问题:\n" "{memories}\n" "==================\n" ) # 多租户配置 user_id: Optional[str] = None # 用户标识 agent_id: Optional[str] = None # Bot 标识 session_id: Optional[str] = None # 会话标识 # LLM 实例(用于 Mem0 的记忆提取和增强) llm_instance: Optional["BaseChatModel"] = None # LangChain LLM 实例 def get_attribution_tuple(self) -> tuple[str, str]: """获取 attribution 所需的元组 (user_id, agent_id) Returns: (user_id, agent_id) 元组 """ if not self.user_id or not self.agent_id: raise ValueError("user_id and agent_id are required for attribution") return (self.user_id, self.agent_id) def is_enabled(self) -> bool: """检查 Mem0 功能是否启用 Returns: bool: 是否启用 """ return self.enabled def get_memory_prompt(self, memories: list[str]) -> str: """根据记忆列表生成注入提示词 Args: memories: 记忆内容列表 Returns: str: 格式化的记忆提示词 """ if not memories: return "" memory_text = "\n".join(f"- {m}" for m in memories) return self.memory_prompt_template.format(memories=memory_text) def get_custom_fact_extraction_prompt(self) -> str: """获取自定义记忆提取提示词 从 prompt/FACT_RETRIEVAL_PROMPT.md 读取模板并替换 {current_time} 为当前日期 Returns: str: 自定义记忆提取提示词 """ template = _load_fact_extraction_prompt() current_date = datetime.now().strftime("%Y-%m-%d") return template.format(current_time=current_date) def with_session(self, session_id: str) -> "Mem0Config": """创建带有新 session_id 的配置副本 Args: session_id: 新的会话 ID Returns: 新的 Mem0Config 实例 """ new_config = Mem0Config(**self.__dict__) new_config.session_id = session_id return new_config