xinmi/internal/application/service/wiki_routing_metadata_test.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/models/embedding"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
)

func routingMetadataFixture() combinedExtraction {
	return combinedExtraction{
		Entities: []extractedItem{{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{"A"},
			Description: "UNSUPPORTED_DESCRIPTION_ALPHA", Details: "UNSUPPORTED_DETAILS_ALPHA", SourceChunks: []string{"source-a"}}},
		Concepts: []extractedItem{{Name: "Beta", Slug: "concept/beta", Aliases: []string{"B"},
			Description: "UNSUPPORTED_DESCRIPTION_BETA", Details: "UNSUPPORTED_DETAILS_BETA", SourceChunks: []string{"source-b"}}},
	}
}

func requireNoGeneratedRoutingProse(t *testing.T, text string) {
	t.Helper()
	require.NotContains(t, text, "UNSUPPORTED_DESCRIPTION_")
	require.NotContains(t, text, "UNSUPPORTED_DETAILS_")
}

func TestWikiRoutingMetadataCitationPreservesIdentitiesAndOriginalEvidence(t *testing.T) {
	ctx := context.Background()
	fixture := routingMetadataFixture()
	original, err := json.Marshal(fixture)
	require.NoError(t, err)
	catalog := renderCandidateSlugsXML(fixture.Entities, fixture.Concepts)
	for _, item := range append(append([]extractedItem(nil), fixture.Entities...), fixture.Concepts...) {
		require.Contains(t, catalog, item.Slug)
		require.Contains(t, catalog, fmt.Sprintf("name: %q", item.Name))
		require.Contains(t, catalog, fmt.Sprintf("aliases=%q", strings.Join(item.Aliases, ", ")))
	}
	requireNoGeneratedRoutingProse(t, catalog)
	chunks := []*types.Chunk{
		{ID: "source-a", Content: "Alpha: first source reports 65/35; condition X is mandatory.\n", ChunkIndex: 0},
		{ID: "source-b", Content: "Beta: another source reports 65/30; no correction of OCR value ?.\n", ChunkIndex: 1},
	}
	batches := splitChunksIntoCitationBatches(chunks)
	require.Len(t, batches, 1)
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		var prompt strings.Builder
		for _, message := range messages {
			prompt.WriteString(message.Content)
		}
		requireNoGeneratedRoutingProse(t, prompt.String())
		for _, chunk := range chunks {
			require.Contains(t, prompt.String(), chunk.Content, "source text is not replaced by generated metadata")
		}
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"],"concept/beta":["c001"]}}`}, nil
	}}
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "routing.sqlite"))}
	result, err := svc.classifyCitationBatch(ctx, model, catalog, batches[0], "English", false)
	require.NoError(t, err)
	require.Equal(t, 1, calls)
	require.Equal(t, []string{"source-a"}, result.Citations["entity/alpha"])
	require.Equal(t, []string{"source-b"}, result.Citations["concept/beta"])
	unchanged, err := json.Marshal(fixture)
	require.NoError(t, err)
	require.JSONEq(t, string(original), string(unchanged))
	scope := &wikiCheckpointScope{7, "kb", "doc", "routing-audit"}
	require.NoError(t, svc.saveWikiCheckpoint(ctx, scope, "candidates", "complete", fixture))
	var restored combinedExtraction
	hit, err := svc.loadWikiCheckpoint(ctx, scope, "candidates", "complete", &restored)
	require.NoError(t, err)
	require.True(t, hit)
	require.Equal(t, fixture, restored, "audit descriptions, aliases and all source references remain persisted")
}

func TestWikiRoutingMetadataTaxonomyReceivesOnlyIdentityInputs(t *testing.T) {
	fixture := routingMetadataFixture()
	updates := map[string][]SlugUpdate{
		"entity/alpha": {{Type: types.WikiPageTypeEntity, Item: fixture.Entities[0]}},
		"concept/beta": {{Type: types.WikiPageTypeConcept, Item: fixture.Concepts[0]}},
	}
	original, err := json.Marshal(updates)
	require.NoError(t, err)
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		for _, message := range messages {
			requireNoGeneratedRoutingProse(t, message.Content)
		}
		require.Contains(t, messages[0].Content, "entity/alpha | title: Alpha | type: entity")
		require.Contains(t, messages[0].Content, "concept/beta | title: Beta | type: concept")
		return &types.ChatResponse{Content: `{"assignments":[{"slug":"entity/alpha","path":["People"]},{"slug":"concept/beta","path":["Topics"]}]}`}, nil
	}}
	svc := &wikiIngestService{}
	got, err := svc.planBatchTaxonomyChecked(context.Background(), model, &types.KnowledgeBase{ID: "kb"}, updates, "English")
	require.NoError(t, err)
	require.Equal(t, 1, calls)
	require.Equal(t, []string{"People"}, got["entity/alpha"])
	require.Equal(t, []string{"Topics"}, got["concept/beta"])
	unchanged, err := json.Marshal(updates)
	require.NoError(t, err)
	require.JSONEq(t, string(original), string(unchanged))
}

type routingMetadataEmbedder struct {
	embedding.Embedder
	inputs [][]string
}

func (e *routingMetadataEmbedder) BatchEmbed(_ context.Context, inputs []string) ([][]float32, error) {
	e.inputs = append(e.inputs, append([]string(nil), inputs...))
	result := make([][]float32, len(inputs))
	for i := range result {
		result[i] = []float32{1, 0}
	}
	return result, nil
}

func TestWikiRoutingMetadataTaxonomyEmbeddingExcludesGeneratedProse(t *testing.T) {
	fixture := routingMetadataFixture()
	items := collectTaxonomyItems(map[string][]SlugUpdate{
		"entity/alpha": {{Type: types.WikiPageTypeEntity, Item: fixture.Entities[0]}},
		"concept/beta": {{Type: types.WikiPageTypeConcept, Item: fixture.Concepts[0]}},
	})
	pool := make([][]string, wikiTaxonomyFeedAllMaxFolders+1)
	for i := range pool {
		pool[i] = []string{"Existing", fmt.Sprintf("Folder %d", i)}
	}
	embedder := &routingMetadataEmbedder{}
	svc := &wikiIngestService{modelService: parentChildModelService{embedder: embedder}}
	_, err := svc.selectRelevantFoldersChecked(context.Background(), &types.KnowledgeBase{EmbeddingModelID: "embedding"}, items, pool)
	require.NoError(t, err)
	require.Len(t, embedder.inputs, 2)
	require.Equal(t, []string{"Beta", "Alpha"}, embedder.inputs[1])
	for _, batch := range embedder.inputs {
		requireNoGeneratedRoutingProse(t, strings.Join(batch, "\n"))
	}
}