xinmi/internal/application/service/wiki_full_read_test.go

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, "<source_segment>") {
			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, "<source_segment>") {
			return &types.ChatResponse{Content: "SUMMARY: too long\n\n" + strings.Repeat("甲", wikiReadSummaryRunes)}, nil
		}
		require.NotContains(t, messages[0].Content, "<chunks>", "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")
}