diff --git a/.env.example b/.env.example index dd8bcf625..c3642166d 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/docker-compose.yml b/docker-compose.yml index 28386f54f..adc1e5fbf 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -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} diff --git a/internal/application/repository/retriever/milvus/repository.go b/internal/application/repository/retriever/milvus/repository.go index 29fb5405f..6e9344eb7 100644 --- a/internal/application/repository/retriever/milvus/repository.go +++ b/internal/application/repository/retriever/milvus/repository.go @@ -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} diff --git a/internal/application/repository/retriever/milvus/structs.go b/internal/application/repository/retriever/milvus/structs.go index 37d404471..c3b44e0f4 100644 --- a/internal/application/repository/retriever/milvus/structs.go +++ b/internal/application/repository/retriever/milvus/structs.go @@ -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 }