85 lines
2.3 KiB
Go
85 lines
2.3 KiB
Go
package vector
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"eino-project/internal/conf"
|
|
|
|
"github.com/go-kratos/kratos/v2/log"
|
|
"github.com/google/wire"
|
|
)
|
|
|
|
// ProviderSet is vector providers.
|
|
var ProviderSet = wire.NewSet(
|
|
NewVectorServiceFromBootstrapConfig,
|
|
NewKnowledgeSearcher,
|
|
NewDocumentProcessor,
|
|
)
|
|
|
|
// VectorServiceType 向量服务类型
|
|
type VectorServiceType string
|
|
|
|
const (
|
|
VectorServiceTypeMemory VectorServiceType = "memory"
|
|
VectorServiceTypeChromaDB VectorServiceType = "chromadb"
|
|
)
|
|
|
|
// NewVectorServiceFromBootstrapConfig 从 Bootstrap 配置创建向量服务
|
|
func NewVectorServiceFromBootstrapConfig(c *conf.Bootstrap, logger log.Logger) (VectorService, error) {
|
|
// 从配置中获取向量数据库配置
|
|
var config *Config
|
|
if c.Vector != nil && c.Vector.Chromadb != nil {
|
|
config = &Config{
|
|
Endpoint: c.Vector.Chromadb.Endpoint,
|
|
Collection: c.Vector.Chromadb.Collection,
|
|
Timeout: c.Vector.Chromadb.Timeout.AsDuration().String(),
|
|
}
|
|
} else {
|
|
// 使用默认配置
|
|
config = &Config{
|
|
Endpoint: "http://127.0.0.1:8000", // ChromaDB 默认端点
|
|
Collection: "customer_knowledge",
|
|
Timeout: "30s",
|
|
}
|
|
}
|
|
|
|
return NewVectorServiceFromConfig(config, logger)
|
|
}
|
|
|
|
// NewVectorServiceFromConfig 根据配置创建向量服务
|
|
func NewVectorServiceFromConfig(config *Config, logger log.Logger) (VectorService, error) {
|
|
if config == nil {
|
|
return nil, fmt.Errorf("vector service configuration is missing")
|
|
}
|
|
|
|
// 根据 endpoint 判断服务类型
|
|
serviceType := determineServiceType(config.Endpoint)
|
|
|
|
switch serviceType {
|
|
case VectorServiceTypeChromaDB:
|
|
log.Context(nil).Info("Creating ChromaDB vector service")
|
|
return NewChromaDBService(config)
|
|
case VectorServiceTypeMemory:
|
|
log.Context(nil).Info("Creating in-memory vector service")
|
|
return NewVectorService(config)
|
|
default:
|
|
return nil, fmt.Errorf("unsupported vector service type: %s", serviceType)
|
|
}
|
|
}
|
|
|
|
// determineServiceType 根据 endpoint 确定服务类型
|
|
func determineServiceType(endpoint string) VectorServiceType {
|
|
if endpoint == "" || strings.Contains(endpoint, "memory") {
|
|
return VectorServiceTypeMemory
|
|
}
|
|
|
|
// 如果包含 ChromaDB 相关的关键词，使用 ChromaDB
|
|
if strings.Contains(strings.ToLower(endpoint), "chroma") ||
|
|
strings.Contains(endpoint, ":8000") {
|
|
return VectorServiceTypeChromaDB
|
|
}
|
|
|
|
// 默认使用内存版本
|
|
return VectorServiceTypeMemory
|
|
} |