xinmi/internal/application/service/wiki_source_identity_repair_test.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/agent"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/require"
)

func TestWikiIdentityRepairPatchesOnlyRejectedNamesUnderExistingBudget(t *testing.T) {
	guard := &candidateRepairTestGuard{}
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	ctx = requestbudget.WithGuard(ctx, guard)
	deadline, _ := ctx.Deadline()
	source := "Alpha and LH are separate. FSH has a different role."
	data := map[string]string{"Content": source, "Language": "English", "PreviousSlugs": "entity/existing"}
	original := combinedExtraction{
		Entities: []extractedItem{
			{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{}, Description: "Valid description", Details: "Keep uncertainty", SourceChunks: []string{"original-1"}},
			{Name: "Luteinizing hormone", Slug: "entity/lh", Description: "Original LH description", Details: "Do not edit this", SourceChunks: []string{"original-2"}},
			{Name: "Expanded FSH", Slug: "entity/fsh", Description: "Original FSH description", Details: "Keep conditions"},
		},
		Concepts: []extractedItem{{Name: "Separate roles", Slug: "concept/roles", Description: "Concept need not be literal"}},
	}
	encoded, err := json.Marshal(original)
	require.NoError(t, err)
	before := string(encoded)
	calls := 0
	var initialPrompt string
	model := &wikiScriptedChat{id: t.Name(), call: func(callCtx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		require.Same(t, guard, requestbudget.FromContext(callCtx))
		require.Equal(t, "wiki_candidate_slug", requestbudget.Phase(callCtx))
		actualDeadline, ok := callCtx.Deadline()
		require.True(t, ok)
		require.Equal(t, deadline, actualDeadline)
		require.Contains(t, messages[0].Content, source)
		if calls == 1 {
			initialPrompt = strings.TrimSuffix(messages[0].Content, "\n\n"+wikiCandidateGenerationInstruction)
			return &types.ChatResponse{Content: before}, nil
		}
		require.Equal(t, 2, calls)
		require.True(t, strings.HasPrefix(messages[0].Content, initialPrompt+"\n\n"+wikiIdentityRepairInstruction))
		require.Contains(t, messages[0].Content, wikiSourceIdentityRepairVersion)
		require.Contains(t, messages[0].Content, `"index":1`)
		require.Contains(t, messages[0].Content, `"index":2`)
		require.NotContains(t, messages[0].Content, "Valid description")
		return &types.ChatResponse{Content: `{"corrections":[{"kind":"entity","index":2,"name":"FSH","source_quote":"FSH has a different role."},{"kind":"entity","index":1,"name":"LH","source_quote":"Alpha and LH are separate."}]}`}, nil
	}}
	got, err := (&wikiIngestService{}).generateWikiFreshExtraction(ctx, model, agent.WikiCandidateSlugPrompt, data)
	require.NoError(t, err)
	require.Equal(t, 2, calls)
	require.Equal(t, "LH", got.Entities[1].Name)
	require.Equal(t, "FSH", got.Entities[2].Name)
	got.Entities[1].Name, got.Entities[2].Name = original.Entities[1].Name, original.Entities[2].Name
	require.Equal(t, original, got, "only two rejected names may change; no metadata or candidate can disappear")
	after, err := json.Marshal(original)
	require.NoError(t, err)
	require.Equal(t, before, string(after))
	require.Equal(t, source, data["Content"])
	require.GreaterOrEqual(t, guard.checks, 4)
}

