xinmi/internal/application/service/wiki_window_merge_pruning_test.go

package service

import (
	"context"
	"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 entity(name, slug string, aliases ...string) extractedItem {
	return extractedItem{Name: name, Slug: slug, Aliases: aliases}
}

// A later window may claim a slug that another identity already owns. The window
// must not discard the whole document: the colliding candidate is dropped and
// counted, and the first identity keeps the slug.
func TestWikiWindowMergeDropsSlugConflictInsteadOfFailing(t *testing.T) {
	accumulated := combinedExtraction{Entities: []extractedItem{entity("口爱技法", "entity/kou-ai")}}
	incoming := combinedExtraction{Entities: []extractedItem{
		entity("口爱前准备", "entity/kou-ai"),
		entity("舌技基础", "entity/she-ji"),
	}}
	require.NoError(t, mergeWikiWindowCandidates(&accumulated, incoming))
	require.Len(t, accumulated.Entities, 2, "the first identity keeps the slug and the supported candidate survives")
	require.Equal(t, "口爱技法", accumulated.Entities[0].Name)
	require.Equal(t, "舌技基础", accumulated.Entities[1].Name)
	require.Equal(t, 1, accumulated.CandidatePruning.SlugConflicts)
}

// The same identity repeated in a later window still merges into one page.
func TestWikiWindowMergeKeepsRepeatedIdentity(t *testing.T) {
	accumulated := combinedExtraction{Entities: []extractedItem{entity("舌技基础", "entity/she-ji")}}
	incoming := combinedExtraction{Entities: []extractedItem{{Name: "舌技基础", Slug: "entity/she-ji", Description: "更完整的描述"}}}
	require.NoError(t, mergeWikiWindowCandidates(&accumulated, incoming))
	require.Len(t, accumulated.Entities, 1)
	require.Equal(t, "更完整的描述", accumulated.Entities[0].Description)
	require.Equal(t, 0, accumulated.CandidatePruning.SlugConflicts)
}

// Candidates that cannot be published because they carry aliases are dropped and
// counted rather than aborting the window.
func TestWikiWindowMergeDropsAliasedCandidates(t *testing.T) {
	accumulated := combinedExtraction{}
	incoming := combinedExtraction{
		Entities: []extractedItem{entity("舌技基础", "entity/she-ji"), entity("指技", "entity/zhi-ji", "指法")},
	}
	require.NoError(t, mergeWikiWindowCandidates(&accumulated, incoming))
	require.Len(t, accumulated.Entities, 1)
	require.Equal(t, "舌技基础", accumulated.Entities[0].Name)
	require.Equal(t, 1, accumulated.CandidatePruning.SlugConflicts)
}

// Pruning must never publish nothing: a window whose candidates were all dropped
// fails closed instead of silently producing an empty set.
func TestWikiWindowMergeFailsWhenEveryCandidateIsUnsupported(t *testing.T) {
	accumulated := combinedExtraction{}
	incoming := combinedExtraction{Entities: []extractedItem{entity("指技", "entity/zhi-ji", "指法")}}
	err := mergeWikiWindowCandidates(&accumulated, incoming)
	require.Error(t, err)
	var identity *wikiSourceIdentityError
	require.ErrorAs(t, err, &identity)
}

// The within-set primitive keeps the first identity for a slug and counts the rest.
func TestPruneIdentityCollisionsCountsLaterClaimants(t *testing.T) {
	pruned, dropped := pruneIdentityCollisions(combinedExtraction{
		Entities: []extractedItem{entity("甲技法", "entity/jia"), entity("乙技法", "entity/jia"), entity("丙技法", "entity/bing")},
		Concepts: []extractedItem{entity("概念甲", "concept/k"), entity("概念乙", "concept/k")},
	})
	require.Equal(t, 2, dropped)
	require.Equal(t, 2, len(pruned.Entities))
	require.Equal(t, "甲技法", pruned.Entities[0].Name)
	require.Equal(t, 1, len(pruned.Concepts))
	require.Equal(t, "概念甲", pruned.Concepts[0].Name)
}

// The candidate stage must have room for a dense JSON answer: window 29 of a real
// document came back truncated at the old 32,768-token output cap, which no retry
// can fix because the input is unchanged. The provider allows far more.
func TestWikiLLMOutputBudgetIsLargeEnoughForDenseCandidateSets(t *testing.T) {
	require.GreaterOrEqual(t, wikiLLMMaxTokens, 65536,
		"a dense candidate list must fit inside the provider output budget")
	_, err := (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(),
		&wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
			return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil
		}}, agent.WikiCandidateSlugPrompt, map[string]string{"Content": "Alpha and Beta."})
	require.NoError(t, err)
}

// The Pass 0 candidate checkpoint is validated on every queue turn. Two names that
// claim one slug must be pruned and counted there too, instead of aborting a
// document that has already read its whole source.
func TestWikiCandidateCheckpointPruningKeepsTheDocumentMoving(t *testing.T) {
	source := "Alpha and Beta are supported."
	collided := combinedExtraction{Entities: []extractedItem{
		entity("Alpha", "entity/shared"), entity("Beta", "entity/shared"), entity("Alpha", "entity/alpha"),
	}}
	pruned, dropped, err := pruneCandidateSetCollisions(collided)
	require.NoError(t, err)
	require.Equal(t, 1, dropped, "the later claimant is dropped")
	require.Len(t, pruned.Entities, 2)
	require.NoError(t, validateWikiCandidateCheckpoint(pruned, source))
	require.Error(t, validateWikiCandidateCheckpoint(collided, source),
		"the strict validator still rejects the unpruned set")

	// A set whose every candidate is unpublishable (unsupported aliases) must not
	// silently become an empty candidate list.
	unpublishable := combinedExtraction{Entities: []extractedItem{entity("Alpha", "entity/alpha", "别名")}}
	_, _, err = pruneCandidateSetCollisions(unpublishable)
	require.Error(t, err, "a set that prunes to nothing still fails closed")
}

// An empty provider response is reported as such instead of as broken JSON, so a
// content filter is diagnosable from one line.
func TestWikiExtractionReportsEmptyProviderResponse(t *testing.T) {
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		return &types.ChatResponse{Content: "", FinishReason: "content_filter"}, nil
	}}
	_, err := (&wikiIngestService{}).generateWikiFreshExtraction(context.Background(), model, agent.WikiCandidateSlugPrompt,
		map[string]string{"Content": "Alpha and Beta are in the source."})
	require.Error(t, err)
	require.Contains(t, err.Error(), "provider_empty_response")
}