190 lines
8.0 KiB
Go
190 lines
8.0 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"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/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestWikiCandidateRepairClassifiesStrictOutputWithoutModelValues(t *testing.T) {
|
||
|
|
for _, fixture := range []struct{ content, code, field string }{
|
||
|
|
{`{"entities":[`, "invalid_json_syntax", ""},
|
||
|
|
{`{"entities":"PRIVATE_VALUE","concepts":[]}`, "json_type_mismatch", ""},
|
||
|
|
{`null`, "missing_entities_array", "entities"},
|
||
|
|
{`{"concepts":[]}`, "missing_entities_array", "entities"},
|
||
|
|
{`{"entities":[],"concepts":null}`, "missing_concepts_array", "concepts"},
|
||
|
|
{`{"entities":[{"name":" ","slug":"entity/PRIVATE_VALUE"}],"concepts":[]}`, "candidate_name_empty", "entities.name"},
|
||
|
|
{`{"entities":[],"concepts":[{"name":"PRIVATE_VALUE","slug":"concept/ "}]}`, "candidate_slug_invalid", "concepts.slug"},
|
||
|
|
{`{"entities":[{"name":"PRIVATE_VALUE","slug":"concept/private"}],"concepts":[]}`, "candidate_slug_invalid", "entities.slug"},
|
||
|
|
{`{"entities":[{"name":"PRIVATE_VALUE","slug":"entity/private","aliases":17}],"concepts":[]}`, "json_type_mismatch", ""},
|
||
|
|
} {
|
||
|
|
t.Run(fixture.code+fixture.field, func(t *testing.T) {
|
||
|
|
result, validation := parseWikiCandidateOutput(fixture.content)
|
||
|
|
require.NotNil(t, validation)
|
||
|
|
require.Empty(t, result.Entities)
|
||
|
|
require.Empty(t, result.Concepts)
|
||
|
|
require.Equal(t, fixture.code, validation.Code)
|
||
|
|
require.Equal(t, fixture.field, validation.Field)
|
||
|
|
require.NotContains(t, validation.String(), "PRIVATE_VALUE")
|
||
|
|
require.NotContains(t, wikiCandidateCorrectionInstruction(*validation), "PRIVATE_VALUE")
|
||
|
|
if fixture.code == "invalid_json_syntax" || fixture.code == "json_type_mismatch" {
|
||
|
|
require.Positive(t, validation.JSONOffset)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateRepairUsesExactSourceSharedContextAndOneFreshRequest(t *testing.T) {
|
||
|
|
guard := &candidateRepairTestGuard{}
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
||
|
|
defer cancel()
|
||
|
|
ctx = requestbudget.WithGuard(context.WithValue(ctx, types.TenantIDContextKey, uint64(71)), guard)
|
||
|
|
deadline, _ := ctx.Deadline()
|
||
|
|
source := "Valid\n" + strings.Repeat("original-source\n", 500)
|
||
|
|
data := map[string]string{"Content": source, "Language": "English", "PreviousSlugs": "entity/known"}
|
||
|
|
wantData := make(map[string]string)
|
||
|
|
for key, value := range data {
|
||
|
|
wantData[key] = value
|
||
|
|
}
|
||
|
|
calls := 0
|
||
|
|
var first 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))
|
||
|
|
actualDeadline, ok := callCtx.Deadline()
|
||
|
|
require.True(t, ok)
|
||
|
|
require.Equal(t, deadline, actualDeadline)
|
||
|
|
require.NoError(t, callCtx.Err())
|
||
|
|
require.Len(t, messages, 1)
|
||
|
|
if calls == 1 {
|
||
|
|
first = messages[0].Content
|
||
|
|
require.Contains(t, first, source)
|
||
|
|
return &types.ChatResponse{Content: `{"entities":"PRIVATE_REJECTED_OUTPUT","concepts":[]}`, FinishReason: "stop"}, nil
|
||
|
|
}
|
||
|
|
require.Equal(t, 2, calls)
|
||
|
|
require.True(t, strings.HasPrefix(messages[0].Content, first+"\n\n"+wikiCandidateRepairInstruction))
|
||
|
|
require.NotContains(t, messages[0].Content, "PRIVATE_REJECTED_OUTPUT")
|
||
|
|
require.Contains(t, messages[0].Content, `"code":"json_type_mismatch"`)
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, agent.WikiCandidateSlugPrompt, data)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 2, calls)
|
||
|
|
require.Len(t, result.Entities, 1)
|
||
|
|
require.Equal(t, wantData, data)
|
||
|
|
require.GreaterOrEqual(t, guard.checks, 4)
|
||
|
|
}
|
||
|
|
|
||
|
|
type candidateRepairTestGuard struct {
|
||
|
|
checks int
|
||
|
|
stopped bool
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *candidateRepairTestGuard) Key() string { return "candidate-repair-offline" }
|
||
|
|
func (g *candidateRepairTestGuard) Check(context.Context) error {
|
||
|
|
g.checks++
|
||
|
|
if g.stopped {
|
||
|
|
return requestbudget.ErrStopped
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
func (g *candidateRepairTestGuard) Reserve(ctx context.Context, _, _ string) error {
|
||
|
|
return g.Check(ctx)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateRepairCannotBypassCancellationOrTransportBudget(t *testing.T) {
|
||
|
|
for _, stop := range []string{"cancel", "budget"} {
|
||
|
|
t.Run(stop, func(t *testing.T) {
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
defer cancel()
|
||
|
|
guard := &candidateRepairTestGuard{}
|
||
|
|
ctx = requestbudget.WithGuard(ctx, guard)
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if stop == "cancel" {
|
||
|
|
cancel()
|
||
|
|
} else {
|
||
|
|
guard.stopped = true
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[],"concepts":null}`, FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Equal(t, 1, calls)
|
||
|
|
require.Empty(t, result.Entities)
|
||
|
|
if stop == "cancel" {
|
||
|
|
require.ErrorIs(t, err, context.Canceled)
|
||
|
|
} else {
|
||
|
|
require.ErrorIs(t, err, requestbudget.ErrStopped)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateRepairProviderFailuresNeverStartAnotherCorrection(t *testing.T) {
|
||
|
|
for _, when := range []string{"initial", "correction"} {
|
||
|
|
for _, code := range []string{"401", "403", "429", "503"} {
|
||
|
|
t.Run(when+code, func(t *testing.T) {
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(72))
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if when == "correction" && calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"concepts":[]}`, FinishReason: "stop"}, nil
|
||
|
|
}
|
||
|
|
return nil, errors.New("API request failed with status " + code + ": PRIVATE_PROVIDER_DETAIL")
|
||
|
|
}}
|
||
|
|
key := wikiProviderCircuitKey(72, model.id)
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
|
||
|
|
result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
|
||
|
|
require.Error(t, err)
|
||
|
|
want := 1
|
||
|
|
if when == "correction" {
|
||
|
|
want = 2
|
||
|
|
} else if code == "503" {
|
||
|
|
want = wikiLLMMaxAttempts
|
||
|
|
}
|
||
|
|
require.Equal(t, want, calls)
|
||
|
|
require.Empty(t, result.Entities)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateRepairPartialFailureCannotAdvanceWindowCheckpoint(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "candidate-correction"}
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
if calls == 2 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[{"name":"WINDOW_000","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"PRIVATE_VALUE","slug":"entity/private"}]}`, FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
entities, concepts, items, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", wikiFullReadTestSource(2), "English", nil, &WikiBatchContext{}, scope)
|
||
|
|
require.ErrorContains(t, err, "partition 2/2")
|
||
|
|
require.ErrorContains(t, err, "candidate_slug_invalid")
|
||
|
|
require.NotContains(t, err.Error(), "PRIVATE_VALUE")
|
||
|
|
require.Equal(t, 4, calls)
|
||
|
|
require.Nil(t, entities)
|
||
|
|
require.Nil(t, concepts)
|
||
|
|
require.Nil(t, items)
|
||
|
|
var count int64
|
||
|
|
require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage IN ?", []string{"read_candidate", "read_progress", "candidates", "map", "page"}).Count(&count).Error)
|
||
|
|
require.Zero(t, count, "partial successful leaf must not become a completed window or progress")
|
||
|
|
}
|