xinmi/internal/application/service/wiki_source_identity_integration_test.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"testing"

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

func TestWikiSourceIdentityCitationMergePreservesInputsAndEmptyAliases(t *testing.T) {
	entities := []extractedItem{{Name: "LH", Slug: "entity/lh", SourceChunks: []string{"old"}}}
	declarations := []newSlugFromCitation{
		{Type: "entity", Name: " LH ", Slug: "entity/lh", SourceChunks: []string{"declared-1"}},
		{Type: "entity", Name: "LH", Slug: "entity/lh", SourceChunks: []string{"declared-2"}},
	}
	before, err := json.Marshal([]any{entities, declarations})
	require.NoError(t, err)
	merged, concepts, uncited, err := mergeWikiSourceCitationItems(entities, nil, map[string][]string{"entity/lh": {"cited"}}, declarations)
	require.NoError(t, err)
	require.Zero(t, uncited)
	require.Empty(t, concepts)
	require.Len(t, merged, 1)
	require.Equal(t, "LH", merged[0].Name)
	require.Empty(t, merged[0].Aliases)
	require.ElementsMatch(t, []string{"cited", "declared-1", "declared-2"}, merged[0].SourceChunks)
	after, err := json.Marshal([]any{entities, declarations})
	require.NoError(t, err)
	require.Equal(t, before, after)
}

func TestWikiSourceIdentityCitationMergeRejectsConflictBeforeMutation(t *testing.T) {
	entities := []extractedItem{{Name: "LH", Slug: "entity/lh", SourceChunks: []string{"original"}}}
	declarations := []newSlugFromCitation{{Type: "entity", Name: "Alpha", Slug: "entity/lh", SourceChunks: []string{"new"}}}
	before, err := json.Marshal([]any{entities, declarations})
	require.NoError(t, err)
	merged, concepts, _, err := mergeWikiSourceCitationItems(entities, nil, nil, declarations)
	requireWikiSourceIdentityError(t, err, "wiki_identity_slug_conflict")
	require.Nil(t, merged)
	require.Nil(t, concepts)
	after, err := json.Marshal([]any{entities, declarations})
	require.NoError(t, err)
	require.Equal(t, before, after)
}

func TestWikiSourceIdentitySplitCollisionPreservesDistinctCandidates(t *testing.T) {
	source := "Alpha " + strings.Repeat("x", 3000) + "Beta " + strings.Repeat("y", 3000)
	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
		}
		name := "Beta"
		if strings.Contains(messages[0].Content, "Alpha") {
			name = "Alpha"
		}
		return &types.ChatResponse{Content: fmt.Sprintf(`{"entities":[{"name":%q,"slug":"entity/shared","aliases":[]}],"concepts":[]}`, name)}, nil
	}}
	result, err := (&wikiIngestService{}).generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": source})
	require.NoError(t, err)
	require.Len(t, result.Entities, 2)
	require.NoError(t, validateWikiIdentityCollisions(result))
	for _, item := range result.Entities {
		require.Contains(t, []string{"Alpha", "Beta"}, item.Name)
		require.True(t, strings.HasPrefix(item.Slug, "entity/shared-"))
	}
	require.Equal(t, 3, calls)
}

func TestWikiSourceIdentityPageRejectionPreservesCommittedPageAndCheckpoint(t *testing.T) {
	for _, test := range []struct {
		name, code string
		change     func(*SlugUpdate)
	}{
		{"different title", "wiki_identity_slug_conflict", func(u *SlugUpdate) { u.Item.Name = "Beta" }},
		{"different type", "wiki_identity_slug_conflict", func(u *SlugUpdate) { u.Type = types.WikiPageTypeConcept }},
		{"new alias", "wiki_identity_alias_unsupported", func(u *SlugUpdate) { u.Item.Aliases = []string{"First"} }},
		{"name absent from cited source", "wiki_identity_name_not_in_source", func(u *SlugUpdate) {}},
	} {
		t.Run(test.name, func(t *testing.T) {
			original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "entity/alpha", Title: "Alpha", PageType: types.WikiPageTypeEntity,
				Status: types.WikiPageStatusPublished, Content: "Original supported content", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}, ChunkRefs: types.StringArray{"old-chunk"}, Aliases: types.StringArray{"Legacy alias"}}
			before, err := json.Marshal(original)
			require.NoError(t, err)
			pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
			db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
			svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}}
			batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
			batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "fingerprint"})
			update := wikiEvidenceAddition()
			update.Type, update.Slug, update.Item.Slug = types.WikiPageTypeEntity, original.Slug, original.Slug
			test.change(&update)
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				t.Fatal("identity rejection must happen before page generation")
				return nil, nil
			}}
			updates := []SlugUpdate{{KnowledgeID: "old-doc", SourceRef: "old-doc", Slug: original.Slug, Type: "retract"}, update}
			changed, _, failed, err := svc.reduceSlugUpdates(context.Background(), model, "kb", original.Slug, updates, 7, batch, nil)
			requireWikiSourceIdentityError(t, err, test.code)
			require.False(t, changed)
			require.True(t, failed)
			require.Zero(t, pages.writes)
			after, err := json.Marshal(pages.pages[original.Slug])
			require.NoError(t, err)
			require.Equal(t, before, after)
			var checkpoints int64
			require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&checkpoints).Error)
			require.Zero(t, checkpoints)
		})
	}
}