func TestWikiIdentityRepairRejectsPartialOrUnsupportedPatchesWithoutMutation(t *testing.T) {
	source := "LH and FSH are separate from Alpha."
	initial := combinedExtraction{Entities: []extractedItem{
		{Name: "Expanded LH", Slug: "entity/lh", Description: "Retain description", SourceChunks: []string{"chunk"}},
		{Name: "Expanded FSH", Slug: "entity/fsh"},
		{Name: "Alpha", Slug: "entity/alpha"},
	}, Concepts: []extractedItem{}}
	valid := `{"corrections":[{"kind":"entity","index":0,"name":"LH","source_quote":"LH and FSH"},{"kind":"entity","index":1,"name":"FSH","source_quote":"LH and FSH"}]}`
	for name, test := range map[string]struct{ raw, code string }{
		"whole extraction":       {`{"entities":[],"concepts":[]}`, "json_decode_failed"},
		"missing correction":     {`{"corrections":[{"kind":"entity","index":0,"name":"LH","source_quote":"LH and FSH"}]}`, "correction_count_mismatch"},
		"empty corrections":      {`{"corrections":[]}`, "correction_count_mismatch"},
		"duplicate position":     {strings.Replace(valid, `"index":1`, `"index":0`, 1), "duplicate_position"},
		"valid candidate edited": {strings.Replace(valid, `"index":1`, `"index":2`, 1), "unknown_position"},
		"new candidate":          {strings.Replace(valid, `"index":1`, `"index":3`, 1), "unknown_position"},
		"missing index":          {strings.Replace(valid, `"index":0,`, "", 1), "missing_index"},
		"negative index":         {strings.Replace(valid, `"index":0`, `"index":-1`, 1), "unknown_position"},
		"wrong kind":             {strings.Replace(valid, `"kind":"entity"`, `"kind":"concept"`, 1), "unknown_position"},
		"metadata field":         {strings.Replace(valid, `"name":"LH"`, `"name":"LH","description":"rewritten"`, 1), "json_decode_failed"},
		"slug field":             {strings.Replace(valid, `"name":"LH"`, `"name":"LH","slug":"entity/new"`, 1), "json_decode_failed"},
		"aliases field":          {strings.Replace(valid, `"name":"LH"`, `"name":"LH","aliases":[]`, 1), "json_decode_failed"},
		"empty name":             {strings.Replace(valid, `"name":"LH"`, `"name":""`, 1), "empty_name"},
		"name whitespace":        {strings.Replace(valid, `"name":"LH"`, `"name":" LH"`, 1), "name_outer_whitespace"},
		"missing quote":          {strings.Replace(valid, `,"source_quote":"LH and FSH"`, "", 1), "empty_source_quote"},
		"fabricated quote":       {strings.Replace(valid, `"source_quote":"LH and FSH"`, `"source_quote":"LH and different text"`, 1), "quote_not_in_source"},
		"name absent from quote": {strings.Replace(valid, `"source_quote":"LH and FSH"`, `"source_quote":"FSH are separate"`, 1), "name_not_in_quote"},
		"expanded name":          {strings.Replace(valid, `"name":"LH"`, `"name":"Expanded LH"`, 1), "name_not_in_quote"},
		"case changed":           {strings.Replace(valid, `"name":"LH"`, `"name":"lh"`, 1), "name_not_in_quote"},
		"trailing output":        {valid + " {}", "trailing_json"},
		"unfinished json":        {strings.TrimSuffix(valid, "}"), "json_decode_failed"},
		"collides with valid":    {strings.Replace(strings.Replace(valid, `"name":"LH"`, `"name":"Alpha"`, 1), `"source_quote":"LH and FSH"`, `"source_quote":"Alpha"`, 1), "identity_collision"},
		"two repairs collide":    {strings.Replace(valid, `"name":"FSH"`, `"name":"LH"`, 1), "identity_collision"},
	} {
		t.Run(name, func(t *testing.T) {
			before, err := json.Marshal(initial)
			require.NoError(t, err)
			got, err := applyWikiSourceIdentityCorrections(initial, source, wikiRejectedSourceIdentities(initial, source), test.raw)
			var diagnostic *wikiIdentityRepairError
			require.ErrorAs(t, err, &diagnostic)
			require.Equal(t, test.code, diagnostic.Code)
			require.Empty(t, got.Entities)
			after, err := json.Marshal(initial)
			require.NoError(t, err)
			require.Equal(t, before, after)
		})
	}
}

