xinmi/internal/application/service/wiki_ingest_cite_test.go

package service

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

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

// TestMergeCitationsIntoItems_PopulatesSourceChunksOnCandidates verifies that
// citations returned by the chunk-classification pass are attached back onto
// the matching candidate items while non-cited candidates are left untouched.
func TestMergeCitationsIntoItems_PopulatesSourceChunksOnCandidates(t *testing.T) {
	entities := []extractedItem{
		{Name: "Acme", Slug: "entity/acme"},
		{Name: "Beta", Slug: "entity/beta"},
	}
	concepts := []extractedItem{
		{Name: "RAG", Slug: "concept/rag"},
	}
	citations := map[string][]string{
		"entity/acme": {"chunk-1", "chunk-3"},
		"concept/rag": {"chunk-2"},
	}

	gotE, gotC, uncited := mergeCitationsIntoItems(entities, concepts, citations, nil)

	if len(gotE) != 2 || len(gotC) != 1 {
		t.Fatalf("expected 2 entities + 1 concept, got %d + %d", len(gotE), len(gotC))
	}
	acme := findBySlug(gotE, "entity/acme")
	if acme == nil {
		t.Fatalf("entity/acme missing from result")
	}
	if !equalStrings(acme.SourceChunks, []string{"chunk-1", "chunk-3"}) {
		t.Errorf("entity/acme source_chunks = %v, want [chunk-1 chunk-3]", acme.SourceChunks)
	}
	beta := findBySlug(gotE, "entity/beta")
	if beta == nil {
		t.Fatalf("entity/beta missing from result")
	}
	if len(beta.SourceChunks) != 0 {
		t.Errorf("entity/beta should have no citations, got %v", beta.SourceChunks)
	}
	rag := findBySlug(gotC, "concept/rag")
	if rag == nil {
		t.Fatalf("concept/rag missing")
	}
	if !equalStrings(rag.SourceChunks, []string{"chunk-2"}) {
		t.Errorf("concept/rag source_chunks = %v, want [chunk-2]", rag.SourceChunks)
	}
	if uncited != 1 {
		t.Errorf("uncited = %d, want 1", uncited)
	}
}

// TestMergeCitationsIntoItems_AddsNewSlugsAndUnionsChunksAcrossBatches checks
// that genuinely new slugs (ones Pass 0 missed) are appended to the right
// type slice, and that a slug surfacing in two batches ends up with the union
// of its source chunks.
func TestMergeCitationsIntoItems_AddsNewSlugsAndUnionsChunksAcrossBatches(t *testing.T) {
	entities := []extractedItem{
		{Name: "Known", Slug: "entity/known"},
	}
	concepts := []extractedItem{}

	newSlugs := []newSlugFromCitation{
		{
			Type:         "entity",
			Name:         "Fresh Entity",
			Slug:         "entity/fresh",
			Description:  "desc",
			Details:      "details",
			SourceChunks: []string{"c001", "c002"},
		},
		{
			// Same slug as above, appears in another batch — must union.
			Type:         "entity",
			Name:         "Fresh Entity",
			Slug:         "entity/fresh",
			SourceChunks: []string{"c002", "c003"},
		},
		{
			Type:         "concept",
			Name:         "New Concept",
			Slug:         "concept/new-concept",
			SourceChunks: []string{"c010"},
		},
		{
			// Duplicate of an existing candidate — should NOT produce a
			// duplicate entry (Known already exists in `entities`).
			Type:         "entity",
			Name:         "Known",
			Slug:         "entity/known",
			SourceChunks: []string{"c020"},
		},
	}

	gotE, gotC, _ := mergeCitationsIntoItems(entities, concepts, nil, newSlugs)

	if len(gotE) != 2 {
		t.Fatalf("expected 2 entities, got %d (%+v)", len(gotE), gotE)
	}
	if len(gotC) != 1 {
		t.Fatalf("expected 1 concept, got %d (%+v)", len(gotC), gotC)
	}
	fresh := findBySlug(gotE, "entity/fresh")
	if fresh == nil {
		t.Fatalf("entity/fresh missing")
	}
	sort.Strings(fresh.SourceChunks)
	if !equalStrings(fresh.SourceChunks, []string{"c001", "c002", "c003"}) {
		t.Errorf("entity/fresh source_chunks = %v, want union [c001 c002 c003]", fresh.SourceChunks)
	}
	newC := findBySlug(gotC, "concept/new-concept")
	if newC == nil || !equalStrings(newC.SourceChunks, []string{"c010"}) {
		t.Errorf("concept/new-concept missing or wrong chunks: %+v", newC)
	}
	known := findBySlug(gotE, "entity/known")
	require.NotNil(t, known)
	require.Equal(t, []string{"c020"}, known.SourceChunks)
}