func TestWikiSourceIdentityExactPageUpdatePreservesLegacyAliases(t *testing.T) {
	original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "entity/alpha", Title: "Alpha", PageType: types.WikiPageTypeEntity,
		Status: types.WikiPageStatusPublished, Content: "Old source", SourceRefs: types.StringArray{"old-doc"}, Aliases: types.StringArray{"Legacy alias"}}
	pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
	source := wikiEvidenceChunk()
	source.Content = "Alpha has new documented details."
	svc := &wikiIngestService{wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{source}}}
	update := wikiEvidenceAddition()
	update.Type, update.Slug, update.Item.Slug = types.WikiPageTypeEntity, original.Slug, original.Slug
	update.Item.Name = " Alpha "
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		var prompt strings.Builder
		for _, message := range messages {
			prompt.WriteString(message.Content)
		}
		require.Contains(t, prompt.String(), source.Content)
		return &types.ChatResponse{Content: "SUMMARY: Updated documented information\nAlpha has new documented details."}, nil
	}}
	batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
	changed, _, failed, err := svc.reduceSlugUpdatesUncheckpointed(context.Background(), model, "kb", original.Slug, []SlugUpdate{update}, 7, batch, nil)
	require.NoError(t, err)
	require.True(t, changed)
	require.False(t, failed)
	require.Equal(t, 1, pages.writes)
	require.Equal(t, "Alpha", pages.pages[original.Slug].Title)
	require.Equal(t, types.StringArray{"Legacy alias"}, pages.pages[original.Slug].Aliases)
	require.ElementsMatch(t, []string{"old-doc", "doc"}, []string(pages.pages[original.Slug].SourceRefs))
	require.Equal(t, "Old source", original.Content)
}

func TestWikiSourceIdentityFailedWindowPreservesPriorCheckpointAndCoverage(t *testing.T) {
	for _, conflict := range []bool{false, true} {
		t.Run(fmt.Sprintf("cross-window-conflict-%t", conflict), func(t *testing.T) {
			svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
			scope := &wikiCheckpointScope{7, "kb", "doc", "identity-window"}
			source := wikiFullReadTestSource(2)
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
				calls++
				if strings.Contains(messages[0].Content, wikiIdentityRepairInstruction) {
					return &types.ChatResponse{Content: `{"corrections":[{"kind":"entity","index":0,"name":"Inferred expansion","source_quote":"Inferred expansion"}]}`}, nil
				}
				if !strings.Contains(messages[0].Content, "WINDOW_000") && !strings.Contains(messages[0].Content, "WINDOW_001") {
					return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil
				}
				name := "WINDOW_000"
				if strings.Contains(messages[0].Content, "WINDOW_001") {
					name = "Inferred expansion"
					if conflict {
						name = "WINDOW_001"
					}
				}
				return &types.ChatResponse{Content: fmt.Sprintf(`{"entities":[{"name":%q,"slug":"entity/shared","aliases":[]}],"concepts":[]}`, name)}, nil
			}}
			entities, concepts, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", source, "English", nil, &WikiBatchContext{}, scope)
			if conflict {
				require.NoError(t, err)
				require.Len(t, entities, 2)
				require.NoError(t, validateWikiIdentityCollisions(combinedExtraction{Entities: entities}))
			} else {
				requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source")
				require.Nil(t, entities)
				require.Nil(t, concepts)
			}
			require.Equal(t, 4, calls, "valid parents read both halves; a rejected half gets one identity patch then stops")
			windows := splitWikiReadWindows(source, maxContentForWiki)
			for i, window := range windows {
				var cached combinedExtraction
				hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_candidate", wikiCheckpointHash([]any{wikiFullReadVersion, window.Text}), &cached)
				require.NoError(t, err)
				require.Equal(t, i == 0 || conflict, hit, "only source-validated fresh windows may be cached")
				if hit {
					require.Equal(t, fmt.Sprintf("WINDOW_%03d", i), cached.Entities[0].Name)
					require.Equal(t, "entity/shared", cached.Entities[0].Slug, "raw routing proposals are retained")
				}
			}
			var progress wikiReadProgress
			hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress)
			require.NoError(t, err)
			require.True(t, hit)
			wantWindows := 1
			if conflict {
				wantWindows = 2
			}
			require.Equal(t, wantWindows, progress.CompletedWindows)
			require.Equal(t, maxContentForWiki*wantWindows, progress.ProcessedRunes)
			require.Equal(t, conflict, progress.StageComplete)
			require.Equal(t, conflict, progress.SourceWindowsComplete)
		})
	}
}