87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
"""
|
|
记忆存储节点模块
|
|
负责将对话历史提交给 Mem0 进行事实提取和存储
|
|
"""
|
|
|
|
from typing import Any, Dict
|
|
|
|
# 本地模块
|
|
from ..graph.state import MessagesState
|
|
from ..memory.mem0_client import Mem0Client
|
|
from ..utils.logging import log_state_change
|
|
from ..logger import debug, info, error, warning
|
|
|
|
def create_summarize_node(mem0_client: Mem0Client):
|
|
"""
|
|
工厂函数：创建记忆存储节点
|
|
|
|
Args:
|
|
mem0_client: Mem0 客户端实例
|
|
|
|
Returns:
|
|
异步节点函数
|
|
"""
|
|
|
|
from langchain_core.runnables.config import RunnableConfig
|
|
|
|
async def summarize_conversation(state: MessagesState, config: RunnableConfig) -> Dict[str, Any]:
|
|
"""
|
|
记忆存储节点 - 使用 Mem0
|
|
|
|
Args:
|
|
state: 当前对话状态
|
|
config: 运行时配置
|
|
|
|
Returns:
|
|
重置计数器的状态更新
|
|
"""
|
|
log_state_change("summarize", state, "进入")
|
|
|
|
messages = state["messages"]
|
|
if len(messages) < 4:
|
|
debug("📝 [记忆添加] 对话过短，跳过")
|
|
return {"turns_since_last_summary": 0}
|
|
|
|
# 从 metadata 中获取 user_id
|
|
user_id = config.get("metadata", {}).get("user_id", "default_user")
|
|
|
|
# 确保 Mem0 已初始化（懒加载）
|
|
if not mem0_client._initialized:
|
|
await mem0_client.initialize()
|
|
|
|
# 将整个对话历史转换为 Mem0 需要的消息格式
|
|
mem0_messages = []
|
|
for msg in messages:
|
|
# 兼容 dict 和对象两种格式
|
|
if isinstance(msg, dict):
|
|
msg_type = msg.get("type", "")
|
|
msg_content = msg.get("content", "")
|
|
else:
|
|
msg_type = getattr(msg, 'type', '')
|
|
msg_content = getattr(msg, 'content', '')
|
|
|
|
if msg_type == "human":
|
|
mem0_messages.append({"role": "user", "content": msg_content})
|
|
elif msg_type == "ai":
|
|
mem0_messages.append({"role": "assistant", "content": msg_content})
|
|
elif msg_type == "tool":
|
|
mem0_messages.append({"role": "system", "content": f"[工具返回] {msg_content}"})
|
|
|
|
if mem0_client.mem0:
|
|
try:
|
|
# 异步调用 Mem0 自动提取并存储事实
|
|
success = await mem0_client.add_memories(
|
|
mem0_messages,
|
|
user_id=user_id
|
|
)
|
|
if success:
|
|
info(f"📝 [记忆添加] 已提交给 Mem0 进行事实提取")
|
|
except Exception as e:
|
|
error(f"❌ Mem0 记忆添加失败: {e}")
|
|
else:
|
|
warning("⚠️ Mem0 未初始化，跳过记忆添加")
|
|
|
|
log_state_change("summarize", state, "离开")
|
|
return {"turns_since_last_summary": 0}
|
|
|
|
return summarize_conversation |