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,240 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWikiCitationCandidateHandlesPreserveExactDurableIdentity(t *testing.T) {
+	entities := []extractedItem{{Name: "Long name", Slug: "entity/a-very-long-uuid-9f2704c9-f7e7-4848-ad9c-44af056c4cf1", Description: "Do not edit entity/inside-description"}}
+	concepts := []extractedItem{{Name: "Concept", Slug: "concept/a"}}
+	xml := renderCandidateSlugsXML(entities, concepts)
+	encoded, handles := encodeWikiCitationCandidates(xml)
+	require.Contains(t, encoded, "- id: s001, type: entity")
+	require.NotContains(t, encoded, "entity/inside-description", "generated descriptions are audit-only, not routing metadata")
+	require.NotContains(t, encoded, entities[0].Slug)
+	_, reversed := encodeWikiCitationCandidates(renderCandidateSlugsXML(nil, []extractedItem{concepts[0], entities[0]}))
+	want, _ := handles.Handle(entities[0].Slug)
+	got, _ := reversed.Handle(entities[0].Slug)
+	require.Equal(t, want, got, "mapping must not depend on candidate order")
+	batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "source", Content: "facts"}})[0]
+	parsed, err := parseWikiCitationBatch(`{"citations":{"s001":["c000"],"s000":["c000"]}}`, batch, wikiCandidateSlugSet(xml), handles)
+	require.NoError(t, err)
+	require.Equal(t, []string{"source"}, parsed.Citations[entities[0].Slug])
+	require.Equal(t, []string{"source"}, parsed.Citations[concepts[0].Slug])
+	for _, raw := range []string{
+		`{"citations":{"s999":["c000"]}}`,
+		`{"citations":{"c000":["c000"]}}`,
+		`{"citations":{"s001":["s000"]}}`,
+		`{"citations":{},"new_slugs":[{"type":"entity","slug":"s001","name":"wrong","source_chunks":["c000"]}]}`,
+	} {
+		_, err := parseWikiCitationBatch(raw, batch, wikiCandidateSlugSet(xml), handles)
+		require.Error(t, err, raw)
+	}
+}
+
+func TestWikiCitationUndeclaredSlugRepairPreservesEveryEdgeAndCheckpoint(t *testing.T) {
+	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
+	scope := &wikiCheckpointScope{7, "kb", "doc", "repair"}
+	chunks := []*types.Chunk{{ID: "first", Content: "Alpha source; New fact one."}, {ID: "second", ChunkIndex: 1, Content: "New fact two."}}
+	calls := 0
+	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
+		calls++
+		if calls == 1 {
+			require.Contains(t, messages[0].Content, "- id: s000")
+			return &types.ChatResponse{Content: `{"citations":{"s000":["c000"],"concept/new":["c000","c001"]}}`}, nil
+		}
+		require.Equal(t, 2, calls, "repair must not become repeated reclassification")
+		require.Contains(t, messages[0].Content, `"concept/new":["c000","c001"]`)
+		return &types.ChatResponse{Content: `{"cannot_repair":false,"new_slugs":[{"type":"concept","slug":"concept/new","name":"New","description":"New source facts","details":"New fact one; New fact two","source_chunks":["c001","c000"],"evidence_quotes":[{"chunk":"c000","quote":"New fact one."},{"chunk":"c001","quote":"New fact two."}]}]}`}, nil
+	}}
+	citations, newSlugs, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
+	require.NoError(t, err)
+	require.Equal(t, 1, completed)
+	require.Equal(t, 2, calls)
+	require.Equal(t, []string{"first", "second"}, citations["concept/new"])
+	require.Equal(t, []string{"first"}, citations["entity/alpha"])
+	require.Len(t, newSlugs, 1)
+	require.ElementsMatch(t, []string{"first", "second"}, newSlugs[0].SourceChunks)
+	_, _, _, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
+	require.NoError(t, err)
+	require.Equal(t, 2, calls, "validated repaired checkpoint must resume without another model call")
+}
+
+func TestWikiCitationDeclarationRepairCannotRewriteOrDropEvidence(t *testing.T) {
+	batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "source", Content: "Fact one."}, {ID: "other", ChunkIndex: 1, Content: "Other fact."}})[0]
+	original := citationBatchResult{Citations: map[string][]string{"entity/known": {"other"}, "concept/new": {"source"}}}
+	missing := map[string][]string{"concept/new": {"c000"}}
+	valid := `{"cannot_repair":false,"new_slugs":[{"type":"concept","slug":"concept/new","name":"New","description":"Fact","details":"Fact one.","source_chunks":["c000"],"evidence_quotes":[{"chunk":"c000","quote":"Fact one."}]}]}`
+	for name, raw := range map[string]string{
+		"missing declaration": `{"new_slugs":[]}`,
+		"declined":            strings.Replace(valid, `"cannot_repair":false`, `"cannot_repair":true`, 1),
+		"renamed":             strings.Replace(valid, "concept/new", "concept/other", 1),
+		"wrong type":          strings.Replace(valid, `"type":"concept"`, `"type":"entity"`, 1),
+		"drop source":         strings.Replace(valid, `"source_chunks":["c000"]`, `"source_chunks":[]`, 1),
+		"changed source":      strings.Replace(valid, `"source_chunks":["c000"]`, `"source_chunks":["c001"]`, 1),
+		"unknown source":      strings.Replace(valid, `"source_chunks":["c000"]`, `"source_chunks":["c999"]`, 1),
+		"fake quote":          strings.Replace(valid, `"quote":"Fact one."`, `"quote":"Made up"`, 1),
+		"missing quote":       strings.Replace(valid, `"evidence_quotes":[{"chunk":"c000","quote":"Fact one."}]`, `"evidence_quotes":[]`, 1),
+		"rewrite citations":   strings.Replace(valid, `"cannot_repair":false`, `"cannot_repair":false,"citations":{}`, 1),
+		"trailing object":     valid + `{}`,
+	} {
+		t.Run(name, func(t *testing.T) {
+			got, err := applyWikiCitationDeclarations(raw, original, missing, batch, map[string]bool{"entity/known": true})
+			require.Error(t, err)
+			require.Nil(t, got.Citations)
+			require.Equal(t, []string{"source"}, original.Citations["concept/new"])
+			require.Empty(t, original.NewSlugs)
+		})
+	}
+}
+
+func TestWikiCitationUnknownSourceAndHandleNeverTriggerDeclarationRepair(t *testing.T) {
+	for _, raw := range []string{
+		`{"citations":{"concept/new":["c999"]}}`,
+		`{"citations":{"s999":["c000"]}}`,
+		`{"citations":{"concept/new":[]}}`,
+	} {
+		calls := 0
+		model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+			calls++
+			return &types.ChatResponse{Content: raw}, nil
+		}}
+		_, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "one", Content: "Source one"}, {ID: "two", Content: "Source two"}}, "English", &WikiBatchContext{})
+		require.Error(t, err)
+		require.Zero(t, completed)
+		if strings.Contains(raw, "c999") {
+			require.Equal(t, 2, calls, "unknown source gets only one independent full regeneration")
+		} else {
+			require.Equal(t, 1, calls)
+		}
+	}
+}
+
+func TestWikiCitationUnknownSourceRegenerationPreservesWholeInputAndValidEdges(t *testing.T) {
+	chunks := []*types.Chunk{{ID: "first", Content: "Alpha fact one."}, {ID: "second", ChunkIndex: 1, Content: "Alpha fact two."}}
+	calls := 0
+	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
+		calls++
+		if calls == 1 {
+			return &types.ChatResponse{Content: `{"citations":{"s000":["c000","c999"]}}`}, nil
+		}
+		prompt := messages[0].Content
+		require.Contains(t, prompt, "complete_citation_regeneration")
+		require.Contains(t, prompt, `The ONLY allowed source IDs are: ["c000","c001"]`)
+		require.Contains(t, prompt, `{"s000":["c000"]}`)
+		require.Contains(t, prompt, chunks[0].Content)
+		require.Contains(t, prompt, chunks[1].Content)
+		require.NotContains(t, prompt, "c999", "invalid source identifier must not be proposed as a mapping")
+		return &types.ChatResponse{Content: `{"citations":{"s000":["c000","c001"]},"new_slugs":[]}`}, nil
+	}}
+	got, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
+	require.NoError(t, err)
+	require.Equal(t, 2, calls)
+	require.Equal(t, 1, completed)
+	require.Equal(t, []string{"first", "second"}, got["entity/alpha"])
+}
+
+func TestWikiCitationSourceRegenerationCannotDropKnownEvidenceOrCascadeRetries(t *testing.T) {
+	for name, reply := range map[string]string{
+		"dropped-known-edge":      `{"citations":{"s000":["c001"]}}`,
+		"still-unknown":           `{"citations":{"s000":["c000","c999"]}}`,
+		"missing-new-declaration": `{"citations":{"s000":["c000"],"concept/new":["c001"]}}`,
+		"malformed":               `{"citations":`,
+	} {
+		t.Run(name, func(t *testing.T) {
+			calls := 0
+			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+				calls++
+				if calls == 1 {
+					return &types.ChatResponse{Content: `{"citations":{"s000":["c000","c999"]}}`}, nil
+				}
+				return &types.ChatResponse{Content: reply}, nil
+			}}
+			got, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Alpha fact one."}, {ID: "second", Content: "Alpha fact two."}}, "English", &WikiBatchContext{})
+			require.Error(t, err)
+			require.Nil(t, got)
+			require.Zero(t, completed)
+			require.Equal(t, 2, calls, "no bisection, declaration repair or provider retry can be stacked")
+		})
+	}
+}
+
+func TestWikiCitationSourceRegenerationKeepsUndeclaredSlugEvidence(t *testing.T) {
+	calls := 0
+	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
+		calls++
+		if calls == 1 {
+			return &types.ChatResponse{Content: `{"citations":{"s000":["c999"],"concept/new":["c000"]}}`}, nil
+		}
+		require.Contains(t, messages[0].Content, `"concept/new":["c000"]`)
+		return &types.ChatResponse{Content: `{"citations":{"s000":["c000"]}}`}, nil
+	}}
+	got, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Alpha source."}}, "English", &WikiBatchContext{})
+	require.ErrorContains(t, err, "dropped previously valid evidence")
+	require.Nil(t, got)
+	require.Zero(t, completed)
+	require.Equal(t, 2, calls)
+}
+
+func TestWikiCitationUnknownSourceCannotRepairInvalidCandidateOrType(t *testing.T) {
+	for _, raw := range []string{
+		`{"citations":{"s999":["c999"],"s000":["c999"]}}`,
+		`{"citations":{"s000":["c999"]},"new_slugs":[{"type":"invalid","slug":"concept/new","name":"New","source_chunks":["c000"]}]}`,
+	} {
+		calls := 0
+		model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+			calls++
+			return &types.ChatResponse{Content: raw}, nil
+		}}
+		_, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Alpha source."}}, "English", &WikiBatchContext{})
+		require.Error(t, err)
+		require.Zero(t, completed)
+		require.Equal(t, 1, calls)
+	}
+}
+
+func TestWikiCitationDeclarationRepairHasOneCallAndNoSplitFallback(t *testing.T) {
+	for _, providerFailure := range []bool{false, true} {
+		calls := 0
+		model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+			calls++
+			if calls == 1 {
+				return &types.ChatResponse{Content: `{"citations":{"concept/new":["c000"]}}`}, nil
+			}
+			if providerFailure {
+				return nil, errors.New("HTTP 429 rate limit")
+			}
+			return &types.ChatResponse{Content: `{"new_slugs":[]}`}, nil
+		}}
+		_, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "one", Content: "Source one"}, {ID: "two", Content: "Source two"}}, "English", &WikiBatchContext{})
+		require.Error(t, err)
+		require.Zero(t, completed)
+		require.Equal(t, 2, calls)
+	}
+}
+
+func TestWikiCitationPersistedUnknownSlugDoesNotCallModel(t *testing.T) {
+	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
+	scope := &wikiCheckpointScope{7, "kb", "doc", "corrupt"}
+	chunks := []*types.Chunk{{ID: "source", Content: "Source text"}}
+	key := wikiCheckpointHash([]any{"entity/alpha", "English", wikiCheckpointChunks(chunks)})
+	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation", key, citationBatchResult{Citations: map[string][]string{"concept/undeclared": {"source"}}}))
+	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+		t.Fatal("persisted corruption must not spend a model request")
+		return nil, nil
+	}}
+	_, _, _, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
+	require.ErrorContains(t, err, "invalid persisted citation")
+	// Diagnostics must not contain private candidate text or original documents.
+	encoded, _ := json.Marshal(err.Error())
+	require.NotContains(t, string(encoded), "undeclared\"")
+}