import json
|
import asyncio
|
import traceback
|
|
from ..base import MemoryProviderBase, logger
|
from mem0 import MemoryClient
|
from core.utils.util import check_model_key
|
|
TAG = __name__
|
|
|
class MemoryProvider(MemoryProviderBase):
|
def __init__(self, config, summary_memory=None):
|
super().__init__(config)
|
self.api_key = config.get("api_key", "")
|
self.api_version = config.get("api_version", "v1.1")
|
model_key_msg = check_model_key("Mem0ai", self.api_key)
|
if model_key_msg:
|
logger.bind(tag=TAG).error(model_key_msg)
|
self.use_mem0 = False
|
return
|
else:
|
self.use_mem0 = True
|
|
try:
|
self.client = MemoryClient(api_key=self.api_key)
|
logger.bind(tag=TAG).info("成功连接到 Mem0ai 服务")
|
except Exception as e:
|
logger.bind(tag=TAG).error(f"连接到 Mem0ai 服务时发生错误: {str(e)}")
|
logger.bind(tag=TAG).error(f"详细错误: {traceback.format_exc()}")
|
self.use_mem0 = False
|
|
async def save_memory(self, msgs, session_id=None):
|
try:
|
if self.use_mem0 and len(msgs) >= 2:
|
# Format the content as a message list for mem0
|
messages = []
|
for message in msgs:
|
if message.role == "system":
|
continue
|
|
if message.role == "tool":
|
continue
|
|
content = message.content
|
|
if content is None:
|
continue
|
|
# Extract content from JSON format if present (for ASR with emotion/language tags)
|
# Same logic as in query_memory method
|
try:
|
if content and content.strip().startswith("{") and content.strip().endswith("}"):
|
data = json.loads(content)
|
if "content" in data:
|
content = data["content"]
|
except (json.JSONDecodeError, KeyError, TypeError):
|
# If parsing fails, use original content
|
pass
|
|
messages.append({"role": message.role, "content": content})
|
|
result = await asyncio.to_thread(
|
self.client.add, messages, user_id=self.role_id
|
)
|
logger.bind(tag=TAG).debug(f"Save memory result: {result}")
|
except Exception as e:
|
logger.bind(tag=TAG).error(f"保存记忆失败: {str(e)}")
|
|
return None
|
|
async def query_memory(self, query: str) -> str:
|
if not self.use_mem0:
|
return ""
|
try:
|
if not getattr(self, "role_id", None):
|
return ""
|
|
filters = {"user_id": self.role_id}
|
|
search_query = query
|
try:
|
if query.strip().startswith("{") and query.strip().endswith("}"):
|
data = json.loads(query)
|
if "content" in data:
|
search_query = data["content"]
|
except (json.JSONDecodeError, KeyError):
|
pass
|
|
results = await asyncio.to_thread(
|
self.client.search, search_query, filters=filters
|
)
|
if not results or "results" not in results:
|
return ""
|
|
# Format each memory entry with its update time up to minutes
|
memories = []
|
for entry in results["results"]:
|
timestamp = entry.get("updated_at", "")
|
if timestamp:
|
try:
|
# Parse and reformat the timestamp
|
dt = timestamp.split(".")[0] # Remove milliseconds
|
formatted_time = dt.replace("T", " ")
|
except:
|
formatted_time = timestamp
|
memory = entry.get("memory", "")
|
if timestamp and memory:
|
# Store tuple of (timestamp, formatted_string) for sorting
|
memories.append((timestamp, f"[{formatted_time}] {memory}"))
|
|
# Sort by timestamp in descending order (newest first)
|
memories.sort(key=lambda x: x[0], reverse=True)
|
|
# Extract only the formatted strings
|
memories_str = "\n".join(f"- {memory[1]}" for memory in memories)
|
logger.bind(tag=TAG).debug(f"Query results: {memories_str}")
|
return memories_str
|
except Exception as e:
|
logger.bind(tag=TAG).error(f"查询记忆失败: {str(e)}")
|
return ""
|