func TestWikiIdentityRepairNeverDeletesUnsupportedAliases(t *testing.T) {
	for _, kind := range []string{"entity", "concept"} {
		for _, aliases := range [][]string{{"Alpha"}, {"LH"}, {""}} {
			result := combinedExtraction{Entities: []extractedItem{}, Concepts: []extractedItem{}}
			item := extractedItem{Name: "LH", Slug: kind + "/lh", Aliases: aliases}
			if kind == "entity" {
				result.Entities = append(result.Entities, item)
			} else {
				result.Concepts = append(result.Concepts, item)
			}
			before, err := json.Marshal(result)
			require.NoError(t, err)
			patch := `{"corrections":[{"kind":"` + kind + `","index":0,"name":"LH","source_quote":"LH and Alpha"}]}`
			got, err := applyWikiSourceIdentityCorrections(result, "LH and Alpha", wikiRejectedSourceIdentities(result, "LH and Alpha"), patch)
			var diagnostic *wikiIdentityRepairError
			require.ErrorAs(t, err, &diagnostic)
			require.Equal(t, "unsupported_aliases", diagnostic.Code)
			require.Equal(t, kind, diagnostic.Kind)
			require.NotNil(t, diagnostic.Index)
			require.Zero(t, *diagnostic.Index)
			require.Empty(t, got.Entities)
			require.Empty(t, got.Concepts)
			after, err := json.Marshal(result)
			require.NoError(t, err)
			require.Equal(t, before, after)
		}
	}
}

func TestWikiIdentityRepairDiagnosticPositionsAndRedaction(t *testing.T) {
	initial := combinedExtraction{Entities: []extractedItem{{Name: "Private expanded LH", Slug: "entity/lh"}}}
	source := "Private source LH"
	valid := `{"corrections":[{"kind":"entity","index":0,"name":"LH","source_quote":"LH"}]}`
	for name, test := range map[string]struct {
		raw, code, kind string
		index           *int
	}{
		"unknown kind":  {strings.Replace(valid, `"kind":"entity"`, `"kind":"PRIVATE_MODEL_KIND"`, 1), "unknown_position", "unknown", new(int)},
		"empty name":    {strings.Replace(valid, `"name":"LH"`, `"name":""`, 1), "empty_name", "entity", new(int)},
		"schema":        {strings.Replace(valid, `"name":"LH"`, `"PRIVATE_MODEL_FIELD":"PRIVATE_MODEL_VALUE"`, 1), "json_decode_failed", "", nil},
		"missing index": {strings.Replace(valid, `"index":0,`, "", 1), "missing_index", "entity", nil},
	} {
		t.Run(name, func(t *testing.T) {
			_, err := applyWikiSourceIdentityCorrections(initial, source, wikiRejectedSourceIdentities(initial, source), test.raw)
			var diagnostic *wikiIdentityRepairError
			require.ErrorAs(t, err, &diagnostic)
			require.Equal(t, test.code, diagnostic.Code)
			require.Equal(t, test.kind, diagnostic.Kind)
			require.Equal(t, test.index, diagnostic.Index)
			var fields map[string]any
			require.NoError(t, json.Unmarshal([]byte(err.Error()), &fields))
			require.NotContains(t, err.Error(), "PRIVATE_MODEL")
			require.NotContains(t, err.Error(), "Private")
			require.NotContains(t, err.Error(), "LH")
		})
	}
}

func TestWikiIdentityRepairRejectsInconsistentTargetsWithSpecificCodes(t *testing.T) {
	initial := combinedExtraction{Entities: []extractedItem{{Name: "Expanded LH", Slug: "entity/lh"}}}
	valid := `{"corrections":[{"kind":"entity","index":0,"name":"LH","source_quote":"LH"}]}`
	for name, test := range map[string]struct {
		result combinedExtraction
		target wikiIdentityRepairTarget
		raw    string
		code   string
	}{
		"unsupported error":          {initial, wikiIdentityRepairTarget{Kind: "entity", Index: 0, Code: "wiki_identity_slug_conflict"}, valid, "unsupported_identity_error"},
		"target aliases":             {initial, wikiIdentityRepairTarget{Kind: "entity", Index: 0, Code: "wiki_identity_name_not_in_source", Candidate: extractedItem{Aliases: []string{"Private alias"}}}, valid, "unsupported_aliases"},
		"unsupported kind":           {initial, wikiIdentityRepairTarget{Kind: "concept", Index: 0, Code: "wiki_identity_name_not_in_source"}, strings.Replace(valid, `"kind":"entity"`, `"kind":"concept"`, 1), "unsupported_kind"},
		"out of range":               {initial, wikiIdentityRepairTarget{Kind: "entity", Index: 1, Code: "wiki_identity_name_not_in_source"}, strings.Replace(valid, `"index":0`, `"index":1`, 1), "candidate_index_out_of_range"},
		"unlisted invalid candidate": {combinedExtraction{Entities: []extractedItem{{Name: "Expanded LH", Slug: "entity/lh"}, {Name: "Private unsupported candidate", Slug: "entity/other"}}}, wikiIdentityRepairTarget{Kind: "entity", Index: 0, Code: "wiki_identity_name_not_in_source"}, valid, "source_revalidation_failed"},
	} {
		t.Run(name, func(t *testing.T) {
			before, err := json.Marshal(test.result)
			require.NoError(t, err)
			got, err := applyWikiSourceIdentityCorrections(test.result, "LH", []wikiIdentityRepairTarget{test.target}, test.raw)
			var diagnostic *wikiIdentityRepairError
			require.ErrorAs(t, err, &diagnostic)
			require.Equal(t, test.code, diagnostic.Code)
			if name == "unlisted invalid candidate" {
				require.Equal(t, "entity", diagnostic.Kind)
				require.NotNil(t, diagnostic.Index)
				require.Equal(t, 1, *diagnostic.Index)
			}
			require.Empty(t, got.Entities)
			after, err := json.Marshal(test.result)
			require.NoError(t, err)
			require.Equal(t, before, after)
		})
	}
	_, err := applyWikiSourceIdentityCorrections(initial, "LH", nil, `{"corrections":[]}`)
	var diagnostic *wikiIdentityRepairError
	require.ErrorAs(t, err, &diagnostic)
	require.Equal(t, "no_rejected_candidates", diagnostic.Code)
}

