xinmi/internal/application/service/wiki_candidate_partition_test.go

package service

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"
	"unicode/utf8"

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

func wikiPartitionRequestSource(t *testing.T, prompt string) string {
	t.Helper()
	_, after, found := strings.Cut(prompt, "<content>\n")
	require.True(t, found)
	source, _, found := strings.Cut(after, "\n</content>")
	require.True(t, found)
	return source
}

func TestWikiCandidatePartitionProactiveHalvesPreserveSourceAndBudget(t *testing.T) {
	guard := &candidateRepairTestGuard{}
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	ctx = requestbudget.WithGuard(ctx, guard)
	deadline, _ := ctx.Deadline()
	left := "Alpha " + strings.Repeat("甲😀", 8100) + "\n"
	right := "Beta " + strings.Repeat("乙😀", 8200) + " \n\t"
	source := left + right
	input := map[string]string{"Content": source, "Language": "English", "PreviousSlugs": "unchanged"}
	var seen []string
	model := &wikiScriptedChat{id: t.Name(), call: func(callCtx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		require.Same(t, guard, requestbudget.FromContext(callCtx))
		actual, ok := callCtx.Deadline()
		require.True(t, ok)
		require.Equal(t, deadline, actual)
		require.True(t, strings.HasSuffix(messages[0].Content, "\n\n"+wikiCandidateGenerationInstruction))
		part := wikiPartitionRequestSource(t, messages[0].Content)
		seen = append(seen, part)
		require.NotEqual(t, source, part, "no oversized parent request precedes the two halves")
		require.LessOrEqual(t, utf8.RuneCountInString(part), (utf8.RuneCountInString(source)+1)/2+500)
		name := "Alpha"
		if strings.Contains(part, "Beta") {
			name = "Beta"
		}
		return &types.ChatResponse{Content: `{"entities":[{"name":"` + name + `","slug":"entity/` + strings.ToLower(name) + `"}],"concepts":[]}`}, nil
	}}
	result, err := (&wikiIngestService{}).generateWikiFreshExtraction(ctx, model, agent.WikiCandidateSlugPrompt, input)
	require.NoError(t, err)
	require.Equal(t, []string{left, right}, seen)
	require.Equal(t, source, strings.Join(seen, ""))
	require.Equal(t, source, input["Content"])
	require.Equal(t, "English", input["Language"])
	require.Equal(t, "unchanged", input["PreviousSlugs"])
	require.Len(t, result.Entities, 2)
	require.Equal(t, "Alpha", result.Entities[0].Name)
	require.Equal(t, "Beta", result.Entities[1].Name)
	require.GreaterOrEqual(t, guard.checks, 4)
}

func TestWikiCandidatePartitionThresholdAndPromptApplyToEveryRequest(t *testing.T) {
	for _, size := range []int{wikiCandidateProactiveSplitRunes, wikiCandidateProactiveSplitRunes + 1} {
		t.Run(fmt.Sprint(size), func(t *testing.T) {
			var parts []string
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
				require.Equal(t, 1, strings.Count(messages[0].Content, wikiCandidateGenerationInstruction))
				parts = append(parts, wikiPartitionRequestSource(t, messages[0].Content))
				return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil
			}}
			source := strings.Repeat("甲", size)
			_, err := (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(), model, agent.WikiCandidateSlugPrompt, map[string]string{"Content": source})
			require.NoError(t, err)
			want := 1
			if size > wikiCandidateProactiveSplitRunes {
				want = 2
			}
			require.Len(t, parts, want)
			require.Equal(t, source, strings.Join(parts, ""))
		})
	}
}

