package service import ( "context" "fmt" "path/filepath" "strings" "testing" "unicode/utf8" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" "github.com/stretchr/testify/require" ) func TestWikiFullReadWindowsPreserveEverySourceCharacter(t *testing.T) { for _, source := range []string{ strings.Repeat("原文😀", 27000) + "最后的限制条件", strings.Repeat("段萜\n", 12000) + "\n\n尟郚", strings.Repeat("x", 32768) + "\n", } { windows := splitWikiReadWindows(source, maxContentForWiki) var joined strings.Builder end := 0 for _, window := range windows { require.Equal(t, end, window.StartRune) require.Equal(t, utf8.RuneCountInString(window.Text), window.EndRune-window.StartRune) require.LessOrEqual(t, window.EndRune-window.StartRune, maxContentForWiki) joined.WriteString(window.Text) end = window.EndRune } require.Equal(t, source, joined.String()) require.Equal(t, utf8.RuneCountInString(source), end) } } func wikiFullReadTestSource(windows int) string { var content strings.Builder for i := 0; i < windows; i++ { label := fmt.Sprintf("WINDOW_%03d", i) content.WriteString(label + strings.Repeat("甲", maxContentForWiki-len(label))) } return content.String() } func TestWikiFullReadCandidatesResumeAndIncludeTail(t *testing.T) { path := filepath.Join(t.TempDir(), "checkpoints.sqlite") svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)} scope := &wikiCheckpointScope{7, "kb", "doc", "full-content-signature"} content := wikiFullReadTestSource(6) seen := make(map[int]int) model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) { for i := 0; i < 6; i++ { if strings.Contains(messages[0].Content, fmt.Sprintf("WINDOW_%03d", i)) { seen[i]++ return &types.ChatResponse{Content: fmt.Sprintf(`{"entities":[{"name":"WINDOW_%03d","slug":"entity/item-%d"}],"concepts":[]}`, i, i)}, nil } } // The second half contains only repeated filler and has no candidate. return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil }} entities, _, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", content, "Chinese", nil, &WikiBatchContext{}, scope) require.ErrorIs(t, err, errWikiCheckpointYield) require.Nil(t, entities, "a partial document must not enter publication") require.Len(t, seen, 4) var progress wikiReadProgress hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress) require.NoError(t, err) require.True(t, hit) require.Equal(t, 4, progress.CompletedWindows) require.False(t, progress.SourceWindowsComplete) sqlDB, _ := svc.checkpointDB.DB() require.NoError(t, sqlDB.Close()) svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)} entities, _, _, err = svc.extractWikiFullCandidates(context.Background(), model, "kb", content, "Chinese", nil, &WikiBatchContext{}, scope) require.NoError(t, err) require.Len(t, entities, 6) for i := 0; i < 6; i++ { require.Equal(t, 1, seen[i], "a completed source window must not be reread") } _, err = svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress) require.NoError(t, err) require.True(t, progress.StageComplete) require.True(t, progress.SourceWindowsComplete) require.False(t, progress.EvidenceVerified) require.Equal(t, 6*maxContentForWiki, progress.ProcessedRunes) } func TestWikiFullReadRepeatedWindowsReuseArtifactButCountAllRanges(t *testing.T) { svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))} scope := &wikiCheckpointScope{7, "kb", "doc", "duplicate-source"} content := strings.Repeat("甲", maxContentForWiki*9) calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { calls++ return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil }} _, _, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", content, "Chinese", nil, &WikiBatchContext{}, scope) require.NoError(t, err) require.Equal(t, 2, calls, "one parent is extracted in two halves; all repeated parent ranges reuse it") var progress wikiReadProgress _, err = svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress) require.NoError(t, err) require.Equal(t, 9, progress.CompletedWindows) require.Equal(t, maxContentForWiki*9, progress.ProcessedRunes) var ranges []wikiReadRange hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_plan", "source", &ranges) require.NoError(t, err) require.True(t, hit) require.Len(t, ranges, 9) for i, segment := range ranges { require.Equal(t, i*maxContentForWiki, segment.StartRune) require.Equal(t, (i+1)*maxContentForWiki, segment.EndRune) require.Equal(t, ranges[0].TextHash, segment.TextHash) } } func TestWikiFullReadInvalidCachedCandidateDoesNotCallModelOrMarkCoverage(t *testing.T) { svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))} scope := &wikiCheckpointScope{7, "kb", "doc", "invalid-cached-candidate"} source := wikiFullReadTestSource(2) windows := splitWikiReadWindows(source, maxContentForWiki) key := wikiCheckpointHash([]any{wikiFullReadVersion, windows[0].Text}) require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "read_candidate", key, combinedExtraction{Entities: []extractedItem{{Name: "private-name", Slug: "concept/private-slug"}}, Concepts: []extractedItem{}})) model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { t.Fatal("invalid persisted output must not trigger another model call") return nil, nil }} _, _, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", source, "Chinese", nil, &WikiBatchContext{}, scope) require.ErrorContains(t, err, "invalid persisted wiki candidate") require.ErrorContains(t, err, "entities[0].slug") require.NotContains(t, err.Error(), "private-") var progress wikiReadProgress hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress) require.NoError(t, err) require.False(t, hit) } func TestWikiFullReadSummaryBoundedHierarchyResumes(t *testing.T) { path := filepath.Join(t.TempDir(), "checkpoints.sqlite") svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)} scope := &wikiCheckpointScope{7, "kb", "doc", "summary-scope"} content := wikiFullReadTestSource(7) seen, mergeCalls := make(map[int]int), 0 model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) { prompt := messages[0].Content require.Less(t, utf8.RuneCountInString(prompt), maxContentForWiki+3000) if strings.Contains(prompt, "") { for i := 0; i < 7; i++ { if strings.Contains(prompt, fmt.Sprintf("WINDOW_%03d", i)) { seen[i]++ return &types.ChatResponse{Content: fmt.Sprintf("SUMMARY: Window %d\n\nNotes from section %d.", i, i)}, nil } } t.Fatal("missing source window") } require.Contains(t, prompt, "NOT original source evidence") mergeCalls++ return &types.ChatResponse{Content: "SUMMARY: Combined overview\n\nLossy navigation notes."}, nil }} for attempt := 0; attempt < 4; attempt++ { output, err := svc.generateWikiFullSummary(context.Background(), model, content, "Chinese", "", &WikiBatchContext{}, scope) if err == nil { require.Contains(t, output, "7 䞪窗口") require.Contains(t, output, "䞍等同完敎原文") break } require.ErrorIs(t, err, errWikiCheckpointYield) require.Empty(t, output) sqlDB, _ := svc.checkpointDB.DB() require.NoError(t, sqlDB.Close()) svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)} } for i := 0; i < 7; i++ { require.Equal(t, 1, seen[i]) } require.Equal(t, 2, mergeCalls) var progress wikiReadProgress _, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "summary", &progress) require.NoError(t, err) require.True(t, progress.StageComplete) require.True(t, progress.SourceWindowsComplete) require.False(t, progress.EvidenceVerified) } func TestWikiFullReadCandidateCatalogCanExceedOneRequest(t *testing.T) { svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))} scope := &wikiCheckpointScope{7, "kb", "doc", "candidate-budget"} content := wikiFullReadTestSource(2) for i, window := range splitWikiReadWindows(content, maxContentForWiki) { part := combinedExtraction{Concepts: []extractedItem{{Name: fmt.Sprintf("Source %d ", i) + strings.Repeat("topic ", 2200), Slug: fmt.Sprintf("concept/source-%d", i)}}} key := wikiCheckpointHash([]any{wikiFullReadVersion, window.Text}) require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "read_candidate", key, part)) } calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { calls++ return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil }} for i := 0; i < 2; i++ { entities, concepts, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", content, "Chinese", nil, &WikiBatchContext{}, scope) require.NoError(t, err) require.Greater(t, utf8.RuneCountInString(renderCandidateSlugsXML(entities, concepts)), wikiReadCandidateRunes) shards, err := buildWikiCandidateShards(entities, concepts, wikiReadCandidateRunes) require.NoError(t, err) require.Len(t, shards, 2) } require.Zero(t, calls) var progress wikiReadProgress _, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress) require.NoError(t, err) require.Equal(t, "complete", progress.Status) require.True(t, progress.StageComplete) require.True(t, progress.SourceWindowsComplete) } func TestWikiFullReadShortSourceWithHugeCatalogUsesBoundedSourceNotes(t *testing.T) { svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))} scope := &wikiCheckpointScope{7, "kb", "doc", "short-large-catalog"} source := "The original source includes a restriction." calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) { calls++ require.Contains(t, messages[0].Content, source) require.NotContains(t, messages[0].Content, "HUGE_CATALOG_SENTINEL") require.Less(t, utf8.RuneCountInString(messages[0].Content), wikiReadCandidateRunes) return &types.ChatResponse{Content: "SUMMARY: Source\n\nA concise overview preserving the restriction."}, nil }} for i := 0; i < 2; i++ { output, err := svc.generateWikiFullSummary(context.Background(), model, source, "English", strings.Repeat("HUGE_CATALOG_SENTINEL", 2000), &WikiBatchContext{}, scope) require.NoError(t, err) require.Contains(t, output, "restriction") } require.Equal(t, 1, calls) } func TestWikiFullReadNeedsDurableStorageAndRejectsPartialSummary(t *testing.T) { svc := &wikiIngestService{} _, _, _, err := svc.extractWikiFullCandidates(context.Background(), nil, "kb", wikiFullReadTestSource(2), "Chinese", nil, &WikiBatchContext{}, nil) require.ErrorContains(t, err, "需芁持久断点") require.Error(t, validateWikiReadSummary("SUMMARY: headline")) require.Error(t, validateWikiReadSummary("SUMMARY: headline\n\n"+strings.Repeat("甲", wikiReadSummaryRunes))) } // Policy: an over-long window summary gets exactly ONE bounded compression retry, // because the observed failure is a model ignoring an explicit character budget // rather than an unsolvable segment. A second overshoot parks the document and the // pause is remembered, so a model that cannot follow the budget never loops. func TestWikiFullReadSummaryBudgetPauseRetriesOnceThenParks(t *testing.T) { svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))} scope := &wikiCheckpointScope{7, "kb", "doc", "summary-budget"} calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { calls++ return &types.ChatResponse{Content: "SUMMARY: overlong\n\n" + strings.Repeat("甲", wikiReadSummaryRunes)}, nil }} for i := 0; i < 2; i++ { output, err := svc.generateWikiFullSummary(context.Background(), model, wikiFullReadTestSource(2), "Chinese", "", &WikiBatchContext{}, scope) require.Empty(t, output) require.ErrorContains(t, err, "党文倄理已暂停") } require.Equal(t, 2, calls, "one initial call plus one compression retry, then the remembered pause") var progress wikiReadProgress _, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "summary", &progress) require.NoError(t, err) require.Equal(t, "paused_budget", progress.Status) require.False(t, progress.SourceWindowsComplete) } func TestWikiFullReadArtifactsCannotCrossTenantDocumentOrRevision(t *testing.T) { svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))} calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { calls++ return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil }} for _, scope := range []*wikiCheckpointScope{ {7, "kb", "doc", "revision-a"}, {8, "kb", "doc", "revision-a"}, {7, "other-kb", "doc", "revision-a"}, {7, "kb", "other-doc", "revision-a"}, {7, "kb", "doc", "revision-b"}, } { _, _, _, err := svc.extractWikiFullCandidates(context.Background(), model, scope.KnowledgeBaseID, wikiFullReadTestSource(2), "Chinese", nil, &WikiBatchContext{}, scope) require.NoError(t, err) } require.Equal(t, 20, calls, "each of two windows has two halves; artifacts remain inside the exact source scope") } func TestWikiFullReadPausedSummaryDoesNotStartCitationWork(t *testing.T) { svc := &wikiIngestService{ checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db")), knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: checkpointChunks(4)}, wikiService: &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}, } calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) { calls++ if strings.Contains(messages[0].Content, "") { return &types.ChatResponse{Content: "SUMMARY: too long\n\n" + strings.Repeat("甲", wikiReadSummaryRunes)}, nil } require.NotContains(t, messages[0].Content, "", "paused summaries must not spend more on citations") return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil }} payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"} op := WikiPendingOp{KnowledgeID: "doc", dbID: 42} for attempt := 0; attempt < 2; attempt++ { result, updates, err := svc.mapOneDocument(context.Background(), model, payload, op, svc.newWikiBatchContext("kb", nil)) require.ErrorContains(t, err, "党文倄理已暂停") require.Nil(t, result) require.Nil(t, updates) } require.Equal(t, 5, calls, "three candidate halves, one over-budget summary and its single compression retry; the second run is cache-only") }