func TestWikiIdentityRepairLogsOnlyBoundedDiagnostics(t *testing.T) {
	const privateText = "PRIVATE_PROVIDER_OR_SOURCE_BODY"
	for name, test := range map[string]struct {
		response    *types.ChatResponse
		failure     error
		stage, code string
	}{
		"accepted":     {&types.ChatResponse{Content: `{"corrections":[{"kind":"entity","index":0,"name":"LH","source_quote":"LH"}]}`}, nil, "none", "none"},
		"validation":   {&types.ChatResponse{Content: `{"corrections":[{"kind":"` + privateText + `","index":0,"name":"LH","source_quote":"LH"}]}`}, nil, "validation", "unknown_position"},
		"malformed":    {&types.ChatResponse{Content: privateText}, nil, "validation", "json_decode_failed"},
		"truncated":    {&types.ChatResponse{Content: privateText, FinishReason: "length"}, nil, "output", "model_output_incomplete"},
		"empty output": {&types.ChatResponse{Content: ""}, nil, "output", "model_output_invalid"},
		"provider":     {nil, errors.New(privateText), "transport", "provider_error"},
		"timeout":      {nil, context.DeadlineExceeded, "control", "deadline_exceeded"},
		"cancelled":    {nil, context.Canceled, "control", "cancelled"},
		"budget":       {nil, requestbudget.ErrStopped, "control", "budget_stopped"},
	} {
		t.Run(name, func(t *testing.T) {
			var log bytes.Buffer
			localLogger := logrus.New()
			localLogger.SetOutput(&log)
			ctx := context.WithValue(context.Background(), types.LoggerContextKey, logrus.NewEntry(localLogger))
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				calls++
				return test.response, test.failure
			}}
			source := "LH " + privateText
			_, err := (&wikiIngestService{}).repairWikiFreshSourceIdentities(ctx, model, "{{.Content}}", map[string]string{"Content": source}, combinedExtraction{Entities: []extractedItem{{Name: "Expanded LH", Slug: "entity/lh"}}})
			require.Equal(t, 1, calls)
			stage, code, _, _ := wikiIdentityRepairFailureDetails(err)
			require.Equal(t, test.stage, stage)
			require.Equal(t, test.code, code)
			if test.failure != nil {
				require.ErrorIs(t, err, test.failure)
			}
			require.Contains(t, log.String(), "stage="+test.stage+" code="+test.code)
			require.Contains(t, log.String(), "source_sha256="+wikiFailureTextHash(source))
			require.Contains(t, log.String(), "rejected=1 entities=1 concepts=0")
			require.NotContains(t, log.String(), privateText)
			require.NotContains(t, log.String(), "Expanded LH")
			if name == "validation" {
				require.Contains(t, log.String(), "kind=unknown index=0")
			}
		})
	}
}

