l_ai_knowledge/docs/使用其他向量数据库.md

6.3 KiB
Raw Permalink Blame History

如何集成新的向量数据库

本文提供了向 WeKnora 项目添加新向量数据库支持的完整指南。通过实现标准化接口和遵循结构化流程,开发者可以高效地集成自定义向量数据库。

集成流程

1. 实现基础检索引擎接口

首先需要实现 interfaces 包中的 RetrieveEngine 接口,定义检索引擎的核心能力:

type RetrieveEngine interface {
    // 返回检索引擎的类型标识
    EngineType() types.RetrieverEngineType

    // 执行检索操作,返回匹配结果
    Retrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error)

    // 返回该引擎支持的检索类型列表
    Support() []types.RetrieverType
}

2. 实现存储层接口

实现 RetrieveEngineRepository 接口,扩展基础检索引擎能力,添加索引管理功能:

type RetrieveEngineRepository interface {
    // 保存单个索引信息
    Save(ctx context.Context, indexInfo *types.IndexInfo, params map[string]any) error
    
    // 批量保存多个索引信息
    BatchSave(ctx context.Context, indexInfoList []*types.IndexInfo, params map[string]any) error
    
    // 估算索引存储所需空间
    EstimateStorageSize(ctx context.Context, indexInfoList []*types.IndexInfo, params map[string]any) int64
    
    // 通过分块ID列表删除索引
    DeleteByChunkIDList(ctx context.Context, indexIDList []string, dimension int) error
    
    // 复制索引数据,避免重新计算嵌入向量
    CopyIndices(
        ctx context.Context,
        sourceKnowledgeBaseID string,
        sourceToTargetKBIDMap map[string]string,
        sourceToTargetChunkIDMap map[string]string,
        targetKnowledgeBaseID string,
        dimension int,
    ) error
    
    // 通过知识ID列表删除索引
    DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, dimension int) error
    
    // 继承RetrieveEngine接口
    RetrieveEngine
}

3. 实现服务层接口

创建实现 RetrieveEngineService 接口的服务,负责处理索引创建和管理的业务逻辑:

type RetrieveEngineService interface {
    // 创建单个索引
    Index(ctx context.Context,
        embedder embedding.Embedder,
        indexInfo *types.IndexInfo,
        retrieverTypes []types.RetrieverType,
    ) error

    // 批量创建索引
    BatchIndex(ctx context.Context,
        embedder embedding.Embedder,
        indexInfoList []*types.IndexInfo,
        retrieverTypes []types.RetrieverType,
    ) error

    // 估算索引存储空间
    EstimateStorageSize(ctx context.Context,
        embedder embedding.Embedder,
        indexInfoList []*types.IndexInfo,
        retrieverTypes []types.RetrieverType,
    ) int64
    
    // 复制索引数据
    CopyIndices(
        ctx context.Context,
        sourceKnowledgeBaseID string,
        sourceToTargetKBIDMap map[string]string,
        sourceToTargetChunkIDMap map[string]string,
        targetKnowledgeBaseID string,
        dimension int,
    ) error

    // 删除索引
    DeleteByChunkIDList(ctx context.Context, indexIDList []string, dimension int) error
    DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, dimension int) error

    // 继承RetrieveEngine接口
    RetrieveEngine
}

4. 添加环境变量配置

在环境配置中添加新数据库的必要连接参数:

# 在RETRIEVE_DRIVER中添加新数据库驱动名称多个驱动用逗号分隔
RETRIEVE_DRIVER=postgres,elasticsearch_v8,your_database

# 新数据库的连接参数
YOUR_DATABASE_ADDR=your_database_host:port
YOUR_DATABASE_USERNAME=username
YOUR_DATABASE_PASSWORD=password
# 其他必要的连接参数...

5. 注册检索引擎

internal/container/container.go 文件的 initRetrieveEngineRegistry 函数中添加新数据库的初始化与注册逻辑:

func initRetrieveEngineRegistry(db *gorm.DB, cfg *config.Config) (interfaces.RetrieveEngineRegistry, error) {
    registry := retriever.NewRetrieveEngineRegistry()
    retrieveDriver := strings.Split(os.Getenv("RETRIEVE_DRIVER"), ",")
    log := logger.GetLogger(context.Background())

    // 已有的PostgreSQL和Elasticsearch初始化代码...
    
    // 添加新向量数据库的初始化代码
    if slices.Contains(retrieveDriver, "your_database") {
        // 初始化数据库客户端
        client, err := your_database.NewClient(your_database.Config{
            Addresses: []string{os.Getenv("YOUR_DATABASE_ADDR")},
            Username:  os.Getenv("YOUR_DATABASE_USERNAME"),
            Password:  os.Getenv("YOUR_DATABASE_PASSWORD"),
            // 其他连接参数...
        })
        
        if err != nil {
            log.Errorf("Create your_database client failed: %v", err)
        } else {
            // 创建检索引擎仓库
            yourDatabaseRepo := your_database.NewYourDatabaseRepository(client, cfg)
            
            // 注册检索引擎
            if err := registry.Register(
                retriever.NewKVHybridRetrieveEngine(
                    yourDatabaseRepo, types.YourDatabaseRetrieverEngineType,
                ),
            ); err != nil {
                log.Errorf("Register your_database retrieve engine failed: %v", err)
            } else {
                log.Infof("Register your_database retrieve engine success")
            }
        }
    }

    return registry, nil
}

6. 定义检索引擎类型常量

internal/types/retriever.go 文件中添加新的检索引擎类型常量:

// RetrieverEngineType 定义检索引擎类型
const (
    ElasticsearchRetrieverEngineType RetrieverEngineType = "elasticsearch"
    PostgresRetrieverEngineType      RetrieverEngineType = "postgres"
    YourDatabaseRetrieverEngineType  RetrieverEngineType = "your_database" // 添加新数据库类型
)

参考实现示例

建议参考现有的 PostgreSQL 和 Elasticsearch 实现作为开发模板。这些实现位于以下目录:

  • PostgreSQL: internal/application/repository/retriever/postgres/
  • ElasticsearchV7: internal/application/repository/retriever/elasticsearch/v7/
  • ElasticsearchV8: internal/application/repository/retriever/elasticsearch/v8/

通过遵循以上步骤和参考现有实现,你可以成功集成新的向量数据库到 WeKnora 系统中,扩展其向量检索能力。