basic_operations.py

#!/usr/bin/env python3
"""
Scintirete SDK 基本操作示例

演示如何使用 Scintirete Python SDK 进行基本的数据库和向量操作。
"""

import random
from typing import List

from scintirete_sdk import (
    ScintireteClient,
    DistanceMetric,
    HnswConfig,
    Vector,
    ScintireteError,
)


def generate_random_vector(dimension: int) -> List[float]:
    """生成随机向量"""
    return [random.random() for _ in range(dimension)]


def main():
    """主函数"""
    # 连接配置
    SERVER_ADDRESS = "localhost:50051"
    PASSWORD = None  # 如果服务器需要密码，请设置
    
    # 测试数据配置
    DB_NAME = "example_db"
    COLLECTION_NAME = "example_collection"
    VECTOR_DIMENSION = 128
    
    try:
        # 创建客户端连接
        print("🔗 连接到 Scintirete 服务器...")
        with ScintireteClient(SERVER_ADDRESS, password=PASSWORD) as client:
            print("✅ 连接成功!")
            
            # 1. 数据库操作
            print("\n📊 数据库操作")
            print("=" * 50)
            
            # 列出现有数据库
            databases = client.list_databases()
            print(f"现有数据库: {databases}")
            
            # 删除测试数据库（如果存在）
            if DB_NAME in databases:
                print(f"删除现有数据库: {DB_NAME}")
                success, dropped_collections = client.drop_database(DB_NAME)
                print(f"删除结果: success={success}, dropped_collections={dropped_collections}")
            
            # 创建新数据库
            print(f"创建数据库: {DB_NAME}")
            success = client.create_database(DB_NAME)
            print(f"创建结果: {success}")
            
            # 验证数据库创建
            databases = client.list_databases()
            print(f"更新后的数据库列表: {databases}")
            
            # 2. 集合操作
            print("\n📁 集合操作")
            print("=" * 50)
            
            # 创建 HNSW 配置
            hnsw_config = HnswConfig(
                m=16,                   # 每个节点最大连接数
                ef_construction=200     # 构建时搜索范围
            )
            
            # 创建集合
            print(f"创建集合: {COLLECTION_NAME}")
            collection_info = client.create_collection(
                db_name=DB_NAME,
                collection_name=COLLECTION_NAME,
                metric_type=DistanceMetric.COSINE,
                hnsw_config=hnsw_config
            )
            
            print(f"集合信息:")
            print(f"  名称: {collection_info.name}")
            print(f"  维度: {collection_info.dimension}")
            print(f"  向量数: {collection_info.vector_count}")
            print(f"  度量类型: {collection_info.metric_type}")
            print(f"  HNSW配置: m={collection_info.hnsw_config.m}, ef_construction={collection_info.hnsw_config.ef_construction}")
            
            # 列出集合
            collections = client.list_collections(DB_NAME)
            print(f"数据库 {DB_NAME} 中的集合: {[c.name for c in collections]}")
            
            # 3. 向量操作
            print("\n🎯 向量操作")
            print("=" * 50)
            
            # 准备测试向量
            print("准备测试向量...")
            vectors = []
            for i in range(10):
                vector = Vector(
                    elements=generate_random_vector(VECTOR_DIMENSION),
                    metadata={
                        "id": i,
                        "category": f"category_{i % 3}",
                        "timestamp": f"2024-01-{i+1:02d}",
                        "description": f"这是第 {i+1} 个测试向量"
                    }
                )
                vectors.append(vector)
            
            print(f"生成了 {len(vectors)} 个 {VECTOR_DIMENSION} 维向量")
            
            # 插入向量
            print("插入向量...")
            inserted_ids, inserted_count = client.insert_vectors(
                db_name=DB_NAME,
                collection_name=COLLECTION_NAME,
                vectors=vectors
            )
            
            print(f"插入结果: {inserted_count} 个向量")
            print(f"分配的ID: {inserted_ids}")
            
            # 获取集合信息（更新后）
            collection_info = client.get_collection_info(DB_NAME, COLLECTION_NAME)
            print(f"更新后的集合信息:")
            print(f"  维度: {collection_info.dimension}")
            print(f"  向量数: {collection_info.vector_count}")
            print(f"  内存使用: {collection_info.memory_bytes} 字节")
            
            # 搜索向量
            print("\n🔍 搜索向量")
            print("-" * 30)
            
            # 使用第一个向量作为查询向量
            query_vector = vectors[0].elements
            print(f"使用查询向量搜索（前5个维度: {query_vector[:5]}...）")
            
            search_results = client.search(
                db_name=DB_NAME,
                collection_name=COLLECTION_NAME,
                query_vector=query_vector,
                top_k=5,
                include_vector=False  # 不返回向量数据以提高性能
            )
            
            print(f"搜索结果 (top {len(search_results)}):")
            for i, result in enumerate(search_results):
                print(f"  {i+1}. ID: {result.id}, 距离: {result.distance:.4f}")
                if result.metadata:
                    print(f"      元数据: {result.metadata}")
            
            # 搜索并返回向量数据
            print("\n搜索并返回向量数据:")
            search_results_with_vectors = client.search(
                db_name=DB_NAME,
                collection_name=COLLECTION_NAME,
                query_vector=query_vector,
                top_k=3,
                include_vector=True
            )
            
            for i, result in enumerate(search_results_with_vectors):
                print(f"  {i+1}. ID: {result.id}, 距离: {result.distance:.4f}")
                if result.vector:
                    print(f"      向量维度: {len(result.vector.elements)}")
                    print(f"      向量前5维: {result.vector.elements[:5]}")
            
            # 删除部分向量
            print("\n🗑️ 删除向量")
            print("-" * 20)
            
            # 删除前3个向量
            ids_to_delete = inserted_ids[:3]
            print(f"删除向量 ID: {ids_to_delete}")
            
            deleted_count = client.delete_vectors(
                db_name=DB_NAME,
                collection_name=COLLECTION_NAME,
                ids=ids_to_delete
            )
            
            print(f"删除了 {deleted_count} 个向量")
            
            # 验证删除结果
            collection_info = client.get_collection_info(DB_NAME, COLLECTION_NAME)
            print(f"删除后的集合信息:")
            print(f"  向量数: {collection_info.vector_count}")
            print(f"  删除数: {collection_info.deleted_count}")
            
            # 4. 持久化操作
            print("\n💾 持久化操作")
            print("=" * 50)
            
            # 同步保存
            print("执行同步保存...")
            success, message, size, duration = client.save()
            print(f"保存结果: success={success}")
            print(f"消息: {message}")
            print(f"快照大小: {size} 字节")
            print(f"耗时: {duration:.3f} 秒")
            
            # 后台保存
            print("\n执行后台保存...")
            success, message, job_id = client.bg_save()
            print(f"后台保存结果: success={success}")
            print(f"消息: {message}")
            print(f"任务ID: {job_id}")
            
            print("\n🎉 基本操作演示完成!")
    
    except ScintireteError as e:
        print(f"❌ Scintirete 错误: {e}")
    except Exception as e:
        print(f"❌ 未预期错误: {e}")
        raise


if __name__ == "__main__":
    main()