func TestMergeCitationsIntoItems_PreservesValidatedEvidenceRepresentations(t *testing.T) {
	batch := splitChunksIntoCitationBatches([]*types.Chunk{
		{ID: "first", Content: "Alpha first", ChunkIndex: 0},
		{ID: "second", Content: "Alpha second", ChunkIndex: 1},
	})[0]
	for _, kind := range []string{"entity", "concept"} {
		for _, existing := range []bool{false, true} {
			for _, test := range []struct {
				name                string
				citations, declared []string
				want                []string
			}{
				{name: "source-chunks-only", declared: []string{"c000"}, want: []string{"first"}},
				{name: "citations-only", citations: []string{"c001"}, want: []string{"second"}},
				{name: "both-with-duplicates", citations: []string{"c000", "c000"}, declared: []string{"c000", "c001", "c001"}, want: []string{"first", "second"}},
			} {
				t.Run(fmt.Sprintf("%s/existing-%t/%s", kind, existing, test.name), func(t *testing.T) {
					slug := kind + "/alpha"
					result := citationBatchResult{
						Citations: map[string][]string{},
						NewSlugs:  []newSlugFromCitation{{Type: kind, Name: "Alpha", Slug: slug, SourceChunks: test.declared}},
					}
					if test.citations != nil {
						result.Citations[slug] = test.citations
					}
					allowed := map[string]bool{}
					var entities, concepts []extractedItem
					if existing {
						allowed[slug] = true
						candidate := extractedItem{Name: "Alpha", Slug: slug}
						if kind == "entity" {
							entities = append(entities, candidate)
						} else {
							concepts = append(concepts, candidate)
						}
					}
					raw, err := json.Marshal(result)
					require.NoError(t, err)
					parsed, err := parseWikiCitationBatch(string(raw), batch, allowed)
					require.NoError(t, err)
					require.NoError(t, validateWikiCitationCheckpoint(&parsed, batch, allowed))
					gotE, gotC, uncited := mergeCitationsIntoItems(entities, concepts, parsed.Citations, parsed.NewSlugs)
					all := append(gotE, gotC...)
					require.Len(t, all, 1)
					require.Equal(t, slug, all[0].Slug)
					require.Equal(t, test.want, all[0].SourceChunks)
					require.Zero(t, uncited)
				})
			}
		}
	}
}

func TestMergeCitationsIntoItems_MergesDeclarationIdentityMetadata(t *testing.T) {
	entities := []extractedItem{{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{"original"}, Description: "short", Details: "short"}}
	newSlugs := []newSlugFromCitation{
		{Type: "entity", Name: "A lpha", Slug: "entity/alpha", Aliases: []string{"first alias"}, Description: "longer description", SourceChunks: []string{"first"}},
		{Type: "entity", Name: "Alpha", Slug: "entity/alpha", Aliases: []string{"second alias", "first alias"}, Details: "longer details", SourceChunks: []string{"second", "first"}},
		{Type: "concept", Name: "B eta", Slug: "concept/beta", Aliases: []string{"first alias"}, Description: "short", Details: "longer details", SourceChunks: []string{"first"}},
		{Type: "concept", Name: "Beta", Slug: "concept/beta", Aliases: []string{"second alias"}, Description: "longer description", Details: "short", SourceChunks: []string{"second"}},
	}
	gotE, gotC, uncited := mergeCitationsIntoItems(entities, nil, nil, newSlugs)
	require.Zero(t, uncited)
	require.Len(t, gotE, 1)
	require.Len(t, gotC, 1)
	require.Equal(t, "Alpha", gotE[0].Name)
	require.ElementsMatch(t, []string{"original", "A lpha", "first alias", "second alias"}, gotE[0].Aliases)
	require.Equal(t, "Beta", gotC[0].Name)
	require.ElementsMatch(t, []string{"B eta", "first alias", "second alias"}, gotC[0].Aliases)
	for _, item := range append(gotE, gotC...) {
		require.Equal(t, "longer description", item.Description)
		require.Equal(t, "longer details", item.Details)
		require.Equal(t, []string{"first", "second"}, item.SourceChunks)
	}
}

func TestMergeCitationsIntoItems_DeduplicatesExistingCandidateCitations(t *testing.T) {
	got, _, uncited := mergeCitationsIntoItems(
		[]extractedItem{{Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"unvalidated-extraction-source"}}}, nil,
		map[string][]string{"entity/alpha": {"first", "first", "second"}}, nil,
	)
	require.Zero(t, uncited)
	require.Equal(t, []string{"first", "second"}, got[0].SourceChunks)
}

