openClaw_agent_dm/skills/mem0-integration/openclaw_interceptor.py

#!/usr/bin/env python3
"""OpenClaw 拦截器：Pre-Hook + Post-Hook"""

import asyncio
import logging
import sys
sys.path.insert(0, '/root/.openclaw/workspace/skills/mem0-integration')
from mem0_client import mem0_client

logger = logging.getLogger(__name__)


class ConversationInterceptor:
    def __init__(self):
        self.enabled = True
    
    async def pre_hook(self, query: str, context: dict) -> str:
        if not self.enabled:
            return None
        try:
            user_id = context.get('user_id', 'default')
            agent_id = context.get('agent_id', 'general')
            memories = await mem0_client.pre_hook_search(query=query, user_id=user_id, agent_id=agent_id)
            if memories:
                return mem0_client.format_memories_for_prompt(memories)
            return None
        except Exception as e:
            logger.error(f"Pre-Hook 失败：{e}")
            return None
    
    async def post_hook(self, user_message: str, assistant_message: str, context: dict):
        if not self.enabled:
            return
        try:
            user_id = context.get('user_id', 'default')
            agent_id = context.get('agent_id', 'general')
            await mem0_client.post_hook_add(user_message, assistant_message, user_id, agent_id)
            logger.debug(f"Post-Hook: 已提交对话")
        except Exception as e:
            logger.error(f"Post-Hook 失败：{e}")


interceptor = ConversationInterceptor()


async def intercept_before_llm(query: str, context: dict):
    return await interceptor.pre_hook(query, context)


async def intercept_after_response(user_msg: str, assistant_msg: str, context: dict):
    await interceptor.post_hook(user_msg, assistant_msg, context)