xinmi/internal/application/service/wiki_candidate_quality_test.go

package service

import (
	"context"
	"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"
)

func extractionModel(t *testing.T, content func(call int) string) (*wikiScriptedChat, *int) {
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		calls++
		return &types.ChatResponse{Content: content(calls)}, nil
	}}
	return model, &calls
}

func freshExtraction(t *testing.T, model *wikiScriptedChat, source string) (combinedExtraction, error) {
	return (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(), model, agent.WikiCandidateSlugPrompt,
		map[string]string{"Content": source})
}

// Duplicated names and the overflow above the per-source cap are dropped and
// counted instead of failing the document.
func TestWikiCandidateQualityPrunesDuplicatesAndOverflow(t *testing.T) {
	names := make([]string, 0, 20)
	for i := 0; i < 20; i++ {
		names = append(names, "Term"+string(rune('A'+i)))
	}
	// 119 runes of names plus 2000 runes of padding put the per-source cap at 17.
	source := strings.Join(names, " ") + " " + strings.Repeat("x ", 1000)

	entities := []string{`{"name":"TermA","slug":"entity/term-a"}`, `{"name":"TermA","slug":"entity/term-a"}`}
	for _, name := range names[1:] {
		entities = append(entities, `{"name":"`+name+`","slug":"entity/`+strings.ToLower(name)+`"}`)
	}
	model, calls := extractionModel(t, func(int) string {
		return `{"entities":[` + strings.Join(entities, ",") + `],"concepts":[]}`
	})

	result, err := freshExtraction(t, model, source)
	require.NoError(t, err)
	require.Equal(t, 1, *calls, "pruning is local: no correction request for duplicates or overflow")
	require.Equal(t, 1, result.CandidatePruning.Duplicates)
	require.Equal(t, 3, result.CandidatePruning.Capped)
	require.Equal(t, 0, result.CandidatePruning.Unmatched)
	require.Equal(t, 17, len(result.Entities), "the cap keeps the earliest candidates")
	require.Equal(t, 4, result.CandidatePruning.Dropped())
	for _, entity := range result.Entities {
		require.True(t, strings.HasPrefix(entity.Slug, "entity/"), "surviving candidates keep their slugs")
	}
}

// A clean candidate set is returned untouched and reports no pruning.
func TestWikiCandidateQualityKeepsSupportedCandidatesUnchanged(t *testing.T) {
	source := "Alpha and Beta are both in the source."
	model, calls := extractionModel(t, func(int) string {
		return `{"entities":[{"name":"Alpha","slug":"entity/alpha"},{"name":"Beta","slug":"entity/beta"}],"concepts":[]}`
	})
	result, err := freshExtraction(t, model, source)
	require.NoError(t, err)
	require.Equal(t, 1, *calls)
	require.Equal(t, 2, len(result.Entities))
	require.Equal(t, 0, result.CandidatePruning.Dropped())
}

// A small number of unsupported names still gets the bounded correction request;
// when the correction cannot be applied the offending candidate is dropped and
// counted instead of failing the whole document.
func TestWikiIdentityRepairFallsBackToDroppingTheCandidate(t *testing.T) {
	source := "Alpha and Beta are both in the source."
	model, calls := extractionModel(t, func(call int) string {
		if call == 1 {
			return `{"entities":[{"name":"Alpha","slug":"entity/alpha"},{"name":"Beta","slug":"entity/beta"},{"name":"Unsupported","slug":"entity/unsupported"}],"concepts":[]}`
		}
		return `{"corrections":[{"kind":"entity","index":2,"name":"","source_quote":""}]}`
	})
	result, err := freshExtraction(t, model, source)
	require.NoError(t, err)
	require.Equal(t, 2, *calls, "a bounded rejection still asks the model once")
	require.Equal(t, 2, len(result.Entities))
	require.Equal(t, 1, result.CandidatePruning.Unmatched)
	for _, entity := range result.Entities {
		require.Contains(t, source, entity.Name)
	}
}

