6.3 KiB
6.3 KiB
如何集成新的向量数据库
本文提供了向 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 系统中,扩展其向量检索能力。