### 如何集成新的向量数据库 本文提供了向 WeKnora 项目添加新向量数据库支持的完整指南。通过实现标准化接口和遵循结构化流程,开发者可以高效地集成自定义向量数据库。 ### 集成流程 #### 1. 实现基础检索引擎接口 首先需要实现 `interfaces` 包中的 `RetrieveEngine` 接口,定义检索引擎的核心能力: ```go type RetrieveEngine interface { // 返回检索引擎的类型标识 EngineType() types.RetrieverEngineType // 执行检索操作,返回匹配结果 Retrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) // 返回该引擎支持的检索类型列表 Support() []types.RetrieverType } ``` #### 2. 实现存储层接口 实现 `RetrieveEngineRepository` 接口,扩展基础检索引擎能力,添加索引管理功能: ```go 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` 接口的服务,负责处理索引创建和管理的业务逻辑: ```go 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` 函数中添加新数据库的初始化与注册逻辑: ```go 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` 文件中添加新的检索引擎类型常量: ```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 系统中,扩展其向量检索能力。