153 lines
5.0 KiB
Python
153 lines
5.0 KiB
Python
"""
|
|
Mem0 记忆层客户端封装模块
|
|
负责 Mem0 的初始化、检索和存储
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Optional, List, Dict, Any
|
|
from mem0 import AsyncMemory
|
|
|
|
from app.config import QDRANT_URL, QDRANT_COLLECTION_NAME, LLAMACPP_EMBEDDING_URL, LLAMACPP_API_KEY
|
|
from app.logger import info, warning, error
|
|
|
|
|
|
class Mem0Client:
|
|
"""Mem0 异步客户端封装类"""
|
|
|
|
def __init__(self, llm_instance):
|
|
"""
|
|
初始化 Mem0 客户端
|
|
|
|
Args:
|
|
llm_instance: LangChain LLM 实例（用于事实提取）
|
|
"""
|
|
self.llm = llm_instance
|
|
self.mem0: Optional[AsyncMemory] = None
|
|
self._initialized = False
|
|
|
|
async def initialize(self):
|
|
"""异步初始化 Mem0 客户端，并进行实际连接测试"""
|
|
if self._initialized:
|
|
return
|
|
|
|
try:
|
|
# Mem0 配置
|
|
config = {
|
|
"vector_store": {
|
|
"provider": "qdrant",
|
|
"config": {
|
|
"url": QDRANT_URL, # 直接使用完整 URL
|
|
"collection_name": QDRANT_COLLECTION_NAME,
|
|
"embedding_model_dims": 768,
|
|
}
|
|
},
|
|
"llm": {
|
|
"provider": "langchain",
|
|
"config": {
|
|
"model": self.llm
|
|
}
|
|
},
|
|
"embedder": {
|
|
"provider": "openai",
|
|
"config": {
|
|
"model": "embeddinggemma-300M-Q8_0",
|
|
"api_key": LLAMACPP_API_KEY,
|
|
"openai_base_url": LLAMACPP_EMBEDDING_URL,
|
|
},
|
|
},
|
|
"version": "v1.1"
|
|
}
|
|
|
|
self.mem0 = AsyncMemory.from_config(config)
|
|
info("✅ Mem0 配置加载成功，开始连接测试...")
|
|
|
|
# 实际连接测试：调用一次 search 确保 Qdrant 和 Embedding 都可达
|
|
await asyncio.wait_for(
|
|
self.mem0.search("ping", user_id="test", limit=1),
|
|
timeout=60.0
|
|
)
|
|
info("✅ Mem0 实际连接测试成功，初始化完成")
|
|
self._initialized = True
|
|
|
|
except asyncio.TimeoutError:
|
|
error("❌ Mem0 连接测试超时 (10s)，请检查 Qdrant 或 Embedding 服务响应")
|
|
self.mem0 = None
|
|
self._initialized = False
|
|
except Exception as e:
|
|
error(f"❌ Mem0 初始化或连接测试失败: {e}")
|
|
import traceback
|
|
error(f"详细错误信息:\n{traceback.format_exc()}")
|
|
self.mem0 = None
|
|
self._initialized = False
|
|
|
|
async def search_memories(self, query: str, user_id: str, limit: int = 5) -> List[str]:
|
|
"""
|
|
检索相关记忆
|
|
|
|
Args:
|
|
query: 查询文本
|
|
user_id: 用户 ID
|
|
limit: 返回结果数量限制
|
|
|
|
Returns:
|
|
List[str]: 记忆事实列表
|
|
"""
|
|
if not self.mem0:
|
|
warning("⚠️ Mem0 未初始化，跳过记忆检索")
|
|
return []
|
|
|
|
try:
|
|
memories = await asyncio.wait_for(
|
|
self.mem0.search(query, user_id=user_id, limit=limit),
|
|
timeout=30.0
|
|
)
|
|
|
|
if memories and "results" in memories:
|
|
facts = [m["memory"] for m in memories["results"] if m.get("memory")]
|
|
if facts:
|
|
info(f"🔍 [记忆检索] Mem0 返回 {len(facts)} 条记忆")
|
|
return facts
|
|
|
|
info("🔍 [记忆检索] 未找到相关记忆")
|
|
return []
|
|
|
|
except asyncio.TimeoutError:
|
|
warning("⚠️ Mem0 检索超时 (30s)，跳过本次记忆检索")
|
|
return []
|
|
except Exception as e:
|
|
warning(f"⚠️ Mem0 检索失败: {e}")
|
|
return []
|
|
|
|
async def add_memories(self, messages: List[Dict[str, str]], user_id: str) -> bool:
|
|
"""
|
|
添加记忆（自动提取事实并存储）
|
|
|
|
Args:
|
|
messages: 消息列表，格式为 [{"role": "user/assistant/system", "content": "..."}]
|
|
user_id: 用户 ID
|
|
|
|
Returns:
|
|
bool: 是否成功
|
|
"""
|
|
if not self.mem0:
|
|
warning("⚠️ Mem0 未初始化，跳过记忆添加")
|
|
return False
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
self.mem0.add(
|
|
messages,
|
|
user_id=user_id,
|
|
metadata={"type": "conversation"}
|
|
),
|
|
timeout=60.0
|
|
)
|
|
info("📝 [记忆添加] 已提交给 Mem0 进行事实提取")
|
|
return True
|
|
|
|
except asyncio.TimeoutError:
|
|
error("❌ Mem0 记忆添加超时 (60s)")
|
|
return False
|
|
except Exception as e:
|
|
error(f"❌ Mem0 记忆添加失败: {e}")
|
|
return False |