本页面提供一个通过 Python SDK 初始化记忆库、写入数据和检索记忆的完整请求示例。在完成安装与初始化SDK的基础上，您可以参考以下示例代码体验 Viking 长期记忆的核心流程。

"""
VikingDB 会话记忆完整示例
包含初始化、写入会话数据和检索会话记忆的完整流程
"""

import os
from datetime import datetime
from vikingdb import IAM
from vikingdb.memory import VikingMem
from vikingdb.exceptions import VikingException

def init_memory_client():
    """初始化记忆库客户端"""
    print("=== 初始化 VikingDB 客户端 ===")
    
    # 认证配置
    _auth = IAM(
        ak=os.getenv("VIKINGDB_AK", "your_ak"),
        sk=os.getenv("VIKINGDB_SK", "your_sk"),
    )
    
    # 创建客户端
    client = VikingMem(
        host="api-knowledgebase.mlp.cn-beijing.volces.com",
        region="cn-beijing",
        auth=_auth,
        scheme="http",
    )
    
    print("客户端初始化成功")
    return client
def get_or_create_collection(client, collection_name, project_name="default"):
    """获取或创建记忆库集合"""
    print(f"\n=== 获取记忆库集合: {collection_name} ===")
    
    try:
        # 尝试获取已存在的集合
        collection = client.get_collection(
            collection_name=collection_name,
            project_name=project_name
        )
        print(f"成功获取现有集合: {collection_name}")
        return collection
        
    except VikingException as e:
        print(f"获取集合失败: {e.message}")
        raise

def add_memory_sessions(collection):
    """添加会话记忆数据"""
    print(f"\n=== 添加会话记忆 ===")
    
    # 模拟对话会话数据
    memory_sessions = [
        {
            "session_id": "session_001",
            "messages": [
                {"role": "user", "content": "今天天气怎么样？"},
                {"role": "assistant", "content": "今天是晴天，适合外出活动！"},
                {"role": "user", "content": "那太好了，我正好想出去走走"},
                {"role": "assistant", "content": "记得涂防晒霜，享受美好的一天！"}
            ],
            "metadata": {
                "default_user_id": "user_001",
                "default_assistant_id": "assistant_001",
                "time": int(datetime.now().timestamp() * 1000),
            }
        },
        {
            "session_id": "session_002",
            "messages": [
                {"role": "user", "content": "你推荐什么咖啡？我喜欢浓郁的口味"},
                {"role": "assistant", "content": "推荐您尝试意式浓缩或美式咖啡，都是浓郁型的！"},
                {"role": "user", "content": "那拿铁和卡布奇诺呢？"},
                {"role": "assistant", "content": "拿铁奶味重一些，卡布奇诺奶泡丰富，都很有特色"}
            ],
            "metadata": {
                "default_user_id": "user_001",
                "default_assistant_id": "assistant_001",
                "time": int(datetime.now().timestamp() * 1000),
            }
        },
        {
            "session_id": "session_003",
            "messages": [
                {"role": "user", "content": "如何设置提醒功能？我是新用户"},
                {"role": "assistant", "content": "点击右上角的铃铛图标，然后选择'添加提醒'即可"},
                {"role": "user", "content": "可以设置重复提醒吗？"},
                {"role": "assistant", "content": "可以的，在设置提醒时选择重复频率即可"}
            ],
            "metadata": {
                "default_user_id": "user_001",
                "default_assistant_id": "assistant_001",
                "time": int(datetime.now().timestamp() * 1000),
            }
        }
    ]
    
    results = []
    for i, session_data in enumerate(memory_sessions, 1):
        try:
            print(f"添加会话 {i}/{len(memory_sessions)}...")
            
            result = collection.add_session(
                session_id=session_data["session_id"],
                messages=session_data["messages"],
                metadata=session_data["metadata"]
            )
            
            results.append(result)
            print(f"会话 {i} 添加成功")
            
        except VikingException as e:
            print(f"会话 {i} 添加失败: {e.message}")
            continue
    
    print(f"成功添加 {len(results)} 个会话记忆")
    return results

def add_user_profile(collection):
    """添加用户画像数据"""
    print(f"\n=== 添加用户画像 ===")
    
    try:
        result = collection.add_profile(
            profile_type="sys_profile_v1",  # 使用系统默认的profile type
            user_id="user_001",  # 使用新的用户ID避免重复
            assistant_id="assistant_001",
            memory_info={
                "user_profile": "### 用户画像\n- 性别: 女\n- 年龄: 28岁\n- 职业: 产品经理\n- 兴趣: 咖啡、阅读、旅行\n- 性格: 开朗、细心、有创造力\n- 饮品偏好: 咖啡\n- 爱好: 阅读、旅行、摄影\n- 沟通风格: 友好直接"
            }
        )
        
        print("用户画像添加成功")
        return result
        
    except VikingException as e:
        if "Profile already exists" in e.message:
            print("用户画像已存在，跳过添加步骤")
            return None
        else:
            print(f"用户画像添加失败: {e.message}")
            raise

def search_memories(collection):
    """搜索会话记忆数据"""
    print(f"\n=== 会话记忆搜索演示 ===")
    
    # 基础搜索 - 按用户ID和关键词
    print("\n--- 搜索用户001的会话记忆 ---")
    try:
        result = collection.search_memory(
            query="咖啡",
            filter={
                "memory_type": ["sys_profile_v1", "sys_event_v1"],
                "user_id": "user_001",
                "assistant_id": "assistant_001"
            },
            limit=10
        )
        
        if result and isinstance(result, dict) and result.get('data'):
            result_data = result['data']
            count = result_data.get('count', 0)
            print(f"找到 {count} 个结果")
            
            if count > 0 and 'result_list' in result_data:
                for item in result_data['result_list'][:3]:
                    print(f"- 结果: {item}")
            else:
                print("结果列表为空")
        else:
            print("未找到相关结果")
            
    except VikingException as e:
        print(f"搜索失败: {e.message}")
    
def main():
    """主函数 - 完整流程演示"""
    print("VikingDB 会话记忆完整示例开始")
    print("=" * 50)
    
    try:
        # 1. 初始化客户端
        client = init_memory_client()
        
        # 2. 获取集合（使用现有集合）
        collection = get_or_create_collection(
            client, 
            collection_name="my_first_memory_collection", #替换为您的记忆库名称
            project_name="default"
        )
        
        # 添加会话记忆
        add_memory_sessions(collection)
        
        # 添加用户画像
        add_user_profile(collection)
        
        # 搜索记忆
        search_memories(collection)
        
        print("完整示例执行成功")
        print("总结:")
        print("- 客户端初始化成功")
        print("- 会话记忆添加成功 (3个会话)")
        print("- 用户画像添加成功")
        print("- 会话记忆搜索功能正常")
        
    except VikingException as e:
        print(f"VikingDB异常: {e.message}")
        
    except Exception as e:
        print(f"异常: {str(e)}")
        
    finally:
        print("=" * 50)
        print("示例执行完成")

if __name__ == "__main__":
    main()