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