fix: make Milvus vector metric type configurable via MILVUS_METRIC_TYPE

This commit is contained in:
Dounx
2026-03-25 21:20:49 +08:00
committed by lyingbug
parent 5af8e4e9f8
commit 8df12aeee2
4 changed files with 25 additions and 1 deletions
+4
View File
@@ -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
+1
View File
@@ -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
}