RDS长期记忆

#!/usr/bin/env python3
"""
Mem0 SDK 基本功能测试脚本
使用 MemoryClient 进行增删改查操作

使用方法:
    python test_mem0_sdk_basic.py
"""
import os
import json
import time
from mem0 import MemoryClient

# ============================================

def main():
    # 从环境变量获取配置
    api_key = os.environ.get("MEM0_API_KEY")
    host = os.environ.get("MEM0_HOST")

    print("=" * 60)
    print("Mem0 SDK 基本功能测试")
    print("=" * 60)
    print(f"服务地址: {host}")

    # 初始化客户端
    client = MemoryClient(host=host, api_key=api_key)

    # 测试用户
    user_id = f"sdk_test_user_{int(time.time())}"
    print(f"测试用户: {user_id}")
    print("=" * 60)

    memory_id = None

    try:
        # =================== 1. 添加记忆 (Add) ===================
        print("\n[1] 添加记忆 (Add)")
        print("-" * 40)

        messages = [
            {"role": "user", "content": "我是张三，在北京工作，喜欢打篮球。"},
            {"role": "assistant", "content": "你好张三！北京是个好地方，篮球也是很棒的运动。"},
        ]

        result = client.add(messages, user_id=user_id)
        print(f"添加结果: {json.dumps(result, ensure_ascii=False, indent=2)}")

        # 等待处理完成
        time.sleep(2)

        # =================== 2. 获取所有记忆 (Get All) ===================
        print("\n[2] 获取所有记忆 (Get All)")
        print("-" * 40)

        all_memories = client.get_all(user_id=user_id)
        print(f"记忆列表: {json.dumps(all_memories, ensure_ascii=False, indent=2)}")

        memories = all_memories.get("results", [])
        print(f"共 {len(memories)} 条记忆")

        if memories:
            memory_id = memories[0]["id"]
            print(f"第一条记忆 ID: {memory_id}")

        # =================== 3. 搜索记忆 (Search) ===================
        print("\n[3] 搜索记忆 (Search)")
        print("-" * 40)

        search_result = client.search("张三在哪里工作", user_id=user_id)
        print(f"搜索结果: {json.dumps(search_result, ensure_ascii=False, indent=2)}")

        # =================== 4. 获取单条记忆 (Get) ===================
        print("\n[4] 获取单条记忆 (Get)")
        print("-" * 40)

        if memory_id:
            single_memory = client.get(memory_id)
            print(f"记忆详情: {json.dumps(single_memory, ensure_ascii=False, indent=2)}")
        else:
            print("无可用记忆 ID")

        # =================== 5. 更新记忆 (Update) ===================
        print("\n[5] 更新记忆 (Update)")
        print("-" * 40)

        if memory_id:
            new_content = "张三现在在上海工作，依然喜欢打篮球"
            update_result = client.update(memory_id, new_content)
            print(f"更新结果: {json.dumps(update_result, ensure_ascii=False, indent=2)}")
        else:
            print("无可用记忆 ID")

        # =================== 6. 获取记忆历史 (History) ===================
        print("\n[6] 获取记忆历史 (History)")
        print("-" * 40)

        if memory_id:
            history = client.history(memory_id)
            print(f"历史记录: {json.dumps(history, ensure_ascii=False, indent=2)}")
        else:
            print("无可用记忆 ID")

        # =================== 7. 删除单条记忆 (Delete) ===================
        print("\n[7] 删除单条记忆 (Delete)")
        print("-" * 40)

        if memory_id:
            delete_result = client.delete(memory_id)
            print(f"删除结果: {json.dumps(delete_result, ensure_ascii=False, indent=2)}")
        else:
            print("无可用记忆 ID")

        # =================== 8. 验证删除 ===================
        print("\n[8] 验证删除")
        print("-" * 40)

        final_memories = client.get_all(user_id=user_id)
        remaining = len(final_memories.get("results", []))
        print(f"剩余记忆数: {remaining}")

        # =================== 9. 删除所有记忆 (Delete All) ===================
        print("\n[9] 删除所有记忆 (Delete All)")
        print("-" * 40)

        delete_all_result = client.delete_all(user_id=user_id)
        print(f"删除全部结果: {json.dumps(delete_all_result, ensure_ascii=False, indent=2)}")

        # =================== 测试完成 ===================
        print("\n" + "=" * 60)
        print("测试完成！")
        print("=" * 60)

    except Exception as e:
        print(f"\n错误: {e}")
        import traceback
        traceback.print_exc()
        return 1

    return 0

