package service import ( "context" "encoding/json" "errors" "strings" "testing" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/alicebob/miniredis/v2" "github.com/redis/go-redis/v9" "github.com/stretchr/testify/require" ) func requireWikiSourceIdentityError(t *testing.T, err error, code string) *wikiSourceIdentityError { t.Helper() var identityErr *wikiSourceIdentityError require.ErrorAs(t, err, &identityErr) require.Equal(t, code, identityErr.Code) require.Equal(t, wikiSourceIdentityVersion, identityErr.Contract) var outputErr *wikiOutputError require.False(t, errors.As(err, &outputErr), "identity rejection must not enter output repair") return identityErr } func TestWikiSourceIdentityAllowsLiteralEntitiesAndDescriptiveConcepts(t *testing.T) { source := "Record LH changes. Retrieve materials, then verify their sources." result := combinedExtraction{ Entities: []extractedItem{{Name: " LH ", Slug: "entity/lh", Aliases: []string{}}}, Concepts: []extractedItem{{Name: "Source verification workflow", Slug: "concept/source-checking", Aliases: []string{}}}, } require.NotContains(t, source, result.Concepts[0].Name) before, err := json.Marshal(result) require.NoError(t, err) require.NoError(t, validateWikiSourceIdentities(result, source)) after, err := json.Marshal(result) require.NoError(t, err) require.Equal(t, before, after, "validation must preserve the original candidates") } func TestWikiSourceIdentityRejectsInferenceAndAllNewAliases(t *testing.T) { source := "LH and Alpha are recorded separately. Alpha is also explicitly called Beta. L H; LH-2." for _, test := range []struct { name, kind, title, code string aliases []string }{ {name: "expanded name", kind: "entity", title: "Luteinizing hormone", code: "wiki_identity_name_not_in_source"}, {name: "case folded name", kind: "entity", title: "lh", code: "wiki_identity_name_not_in_source"}, {name: "internal spacing changed", kind: "entity", title: "L H", code: "wiki_identity_name_not_in_source"}, {name: "punctuation removed", kind: "entity", title: "LH2", code: "wiki_identity_name_not_in_source"}, {name: "inferred alias", kind: "entity", title: "LH", aliases: []string{"Luteinizing hormone"}, code: "wiki_identity_alias_unsupported"}, {name: "both names occur", kind: "entity", title: "LH", aliases: []string{"Alpha"}, code: "wiki_identity_alias_unsupported"}, {name: "explicit semantic alias deferred", kind: "entity", title: "Alpha", aliases: []string{"Beta"}, code: "wiki_identity_alias_unsupported"}, {name: "formatting alias deferred", kind: "entity", title: "LH", aliases: []string{"L H"}, code: "wiki_identity_alias_unsupported"}, {name: "empty alias entry", kind: "entity", title: "LH", aliases: []string{""}, code: "wiki_identity_alias_unsupported"}, {name: "concept alias", kind: "concept", title: "Source verification workflow", aliases: []string{"Verification"}, code: "wiki_identity_alias_unsupported"}, } { t.Run(test.name, func(t *testing.T) { item := extractedItem{Name: test.title, Slug: test.kind + "/test", Aliases: test.aliases} err := validateWikiSourceItemIdentity(item, test.kind, 3, []string{source}) identityErr := requireWikiSourceIdentityError(t, err, test.code) require.Equal(t, test.kind, identityErr.Kind) require.Equal(t, 3, identityErr.Index) require.NotEmpty(t, identityErr.SourceHash) }) } } func TestWikiSourceIdentityErrorContainsNoPrivateText(t *testing.T) { source := "PRIVATE-SOURCE: secret notes contain LH." item := extractedItem{Name: "PRIVATE-GENERATED-NAME", Slug: "entity/private", Aliases: []string{"PRIVATE-ALIAS"}} err := validateWikiSourceItemIdentity(item, "entity", 7, []string{source}) identityErr := requireWikiSourceIdentityError(t, err, "wiki_identity_alias_unsupported") require.Equal(t, wikiCheckpointHash([]string{source}), identityErr.SourceHash) for _, private := range []string{source, item.Name, item.Slug, item.Aliases[0], "secret notes"} { require.NotContains(t, err.Error(), private) } var receipt map[string]any require.NoError(t, json.Unmarshal([]byte(err.Error()), &receipt)) require.Len(t, receipt, 5) require.Equal(t, "entity", receipt["kind"]) require.Equal(t, float64(7), receipt["index"]) } func TestWikiSourceIdentityExtractionFailureGetsOnePatchWithoutRegenerationOrSplit(t *testing.T) { for _, test := range []struct { name, output, code string corrected bool }{ {"initial expanded name", `{"entities":[{"name":"Luteinizing hormone","slug":"entity/lh","aliases":[]}],"concepts":[]}`, "wiki_identity_name_not_in_source", false}, {"initial alias", `{"entities":[{"name":"LH","slug":"entity/lh","aliases":["Alpha"]}],"concepts":[]}`, "wiki_identity_alias_unsupported", false}, {"corrected expanded name", `{"entities":[{"name":"Luteinizing hormone","slug":"entity/lh","aliases":[]}],"concepts":[]}`, "wiki_identity_name_not_in_source", true}, {"corrected alias", `{"entities":[{"name":"LH","slug":"entity/lh","aliases":["Alpha"]}],"concepts":[]}`, "wiki_identity_alias_unsupported", true}, } { t.Run(test.name, func(t *testing.T) { source := strings.Repeat("LH and Alpha are different. ", 800) input := map[string]string{"Content": source} 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) if test.corrected && calls == 1 { return &types.ChatResponse{Content: `{"entities":[`}, nil } identityCall := 2 if test.corrected { identityCall++ } if calls == identityCall { require.Contains(t, messages[0].Content, wikiIdentityRepairInstruction) } return &types.ChatResponse{Content: test.output}, nil }} result, err := (&wikiIngestService{}).generateWikiExtraction(context.Background(), model, "{{.Content}}", input) requireWikiSourceIdentityError(t, err, test.code) require.Empty(t, result.Entities) require.Empty(t, result.Concepts) wantCalls := 2 if test.corrected { wantCalls++ } require.Equal(t, wantCalls, calls) require.Equal(t, source, input["Content"]) }) } } func TestWikiSourceIdentitySplitCannotBorrowNameFromOtherHalf(t *testing.T) { source := "Alpha " + strings.Repeat("x", 6000) calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) { calls++ if calls == 1 { return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil } if calls == 3 { require.NotContains(t, messages[0].Content, "Alpha") } if calls == 4 { require.Contains(t, messages[0].Content, wikiIdentityRepairInstruction) return &types.ChatResponse{Content: `{"corrections":[{"kind":"entity","index":0,"name":"Alpha","source_quote":"Alpha"}]}`}, nil } return &types.ChatResponse{Content: `{"entities":[{"name":"Alpha","slug":"entity/alpha","aliases":[]}],"concepts":[]}`}, nil }} result, err := (&wikiIngestService{}).generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": source}) requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source") require.Equal(t, 4, calls) require.Empty(t, result.Entities, "a successful first partition cannot be published after the second fails") } func TestWikiSourceIdentityCitationUsesOnlyDeclaredSources(t *testing.T) { chunks := []*types.Chunk{{ID: "cited", Content: "Alpha and Beta are recorded separately."}, {ID: "uncited", ChunkIndex: 1, Content: "Gamma is only in this chunk."}} batch := splitChunksIntoCitationBatches(chunks)[0] for _, test := range []struct { name, kind, title, code string aliases []string }{ {name: "literal entity", kind: "entity", title: "Alpha"}, {name: "descriptive concept", kind: "concept", title: "Separate measurement workflow"}, {name: "name only in uncited chunk", kind: "entity", title: "Gamma", code: "wiki_identity_name_not_in_source"}, {name: "expanded entity", kind: "entity", title: "Alpha metric", code: "wiki_identity_name_not_in_source"}, {name: "both names in cited source", kind: "entity", title: "Alpha", aliases: []string{"Beta"}, code: "wiki_identity_alias_unsupported"}, } { t.Run(test.name, func(t *testing.T) { slug := test.kind + "/test" raw, err := json.Marshal(citationBatchResult{ Citations: map[string][]string{slug: {"c000"}}, NewSlugs: []newSlugFromCitation{{Type: test.kind, Name: test.title, Slug: slug, Aliases: test.aliases, SourceChunks: []string{"c000"}}}, }) require.NoError(t, err) result, err := parseWikiCitationBatch(string(raw), batch, map[string]bool{}) if test.code != "" { requireWikiSourceIdentityError(t, err, test.code) return } require.NoError(t, err) require.Equal(t, []string{"cited"}, result.NewSlugs[0].SourceChunks) require.Equal(t, []string{"cited"}, result.Citations[slug]) }) } } func TestWikiSourceIdentityRepairedDeclarationsCannotBypassContract(t *testing.T) { batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "cited", Content: "Alpha and Beta are recorded separately."}, {ID: "uncited", ChunkIndex: 1, Content: "Gamma only occurs here."}})[0] for _, test := range []struct{ name, title, aliases, code string }{ {"literal entity", "Alpha", "[]", ""}, {"name only in uncited chunk", "Gamma", "[]", "wiki_identity_name_not_in_source"}, {"expanded entity", "Alpha metric", "[]", "wiki_identity_name_not_in_source"}, {"alias despite exact evidence quote", "Alpha", `["Beta"]`, "wiki_identity_alias_unsupported"}, } { t.Run(test.name, func(t *testing.T) { original := citationBatchResult{Citations: map[string][]string{"entity/alpha": {"cited"}}} before, err := json.Marshal(original) require.NoError(t, err) raw := `{"cannot_repair":false,"new_slugs":[{"type":"entity","name":"` + test.title + `","slug":"entity/alpha","aliases":` + test.aliases + `,"description":"Separate records","details":"","source_chunks":["c000"],"evidence_quotes":[{"chunk":"c000","quote":"Alpha and Beta are recorded separately."}]}]}` got, err := applyWikiCitationDeclarations(raw, original, map[string][]string{"entity/alpha": {"c000"}}, batch, map[string]bool{}) if test.code == "" { require.NoError(t, err) require.Equal(t, original.Citations, got.Citations) require.Len(t, got.NewSlugs, 1) } else { requireWikiSourceIdentityError(t, err, test.code) require.Empty(t, got.Citations) require.Empty(t, got.NewSlugs) } after, marshalErr := json.Marshal(original) require.NoError(t, marshalErr) require.Equal(t, before, after, "repair cannot alter original evidence on success or rejection") }) } } func TestWikiSourceIdentityCitationFailureDoesNotRetryOrSplit(t *testing.T) { calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { calls++ return &types.ChatResponse{Content: `{"citations":{"entity/lh":["c000"]},"new_slugs":[{"type":"entity","name":"Luteinizing hormone","slug":"entity/lh","aliases":[],"source_chunks":["c000"]}]}`}, nil }} batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "first", Content: "LH is recorded."}, {ID: "second", ChunkIndex: 1, Content: "Other information."}})[0] got, err := (&wikiIngestService{}).classifyCitationBatch(context.Background(), model, "", batch, "English", true) requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source") require.Equal(t, 1, calls) require.Empty(t, got.Citations) require.Empty(t, got.NewSlugs) } type wikiSourceIdentityPages struct { interfaces.WikiPageService pages map[string]*types.WikiPageLite similarCalls int lookupFailure error } func (s *wikiSourceIdentityPages) ListBySlugs(_ context.Context, _ string, slugs []string) (map[string]*types.WikiPageLite, error) { if s.lookupFailure != nil { return nil, s.lookupFailure } result := make(map[string]*types.WikiPageLite) for _, slug := range slugs { if page := s.pages[slug]; page != nil { result[slug] = page } } return result, nil } func (s *wikiSourceIdentityPages) FindPagesByNormalizedTitles(_ context.Context, _ string, kind string, names []string) ([]*types.WikiPageLite, error) { if s.lookupFailure != nil { return nil, s.lookupFailure } var result []*types.WikiPageLite for _, page := range s.pages { if page.PageType != kind { continue } for _, name := range names { if normalizeWikiIdentityTitle(page.Title) == name { result = append(result, page) break } } } return result, nil } func (s *wikiSourceIdentityPages) FindSimilarPages(context.Context, string, string, []string, int) ([]*types.WikiPageLite, error) { s.similarCalls++ return nil, errors.New("strict identity must not use semantic candidates") } func TestWikiSourceIdentityResolverUsesExactTitleAndPreservesLegacyAliases(t *testing.T) { for _, test := range []struct{ name, existingTitle, proposed, want, code string }{ {"exact title", "LH", "entity/proposed", "entity/existing", ""}, {"outer whitespace", " LH ", "entity/proposed", "entity/existing", ""}, {"case normalization only", "lh", "entity/proposed", "entity/proposed", ""}, {"spacing normalization only", "L H", "entity/proposed", "entity/proposed", ""}, {"legacy alias only", "Luteinizing hormone", "entity/proposed", "entity/proposed", ""}, {"same slug different title", "Luteinizing hormone", "entity/existing", "", "wiki_identity_slug_conflict"}, } { t.Run(test.name, func(t *testing.T) { page := &types.WikiPageLite{Slug: "entity/existing", Title: test.existingTitle, PageType: types.WikiPageTypeEntity, Aliases: types.StringArray{"LH", "Legacy alias"}} pages := &wikiSourceIdentityPages{pages: map[string]*types.WikiPageLite{page.Slug: page}} before, err := json.Marshal(page) require.NoError(t, err) input := combinedExtraction{Entities: []extractedItem{{Name: "LH", Slug: test.proposed}}} got, err := (&wikiIngestService{wikiService: pages}).resolveWikiSourceIdentities(context.Background(), "kb", input, &WikiBatchContext{}) if test.code != "" { requireWikiSourceIdentityError(t, err, test.code) require.Empty(t, got.Entities) } else { require.NoError(t, err) require.Len(t, got.Entities, 1) require.Equal(t, test.want, got.Entities[0].Slug) require.Empty(t, got.Entities[0].Aliases) } require.Equal(t, test.proposed, input.Entities[0].Slug) require.Zero(t, pages.similarCalls) after, marshalErr := json.Marshal(page) require.NoError(t, marshalErr) require.Equal(t, before, after) }) } } func TestWikiSourceIdentityExtractionDoesNotInvokeSemanticDedup(t *testing.T) { page := &types.WikiPageLite{Slug: "entity/existing", Title: "LH", PageType: types.WikiPageTypeEntity, Aliases: types.StringArray{"Legacy alias"}} pages := &wikiSourceIdentityPages{pages: map[string]*types.WikiPageLite{page.Slug: page}} calls := 0 model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) { calls++ return &types.ChatResponse{Content: `{"entities":[{"name":"LH","slug":"entity/proposed","aliases":[]}],"concepts":[]}`}, nil }} entities, _, items, err := (&wikiIngestService{wikiService: pages}).extractCandidateSlugs(context.Background(), model, "kb", "LH is recorded.", "English", nil, &WikiBatchContext{}) require.NoError(t, err) require.Equal(t, 1, calls) require.Zero(t, pages.similarCalls) require.Len(t, entities, 1) require.Equal(t, "entity/existing", entities[0].Slug) require.Contains(t, items, "entity/existing") require.Equal(t, types.StringArray{"Legacy alias"}, page.Aliases) } func TestWikiSourceIdentityClaimsIgnoreLegacyRedisAndLocalMappings(t *testing.T) { mr := miniredis.RunT(t) rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) t.Cleanup(func() { _ = rdb.Close() }) svc := &wikiIngestService{redisClient: rdb} ctx := context.Background() batch := &WikiBatchContext{} legacy := svc.claimWikiIdentitySlug(ctx, "kb", "entity", "LH", "entity/legacy-expanded", false, batch) require.Equal(t, "entity/legacy-expanded", legacy) for _, test := range []struct{ title, proposed, want string }{ {"LH", "entity/lh", "entity/lh"}, {"lh", "entity/lower", "entity/lower"}, {"L H", "entity/spaced", "entity/spaced"}, {"LH", "entity/another", "entity/lh"}, } { got, err := svc.claimWikiSourceIdentity(ctx, "kb", "entity", test.title, test.proposed, false, batch) require.NoError(t, err) require.Equal(t, test.want, got) } legacyKey := wikiIdentityClaimPrefix + "kb:entity:" + normalizeWikiIdentityTitle("LH") value, err := rdb.Get(ctx, legacyKey).Result() require.NoError(t, err) require.Equal(t, "entity/legacy-expanded", value) local, ok := batch.identityClaims.Load("entity\x00" + normalizeWikiIdentityTitle("LH")) require.True(t, ok) require.Equal(t, "entity/legacy-expanded", local) } func TestWikiSourceIdentityRevalidatesClaimedTarget(t *testing.T) { page := &types.WikiPageLite{Slug: "entity/occupied", Title: "Other", PageType: types.WikiPageTypeEntity} pages := &wikiSourceIdentityPages{pages: map[string]*types.WikiPageLite{page.Slug: page}} svc := &wikiIngestService{wikiService: pages} batch := &WikiBatchContext{} _, err := svc.claimWikiSourceIdentity(context.Background(), "kb", "entity", "LH", page.Slug, false, batch) require.NoError(t, err) got, err := svc.resolveWikiSourceIdentities(context.Background(), "kb", combinedExtraction{Entities: []extractedItem{{Name: "LH", Slug: "entity/proposed"}}}, batch) requireWikiSourceIdentityError(t, err, "wiki_identity_slug_conflict") require.Empty(t, got.Entities) require.Equal(t, "Other", page.Title) } func TestWikiSourceIdentityWindowCollisionPreservesEarlierCandidates(t *testing.T) { all := combinedExtraction{Entities: []extractedItem{{Name: "LH", Slug: "entity/shared", SourceChunks: []string{"first"}, Description: "First source"}}} before, err := json.Marshal(all) require.NoError(t, err) part := combinedExtraction{Entities: []extractedItem{{Name: "Other", Slug: "entity/shared", SourceChunks: []string{"second"}, Description: "Second source"}}} err = mergeWikiWindowCandidates(&all, part) requireWikiSourceIdentityError(t, err, "wiki_identity_slug_conflict") after, err := json.Marshal(all) require.NoError(t, err) require.Equal(t, before, after) require.Empty(t, all.Entities[0].Aliases) part.Entities[0].Name = "LH" require.NoError(t, mergeWikiWindowCandidates(&all, part)) require.Len(t, all.Entities, 1) require.ElementsMatch(t, []string{"first", "second"}, all.Entities[0].SourceChunks) require.Empty(t, all.Entities[0].Aliases) } func TestWikiSourceIdentityWindowOuterWhitespaceDoesNotCreateAlias(t *testing.T) { all := combinedExtraction{Entities: []extractedItem{{Name: "LH", Slug: "entity/lh", SourceChunks: []string{"first"}}}} part := combinedExtraction{Entities: []extractedItem{{Name: " LH ", Slug: "entity/lh", SourceChunks: []string{"second"}}}} require.NoError(t, validateWikiSourceIdentities(part, "LH occurs in the second source.")) require.NoError(t, mergeWikiWindowCandidates(&all, part)) require.Len(t, all.Entities, 1) require.Equal(t, "LH", all.Entities[0].Name) require.Empty(t, all.Entities[0].Aliases, "outer whitespace trimming must not manufacture a new alias") require.ElementsMatch(t, []string{"first", "second"}, all.Entities[0].SourceChunks) } func TestWikiSourceIdentityCandidateCheckpointRejectsUnsupportedIdentity(t *testing.T) { source := "LH and Alpha are recorded separately." for _, test := range []struct { name, title, code string aliases []string }{ {name: "faithful", title: "LH"}, {name: "expanded name", title: "Luteinizing hormone", code: "wiki_identity_name_not_in_source"}, {name: "alias", title: "LH", aliases: []string{"Alpha"}, code: "wiki_identity_alias_unsupported"}, } { t.Run(test.name, func(t *testing.T) { result := combinedExtraction{Entities: []extractedItem{{Name: test.title, Slug: "entity/lh", Aliases: test.aliases}}} err := validateWikiCandidateCheckpoint(result, source) if test.code == "" { require.NoError(t, err) } else { requireWikiSourceIdentityError(t, err, test.code) require.ErrorContains(t, err, "invalid persisted wiki candidate") } }) } } func wikiSourceIdentityMapFixture() wikiMapCheckpoint { return wikiMapCheckpoint{ Result: docIngestResult{KnowledgeID: "doc", DocTitle: "Record", Summary: "Complete summary", Pages: []wikiIngestPageRef{{Slug: "summary/doc", Title: "Record"}, {Slug: "entity/lh", Title: "LH"}}}, Updates: []SlugUpdate{ {KnowledgeID: "doc", SourceRef: "doc", Slug: "summary/doc", Type: types.WikiPageTypeSummary, DocTitle: "Record", SummaryBody: "Complete summary"}, {KnowledgeID: "doc", SourceRef: "doc", Slug: "entity/lh", Type: types.WikiPageTypeEntity, DocTitle: "Record", SourceChunks: []string{"cited"}, Item: extractedItem{Name: "LH", Slug: "entity/lh", SourceChunks: []string{"cited"}}}, }, } } func TestWikiSourceIdentityCompletedMapCheckpointRejectsUnsupportedIdentity(t *testing.T) { chunks := []*types.Chunk{{ID: "cited", Content: "LH and Alpha are recorded separately."}, {ID: "uncited", Content: "Gamma occurs only here."}} scope := &wikiCheckpointScope{7, "kb", "doc", wikiSourceIdentityVersion} require.NoError(t, validateWikiMapCheckpoint(wikiSourceIdentityMapFixture(), scope, chunks)) for _, test := range []struct { name, title, code string aliases []string }{ {name: "expanded name", title: "Luteinizing hormone", code: "wiki_identity_name_not_in_source"}, {name: "name in uncited chunk", title: "Gamma", code: "wiki_identity_name_not_in_source"}, {name: "alias", title: "LH", aliases: []string{"Alpha"}, code: "wiki_identity_alias_unsupported"}, } { t.Run(test.name, func(t *testing.T) { saved := wikiSourceIdentityMapFixture() saved.Updates[1].Item.Name = test.title saved.Updates[1].Item.Aliases = test.aliases before, err := json.Marshal(saved) require.NoError(t, err) err = validateWikiMapCheckpoint(saved, scope, chunks) requireWikiSourceIdentityError(t, err, test.code) after, err := json.Marshal(saved) require.NoError(t, err) require.Equal(t, before, after) }) } }