// More unsupported names than the repair budget are dropped without spending a
// correction request, and the supported candidates survive.
func TestWikiIdentityPruningDropsUnsupportedCandidatesWithoutRequest(t *testing.T) {
	supported := []string{"Alpha", "Beta", "Gamma"}
	source := strings.Join(supported, " ") + " are supported."
	entities := make([]string, 0, 13)
	for _, name := range supported {
		entities = append(entities, `{"name":"`+name+`","slug":"entity/`+strings.ToLower(name)+`"}`)
	}
	for i := 0; i < 10; i++ {
		name := "Missing" + string(rune('A'+i))
		entities = append(entities, `{"name":"`+name+`","slug":"entity/missing-`+strings.ToLower(name)+`"}`)
	}
	model, calls := extractionModel(t, func(int) string {
		return `{"entities":[` + strings.Join(entities, ",") + `],"concepts":[]}`
	})
	result, err := freshExtraction(t, model, source)
	require.NoError(t, err)
	require.Equal(t, 1, *calls, "an over-budget rejection must not spend a correction request")
	require.Equal(t, 3, len(result.Entities))
	require.Equal(t, 10, result.CandidatePruning.Unmatched)
	for _, entity := range result.Entities {
		require.Contains(t, source, entity.Name)
	}
}

// Aliased identities still cannot be published, but an over-budget alias
// rejection drops only the offending candidates.
func TestWikiIdentityRepairDropsUnsupportedAliasesBeyondTheBudget(t *testing.T) {
	source := "Alpha and Beta are supported. GammaA GammaB GammaC GammaD GammaE GammaF GammaG GammaH GammaI GammaJ"
	entities := []string{`{"name":"Alpha","slug":"entity/alpha"}`, `{"name":"Beta","slug":"entity/beta"}`}
	for i := 0; i < 10; i++ {
		name := "Gamma" + string(rune('A'+i))
		entities = append(entities, `{"name":"`+name+`","slug":"entity/alias-`+strings.ToLower(name)+`","aliases":["`+name+` alias"]}`)
	}
	model, calls := extractionModel(t, func(int) string {
		return `{"entities":[` + strings.Join(entities, ",") + `],"concepts":[]}`
	})
	result, err := freshExtraction(t, model, source)
	require.NoError(t, err)
	require.Equal(t, 1, *calls, "more rejected identities than the repair budget must not spend a correction request")
	require.Equal(t, 2, len(result.Entities))
	require.Equal(t, 10, result.CandidatePruning.Unmatched)
}

// Pruning must never publish an empty candidate set: when nothing survives, the
// document still fails, and the reported error is the identity failure that
// caused the drops.
func TestWikiCandidateQualityStillFailsWhenNothingSurvives(t *testing.T) {
	source := "Alpha and Beta are in the source."
	model, _ := extractionModel(t, func(call int) string {
		if call == 1 {
			return `{"entities":[{"name":"Unsupported","slug":"entity/unsupported"}],"concepts":[]}`
		}
		return `{"corrections":[{"kind":"entity","index":0,"name":"","source_quote":""}]}`
	})
	_, err := freshExtraction(t, model, source)
	require.Error(t, err)
	var identity *wikiSourceIdentityError
	require.ErrorAs(t, err, &identity)
}

// A bounded, supported repair path keeps working: a small number of rejected
// identities is still corrected by the model.
func TestWikiCandidateQualityKeepsSmallIdentityRepairPath(t *testing.T) {
	source := "Alpha, LH and FSH are separate."
	model, calls := extractionModel(t, func(call int) string {
		if call == 1 {
			return `{"entities":[{"name":"Alpha","slug":"entity/alpha"},{"name":"Expanded LH","slug":"entity/lh"}],"concepts":[]}`
		}
		return `{"corrections":[{"kind":"entity","index":1,"name":"LH","source_quote":"Alpha, LH and FSH"}]}`
	})
	result, err := freshExtraction(t, model, source)
	require.NoError(t, err)
	require.Equal(t, 2, *calls)
	require.Equal(t, "LH", result.Entities[1].Name)
	require.Equal(t, 0, result.CandidatePruning.Unmatched)
}