if __name__ == "__main__":
    exit(main())

#!/usr/bin/env python3
"""
Mem0 SDK 图搜索功能测试脚本
使用 MemoryClient 进行带图搜索的增删改查操作

功能特点:
- 启用 enable_graph=True 进行图搜索
- 展示 relations（图关系）结果
- 对比普通搜索和图搜索的差异

使用方法:
    python 当前文件
"""

import os
import json
import time
from mem0 import MemoryClient

# ============================================
# 配置
# ============================================
HOST = "http://<your-host>:80/memory"  # 服务地址
API_KEY = "<your-api-key>"              # API 密钥

# ============================================


def print_json(data, title=None):
    """格式化打印 JSON 数据"""
    if title:
        print(f"\n{title}:")
    print(json.dumps(data, ensure_ascii=False, indent=2))


def print_relations(relations):
    """格式化打印图关系"""
    if not relations:
        print("  (无图关系)")
        return
    
    print(f"  共 {len(relations)} 条关系:")
    for i, rel in enumerate(relations, 1):
        source = rel.get("source", "?")
        relationship = rel.get("relationship", "?")
        destination = rel.get("destination") or rel.get("target", "?")
        print(f"    [{i}] {source} --[{relationship}]--> {destination}")


def main():
    # 从环境变量获取配置
    api_key = os.environ.get("MEM0_API_KEY", API_KEY)
    host = os.environ.get("MEM0_HOST", HOST)
    
    print("=" * 70)
    print("Mem0 SDK 图搜索功能测试")
    print("=" * 70)
    print(f"服务地址: {host}")
    
    # 初始化客户端
    client = MemoryClient(host=host, api_key=api_key)
    
    # 测试用户
    user_id = "user_001"
    print(f"测试用户: {user_id}")
    print("=" * 70)
    
    try:
        # =================== 1. 添加多条记忆（构建图关系） ===================
        print("\n[1] 添加多条记忆（构建图关系）")
        print("-" * 50)
        
        # 第一轮对话：个人信息
        messages1 = [
            {"role": "user", "content": "我是张三，在阿里巴巴工作，负责AI平台研发。"},
            {"role": "assistant", "content": "你好张三！阿里巴巴的AI平台是很有影响力的。"},
        ]
        result1 = client.add(messages1, user_id=user_id)
        print(f"添加结果1: {json.dumps(result1, ensure_ascii=False)[:200]}...")
        
        time.sleep(2)
        
        # 第二轮对话：兴趣爱好
        messages2 = [
            {"role": "user", "content": "我喜欢打篮球，周末经常和李四一起打球。"},
            {"role": "assistant", "content": "篮球是很棒的运动，和朋友一起更有趣！"},
        ]
        result2 = client.add(messages2, user_id=user_id)
        print(f"添加结果2: {json.dumps(result2, ensure_ascii=False)[:200]}...")
        
        time.sleep(2)
        
        # 第三轮对话：更多关系
        messages3 = [
            {"role": "user", "content": "我的老板是王总，他是数据智能部门的负责人。"},
            {"role": "assistant", "content": "了解，王总是数据智能部门的领导。"},
        ]
        result3 = client.add(messages3, user_id=user_id)
        print(f"添加结果3: {json.dumps(result3, ensure_ascii=False)[:200]}...")
        
        # 等待图构建完成
        print("\n等待图关系构建完成...")
        time.sleep(5)
        
        # =================== 2. 获取所有记忆（带图） ===================
        print("\n[2] 获取所有记忆（enable_graph=True）")
        print("-" * 50)
        
        all_memories = client.get_all(user_id=user_id, enable_graph=True)
        
        results = all_memories.get("results", [])
        relations = all_memories.get("relations", [])
        
        print(f"\n向量搜索结果: 共 {len(results)} 条记忆")
        for i, mem in enumerate(results[:5], 1):
            print(f"  [{i}] {mem.get('memory', '')[:50]}...")
        
        print(f"\n图关系结果:")
        print_relations(relations)
        
        # =================== 3. 图搜索 vs 普通搜索 ===================
        print("\n[3] 图搜索 vs 普通搜索 对比")
        print("-" * 50)
        
        query = "张三在哪里工作？"
        print(f"搜索查询: {query}")
        
        # 3.1 普通搜索（不启用图）
        print("\n【普通搜索】enable_graph=False:")
        search_normal = client.search(query, user_id=user_id, enable_graph=False)
        normal_results = search_normal.get("results", [])
        normal_relations = search_normal.get("relations")
        
        print(f"  向量结果数: {len(normal_results)}")
        for i, mem in enumerate(normal_results[:3], 1):
            score = mem.get("score", "N/A")
            memory = mem.get("memory", "")[:40]
            print(f"    [{i}] score={score:.4f} | {memory}...")
        print(f"  图关系: {normal_relations}")
        
        # 3.2 图搜索（启用图）
        print("\n【图搜索】enable_graph=True:")
        search_graph = client.search(query, user_id=user_id, enable_graph=True)
        graph_results = search_graph.get("results", [])
        graph_relations = search_graph.get("relations", [])
        
        print(f"  向量结果数: {len(graph_results)}")
        for i, mem in enumerate(graph_results[:3], 1):
            score = mem.get("score", "N/A")
            memory = mem.get("memory", "")[:40]
            print(f"    [{i}] score={score:.4f} | {memory}...")
        
        print(f"  图关系结果:")
        print_relations(graph_relations)
        
        # =================== 4. 多关键词图搜索 ===================
        print("\n[4] 多关键词图搜索")
        print("-" * 50)
        
        queries = [
            "张三和李四的关系",
            "谁是王总",
            "篮球相关的活动",
        ]
        
        for query in queries:
            print(f"\n查询: {query}")
            result = client.search(query, user_id=user_id, enable_graph=True)
            
            results = result.get("results", [])
            relations = result.get("relations", [])
            
            print(f"  向量结果数: {len(results)}")
            if results:
                print(f"    Top1: {results[0].get('memory', '')[:50]}...")
            
            print(f"  图关系:")
            print_relations(relations)
        
        # =================== 5. 清理测试数据 ===================
        print("\n[5] 清理测试数据")
        print("-" * 50)
        
        delete_result = client.delete_all(user_id=user_id)
        print(f"删除结果: {json.dumps(delete_result, ensure_ascii=False)}")
        
        # =================== 测试完成 ===================
        print("\n" + "=" * 70)
        print("图搜索测试完成！")
        print("=" * 70)
        
        print("\n【总结】")
        print("- enable_graph=False: 仅返回向量搜索结果 (results)")
        print("- enable_graph=True:  返回向量结果 (results) + 图关系 (relations)")
        print("- 图关系格式: source --[relationship]--> destination")
        
    except Exception as e:
        print(f"\n错误: {e}")
        import traceback
        traceback.print_exc()
        
        # 清理测试数据
        try:
            client.delete_all(user_id=user_id)
            print("已清理测试数据")
        except Exception:
            pass
        
        return 1
    
    return 0


if __name__ == "__main__":
    exit(main())