mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-09-19 02:18:25 +08:00
fix: make Milvus vector metric type configurable via MILVUS_METRIC_TYPE
This commit is contained in:
@@ -267,6 +267,10 @@ APK_MIRROR_ARG=mirrors.tencent.com
|
||||
# Milvus集合名称,用于存储向量数据
|
||||
# MILVUS_COLLECTION=weknora_embeddings
|
||||
|
||||
# Milvus向量搜索度量类型,支持 IP(默认)、COSINE、L2
|
||||
# 注意:修改度量类型后需要重建collection才能生效
|
||||
# MILVUS_METRIC_TYPE=IP
|
||||
|
||||
# Milvus 用户名(可选)
|
||||
# MILVUS_USERNAME=your_milvus_username
|
||||
|
||||
|
||||
@@ -84,6 +84,7 @@ services:
|
||||
- QDRANT_USE_TLS=${QDRANT_USE_TLS:-false}
|
||||
- MILVUS_ADDRESS=milvus:19530
|
||||
- MILVUS_COLLECTION=${MILVUS_COLLECTION:-weknora_embeddings}
|
||||
- MILVUS_METRIC_TYPE=${MILVUS_METRIC_TYPE:-IP}
|
||||
- DOCREADER_ADDR=${DOCREADER_ADDR:-docreader:50051}
|
||||
- DOCREADER_TRANSPORT=${DOCREADER_TRANSPORT:-grpc}
|
||||
- WEAVIATE_HOST=${WEAVIATE_HOST:-weaviate:8080}
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
|
||||
const (
|
||||
envMilvusCollection = "MILVUS_COLLECTION"
|
||||
envMilvusMetricType = "MILVUS_METRIC_TYPE"
|
||||
defaultCollectionName = "weknora_embeddings"
|
||||
fieldContent = "content"
|
||||
fieldSourceID = "source_id"
|
||||
@@ -51,10 +52,26 @@ func NewMilvusRetrieveEngineRepository(client *client.Client) interfaces.Retriev
|
||||
collectionBaseName = defaultCollectionName
|
||||
}
|
||||
|
||||
metricType := entity.IP
|
||||
if mt := os.Getenv(envMilvusMetricType); mt != "" {
|
||||
switch strings.ToUpper(mt) {
|
||||
case "COSINE":
|
||||
metricType = entity.COSINE
|
||||
case "L2":
|
||||
metricType = entity.L2
|
||||
case "IP":
|
||||
metricType = entity.IP
|
||||
default:
|
||||
log.Warnf("[Milvus] Unknown MILVUS_METRIC_TYPE '%s', using default IP", mt)
|
||||
}
|
||||
}
|
||||
log.Infof("[Milvus] Using metric type: %s", metricType)
|
||||
|
||||
res := &milvusRepository{
|
||||
filter: filter{},
|
||||
client: client,
|
||||
collectionBaseName: collectionBaseName,
|
||||
metricType: metricType,
|
||||
}
|
||||
|
||||
log.Info("[Milvus] Successfully initialized repository")
|
||||
@@ -150,7 +167,7 @@ func (m *milvusRepository) ensureCollection(ctx context.Context, dimension int)
|
||||
|
||||
indexOpts := make([]client.CreateIndexOption, 0)
|
||||
// hnsw index for embedding field
|
||||
indexOpts = append(indexOpts, client.NewCreateIndexOption(collectionName, fieldEmbedding, index.NewHNSWIndex(entity.IP, 16, 128)))
|
||||
indexOpts = append(indexOpts, client.NewCreateIndexOption(collectionName, fieldEmbedding, index.NewHNSWIndex(m.metricType, 16, 128)))
|
||||
indexOpts = append(indexOpts, client.NewCreateIndexOption(collectionName, fieldContentSparse, index.NewAutoIndex(entity.BM25)))
|
||||
// Create payload indexes for filtering
|
||||
indexFields := []string{fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldSourceID, fieldIsEnabled}
|
||||
|
||||
@@ -3,6 +3,7 @@ package milvus
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/milvus-io/milvus/client/v2/entity"
|
||||
client "github.com/milvus-io/milvus/client/v2/milvusclient"
|
||||
)
|
||||
|
||||
@@ -10,6 +11,7 @@ type milvusRepository struct {
|
||||
filter
|
||||
client *client.Client
|
||||
collectionBaseName string
|
||||
metricType entity.MetricType
|
||||
// Cache for initialized collections (dimension -> true)
|
||||
initializedCollections sync.Map
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user