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

internal/application/service/wiki_candidate_catalog_test.go

@@ -0,0 +1,310 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"path/filepath"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+	"unicode/utf8"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/require"
+)
+
+func wikiCatalogTestCandidates() []extractedItem {
+	return []extractedItem{
+		{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{strings.Repeat("甲", 13000)}},
+		{Name: "Beta", Slug: "entity/beta", Aliases: []string{strings.Repeat("乙", 13000)}},
+	}
+}
+
+func wikiCatalogTestChunks(n int) []*types.Chunk {
+	var chunks []*types.Chunk
+	for i := 0; i < n; i++ {
+		chunks = append(chunks, &types.Chunk{ID: fmt.Sprintf("source-%02d", i), ChunkIndex: i,
+			Content: fmt.Sprintf("SOURCE_%02d\n", i) + strings.Repeat("原", 7000)})
+	}
+	return chunks
+}
+
+func TestWikiCandidateCatalogWholeEntriesStableAndInvalidated(t *testing.T) {
+	entities := wikiCatalogTestCandidates()
+	entities[0].Details = "Audit detail is retained but not rendered"
+	entities[0].Description = "Audit description is retained but not rendered"
+	entities[0].SourceChunks = []string{"original-source"}
+	concepts := []extractedItem{{Name: "Topic", Slug: "concept/topic", Aliases: []string{"one\ntwo", "three"}}}
+	shards, err := buildWikiCandidateShards(entities, concepts, wikiReadCandidateRunes)
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, len(shards), 2)
+	var got []wikiCandidateEntry
+	for _, shard := range shards {
+		require.Equal(t, shard.XML, renderWikiCandidateEntries(shard.Entries))
+		require.LessOrEqual(t, utf8.RuneCountInString(shard.XML), wikiReadCandidateRunes)
+		got = append(got, shard.Entries...)
+	}
+	require.ElementsMatch(t, []wikiCandidateEntry{{"entity", entities[0]}, {"entity", entities[1]}, {"concept", concepts[0]}}, got)
+	reordered, err := buildWikiCandidateShards([]extractedItem{entities[1], entities[0]}, concepts, wikiReadCandidateRunes)
+	require.NoError(t, err)
+	require.Equal(t, shards, reordered)
+	duplicates, err := buildWikiCandidateShards([]extractedItem{entities[0], entities[0]}, nil, wikiReadCandidateRunes)
+	require.NoError(t, err)
+	require.Len(t, duplicates, 2, "even duplicate input entries are not silently discarded by planning")
+	require.NotEqual(t, duplicates[0].Hash, duplicates[1].Hash, "identical shards retain distinct occurrence checkpoints")
+	chunks := wikiCatalogTestChunks(3)
+	plan, err := buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
+	require.NoError(t, err)
+	reorderedPlan, err := buildWikiCitationPlan([]extractedItem{entities[1], entities[0]}, concepts, []*types.Chunk{chunks[2], chunks[0], chunks[1]}, "Chinese")
+	require.NoError(t, err)
+	require.Equal(t, plan, reorderedPlan)
+	changed, err := buildWikiCitationPlan(entities, concepts, chunks, "English")
+	require.NoError(t, err)
+	require.NotEqual(t, plan.Hash, changed.Hash)
+	entities[0].Details += " changed audit field"
+	changed, err = buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
+	require.NoError(t, err)
+	require.NotEqual(t, plan.Hash, changed.Hash)
+	chunks[2].Content += "tail condition"
+	changedSource, err := buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
+	require.NoError(t, err)
+	require.NotEqual(t, changed.Hash, changedSource.Hash)
+	_, err = buildWikiCandidateShards([]extractedItem{{Name: "Oversize", Slug: "entity/oversize", Aliases: []string{strings.Repeat("名", wikiReadCandidateRunes)}}}, nil, wikiReadCandidateRunes)
+	require.ErrorContains(t, err, "single-entry budget")
+	_, err = buildWikiCandidateShards([]extractedItem{{Slug: "entity/no-name"}}, nil, wikiReadCandidateRunes)
+	require.ErrorContains(t, err, "empty identity")
+}
+
+func TestWikiCandidateCatalogSourcesPreserveExactRangesAndWrappers(t *testing.T) {
+	chunks := []*types.Chunk{
+		{ID: "second", ChunkIndex: 1, Content: "\nsecond\n\n"},
+		{ID: "first", ChunkIndex: 0, Content: strings.Repeat("原文😀\n", 9100) + " trailing qualification \n"},
+	}
+	sources, err := buildWikiCitationSources(chunks, maxRunesPerCitationBatch)
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, len(sources), 3)
+	joined := make(map[string]string)
+	ends := make(map[string]int)
+	for _, source := range sources {
+		require.LessOrEqual(t, utf8.RuneCountInString(renderChunksXML(source.Batch)), maxRunesPerCitationBatch)
+		for i, span := range source.Ranges {
+			text := source.Batch.chunks[i].Content
+			require.Equal(t, ends[span.ChunkID], span.StartRune)
+			require.Equal(t, span.EndRune-span.StartRune, utf8.RuneCountInString(text))
+			require.Equal(t, wikiCheckpointHash(text), span.TextHash)
+			joined[span.ChunkID] += text
+			ends[span.ChunkID] = span.EndRune
+		}
+	}
+	for _, chunk := range chunks {
+		require.Equal(t, chunk.Content, joined[chunk.ID])
+	}
+	require.NotEqual(t, sources[0].Hash, sources[1].Hash, "source ranges remain distinct even for repeating content")
+	_, err = buildWikiCitationSources([]*types.Chunk{chunks[0], chunks[0]}, maxRunesPerCitationBatch)
+	require.ErrorContains(t, err, "duplicate")
+	_, err = buildWikiCitationSources(chunks, 1)
+	require.ErrorContains(t, err, "wrapper")
+}
+
+func TestWikiCandidateCatalogAllPairsResumeAndLateEvidence(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "db")
+	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
+	scope := &wikiCheckpointScope{7, "kb", "doc", "pairs"}
+	chunks := wikiCatalogTestChunks(9)
+	plan, err := buildWikiCitationPlan(wikiCatalogTestCandidates(), nil, chunks, "Chinese")
+	require.NoError(t, err)
+	require.Len(t, plan.Shards, 2)
+	require.Len(t, plan.Sources, 9)
+	var calls, active, peak atomic.Int32
+	var mu sync.Mutex
+	seen := map[string]int{}
+	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
+		calls.Add(1)
+		concurrent := active.Add(1)
+		defer active.Add(-1)
+		for old := peak.Load(); concurrent > old; old = peak.Load() {
+			if peak.CompareAndSwap(old, concurrent) {
+				break
+			}
+		}
+		prompt := messages[0].Content
+		if utf8.RuneCountInString(prompt) > wikiCitationRequestRunes {
+			return nil, errors.New("unbounded request")
+		}
+		slug := "entity/beta"
+		if strings.Contains(prompt, `name: "Alpha"`) {
+			slug = "entity/alpha"
+		}
+		for i := range chunks {
+			if strings.Contains(prompt, fmt.Sprintf("SOURCE_%02d", i)) {
+				mu.Lock()
+				seen[fmt.Sprintf("%s/%d", slug, i)]++
+				mu.Unlock()
+			}
+		}
+		if slug == "entity/alpha" && !strings.Contains(prompt, "SOURCE_08") {
+			return &types.ChatResponse{Content: `{"citations":{}}`}, nil
+		}
+		return &types.ChatResponse{Content: fmt.Sprintf(`{"citations":{%q:["c000"]}}`, slug)}, nil
+	}}
+	got, _, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
+	require.ErrorIs(t, err, errWikiCheckpointYield)
+	require.Nil(t, got)
+	require.Equal(t, 16, completed)
+	require.Equal(t, int32(16), calls.Load())
+	var progress wikiCitationProgress
+	_, err = svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
+	require.NoError(t, err)
+	require.False(t, progress.StageComplete)
+	require.Equal(t, 18, progress.TotalPairs)
+	sqlDB, err := svc.checkpointDB.DB()
+	require.NoError(t, err)
+	require.NoError(t, sqlDB.Close())
+	svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
+	got, _, completed, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
+	require.NoError(t, err)
+	require.Equal(t, 18, completed)
+	require.Equal(t, []string{chunks[8].ID}, got["entity/alpha"], "an early shard must inspect the last source batch")
+	require.Len(t, got["entity/beta"], 9)
+	require.Equal(t, int32(18), calls.Load())
+	require.LessOrEqual(t, peak.Load(), int32(maxCitationBatchConcurrency))
+	for key, count := range seen {
+		require.Equal(t, 1, count, key)
+	}
+	require.Len(t, seen, 18)
+	_, err = svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
+	require.NoError(t, err)
+	require.True(t, progress.StageComplete)
+	_, _, _, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
+	require.NoError(t, err)
+	require.Equal(t, int32(18), calls.Load())
+	_, _, _, err = (&wikiIngestService{}).classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", nil)
+	require.ErrorContains(t, err, "需要持久断点")
+	require.Equal(t, int32(18), calls.Load())
+}
+
+func TestWikiCandidateCatalogFailedPairKeepsSuccessfulSiblings(t *testing.T) {
+	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
+	scope := &wikiCheckpointScope{7, "kb", "doc", "failure"}
+	chunks := wikiCatalogTestChunks(3)
+	plan, err := buildWikiCitationPlan(wikiCatalogTestCandidates(), nil, chunks, "English")
+	require.NoError(t, err)
+	var failed atomic.Bool
+	var mu sync.Mutex
+	calls := map[string]int{}
+	model := &wikiScriptedChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
+		prompt := messages[0].Content
+		slug := "entity/beta"
+		if strings.Contains(prompt, `name: "Alpha"`) {
+			slug = "entity/alpha"
+		}
+		index := 0
+		for i := range chunks {
+			if strings.Contains(prompt, fmt.Sprintf("SOURCE_%02d", i)) {
+				index = i
+			}
+		}
+		key := fmt.Sprintf("%s/%d", slug, index)
+		mu.Lock()
+		calls[key]++
+		mu.Unlock()
+		if slug == "entity/alpha" && index == 1 && failed.CompareAndSwap(false, true) {
+			deadline := time.Now().Add(3 * time.Second)
+			for time.Now().Before(deadline) {
+				var count int64
+				if err := svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "citation_pair").Count(&count).Error; err != nil {
+					return nil, err
+				}
+				if count > 0 {
+					return nil, errors.New("invalid request")
+				}
+				select {
+				case <-ctx.Done():
+					return nil, ctx.Err()
+				case <-time.After(time.Millisecond):
+				}
+			}
+			return nil, errors.New("sibling checkpoint did not persist")
+		}
+		return &types.ChatResponse{Content: fmt.Sprintf(`{"citations":{%q:["c000"]}}`, slug)}, nil
+	}}
+	got, _, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
+	require.Error(t, err)
+	require.Nil(t, got)
+	require.GreaterOrEqual(t, completed, 1)
+	require.Less(t, completed, 6)
+	var progress wikiCitationProgress
+	_, loadErr := svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
+	require.NoError(t, loadErr)
+	require.False(t, progress.StageComplete)
+	got, _, completed, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
+	require.NoError(t, err)
+	require.Equal(t, 6, completed)
+	require.Len(t, got["entity/alpha"], 3)
+	for key, count := range calls {
+		want := 1
+		if key == "entity/alpha/1" {
+			want = 2
+		}
+		require.Equal(t, want, count, key)
+	}
+}
+
+func TestWikiCandidateCatalogPairIsolationAndEvidenceUnion(t *testing.T) {
+	for _, corrupt := range []bool{false, true} {
+		t.Run(fmt.Sprintf("corrupt-%t", corrupt), func(t *testing.T) {
+			svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
+			scope := &wikiCheckpointScope{7, "kb", "doc", "isolation"}
+			entities := wikiCatalogTestCandidates()
+			chunks := []*types.Chunk{{ID: "first", Content: "Alpha Beta first", ChunkIndex: 0}, {ID: "second", Content: "Alpha Beta second", ChunkIndex: 1}}
+			plan, err := buildWikiCitationPlan(entities, nil, chunks, "English")
+			require.NoError(t, err)
+			require.Len(t, plan.Shards, 2)
+			require.Len(t, plan.Sources, 1)
+			for i, shard := range plan.Shards {
+				result := citationBatchResult{Citations: map[string][]string{"entity/alpha": {"first"}}}
+				if i == 1 && !corrupt {
+					result.NewSlugs = []newSlugFromCitation{{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"second"}}}
+					result.Citations["entity/beta"] = []string{"second"}
+				}
+				key := wikiCheckpointHash([]any{wikiCandidateCatalogVersion, plan.Hash, shard.Hash, plan.Sources[0].Hash, "English"})
+				require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation_pair", key, result))
+			}
+			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+				return nil, errors.New("cached validation must not invoke model repair")
+			}}
+			citations, declarations, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
+			if corrupt {
+				require.ErrorContains(t, err, "undeclared slug")
+				require.Nil(t, citations)
+				require.Less(t, completed, 2)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, 2, completed)
+			got, _, _ := mergeCitationsIntoItems(entities, nil, citations, declarations)
+			require.Equal(t, []string{"first", "second"}, findBySlug(got, "entity/alpha").SourceChunks)
+			require.Equal(t, []string{"second"}, findBySlug(got, "entity/beta").SourceChunks)
+		})
+	}
+}
+
+func TestWikiCandidateCatalogCompleteRequestGuard(t *testing.T) {
+	var calls atomic.Int32
+	inner := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+		calls.Add(1)
+		return &types.ChatResponse{}, nil
+	}}
+	model := wikiCitationBoundedChat{inner}
+	_, err := model.Chat(context.Background(), []chat.Message{{Content: strings.Repeat("x", wikiCitationRequestRunes-1)}, {Content: "repair instruction"}}, nil)
+	require.ErrorContains(t, err, "complete-request budget")
+	require.Zero(t, calls.Load())
+	require.Equal(t, inner.GetModelID(), model.GetModelID())
+	_, err = model.Chat(context.Background(), []chat.Message{{Content: "bounded request"}}, nil)
+	require.NoError(t, err)
+	require.Equal(t, int32(1), calls.Load())
+}