xinmi/internal/application/repository/retriever/milvus/repository_test.go

package milvus

import (
	"context"
	"errors"
	"testing"

	"github.com/milvus-io/milvus/client/v2/column"
	"github.com/milvus-io/milvus/client/v2/entity"
	client "github.com/milvus-io/milvus/client/v2/milvusclient"
	"github.com/stretchr/testify/require"

	"github.com/Tencent/WeKnora/internal/types"
)

func TestContentFieldUsesChineseAnalyzer(t *testing.T) {
	field := newContentField().ProtoMessage()
	require.Equal(t, fieldContent, field.GetName())
	require.Equal(t, int32(entity.FieldTypeVarChar), int32(field.GetDataType()))
	params := make(map[string]string)
	for _, param := range field.GetTypeParams() {
		params[param.GetKey()] = param.GetValue()
	}
	require.Equal(t, "65535", params["max_length"])
	require.Equal(t, "true", params["enable_analyzer"])
	require.Equal(t, "true", params["enable_match"])
	require.JSONEq(t, `{"type":"chinese"}`, params["analyzer_params"])
}

func TestUpdateChunkEnabledStatusInCollectionSkipsEmptyChunkIDs(t *testing.T) {
	repo := &milvusRepository{}

	require.NoError(t, repo.updateChunkEnabledStatusInCollection(
		context.Background(),
		"weknora_embeddings_1024",
		nil,
		false,
	))
	require.NoError(t, repo.updateChunkEnabledStatusInCollection(
		context.Background(),
		"weknora_embeddings_1024",
		[]string{},
		true,
	))
}

func TestUpdateChunkEnabledStatusInCollectionsPropagatesFailure(t *testing.T) {
	wantErr := errors.New("upsert failed")
	err := updateChunkEnabledStatusInCollections(
		context.Background(),
		[]string{"other_collection", "weknora_embeddings_1024"},
		"weknora_embeddings",
		nil,
		[]string{"chunk-1"},
		func(_ context.Context, collection string, _ []string, enabled bool) error {
			if collection == "weknora_embeddings_1024" && !enabled {
				return wantErr
			}
			return nil
		},
	)
	require.ErrorIs(t, err, wantErr)
}

func TestMatchesCollection(t *testing.T) {
	for _, name := range []string{"weknora_embeddings_1", "weknora_embeddings_1024"} {
		require.True(t, matchesCollection(name, "weknora_embeddings"), name)
	}
	for _, name := range []string{
		"weknora_embeddings", "weknora_embeddings_", "weknora_embeddings_0",
		"weknora_embeddings_-1", "weknora_embeddings_+1", "weknora_embeddings_01024",
		"weknora_embeddings_1024_backup", "weknora_embeddings_backup_1024",
		"weknora_embeddings_other_tenant_1024", "weknora_embeddings2_1024",
		"weknora_embeddings_9999999999999999999999999", "other_1024", "_1024",
	} {
		require.False(t, matchesCollection(name, "weknora_embeddings"), name)
	}
	require.False(t, matchesCollection("_1024", ""))
}

func TestUpdateChunkEnabledStatusDoesNotCrossNamespaces(t *testing.T) {
	var touched []string
	err := updateChunkEnabledStatusInCollections(
		context.Background(),
		[]string{"weknora_embeddings_1024_backup", "weknora_embeddings_other_1024", "weknora_embeddings_1024"},
		"weknora_embeddings", nil, []string{"chunk-1"},
		func(_ context.Context, collection string, _ []string, _ bool) error {
			touched = append(touched, collection)
			return nil
		},
	)
	require.NoError(t, err)
	require.Equal(t, []string{"weknora_embeddings_1024", "weknora_embeddings_1024"}, touched)
}

func keywordResult(chunkID string, score float32) []client.ResultSet {
	return []client.ResultSet{{
		ResultCount: 1,
		Scores:      []float32{score},
		Fields: client.DataSet{
			column.NewColumnVarChar(fieldID, []string{"index-" + chunkID}),
			column.NewColumnVarChar(fieldChunkID, []string{chunkID}),
			column.NewColumnVarChar(fieldContent, []string{"test content"}),
		},
	}}
}