func TestWikiCandidatePartitionStopsAfterFailedOrCancelledHalf(t *testing.T) {
	for _, failure := range []string{"provider", "identity", "cancel", "budget"} {
		t.Run(failure, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			guard := &candidateRepairTestGuard{}
			ctx = requestbudget.WithGuard(ctx, guard)
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
				calls++
				if calls == 1 {
					if failure == "cancel" {
						cancel()
					}
					if failure == "budget" {
						guard.stopped = true
					}
					return &types.ChatResponse{Content: `{"entities":[{"name":"Alpha","slug":"entity/alpha"}],"concepts":[]}`}, nil
				}
				if failure == "identity" {
					require.NotContains(t, messages[0].Content, "Alpha ")
					if calls == 3 {
						require.Contains(t, messages[0].Content, wikiIdentityRepairInstruction)
						return &types.ChatResponse{Content: `{"corrections":[]}`}, nil
					}
					return &types.ChatResponse{Content: `{"entities":[{"name":"Unsupported","slug":"entity/unsupported"}],"concepts":[]}`}, nil
				}
				return nil, errors.New("nontransient provider failure")
			}}
			source := "Alpha " + strings.Repeat("x", 15994) + "Beta " + strings.Repeat("y", 15995)
			result, err := (&wikiIngestService{}).generateWikiFreshExtraction(ctx, model, agent.WikiCandidateSlugPrompt, map[string]string{"Content": source})
			require.Error(t, err)
			require.Empty(t, result.Entities, "a completed first half must not escape after later failure")
			want := 2
			switch failure {
			case "identity":
				want = 3
				requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source")
			case "cancel":
				want = 1
				require.ErrorIs(t, err, context.Canceled)
			case "budget":
				want = 1
				require.ErrorIs(t, err, requestbudget.ErrStopped)
			}
			require.Equal(t, want, calls)
		})
	}
}

func TestWikiCandidatePartitionKeepsTwoLevelDepthLimit(t *testing.T) {
	for _, persistent := range []bool{false, true} {
		t.Run(fmt.Sprint(persistent), func(t *testing.T) {
			calls := 0
			var leaves []string
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
				calls++
				part := wikiPartitionRequestSource(t, messages[0].Content)
				require.LessOrEqual(t, utf8.RuneCountInString(part), 16000)
				if utf8.RuneCountInString(part) > 8000 || persistent {
					return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
				}
				leaves = append(leaves, part)
				return &types.ChatResponse{Content: `{"entities":[{"name":"x","slug":"entity/x"}],"concepts":[]}`}, nil
			}}
			source := strings.Repeat("x", 32000)
			result, err := (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(), model, agent.WikiCandidateSlugPrompt, map[string]string{"Content": source})
			if persistent {
				require.Error(t, err)
				require.Equal(t, 2, calls, "first half and first quarter exhaust the existing depth bound")
				require.Empty(t, result.Entities)
			} else {
				require.NoError(t, err)
				require.Equal(t, 6, calls, "two truncated halves and four complete quarters, without a parent call")
				require.Len(t, leaves, 4)
				require.Equal(t, source, strings.Join(leaves, ""))
				require.Len(t, result.Entities, 4, "raw leaf proposals survive until all source checks finish")
			}
		})
	}
}

func TestWikiCandidatePartitionRegenerationRetainsUniquenessInstruction(t *testing.T) {
	calls := 0
	var first string
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		require.Equal(t, 1, strings.Count(messages[0].Content, wikiCandidateGenerationInstruction))
		if calls == 1 {
			first = messages[0].Content
			return &types.ChatResponse{Content: `{"concepts":[]}`}, nil
		}
		require.Equal(t, 2, calls)
		require.True(t, strings.HasPrefix(messages[0].Content, first+"\n\n"+wikiCandidateRepairInstruction))
		return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil
	}}
	_, err := (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(), model, agent.WikiCandidateSlugPrompt, map[string]string{"Content": "A short source"})
	require.NoError(t, err)
	require.Equal(t, 2, calls)
}

func TestWikiCandidatePartitionLeavesLegacyExtractionUnchanged(t *testing.T) {
	source := "Alpha " + strings.Repeat("甲", 32000)
	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, wikiCandidateGenerationInstruction)
		return &types.ChatResponse{Content: `{"entities":[{"name":"Alpha","slug":"entity/alpha"}],"concepts":[]}`}, nil
	}}
	result, err := (&wikiIngestService{}).generateWikiExtraction(context.Background(), model, agent.WikiKnowledgeExtractPrompt, map[string]string{"Content": source})
	require.NoError(t, err)
	require.Equal(t, 1, calls)
	require.Len(t, result.Entities, 1)
}