|
|
#!/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())
|
|
|
|
