|
|
|
|
|
#!/usr/bin/env python3
|
|
|
|
|
|
"""
|
|
|
|
|
|
Session Initialization Hook - 冷启动记忆预加载
|
|
|
|
|
|
在每次 /new 或新会话启动时自动检索最近活跃上下文
|
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
import asyncio
|
|
|
|
|
|
import sys
|
|
|
|
|
|
import os
|
|
|
|
|
|
import logging
|
|
|
|
|
|
|
|
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
|
|
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', '') or os.getenv('DASHSCOPE_API_KEY', '')
|
|
|
|
|
|
if _dashscope_key:
|
|
|
|
|
|
os.environ['OPENAI_API_KEY'] = _dashscope_key
|
|
|
|
|
|
|
|
|
|
|
|
from mem0_client import Mem0Client
|
|
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
# 冷启动检索查询列表
|
|
|
|
|
|
COLD_START_QUERIES = [
|
|
|
|
|
|
"current active tasks deployment progress",
|
|
|
|
|
|
"pending work items todos",
|
|
|
|
|
|
"recent important decisions configurations",
|
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
async def cold_start_retrieval(agent_id: str = "main", user_id: str = "default", top_k: int = 3) -> str:
|
|
|
|
|
|
"""
|
|
|
|
|
|
冷启动记忆检索 - 在会话初始化时调用
|
|
|
|
|
|
返回格式化的记忆片段，注入到 System Prompt
|
|
|
|
|
|
"""
|
|
|
|
|
|
client = Mem0Client()
|
|
|
|
|
|
await client.start()
|
|
|
|
|
|
|
|
|
|
|
|
all_memories = []
|
|
|
|
|
|
|
|
|
|
|
|
for query in COLD_START_QUERIES:
|
|
|
|
|
|
try:
|
|
|
|
|
|
memories = await client.pre_hook_search(
|
|
|
|
|
|
query=query,
|
|
|
|
|
|
user_id=user_id,
|
|
|
|
|
|
agent_id=agent_id,
|
|
|
|
|
|
top_k=top_k
|
|
|
|
|
|
)
|
|
|
|
|
|
all_memories.extend(memories)
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
|
logger.debug(f"检索失败 {query}: {e}")
|
|
|
|
|
|
|
|
|
|
|
|
# 去重（按 memory 内容）
|
|
|
|
|
|
seen = set()
|
|
|
|
|
|
unique_memories = []
|
|
|
|
|
|
for m in all_memories:
|
|
|
|
|
|
mem_text = m.get('memory', '')
|
|
|
|
|
|
if mem_text and mem_text not in seen:
|
|
|
|
|
|
seen.add(mem_text)
|
|
|
|
|
|
unique_memories.append(m)
|
|
|
|
|
|
|
|
|
|
|
|
# 按分数排序，取前 top_k 条
|
|
|
|
|
|
unique_memories.sort(key=lambda x: x.get('score', 0), reverse=True)
|
|
|
|
|
|
unique_memories = unique_memories[:top_k]
|
|
|
|
|
|
|
|
|
|
|
|
await client.shutdown()
|
|
|
|
|
|
|
|
|
|
|
|
# 格式化为 Prompt 片段
|
|
|
|
|
|
if not unique_memories:
|
|
|
|
|
|
return ""
|
|
|
|
|
|
|
|
|
|
|
|
prompt = "\n\n=== 最近活跃上下文（自动加载） ===\n"
|
|
|
|
|
|
for i, mem in enumerate(unique_memories, 1):
|
|
|
|
|
|
mem_text = mem.get('memory', '')
|
|
|
|
|
|
metadata = mem.get('metadata', {})
|
|
|
|
|
|
agent = metadata.get('agent_id', 'unknown')
|
|
|
|
|
|
prompt += f"{i}. [{agent}] {mem_text}\n"
|
|
|
|
|
|
prompt += "=================================\n"
|
|
|
|
|
|
|
|
|
|
|
|
logger.info(f"冷启动检索完成：{len(unique_memories)} 条记忆")
|
|
|
|
|
|
return prompt
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def main():
|
|
|
|
|
|
"""测试冷启动检索"""
|
|
|
|
|
|
result = await cold_start_retrieval(
|
|
|
|
|
|
agent_id="main",
|
|
|
|
|
|
user_id="wang 院长",
|
|
|
|
|
|
top_k=5
|
|
|
|
|
|
)
|
|
|
|
|
|
if result:
|
|
|
|
|
|
print("✅ 冷启动检索成功:")
|
|
|
|
|
|
print(result)
|
|
|
|
|
|
else:
|
|
|
|
|
|
print("⚠️ 无记忆记录（Qdrant 为空）")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
asyncio.run(main())
|
