#!/usr/bin/env python3 """ mem0 Client for OpenClaw - 生产级纯异步架构 Pre-Hook 检索注入 + Post-Hook 异步写入 三级可见性隔离 (public / project / private) 记忆衰减 (expiration_date) + 智能写入过滤 """ import os import re import asyncio import logging import time import yaml from typing import List, Dict, Any from collections import deque from datetime import datetime, timedelta from pathlib import Path # ========== DashScope 环境变量配置 ========== # Gemini Pro Embedding 模型：text-embedding-v4 (1024 维度) os.environ['OPENAI_API_BASE'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1' os.environ['OPENAI_BASE_URL'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1' # 关键：兼容模式需要此变量 _dashscope_key = os.getenv('MEM0_DASHSCOPE_API_KEY', '') if not _dashscope_key: _dashscope_key = os.getenv('DASHSCOPE_API_KEY', '') if _dashscope_key: os.environ['OPENAI_API_KEY'] = _dashscope_key elif not os.environ.get('OPENAI_API_KEY'): logging.warning("MEM0_DASHSCOPE_API_KEY not set; mem0 embedding/LLM calls will fail") try: from mem0 import Memory from mem0.configs.base import MemoryConfig, VectorStoreConfig, LlmConfig, EmbedderConfig except ImportError as e: print(f"⚠️ mem0ai 导入失败：{e}") Memory = None logger = logging.getLogger(__name__) class AsyncMemoryQueue: """纯异步记忆写入队列""" def __init__(self, max_size: int = 100): self.queue = deque(maxlen=max_size) self.lock = asyncio.Lock() self.running = False self._worker_task = None self.callback = None self.batch_size = 10 self.flush_interval = 60 def add(self, item: Dict[str, Any]): """添加任务到队列（同步方法）""" try: if len(self.queue) < self.queue.maxlen: self.queue.append({ 'messages': item['messages'], 'user_id': item['user_id'], 'agent_id': item['agent_id'], 'visibility': item.get('visibility', 'private'), 'project_id': item.get('project_id'), 'memory_type': item.get('memory_type', 'session'), 'timestamp': item.get('timestamp', datetime.now().isoformat()) }) else: logger.warning("异步队列已满，丢弃旧任务") except Exception as e: logger.error(f"队列添加失败：{e}") async def get_batch(self, batch_size: int) -> List[Dict]: """获取批量任务（异步方法）""" async with self.lock: batch = [] while len(batch) < batch_size and self.queue: batch.append(self.queue.popleft()) return batch def start_worker(self, callback, batch_size: int, flush_interval: int): """启动异步后台任务""" self.running = True self.callback = callback self.batch_size = batch_size self.flush_interval = flush_interval self._worker_task = asyncio.create_task(self._worker_loop()) logger.info(f"✅ 异步工作线程已启动 (batch_size={batch_size}, interval={flush_interval}s)") async def _worker_loop(self): """异步工作循环""" last_flush = time.time() while self.running: try: if self.queue: batch = await self.get_batch(self.batch_size) if batch: asyncio.create_task(self._process_batch(batch)) if time.time() - last_flush > self.flush_interval: if self.queue: batch = await self.get_batch(self.batch_size) if batch: asyncio.create_task(self._process_batch(batch)) last_flush = time.time() await asyncio.sleep(1) except asyncio.CancelledError: logger.info("异步工作线程已取消") break except Exception as e: logger.error(f"异步工作线程错误：{e}") await asyncio.sleep(5) async def _process_batch(self, batch: List[Dict]): """处理批量任务""" try: logger.debug(f"开始处理批量任务：{len(batch)} 条") for item in batch: await self.callback(item) logger.debug(f"批量任务处理完成") except Exception as e: logger.error(f"批量处理失败：{e}") async def stop(self): """优雅关闭""" self.running = False if self._worker_task: self._worker_task.cancel() try: await self._worker_task except asyncio.CancelledError: pass logger.info("异步工作线程已关闭") EXPIRATION_MAP = { 'session': timedelta(days=7), 'chat_summary': timedelta(days=30), 'preference': None, 'knowledge': None, } SKIP_PATTERNS = re.compile( r'^(好的|收到|OK|ok|嗯|行|没问题|感谢|谢谢|了解|明白|知道了|👍|✅|❌)$', re.IGNORECASE ) SYSTEM_CMD_PATTERN = re.compile(r'^/') MEMORY_KEYWORDS = re.compile( r'(记住|以后|偏好|配置|设置|规则|永远|始终|总是|不要|禁止)', ) PUBLIC_KEYWORDS = re.compile( r'(所有人|通知|全局|公告|大家|集群)', ) def _load_project_registry() -> Dict: """从 project_registry.yaml 加载项目注册表""" registry_path = Path(__file__).parent / 'project_registry.yaml' if registry_path.exists(): try: with open(registry_path, 'r', encoding='utf-8') as f: return yaml.safe_load(f) or {} except Exception: pass return {} def get_agent_projects(agent_id: str) -> List[str]: """查询一个 agent 所属的所有 project_id""" registry = _load_project_registry() projects = registry.get('projects', {}) result = [] for pid, pconf in projects.items(): members = pconf.get('members', []) if '*' in members or agent_id in members: result.append(pid) return result class Mem0Client: """ 生产级 mem0 客户端 纯异步架构 + 三级可见性 + 记忆衰减 + 智能写入过滤 """ def __init__(self, config: Dict = None): self.config = config or self._load_default_config() self.local_memory = None self.async_queue = None self.cache = {} self._started = False self._init_memory() def _load_default_config(self) -> Dict: """加载默认配置 - 单库融合架构""" return { "qdrant": { "host": os.getenv('MEM0_QDRANT_HOST', 'localhost'), "port": int(os.getenv('MEM0_QDRANT_PORT', '6333')), "collection_name": "mem0_v4_shared" # 统一共享 Collection（陈医生/张大师共用） }, "llm": { "provider": "openai", "config": { "model": os.getenv('MEM0_LLM_MODEL', 'qwen-plus') } }, "embedder": { "provider": "openai", "config": { "model": os.getenv('MEM0_EMBEDDER_MODEL', 'text-embedding-v4'), "dimensions": 1024 # DashScope text-embedding-v4 支持的最大维度 } }, "retrieval": { "enabled": True, "top_k": 5, "min_confidence": 0.7, "timeout_ms": 5000 }, "async_write": { "enabled": True, "queue_size": 100, "batch_size": 10, "flush_interval": 60, "timeout_ms": 5000 }, "cache": { "enabled": True, "ttl": 300, "max_size": 1000 }, "fallback": { "enabled": True, "log_level": "WARNING", "retry_attempts": 2 }, "metadata": { "default_user_id": "default", "default_agent_id": "general" }, "write_filter": { "enabled": True, "min_user_message_length": 5, } } def _init_memory(self): """初始化 mem0（同步操作）- 三位一体完整配置""" if Memory is None: logger.warning("mem0ai 未安装") return try: config = MemoryConfig( vector_store=VectorStoreConfig( provider="qdrant", config={ "host": self.config['qdrant']['host'], "port": self.config['qdrant']['port'], "collection_name": self.config['qdrant']['collection_name'], "on_disk": True, "embedding_model_dims": 1024 # 强制同步 Qdrant 集合维度 } ), llm=LlmConfig( provider="openai", config=self.config['llm']['config'] ), embedder=EmbedderConfig( provider="openai", config={ "model": "text-embedding-v4", "embedding_dims": 1024 # 核心修复：强制覆盖默认的 1536 维度 # api_base 和 api_key 通过环境变量 OPENAI_BASE_URL 和 OPENAI_API_KEY 读取 } ) ) self.local_memory = Memory(config=config) logger.info("✅ mem0 初始化成功（含 Embedder，1024 维度）") except Exception as e: logger.error(f"❌ mem0 初始化失败：{e}") self.local_memory = None async def start(self): """ 显式启动异步工作线程 必须在事件循环中调用：await mem0_client.start() """ if self._started: logger.debug("mem0 Client 已启动") return if self.config['async_write']['enabled']: self.async_queue = AsyncMemoryQueue( max_size=self.config['async_write']['queue_size'] ) self.async_queue.start_worker( callback=self._async_write_memory, batch_size=self.config['async_write']['batch_size'], flush_interval=self.config['async_write']['flush_interval'] ) self._started = True logger.info("✅ mem0 Client 异步工作线程已启动") # ========== Pre-Hook: 智能检索 ========== async def pre_hook_search(self, query: str, user_id: str = None, agent_id: str = None, top_k: int = None) -> List[Dict]: """Pre-Hook: 对话前智能检索""" if not self.config['retrieval']['enabled'] or self.local_memory is None: return [] cache_key = f"{user_id}:{agent_id}:{query}" if self.config['cache']['enabled'] and cache_key in self.cache: cached = self.cache[cache_key] if time.time() - cached['time'] < self.config['cache']['ttl']: logger.debug(f"Cache hit: {cache_key}") return cached['results'] timeout_ms = self.config['retrieval']['timeout_ms'] try: memories = await asyncio.wait_for( self._execute_search(query, user_id, agent_id, top_k or self.config['retrieval']['top_k']), timeout=timeout_ms / 1000 ) if self.config['cache']['enabled'] and memories: self.cache[cache_key] = {'results': memories, 'time': time.time()} self._cleanup_cache() logger.info(f"Pre-Hook 检索完成：{len(memories)} 条记忆") return memories except asyncio.TimeoutError: logger.warning(f"Pre-Hook 检索超时 ({timeout_ms}ms)") return [] except Exception as e: logger.warning(f"Pre-Hook 检索失败：{e}") return [] async def _execute_search(self, query: str, user_id: str, agent_id: str, top_k: int) -> List[Dict]: """ 三阶段检索 — 按可见性分层，合并去重 Phase 1: public (所有 agent 可见) Phase 2: project (同 project_id 成员可见) Phase 3: private (仅 agent_id 本人可见) """ if self.local_memory is None: return [] all_memories: Dict[str, Dict] = {} per_phase = max(top_k, 3) # Phase 1: 检索 public 记忆 try: public_mems = await asyncio.to_thread( self.local_memory.search, query, user_id=user_id, filters={"visibility": "public"}, limit=per_phase ) for mem in (public_mems or []): mid = mem.get('id') if isinstance(mem, dict) else None if mid and mid not in all_memories: all_memories[mid] = mem except Exception as e: logger.debug(f"Public 记忆检索失败：{e}") # Phase 2: 检索 project 记忆 (agent 所属的所有项目) if agent_id and agent_id != 'general': agent_projects = get_agent_projects(agent_id) for project_id in agent_projects: if project_id == 'global': continue try: proj_mems = await asyncio.to_thread( self.local_memory.search, query, user_id=user_id, filters={ "visibility": "project", "project_id": project_id, }, limit=per_phase ) for mem in (proj_mems or []): mid = mem.get('id') if isinstance(mem, dict) else None if mid and mid not in all_memories: all_memories[mid] = mem except Exception as e: logger.debug(f"Project({project_id}) 记忆检索失败：{e}") # Phase 3: 检索 private 记忆 if agent_id and agent_id != 'general': try: private_mems = await asyncio.to_thread( self.local_memory.search, query, user_id=user_id, filters={ "visibility": "private", "agent_id": agent_id, }, limit=per_phase ) for mem in (private_mems or []): mid = mem.get('id') if isinstance(mem, dict) else None if mid and mid not in all_memories: all_memories[mid] = mem except Exception as e: logger.debug(f"Private 记忆检索失败：{e}") # Fallback: 兼容旧数据（无 visibility 字段） if user_id: try: legacy_mems = await asyncio.to_thread( self.local_memory.search, query, user_id=user_id, filters={"agent_id": agent_id} if agent_id and agent_id != 'general' else None, limit=per_phase ) for mem in (legacy_mems or []): mid = mem.get('id') if isinstance(mem, dict) else None if mid and mid not in all_memories: all_memories[mid] = mem except Exception as e: logger.debug(f"Legacy 记忆检索失败：{e}") min_confidence = self.config['retrieval']['min_confidence'] filtered = [ m for m in all_memories.values() if m.get('score', 1.0) >= min_confidence ] filtered.sort(key=lambda x: x.get('score', 0), reverse=True) return filtered[:top_k] def format_memories_for_prompt(self, memories: List[Dict]) -> str: """格式化记忆为 Prompt 片段""" if not memories: return "" prompt = "\n\n=== 相关记忆 ===\n" for i, mem in enumerate(memories, 1): memory_text = mem.get('memory', '') if isinstance(mem, dict) else str(mem) created_at = mem.get('created_at', '') if isinstance(mem, dict) else '' prompt += f"{i}. {memory_text}" if created_at: prompt += f" (记录于：{created_at})" prompt += "\n" prompt += "===============\n" return prompt # ========== Post-Hook: 异步写入 ========== def _should_skip_memory(self, user_message: str, assistant_message: str) -> bool: """判断是否应跳过此对话的记忆写入""" if not self.config.get('write_filter', {}).get('enabled', True): return False min_len = self.config.get('write_filter', {}).get('min_user_message_length', 5) if len(user_message.strip()) < min_len: return True if SKIP_PATTERNS.match(user_message.strip()): return True if SYSTEM_CMD_PATTERN.match(user_message.strip()): return True return False def _classify_memory_type(self, user_message: str, assistant_message: str) -> str: """自动分类记忆类型，决定过期策略""" combined = user_message + ' ' + assistant_message if MEMORY_KEYWORDS.search(combined): return 'preference' if any(kw in combined for kw in ('部署', '配置', '架构', '端口', '安装', '版本')): return 'knowledge' return 'session' def _classify_visibility(self, user_message: str, assistant_message: str, agent_id: str = None) -> str: """自动分类记忆可见性""" combined = user_message + ' ' + assistant_message if PUBLIC_KEYWORDS.search(combined): return 'public' return 'private' def post_hook_add(self, user_message: str, assistant_message: str, user_id: str = None, agent_id: str = None, visibility: str = None, project_id: str = None, memory_type: str = None): """Post-Hook: 对话后异步写入（同步方法，仅添加到队列） 支持三级可见性 (public/project/private) 和记忆衰减 (expiration_date)。 内置智能写入过滤，跳过无价值对话。 """ if not self.config['async_write']['enabled']: return if self._should_skip_memory(user_message, assistant_message): logger.debug(f"Post-Hook 跳过（写入过滤）：{user_message[:30]}") return if not user_id: user_id = self.config['metadata']['default_user_id'] if not agent_id: agent_id = self.config['metadata']['default_agent_id'] if not memory_type: memory_type = self._classify_memory_type(user_message, assistant_message) if not visibility: visibility = self._classify_visibility(user_message, assistant_message, agent_id) messages = [ {"role": "user", "content": user_message}, {"role": "assistant", "content": assistant_message} ] if self.async_queue: self.async_queue.add({ 'messages': messages, 'user_id': user_id, 'agent_id': agent_id, 'visibility': visibility, 'project_id': project_id, 'memory_type': memory_type, 'timestamp': datetime.now().isoformat() }) logger.debug(f"Post-Hook 已提交：user={user_id}, agent={agent_id}, " f"visibility={visibility}, type={memory_type}") else: logger.warning("异步队列未初始化") async def _async_write_memory(self, item: Dict): """异步写入记忆（后台任务）""" if self.local_memory is None: return timeout_ms = self.config['async_write']['timeout_ms'] try: await asyncio.wait_for(self._execute_write(item), timeout=timeout_ms / 1000) logger.debug(f"异步写入成功：user={item['user_id']}, agent={item['agent_id']}") except asyncio.TimeoutError: logger.warning(f"异步写入超时 ({timeout_ms}ms)") except Exception as e: logger.warning(f"异步写入失败：{e}") async def _execute_write(self, item: Dict): """ 执行写入 — 三级可见性 + 记忆衰减 metadata 携带 visibility / project_id / agent_id expiration_date 根据 memory_type 自动设置 """ if self.local_memory is None: return visibility = item.get('visibility', 'private') memory_type = item.get('memory_type', 'session') custom_metadata = { "agent_id": item['agent_id'], "visibility": visibility, "project_id": item.get('project_id') or '', "business_type": item.get('business_type', item['agent_id']), "memory_type": memory_type, "source": "openclaw", "timestamp": item.get('timestamp'), } ttl = EXPIRATION_MAP.get(memory_type) expiration_date = (datetime.now() + ttl).isoformat() if ttl else None add_kwargs = dict( messages=item['messages'], user_id=item['user_id'], agent_id=item['agent_id'], metadata=custom_metadata, ) if expiration_date: add_kwargs['expiration_date'] = expiration_date await asyncio.to_thread(self.local_memory.add, **add_kwargs) def _cleanup_cache(self): """清理过期缓存""" if not self.config['cache']['enabled']: return current_time = time.time() ttl = self.config['cache']['ttl'] expired_keys = [k for k, v in self.cache.items() if current_time - v['time'] > ttl] for key in expired_keys: del self.cache[key] if len(self.cache) > self.config['cache']['max_size']: oldest_keys = sorted(self.cache.keys(), key=lambda k: self.cache[k]['time'])[:len(self.cache) - self.config['cache']['max_size']] for key in oldest_keys: del self.cache[key] def get_status(self) -> Dict: """获取状态""" return { "initialized": self.local_memory is not None, "started": self._started, "async_queue_enabled": self.config['async_write']['enabled'], "queue_size": len(self.async_queue.queue) if self.async_queue else 0, "cache_size": len(self.cache), "qdrant": f"{self.config['qdrant']['host']}:{self.config['qdrant']['port']}" } # ========== Knowledge Publishing (Hub Agent) ========== async def publish_knowledge(self, content, category='knowledge', visibility='public', project_id=None, agent_id='main'): """Publish knowledge/best practices to shared memory. Used by hub agent to share with all agents (public) or project teams.""" if visibility == 'project' and not project_id: raise ValueError("project visibility requires project_id") item = { 'messages': [{'role': 'system', 'content': content}], 'user_id': 'system', 'agent_id': agent_id, 'visibility': visibility, 'project_id': project_id, 'memory_type': category, 'timestamp': datetime.now().isoformat(), } await self._execute_write(item) logger.info(f"Published {category} ({visibility}): {content[:80]}...") # ========== Cold Start (Three-Phase) ========== async def cold_start_search(self, agent_id='main', user_id='default', top_k=5): """Three-phase cold start: public -> project -> private. Returns merged memories ordered by phase then score.""" if self.local_memory is None: return [] all_mems = {} phase = {} for q in ["system best practices and conventions", "shared configuration and architecture decisions"]: try: r = await asyncio.to_thread( self.local_memory.search, q, user_id=user_id, filters={"visibility": "public"}, limit=top_k) for m in (r or []): mid = m.get('id') if isinstance(m, dict) else None if mid and mid not in all_mems: all_mems[mid] = m phase[mid] = 0 except Exception as e: logger.debug(f"Cold start public failed: {e}") for pid in get_agent_projects(agent_id): if pid == 'global': continue for q in ["project guidelines and shared knowledge", "recent project decisions and updates"]: try: r = await asyncio.to_thread( self.local_memory.search, q, user_id=user_id, filters={"visibility": "project", "project_id": pid}, limit=top_k) for m in (r or []): mid = m.get('id') if isinstance(m, dict) else None if mid and mid not in all_mems: all_mems[mid] = m phase[mid] = 1 except Exception as e: logger.debug(f"Cold start project({pid}) failed: {e}") for q in ["current active tasks deployment progress", "pending work items todos", "recent important decisions configurations"]: try: r = await asyncio.to_thread( self.local_memory.search, q, user_id=user_id, filters={"visibility": "private", "agent_id": agent_id}, limit=top_k) for m in (r or []): mid = m.get('id') if isinstance(m, dict) else None if mid and mid not in all_mems: all_mems[mid] = m phase[mid] = 2 except Exception as e: logger.debug(f"Cold start private failed: {e}") mc = self.config['retrieval']['min_confidence'] out = [m for m in all_mems.values() if m.get('score', 1.0) >= mc] out.sort(key=lambda m: (phase.get(m.get('id', ''), 9), -m.get('score', 0))) return out[:top_k] async def shutdown(self): if self.async_queue: await self.async_queue.stop() logger.info("mem0 Client shutdown complete") # 全局客户端实例 mem0_client = Mem0Client()