核心流程--向量数据库VikingDB-火山引擎

文档中心
向量数据库VikingDB
SDK参考
核心流程
本页面提供一个通过 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()
最近更新时间：2026.04.21 20:42:18
这个页面对您有帮助吗？
有用
无用