xinmi/internal/application/service/wiki_fresh_candidates_test.go

package service

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strings"
	"testing"

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

func TestWikiSourceIdentitySingleCJKLayoutNewline(t *testing.T) {
	for _, test := range []struct {
		name, title string
		sources     []string
		allowed     bool
	}{
		{"literal", "高丝汀博士", []string{"原文高丝汀博士描述"}, true},
		{"LF", "高丝汀博士", []string{"原文高丝\n汀博士描述"}, true},
		{"CRLF", "高丝汀博士", []string{"原文高丝\r\n汀博士描述"}, true},
		{"paragraph LF", "高丝汀博士", []string{"高丝\n\n汀博士"}, false},
		{"paragraph CRLF", "高丝汀博士", []string{"高丝\r\n\r\n汀博士"}, false},
		{"two wraps", "高丝汀博士", []string{"高\n丝汀\n博士"}, false},
		{"space before", "高丝汀博士", []string{"高丝 \n汀博士"}, false},
		{"space after", "高丝汀博士", []string{"高丝\n 汀博士"}, false},
		{"bare CR", "高丝汀博士", []string{"高丝\r汀博士"}, false},
		{"latin", "AlphaBeta", []string{"Alpha\nBeta"}, false},
		{"mixed left", "A博士", []string{"A\n博士"}, false},
		{"mixed right", "博士A", []string{"博士\nA"}, false},
		{"punctuation", "高丝汀博士", []string{"高丝、\n汀博士"}, false},
		{"altered character", "高思汀博士", []string{"高丝\n汀博士"}, false},
		{"different chunks", "高丝汀博士", []string{"高丝", "汀博士"}, false},
		{"newline split chunks", "高丝汀博士", []string{"高丝\n", "汀博士"}, false},
	} {
		t.Run(test.name, func(t *testing.T) {
			before := append([]string(nil), test.sources...)
			err := validateWikiSourceItemIdentity(extractedItem{Name: test.title, Slug: "entity/name"}, "entity", 0, test.sources)
			if test.allowed {
				require.NoError(t, err)
			} else {
				requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source")
			}
			require.Equal(t, before, test.sources)
		})
	}
	batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "cited", Content: "高丝"}, {ID: "uncited", ChunkIndex: 1, Content: "高丝\n汀博士"}})[0]
	_, err := parseWikiCitationBatch(`{"citations":{"entity/name":["c000"]},"new_slugs":[{"type":"entity","name":"高丝汀博士","slug":"entity/name","aliases":[],"source_chunks":["c000"]}]}`, batch, map[string]bool{})
	requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source")
}

func TestWikiFreshCandidatesDeterministicIdempotentPreservesAudit(t *testing.T) {
	alphaDigest := fmt.Sprintf("%x", sha256.Sum256([]byte(wikiSourceIdentityKey("concept", "Alpha"))))
	reserved := "concept/shared-" + alphaDigest[:12]
	input := combinedExtraction{Entities: []extractedItem{{Name: "Alpha", Slug: "entity/alpha"}}, Concepts: []extractedItem{
		{Name: "Alpha", Slug: "concept/shared", Description: "first", SourceChunks: []string{"z", "a"}},
		{Name: "Beta", Slug: "concept/shared", Details: "Beta source"},
		{Name: "Alpha", Slug: "concept/shared", Description: "first", SourceChunks: []string{"z", "a"}},
		{Name: "Alpha", Slug: "concept/zz", Description: "second", SourceChunks: []string{"b"}},
		{Name: "Reserved", Slug: reserved},
	}}
	before, err := json.Marshal(input)
	require.NoError(t, err)
	want, err := normalizeWikiFreshCandidates(input)
	require.NoError(t, err)
	require.Len(t, want.Entities, 1)
	require.Len(t, want.Concepts, 3)
	require.Equal(t, "Alpha", want.Concepts[0].Name)
	require.Equal(t, "concept/shared-"+alphaDigest[:16], want.Concepts[0].Slug)
	require.Equal(t, "second", want.Concepts[0].Description)
	require.Equal(t, []string{"a", "b", "z"}, want.Concepts[0].SourceChunks)
	require.Equal(t, reserved, want.Concepts[2].Slug)
	again, err := normalizeWikiFreshCandidates(want)
	require.NoError(t, err)
	require.Equal(t, want, again)
	after, err := json.Marshal(input)
	require.NoError(t, err)
	require.Equal(t, before, after, "raw audit candidates must not be mutated")
	random := rand.New(rand.NewSource(42))
	for i := 0; i < 30; i++ {
		shuffled := combinedExtraction{Entities: append([]extractedItem(nil), input.Entities...), Concepts: append([]extractedItem(nil), input.Concepts...)}
		random.Shuffle(len(shuffled.Concepts), func(i, j int) {
			shuffled.Concepts[i], shuffled.Concepts[j] = shuffled.Concepts[j], shuffled.Concepts[i]
		})
		got, err := normalizeWikiFreshCandidates(shuffled)
		require.NoError(t, err)
		require.Equal(t, want, got)
	}
}

