WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,350 @@
+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)
+}