xinmi/internal/application/service/wiki_read_summary_repair_test.go

package service

import (
	"context"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/agent"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
)

// Every prompt that can produce a window summary must state the character budget the
// validator enforces, because a model that is never told the limit overshoots it and
// parks the document.
func TestWikiSummaryPromptsStateTheEnforcedCharacterBudget(t *testing.T) {
	limit := strconv.Itoa(wikiReadSummaryRunes)
	for name, prompt := range map[string]string{
		"read_section":  wikiReadSectionPrompt,
		"read_merge":    wikiReadMergePrompt,
		"document_page": agent.WikiSummaryPrompt,
	} {
		require.Contains(t, prompt, "3000", name+": a soft target below the hard limit")
		require.Contains(t, prompt, limit, name+": the hard limit the validator enforces")
		require.NotContains(t, prompt, "500-1500 words", name+": word counts do not bound CJK output")
	}
}

func wikiSummaryShortNote() string {
	return "SUMMARY: short note\n\n## Topic\nOne line per topic.\n\n## Takeaways\n- Alpha\n"
}

func wikiSummaryOverlongNote() string {
	return "SUMMARY: overlong note\n\n" + strings.Repeat("Alpha Beta Gamma detail. ", 400)
}

// The window summary path must recover from a single overshoot by compressing, so the
// document is not parked for a budget the model was already told.
func TestWikiFullReadSummaryCompressionRetryCompletesTheDocument(t *testing.T) {
	service := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "summary-compress"}
	batch := &WikiBatchContext{}

	compressions, calls := 0, 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		input := messages[len(messages)-1].Content
		if strings.Contains(input, "<compression_request>") {
			compressions++
			return &types.ChatResponse{Content: wikiSummaryShortNote()}, nil
		}
		if calls == 1 {
			return &types.ChatResponse{Content: wikiSummaryOverlongNote()}, nil
		}
		return &types.ChatResponse{Content: wikiSummaryShortNote()}, nil
	}}

	var summary string
	var err error
	for attempt := 0; attempt < 12; attempt++ {
		summary, err = service.generateWikiFullSummary(context.Background(), model, wikiFullReadTestSource(2), "English", "", batch, scope)
		if err == nil || !strings.Contains(err.Error(), "continuation required") {
			break
		}
	}
	require.NoError(t, err, "one overshoot must not park the document")
	require.Contains(t, summary, "short note")
	require.Equal(t, 1, compressions, "exactly one compression request for the over-long segment")
	require.Less(t, calls, 12, "the retry must be bounded by the per-turn operation budget")
}