func TestWikiFreshCandidatesRetainsPublishedAndCitationCollisionGuards(t *testing.T) {
	input := combinedExtraction{Entities: []extractedItem{{Name: "Alpha", Slug: "entity/shared"}, {Name: "Beta", Slug: "entity/shared"}}}
	got, err := normalizeWikiFreshCandidates(input)
	require.NoError(t, err)
	page := &types.WikiPageLite{Slug: got.Entities[0].Slug, Title: "Published different title", PageType: "entity"}
	svc := &wikiIngestService{wikiService: &wikiSourceIdentityPages{pages: map[string]*types.WikiPageLite{page.Slug: page}}}
	// A published page keeps its slug: the intruding candidate is dropped and counted
	// so a document that already read its whole source keeps going, while the raw
	// validators below stay strict for the paths that merge citations into pages.
	resolved, err := svc.resolveWikiSourceIdentities(context.Background(), "kb", got, &WikiBatchContext{})
	require.NoError(t, err)
	require.Equal(t, 1, resolved.CandidatePruning.SlugConflicts)
	require.Len(t, resolved.Entities, 1)
	require.NotEqual(t, page.Slug, resolved.Entities[0].Slug)
	requireWikiSourceIdentityError(t, validateWikiIdentityCollisions(input), "wiki_identity_slug_conflict")
	_, _, _, err = mergeWikiSourceCitationItems(got.Entities, nil, nil, []newSlugFromCitation{{Type: "entity", Name: "Gamma", Slug: got.Entities[0].Slug}})
	requireWikiSourceIdentityError(t, err, "wiki_identity_slug_conflict")
	input.Entities[0].Aliases = []string{"Alias"}
	_, err = normalizeWikiFreshCandidates(input)
	requireWikiSourceIdentityError(t, err, "wiki_identity_alias_unsupported")
}

func TestWikiFreshCandidatesCrossWindowResumeDefersClaims(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "fresh-collision-resume"}
	source := wikiFullReadTestSource(6)
	batch, calls := &WikiBatchContext{}, 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		for i := 0; i < 6; i++ {
			name := fmt.Sprintf("WINDOW_%03d", i)
			if strings.Contains(messages[0].Content, name) {
				return &types.ChatResponse{Content: fmt.Sprintf(`{"entities":[{"name":%q,"slug":"entity/shared","aliases":[]}],"concepts":[]}`, name)}, nil
			}
		}
		return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil
	}}
	_, _, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", source, "English", nil, batch, scope)
	require.ErrorIs(t, err, errWikiCheckpointYield)
	batch.identityClaims.Range(func(_, _ any) bool { t.Error("partial fresh windows must not claim slugs"); return true })
	entities, _, items, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", source, "English", nil, batch, scope)
	require.NoError(t, err)
	require.Len(t, entities, 6)
	require.Len(t, items, 6)
	require.Equal(t, 12, calls)
	for _, item := range entities {
		require.True(t, strings.HasPrefix(item.Slug, "entity/shared-"))
		claim, ok := batch.identityClaims.Load(wikiSourceIdentityVersion + ":" + wikiSourceIdentityKey("entity", item.Name))
		require.True(t, ok)
		require.Equal(t, item.Slug, claim)
	}
	again, _, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", source, "English", nil, batch, scope)
	require.NoError(t, err)
	require.Equal(t, entities, again)
	require.Equal(t, 12, calls)
}