import asyncio from typing import List, Dict, Any, Optional from mem0 import AsyncMemory from appPublic.log import debug, error, info, exception import atexit # 配置 CONFIG = { "llm": { "provider": "openai", "config": { "model": "Qwen3-14B-FP8", "temperature": 0.3, "max_tokens": 2000, "openai_base_url": "https://t2t.opencomputing.net:10443/qwen3-14b-fp8/v1", "api_key": "any-key" } }, "embedder": { "provider": "openai", "config": { "model": "bge-m3", "embedding_dims": 1024, "openai_base_url": "https://embedding.opencomputing.net:10443/v1", "api_key": "any-key" } }, "vector_store": { "provider": "milvus", "config": { "collection_name": "mem0", "url": "/home/wangmeihua/mem0/milvus_demo.db", "embedding_model_dims": 1024, "metric_type": "COSINE", "db_name": "milvus" } }, "version": "v1.1" } class MemoryManager: """管理用户记忆的类，基于 AsyncMemory 提供添加和检索功能""" _instance = None def __new__(cls): if cls._instance is None: cls._instance = super(MemoryManager, cls).__new__(cls) return cls._instance def __init__(self): if not hasattr(self, '_initialized'): self.memory = None self._initialized = False atexit.register(self._close) async def initialize(self): """异步初始化 AsyncMemory""" if self._initialized: debug(f"AsyncMemory 已初始化，跳过") try: self.memory = await AsyncMemory.from_config(CONFIG) self._initialized = True debug("AsyncMemory 初始化成功") except Exception as e: error(f"AsyncMemory 初始化失败: {e}") raise def _close(self): if self.memory and hasattr(self.memory.vector_store, 'client'): try: self.memory.vector_store.client.close() debug(f"Milvus连接已关闭") except Exception as e: error(f"关闭milvus连接失败：{e}") if self.memory and hasattr(self.memory.db, 'connection'): try: self.memory.db.connection.close() debug("SQLite 连接已关闭") except Exception as e: error(f"关闭 SQLite 连接失败: {e}") async def _ensure_initialized(self): """确保 AsyncMemory 已初始化（懒加载）""" if self.memory is None: await self.initialize() async def add_messages_to_memory(self, messages: List[Dict[str, str]], user_id: str) -> Dict[str, Any]: await self._ensure_initialized() try: result = await self.memory.add(messages, user_id=user_id) debug(f"用户 {user_id} 的消息添加成功，结果: {result}") return {"status": "success", "result": result} except Exception as e: error(f"用户 {user_id} 的消息添加失败: {e}") return {"status": "error", "message": str(e)} async def search_user_memories(self, query: str, user_id: str, limit: int = 5) -> List[Dict[str, Any]]: await self._ensure_initialized() try: result = await self.memory.search(query=query, user_id=user_id, limit=limit) existing_memories = result["results"] debug(f"用户 {user_id} 的记忆检索成功，找到 {len(existing_memories)} 条记录") return existing_memories except Exception as e: error(f"用户 {user_id} 的记忆检索失败: {e}") return [] async def get_all_memories(self, user_id: Optional[str] = None, agent_id: Optional[str] = None, run_id: Optional[str] = None, filters: Optional[Dict[str, Any]] = None, limit: int = 100) -> Dict[str, Any]: await self._ensure_initialized() try: result = await self.memory.get_all(user_id=user_id, agent_id=agent_id, run_id=run_id, filters=filters, limit=limit) debug(f"检索所有记忆成功，找到 {len(result.get('results', []))} 条记录") return {"status": "success", "result": result} except Exception as e: error(f"检索所有记忆失败: {e}") return {"status": "error", "message": str(e)} async def delete_all_memories(self, user_id: Optional[str] = None, agent_id: Optional[str] = None, run_id: Optional[str] = None) -> Dict[str, Any]: await self._ensure_initialized() try: result = await self.memory.delete_all(user_id=user_id, agent_id=agent_id, run_id=run_id) debug(f"删除所有记忆成功: {result}") return {"status": "success", "result": result} except Exception as e: error(f"删除所有记忆失败: {e}") return {"status": "error", "message": str(e)} async def get_memory_history(self, memory_id: str) -> Dict[str, Any]: await self._ensure_initialized() try: result = await self.memory.history(memory_id) debug(f"获取记忆 ID {memory_id} 的历史成功，找到 {len(result)} 条记录") return {"status": "success", "result": result} except Exception as e: error(f"获取记忆 ID {memory_id} 的历史失败: {e}") return {"status": "error", "message": str(e)} async def reset_memory(self) -> Dict[str, Any]: await self._ensure_initialized() try: await self.memory.reset() debug("记忆存储重置成功") return {"status": "success", "result": {"message": "Memory store reset successfully"}} except Exception as e: error(f"记忆存储重置失败: {e}") return {"status": "error", "message": str(e)} async def test_memory_functions(): """异步测试函数，验证 MemoryManager 的 add_messages_to_memory 和 search_user_memories""" # 初始化 MemoryManager manager = MemoryManager() #测试获取所有记忆 test_user_id = "test_user_123" try: result = await manager.get_all_memories(user_id=test_user_id, limit=10) if result["status"] != "success": error(f"测试获取所有记忆失败: {result['message']}") return debug(f"测试获取所有记忆成功: {result}") except Exception as e: error(f"测试获取所有记忆失败: {e}") return # # 测试数据 # test_user_id = "test_user_123" # test_messages = [ # {"role": "user", "content": "您好，请您为我推荐一款饮料"}, # {"role": "assistant", "content": "好的，您可以试试元气森林、冰红茶、水溶C等饮料"}, # {"role": "user", "content": "感谢您的推荐，我喜欢喝元气森林。"}, # {"role": "assistant", "content": "好的"} # ] # # # 测试添加消息 # try: # result = await manager.add_messages_to_memory(test_messages, test_user_id) # if result["status"] != "success": # error(f"测试添加消息失败: {result['message']}") # return # if "result" not in result: # error("测试添加消息失败: 返回结果中缺少 'result' 键") # return # debug(f"测试添加消息成功: {result}") # except Exception as e: # error(f"测试添加消息失败: {e}") # return # # # 测试检索记忆 # query = "今天想吃水果" # try: # memories = await manager.search_user_memories(query, test_user_id, limit=3) # if not isinstance(memories, list): # error("测试检索记忆失败: 返回结果不是列表") # return # if memories and ("user_id" not in memories[0] or memories[0]["user_id"] != test_user_id): # error("测试检索记忆失败: 返回结果的 user_id 不匹配") # return # debug(f"测试检索记忆成功，返回 {len(memories)} 条记录") # debug(f"检索结果: {memories}") # except Exception as e: # error(f"测试检索记忆失败: {e}") if __name__ == "__main__": asyncio.run(test_memory_functions())