370 lines
20 KiB
Go
370 lines
20 KiB
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)
|
||
|
|
}
|