func TestWikiIdentityRepairFailureKeepsOriginalErrorAndNeverRetriesProvider(t *testing.T) {
	for _, failure := range []string{"timeout", "provider", "truncated", "invalid", "cancel", "budget"} {
		t.Run(failure, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			guard := &candidateRepairTestGuard{}
			ctx = requestbudget.WithGuard(ctx, guard)
			calls := 0
			source := strings.Repeat("LH is recorded. ", 1000)
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
				calls++
				if calls == 1 {
					if failure == "cancel" {
						cancel()
					} else if failure == "budget" {
						guard.stopped = true
					}
					return &types.ChatResponse{Content: `{"entities":[{"name":"Unsupported expansion","slug":"entity/lh"}],"concepts":[]}`}, nil
				}
				require.Equal(t, 2, calls, "no provider retry, whole output regeneration or splitting after correction")
				require.Contains(t, messages[0].Content, wikiIdentityRepairInstruction)
				switch failure {
				case "timeout":
					return nil, context.DeadlineExceeded
				case "provider":
					return nil, errors.New("API request failed with status 503: private provider detail")
				case "truncated":
					return &types.ChatResponse{Content: `{"corrections":[`, FinishReason: "length"}, nil
				default:
					return &types.ChatResponse{Content: `{"corrections":[]}`}, nil
				}
			}}
			got, err := (&wikiIngestService{}).generateWikiFreshExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": source})
			identityErr := requireWikiSourceIdentityError(t, err, "wiki_identity_name_not_in_source")
			require.Equal(t, wikiCheckpointHash([]string{source}), identityErr.SourceHash)
			require.Equal(t, "entity", identityErr.Kind)
			require.Zero(t, identityErr.Index)
			require.NotContains(t, err.Error(), "private provider detail")
			require.Empty(t, got.Entities)
			switch failure {
			case "timeout":
				require.ErrorIs(t, err, context.DeadlineExceeded)
			case "cancel":
				require.ErrorIs(t, err, context.Canceled)
			case "budget":
				require.ErrorIs(t, err, requestbudget.ErrStopped)
			}
			wantCalls := 2
			if failure == "cancel" || failure == "budget" {
				wantCalls = 1
			}
			require.Equal(t, wantCalls, calls)
		})
	}
}

func TestWikiIdentityRepairCompletesWindowAndPreservesEarlierCheckpoint(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "identity-repair"}
	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, "WINDOW_000") {
			return &types.ChatResponse{Content: `{"entities":[{"name":"WINDOW_000","slug":"entity/first"}],"concepts":[]}`}, nil
		}
		if !strings.Contains(messages[0].Content, "WINDOW_001") {
			return &types.ChatResponse{Content: `{"entities":[],"concepts":[]}`}, nil
		}
		if !strings.Contains(messages[0].Content, wikiIdentityRepairInstruction) {
			return &types.ChatResponse{Content: `{"entities":[{"name":"Expanded WINDOW_001","slug":"entity/second","description":"Do not replace metadata"}],"concepts":[]}`}, nil
		}
		require.Equal(t, 4, calls)
		require.NotContains(t, messages[0].Content, "WINDOW_000", "repair uses exactly the failed source window")
		return &types.ChatResponse{Content: `{"corrections":[{"kind":"entity","index":0,"name":"WINDOW_001","source_quote":"WINDOW_001"}]}`}, nil
	}}
	got, concepts, _, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", source, "English", nil, &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Len(t, got, 2)
	require.Empty(t, concepts)
	require.Equal(t, 5, calls)
	for index, window := range splitWikiReadWindows(source, maxContentForWiki) {
		var cached combinedExtraction
		hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_candidate", wikiCheckpointHash([]any{wikiFullReadVersion, window.Text}), &cached)
		require.NoError(t, err)
		require.True(t, hit)
		require.Len(t, cached.Entities, 1)
		if index == 1 {
			require.Equal(t, "WINDOW_001", cached.Entities[0].Name)
			require.Equal(t, "Do not replace metadata", cached.Entities[0].Description)
		}
	}
	var progress wikiReadProgress
	hit, err := svc.loadWikiCheckpoint(context.Background(), scope, "read_progress", "candidates", &progress)
	require.NoError(t, err)
	require.True(t, hit)
	require.Equal(t, 2, progress.CompletedWindows)
	require.Equal(t, maxContentForWiki*2, progress.ProcessedRunes)
	require.True(t, progress.StageComplete)
	require.True(t, progress.SourceWindowsComplete)
}