From 02259e18eceaedf5996a1a7626c057007b9cedeb Mon Sep 17 00:00:00 2001 From: Gary Date: Mon, 5 Aug 2024 17:32:02 +0800 Subject: [PATCH] feat(catelog): make topK default 5 (#62) Because user may not know how to use topk parameter. This commit mke topK default 5 when it is not provided. --- pkg/milvus/milvus.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pkg/milvus/milvus.go b/pkg/milvus/milvus.go index 5d4e8c9..9ad5f45 100644 --- a/pkg/milvus/milvus.go +++ b/pkg/milvus/milvus.go @@ -23,6 +23,8 @@ type MilvusClientI interface { DropKnowledgeBaseCollection(ctx context.Context, kbUID string) error ListEmbeddings(ctx context.Context, collectionName string) ([]Embedding, error) SearchSimilarEmbeddings(ctx context.Context, collectionName string, vectors [][]float32, topK int) ([][]SimilarEmbedding, error) + // SearchSimilarEmbeddingsInKB search similar embeddings in knowledge base. + // The topK has default value 5 SearchSimilarEmbeddingsInKB(ctx context.Context, kbUID string, vectors [][]float32, topK int) ([][]SimilarEmbedding, error) DeleteEmbedding(ctx context.Context, collectionName string, embeddingUID []string) error DeleteEmbeddingsInKb(ctx context.Context, kbUID string, embeddingUID []string) error @@ -319,7 +321,12 @@ type SimilarEmbedding struct { } // SearchSimilarEmbeddings searches for embeddings similar to the input vector +// topk has default value 5, when topk <= 0, it will be set to 5. func (m *MilvusClient) SearchSimilarEmbeddings(ctx context.Context, collectionName string, vectors [][]float32, topK int) ([][]SimilarEmbedding, error) { + // set default topK + if topK <= 0 { + topK = 5 + } log, err := logger.GetZapLogger(ctx) if err != nil { log.Error("failed to get logger", zap.Error(err))