func TestKeywordsRetrieveReturnsAllCollectionFailures(t *testing.T) {
	repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
	missingSparseField := errors.New("field content_sparse does not exist")
	conversionFailure := errors.New("cannot decode search result")
	var searched []string
	results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
		Query: "test query", TopK: 10,
	}, []string{"weknora_embeddings_384", "weknora_embeddings_1024", "weknora_embeddings_backup_1024"},
		func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
			request, err := option.Request()
			require.NoError(t, err)
			searched = append(searched, request.CollectionName)
			if request.CollectionName == "weknora_embeddings_384" {
				return nil, missingSparseField
			}
			return []client.ResultSet{{Err: conversionFailure}}, nil
		})
	require.Nil(t, results)
	require.ErrorIs(t, err, missingSparseField)
	require.ErrorIs(t, err, conversionFailure)
	require.ErrorContains(t, err, "all 2 matching Milvus collections")
	require.ErrorContains(t, err, "weknora_embeddings_384")
	require.ErrorContains(t, err, "weknora_embeddings_1024")
	require.Equal(t, []string{"weknora_embeddings_384", "weknora_embeddings_1024"}, searched)
}

func TestKeywordsRetrievePreservesBM25AndRanksAcrossCollections(t *testing.T) {
	repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
	results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
		Query: "test query", TopK: 1, KnowledgeBaseIDs: []string{"kb-allowed"},
		TagIDs: []string{"tag-allowed"}, ExcludeChunkIDs: []string{"chunk-excluded"},
	}, []string{"weknora_embeddings_384", "weknora_embeddings_1024", "weknora_embeddings_other_1536"},
		func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
			request, err := option.Request()
			require.NoError(t, err)
			require.Contains(t, request.Dsl, fieldKnowledgeBaseID)
			require.Contains(t, request.Dsl, fieldTagID)
			require.Contains(t, request.Dsl, fieldChunkID)
			require.Contains(t, request.Dsl, fieldIsEnabled)
			require.NotEmpty(t, request.ExprTemplateValues)
			for _, param := range request.SearchParams {
				if param.Key == "anns_field" {
					require.Equal(t, fieldContentSparse, param.Value)
				}
			}
			switch request.CollectionName {
			case "weknora_embeddings_384":
				return keywordResult("low-score", 2.5), nil
			case "weknora_embeddings_1024":
				return keywordResult("high-score", 17.25), nil
			default:
				t.Fatalf("searched unrelated collection %s", request.CollectionName)
				return nil, nil
			}
		})
	require.NoError(t, err)
	require.Len(t, results, 1)
	require.Equal(t, types.KeywordsRetrieverType, results[0].RetrieverType)
	require.Len(t, results[0].Results, 1)
	require.Equal(t, "high-score", results[0].Results[0].ChunkID)
	require.Equal(t, 17.25, results[0].Results[0].Score)
	require.Equal(t, types.MatchTypeKeywords, results[0].Results[0].MatchType)
}

func TestKeywordsRetrieveKeepsSuccessfulCollections(t *testing.T) {
	for _, empty := range []bool{false, true} {
		t.Run(map[bool]string{false: "results", true: "empty"}[empty], func(t *testing.T) {
			repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
			results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
				Query: "query", TopK: 5,
			}, []string{"weknora_embeddings_384", "weknora_embeddings_1024"},
				func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
					request, err := option.Request()
					require.NoError(t, err)
					if request.CollectionName == "weknora_embeddings_384" {
						return nil, errors.New("missing sparse field")
					}
					if empty {
						return nil, nil
					}
					return keywordResult("available", 8), nil
				})
			require.NoError(t, err)
			require.Len(t, results, 1)
			if empty {
				require.Empty(t, results[0].Results)
			} else {
				require.Len(t, results[0].Results, 1)
				require.Equal(t, "available", results[0].Results[0].ChunkID)
			}
		})
	}
}

func TestKeywordsRetrieveRejectsMissingScores(t *testing.T) {
	repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
	results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
		Query: "query", TopK: 5,
	}, []string{"weknora_embeddings_384"},
		func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
			result := keywordResult("missing-score", 0)
			result[0].Scores = nil
			return result, nil
		})
	require.Nil(t, results)
	require.ErrorContains(t, err, "expected 1 BM25 scores, received 0")
}

func TestKeywordsRetrieveNoMatchingCollections(t *testing.T) {
	repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
	results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
		Query: "query", TopK: 5,
	}, []string{"weknora_embeddings_backup_384"},
		func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
			t.Fatal("must not search unrelated collection")
			return nil, nil
		})
	require.NoError(t, err)
	require.Len(t, results, 1)
	require.Empty(t, results[0].Results)
}

func TestKeywordsRetrieveHonorsCancellation(t *testing.T) {
	repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	results, err := repo.keywordsRetrieveInCollections(ctx, types.RetrieveParams{
		Query: "query", TopK: 5,
	}, []string{"weknora_embeddings_384", "weknora_embeddings_1024"},
		func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
			cancel()
			return keywordResult("cancelled", 8), nil
		})
	require.Nil(t, results)
	require.ErrorIs(t, err, context.Canceled)
}