218 lines
7.5 KiB
Python
218 lines
7.5 KiB
Python
"""
|
||
Mem0 记忆层客户端封装模块
|
||
负责 Mem0 的初始化、检索和存储
|
||
"""
|
||
|
||
import asyncio
|
||
import time
|
||
from typing import Optional, List
|
||
|
||
from mem0 import AsyncMemory
|
||
|
||
from app.config import (
|
||
LLM_API_KEY,
|
||
ZHIPUAI_API_KEY,
|
||
VLLM_BASE_URL,
|
||
QDRANT_URL,
|
||
QDRANT_COLLECTION_NAME,
|
||
QDRANT_API_KEY,
|
||
LLAMACPP_EMBEDDING_URL,
|
||
LLAMACPP_API_KEY,
|
||
ZHIPU_EMBEDDING_MODEL,
|
||
ZHIPU_API_BASE,
|
||
)
|
||
from app.logger import info, warning, error
|
||
from app.model_services import get_embedding_service
|
||
from app.model_services.chat_services import get_chat_service
|
||
|
||
|
||
class Mem0Client:
|
||
"""Mem0 异步客户端封装类"""
|
||
|
||
def __init__(self):
|
||
"""初始化 Mem0 客户端（内部获取所需服务）"""
|
||
self.mem0: Optional[AsyncMemory] = None
|
||
self._initialized = False
|
||
|
||
async def initialize(self):
|
||
"""异步初始化 Mem0 客户端，并进行实际连接测试"""
|
||
if self._initialized:
|
||
return
|
||
|
||
try:
|
||
# 获取可用的 embedding 服务并确定维度
|
||
info("🔄 正在获取嵌入服务...")
|
||
embeddings = get_embedding_service()
|
||
test_embedding = embeddings.embed_query("test")
|
||
embedding_dim = len(test_embedding)
|
||
info(f"✅ 嵌入服务可用，向量维度: {embedding_dim}")
|
||
|
||
# 构建 embedder 配置
|
||
from app.model_services.embedding_services import (
|
||
LocalLlamaCppEmbeddingProvider,
|
||
ZhipuEmbeddingProvider,
|
||
)
|
||
|
||
embedder_config = None
|
||
local_provider = LocalLlamaCppEmbeddingProvider()
|
||
|
||
if local_provider.is_available():
|
||
info("✅ 使用本地 llama.cpp 作为 mem0 embedder")
|
||
embedder_config = {
|
||
"provider": "openai",
|
||
"config": {
|
||
"model": "Qwen3-Embedding-0.6B-Q8_0",
|
||
"api_key": LLAMACPP_API_KEY or "dummy-key",
|
||
"openai_base_url": LLAMACPP_EMBEDDING_URL,
|
||
},
|
||
}
|
||
else:
|
||
# 检查智谱
|
||
zhipu_provider = ZhipuEmbeddingProvider()
|
||
if zhipu_provider.is_available():
|
||
info("✅ 使用智谱 API 作为 mem0 embedder")
|
||
embedder_config = {
|
||
"provider": "openai",
|
||
"config": {
|
||
"model": ZHIPU_EMBEDDING_MODEL,
|
||
"api_key": ZHIPUAI_API_KEY,
|
||
"openai_base_url": ZHIPU_API_BASE,
|
||
},
|
||
}
|
||
else:
|
||
# 都不可用，使用 dummy 配置并警告
|
||
warning("⚠️ 没有可用的 embedder，使用 dummy 配置")
|
||
embedder_config = {
|
||
"provider": "openai",
|
||
"config": {
|
||
"model": "text-embedding-ada-002",
|
||
"api_key": "***",
|
||
"openai_base_url": "http://localhost:8080/v1",
|
||
},
|
||
}
|
||
|
||
# 获取 LLM 服务（内部获取）
|
||
info("🔄 正在获取 LLM 服务...")
|
||
chat_llm = get_chat_service()
|
||
info("✅ LLM 服务获取成功")
|
||
|
||
# Mem0 配置
|
||
info("🔄 正在构建 Mem0 配置...")
|
||
config = {
|
||
"vector_store": {
|
||
"provider": "qdrant",
|
||
"config": {
|
||
"url": QDRANT_URL,
|
||
"api_key": QDRANT_API_KEY,
|
||
"collection_name": QDRANT_COLLECTION_NAME,
|
||
"embedding_model_dims": embedding_dim,
|
||
},
|
||
},
|
||
"llm": {
|
||
"provider": "openai",
|
||
"config": {
|
||
"model": "gpt-3.5-turbo", # 使用一个通用的模型名
|
||
"api_key": LLM_API_KEY or ZHIPUAI_API_KEY or "dummy-key",
|
||
"openai_base_url": VLLM_BASE_URL or ZHIPU_API_BASE,
|
||
"temperature": 0.1,
|
||
"max_tokens": 2000,
|
||
},
|
||
},
|
||
"embedder": embedder_config,
|
||
"version": "v1.1",
|
||
}
|
||
|
||
info("🔄 正在初始化 Mem0 实例...")
|
||
self.mem0 = AsyncMemory.from_config(config)
|
||
info("✅ Mem0 配置加载成功")
|
||
|
||
# 尝试进行连接测试，但失败不会阻止初始化
|
||
try:
|
||
info("🔄 正在测试 Mem0 连接...")
|
||
await asyncio.wait_for(
|
||
self.mem0.search("ping", user_id="test", limit=1),
|
||
timeout=10.0,
|
||
)
|
||
info("✅ Mem0 连接测试成功")
|
||
except Exception as e:
|
||
warning(f"⚠️ Mem0 连接测试遇到问题（但继续使用）: {e}")
|
||
|
||
self._initialized = True
|
||
info("🎉 Mem0 初始化完成")
|
||
|
||
except asyncio.TimeoutError:
|
||
error("❌ Mem0 初始化超时")
|
||
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, user_id):
|
||
"""添加记忆"""
|
||
if not self.mem0:
|
||
return False
|
||
try:
|
||
start = time.time()
|
||
info(f"📝 开始 Mem0 add，消息数: {len(messages)}")
|
||
await asyncio.wait_for(
|
||
self.mem0.add(
|
||
messages, user_id=user_id, metadata={"type": "conversation"}
|
||
),
|
||
timeout=60.0,
|
||
)
|
||
info(f"✅ Mem0 add 完成，耗时: {time.time() - start:.2f}s")
|
||
return True
|
||
except asyncio.TimeoutError:
|
||
error(
|
||
f"❌ Mem0 记忆添加超时 (60s)，已等待 {time.time() - start:.2f}s"
|
||
)
|
||
return False
|
||
except Exception as e:
|
||
error(f"❌ Mem0 add 失败: {e}")
|
||
return False
|