func TestWikiCitationEvidenceUnionStillRejectsInvalidBatchEvidence(t *testing.T) {
	batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "first", Content: "first"}})[0]
	allowed := map[string]bool{"entity/alpha": true}
	for _, test := range []struct{ name, output string }{
		{"undeclared-slug", `{"citations":{"concept/new":["c000"]}}`},
		{"invalid-citation-source", `{"citations":{"entity/alpha":["c999"]}}`},
		{"invalid-declaration-source", `{"citations":{},"new_slugs":[{"type":"entity","name":"Alpha","slug":"entity/alpha","source_chunks":["c999"]}]}`},
		{"invalid-new-slug-type", `{"citations":{},"new_slugs":[{"type":"entity","name":"New","slug":"concept/new","source_chunks":["c000"]}]}`},
	} {
		t.Run(test.name, func(t *testing.T) {
			_, err := parseWikiCitationBatch(test.output, batch, allowed)
			require.Error(t, err)
		})
	}
	for _, result := range []citationBatchResult{
		{Citations: map[string][]string{"entity/alpha": {"foreign-source"}}},
		{Citations: map[string][]string{}, NewSlugs: []newSlugFromCitation{{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"foreign-source"}}}},
	} {
		require.Error(t, validateWikiCitationCheckpoint(&result, batch, allowed))
	}
}

func TestWikiCitationEvidenceUnionSurvivesCheckpointReload(t *testing.T) {
	path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
	db := openWikiCheckpointDB(t, path)
	svc := &wikiIngestService{checkpointDB: db}
	scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
	chunks := []*types.Chunk{{ID: "first", Content: "first Alpha observation", ChunkIndex: 0}, {ID: "second", Content: "second Alpha observation", ChunkIndex: 1}}
	batch := splitChunksIntoCitationBatches(chunks)[0]
	entities := []extractedItem{{Name: "Alpha", Slug: "entity/alpha"}}
	candidates := renderCandidateSlugsXML(entities, nil)
	key := wikiCheckpointHash([]any{candidates, "English", wikiCheckpointChunks(batch.chunks)})
	result := citationBatchResult{
		Citations: map[string][]string{"entity/alpha": {"first"}, "concept/beta": {"second"}},
		NewSlugs: []newSlugFromCitation{
			{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"second"}},
			{Type: "concept", Name: "Beta", Slug: "concept/beta"},
		},
	}
	require.NoError(t, validateWikiCitationCheckpoint(&result, batch, wikiCandidateSlugSet(candidates)))
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation", key, result))
	sqlDB, err := db.DB()
	require.NoError(t, err)
	require.NoError(t, sqlDB.Close())
	svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		t.Fatal("validated citation checkpoint must not call the model again")
		return nil, nil
	}}
	citations, declarations, completed, err := svc.classifyChunkCitations(context.Background(), model, candidates, chunks, "English", &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Equal(t, 1, completed)
	gotE, gotC, uncited := mergeCitationsIntoItems(entities, nil, citations, declarations)
	require.Zero(t, uncited)
	require.Len(t, gotE, 1)
	require.Len(t, gotC, 1)
	require.Equal(t, []string{"first", "second"}, gotE[0].SourceChunks)
	require.Equal(t, []string{"second"}, gotC[0].SourceChunks)
}

// TestSplitChunksIntoCitationBatches_RespectsBudgetAndOrder verifies that the
// batcher never puts too many runes in one batch, preserves document order,
// and that an oversized chunk gets its own batch.
func TestSplitChunksIntoCitationBatches_RespectsBudgetAndOrder(t *testing.T) {
	// Each small chunk is 5k runes → 3 of them should fit in one batch
	// (15k > 12k limit would spill to a second batch).
	mk := func(idx int, runes int, id string) *types.Chunk {
		return &types.Chunk{
			ID:         id,
			ChunkIndex: idx,
			Content:    repeatRune('a', runes),
			ChunkType:  types.ChunkTypeText,
		}
	}
	chunks := []*types.Chunk{
		mk(0, 5000, "id-0"),
		mk(1, 5000, "id-1"),
		mk(2, 5000, "id-2"), // this should start a new batch (15k > 12k)
		// An oversized chunk gets a dedicated batch.
		mk(3, 20000, "id-big"),
		mk(4, 1000, "id-small"),
	}
	batches := splitChunksIntoCitationBatches(chunks)
	if len(batches) < 3 {
		t.Fatalf("expected at least 3 batches, got %d", len(batches))
	}
	// All input IDs should show up in some batch, exactly once, in order.
	seen := []string{}
	for _, b := range batches {
		for _, c := range b.chunks {
			seen = append(seen, c.ID)
		}
	}
	wantOrder := []string{"id-0", "id-1", "id-2", "id-big", "id-small"}
	if !equalStrings(seen, wantOrder) {
		t.Errorf("batch order = %v, want %v", seen, wantOrder)
	}

	// Verify the typed handle table is populated per batch.
	for bi, b := range batches {
		if b.handles.Len() != len(b.chunks) {
			t.Errorf("batch %d handle count %d != chunk count %d", bi, b.handles.Len(), len(b.chunks))
		}
	}
}

// --- helpers ---

func findBySlug(items []extractedItem, slug string) *extractedItem {
	for i := range items {
		if items[i].Slug == slug {
			return &items[i]
		}
	}
	return nil
}

func equalStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func repeatRune(r rune, n int) string {
	out := make([]rune, n)
	for i := range out {
		out[i] = r
	}
	return string(out)
}