You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
69 lines
2.1 KiB
69 lines
2.1 KiB
#!/usr/bin/env python3 |
|
""" |
|
Mem0 Python 集成脚本 |
|
被 Node.js 插件调用，执行实际的记忆操作 |
|
""" |
|
|
|
import sys |
|
import json |
|
import os |
|
import asyncio |
|
|
|
# 设置环境变量 |
|
os.environ['OPENAI_API_BASE'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1' |
|
os.environ['OPENAI_BASE_URL'] = 'https://dashscope.aliyuncs.com/compatible-mode/v1' |
|
os.environ['OPENAI_API_KEY'] = os.getenv('MEM0_DASHSCOPE_API_KEY', 'sk-c1715ee0479841399fd359c574647648') |
|
|
|
sys.path.insert(0, os.path.dirname(__file__)) |
|
|
|
from mem0_client import mem0_client |
|
|
|
|
|
async def main(): |
|
if len(sys.argv) < 2: |
|
print(json.dumps({"error": "No action specified"})) |
|
return |
|
|
|
action = sys.argv[1] |
|
data = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {} |
|
|
|
try: |
|
if action == 'init': |
|
# 初始化 mem0 |
|
await mem0_client.start() |
|
print(json.dumps({ |
|
"status": "initialized", |
|
"qdrant": f"{mem0_client.config['qdrant']['host']}:{mem0_client.config['qdrant']['port']}" |
|
})) |
|
|
|
elif action == 'search': |
|
# 检索记忆 |
|
memories = await mem0_client.pre_hook_search( |
|
query=data.get('query', ''), |
|
user_id=data.get('user_id', 'default'), |
|
agent_id=data.get('agent_id', 'general') |
|
) |
|
print(json.dumps({ |
|
"memories": memories, |
|
"count": len(memories) |
|
})) |
|
|
|
elif action == 'add': |
|
# 添加记忆（异步，不等待） |
|
mem0_client.post_hook_add( |
|
user_message=data.get('user_message', ''), |
|
assistant_message=data.get('assistant_message', ''), |
|
user_id=data.get('user_id', 'default'), |
|
agent_id=data.get('agent_id', 'general') |
|
) |
|
print(json.dumps({"status": "queued"})) |
|
|
|
else: |
|
print(json.dumps({"error": f"Unknown action: {action}"})) |
|
|
|
except Exception as e: |
|
print(json.dumps({"error": str(e)})) |
|
|
|
|
|
if __name__ == '__main__': |
|
asyncio.run(main())
|
|
|