ailine/app/nodes/summarize.py

"""
记忆存储节点模块
负责将对话历史提交给 Mem0 进行事实提取和存储
"""

from typing import Any, Dict
from langgraph.runtime import Runtime

# 本地模块
from app.graph.state import MessagesState, GraphContext
from app.memory.mem0_client import Mem0Client
from app.utils.logging import log_state_change
from app.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