fea6b35779
通过插件 hook 机制允许自定义 Mem0 的 FACT_RETRIEVAL_PROMPT， 包含联系人关系追踪和同姓消歧等增强能力。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
150 lines
4.6 KiB
Python
150 lines
4.6 KiB
Python
"""
|
||
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 # 会话标识
|
||
|
||
@property
|
||
def bot_id(self) -> str:
|
||
"""兼容 execute_hooks 所需的 bot_id 属性"""
|
||
return self.agent_id or ""
|
||
|
||
# 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)
|
||
|
||
async def get_custom_fact_extraction_prompt_async(self) -> str:
|
||
"""异步获取自定义记忆提取提示词（支持 PreMemoryPrompt hook 注入）
|
||
|
||
从 prompt/FACT_RETRIEVAL_PROMPT.md 读取默认模板，
|
||
然后执行 PreMemoryPrompt hooks：如果 hook 返回内容则替换整个模板，
|
||
最后替换 {current_time} 为当前日期。
|
||
|
||
Returns:
|
||
str: 自定义记忆提取提示词
|
||
"""
|
||
from agent.plugin_hook_loader import execute_hooks
|
||
|
||
template = await execute_hooks('PreMemoryPrompt', self) or _load_fact_extraction_prompt()
|
||
return template.format(current_time=datetime.now().strftime("%Y-%m-%d"))
|
||
|
||
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
|