428 lines
21 KiB
Go
428 lines
21 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/xml"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"strings"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
type wikiScriptedChat struct {
|
||
|
|
templateCaptureChatModel
|
||
|
|
id string
|
||
|
|
call func(context.Context, []chat.Message) (*types.ChatResponse, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationTruncationRetriesTwoSmallerBatchesWithLocalHandles(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var inputs []string
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
prompt := messages[0].Content
|
||
|
|
inputs = append(inputs, prompt)
|
||
|
|
if len(inputs) == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{}}`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
if strings.Contains(prompt, `<c id="c001"`) {
|
||
|
|
t.Fatal("retry did not rebuild the local handle table")
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[{"type":"concept","slug":"concept/new","name":"New","source_chunks":["c000"]}]}`}, nil
|
||
|
|
}}
|
||
|
|
chunks := []*types.Chunk{{ID: "first", Content: strings.Repeat("first", 600), ChunkIndex: 0}, {ID: "second", Content: strings.Repeat("second", 500), ChunkIndex: 1}}
|
||
|
|
citations, newSlugs, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err != nil || completed != 1 || len(inputs) != 3 || strings.Join(citations["entity/alpha"], ",") != "first,second" {
|
||
|
|
t.Fatalf("failed bounded citation recovery: completed=%d calls=%d citations=%v err=%v", completed, len(inputs), citations, err)
|
||
|
|
}
|
||
|
|
if len(newSlugs) != 2 || strings.Join(newSlugs[0].SourceChunks, ",") != "first" || strings.Join(newSlugs[1].SourceChunks, ",") != "second" {
|
||
|
|
t.Fatalf("new slug provenance lost: %+v", newSlugs)
|
||
|
|
}
|
||
|
|
if !strings.Contains(inputs[1], chunks[0].Content) || strings.Contains(inputs[1], chunks[1].Content) || !strings.Contains(inputs[2], chunks[1].Content) || strings.Contains(inputs[2], chunks[0].Content) {
|
||
|
|
t.Fatal("retry did not partition all source chunks")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationOversizedChunkRetryPreservesEveryRuneAndSourceID(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
content := strings.Repeat("前", 1200) + "\n" + strings.Repeat("后", 1200)
|
||
|
|
var seen []string
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
prompt := messages[0].Content
|
||
|
|
start, end := strings.Index(prompt, "<chunks>\n"), strings.Index(prompt, "\n</chunks>")
|
||
|
|
var rendered struct {
|
||
|
|
Chunks []struct {
|
||
|
|
Content string `xml:",chardata"`
|
||
|
|
} `xml:"c"`
|
||
|
|
}
|
||
|
|
if start < 0 || end < 0 {
|
||
|
|
t.Fatal("missing chunks block")
|
||
|
|
}
|
||
|
|
if err := xml.Unmarshal([]byte(prompt[start:end+len("\n</chunks>")]), &rendered); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(rendered.Chunks) != 1 {
|
||
|
|
t.Fatalf("unexpected retry chunks: %+v", rendered)
|
||
|
|
}
|
||
|
|
seen = append(seen, strings.TrimSuffix(strings.TrimPrefix(rendered.Chunks[0].Content, "\n"), "\n"))
|
||
|
|
if len(seen) == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":`, FinishReason: "max_tokens"}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
|
||
|
|
}}
|
||
|
|
chunk := &types.Chunk{ID: "source", Content: content}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{chunk}, "Chinese", &WikiBatchContext{})
|
||
|
|
if err != nil || completed != 1 || len(seen) != 3 || seen[1]+seen[2] != content || chunk.Content != content || strings.Join(citations["entity/alpha"], ",") != "source" {
|
||
|
|
t.Fatalf("singleton split lost source or duplicated citation: calls=%d result=%v err=%v", len(seen), citations, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationPartitionFailureCannotReturnPartialCoverage(t *testing.T) {
|
||
|
|
for _, failure := range []struct{ name, output, reason string }{
|
||
|
|
{"truncated", `{"citations":{}}`, "max_output_tokens"},
|
||
|
|
{"malformed", `{"citations":`, "stop"},
|
||
|
|
{"missing-schema", `{}`, "stop"},
|
||
|
|
{"unknown-handle", `{"citations":{"entity/alpha":["c999"]}}`, "stop"},
|
||
|
|
} {
|
||
|
|
t.Run(failure.name, func(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
if calls == 2 {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: failure.output, FinishReason: failure.reason}, nil
|
||
|
|
}}
|
||
|
|
chunks := []*types.Chunk{{ID: "first", Content: "first", ChunkIndex: 0}, {ID: "second", Content: "second", ChunkIndex: 1}}
|
||
|
|
citations, slugs, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err == nil || calls != 3 || completed != 0 || citations != nil || slugs != nil || !strings.Contains(err.Error(), "partition 2/2") {
|
||
|
|
t.Fatalf("partial citation accepted or retry budget exceeded: calls=%d completed=%d err=%v", calls, completed, err)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationDoesNotPartitionProviderFailures(t *testing.T) {
|
||
|
|
for _, status := range []int{401, 402, 403, 429} {
|
||
|
|
t.Run(fmt.Sprint(status), func(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(77))
|
||
|
|
var calls atomic.Int32
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
return nil, fmt.Errorf("API request failed with status %d: private text", status)
|
||
|
|
}}
|
||
|
|
key := wikiProviderCircuitKey(77, model.id)
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
|
||
|
|
chunks := []*types.Chunk{{ID: "first", Content: "first"}, {ID: "second", Content: "second"}}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(ctx, model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err == nil || calls.Load() != 1 || citations != nil || completed != 0 || strings.Contains(err.Error(), "private text") {
|
||
|
|
t.Fatalf("provider failure split or leaked: calls=%d err=%v", calls.Load(), err)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationCancellationStopsBeforeRetryPartition(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
defer cancel()
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
cancel()
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
|
||
|
|
}}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(ctx, model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "first"}, {ID: "second", Content: "second"}}, "English", &WikiBatchContext{})
|
||
|
|
if err == nil || !errors.Is(err, context.Canceled) || calls != 2 || completed != 0 || citations != nil {
|
||
|
|
t.Fatalf("cancelled partition continued: calls=%d completed=%d err=%v", calls, completed, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationTwelveBatchesCompleteOnlyAfterTruncatedBatchRecovers(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var calls atomic.Int32
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
prompt := messages[0].Content
|
||
|
|
if strings.Contains(prompt, "target-left") && strings.Contains(prompt, "target-right") {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
handles := `["c000"]`
|
||
|
|
if strings.Contains(prompt, `<c id="c001"`) {
|
||
|
|
handles = `["c000","c001"]`
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":` + handles + `},"new_slugs":[]}`}, nil
|
||
|
|
}}
|
||
|
|
var chunks []*types.Chunk
|
||
|
|
var expected []string
|
||
|
|
for i := 0; i < 12; i++ {
|
||
|
|
for j := 0; j < 2; j++ {
|
||
|
|
id := fmt.Sprintf("chunk-%02d-%d", i, j)
|
||
|
|
prefix := id
|
||
|
|
if i == 9 {
|
||
|
|
prefix = []string{"target-left", "target-right"}[j]
|
||
|
|
}
|
||
|
|
chunks = append(chunks, &types.Chunk{ID: id, Content: prefix + strings.Repeat("x", 6000-len(prefix)), ChunkIndex: 2*i + j})
|
||
|
|
expected = append(expected, id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err != nil || calls.Load() != 14 || completed != 12 || strings.Join(citations["entity/alpha"], ",") != strings.Join(expected, ",") {
|
||
|
|
t.Fatalf("batch 10/12 recovery lost coverage: calls=%d completed=%d citations=%v err=%v", calls.Load(), completed, citations, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationRepeatedTruncationStopsAtFirstFailedPartition(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var calls atomic.Int32
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
|
||
|
|
}}
|
||
|
|
chunks := []*types.Chunk{{ID: "first", Content: strings.Repeat("x", 3000)}, {ID: "second", Content: strings.Repeat("y", 3000)}}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err == nil || calls.Load() != 2 || completed != 0 || citations != nil || !strings.Contains(err.Error(), "partition 1/2") {
|
||
|
|
t.Fatalf("persistent truncation exceeded retry bound: calls=%d completed=%d err=%v", calls.Load(), completed, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *wikiScriptedChat) GetModelID() string { return m.id }
|
||
|
|
func (m *wikiScriptedChat) Chat(ctx context.Context, messages []chat.Message, _ *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
return m.call(ctx, messages)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionSplitsTruncatedOutputWithoutDroppingEitherHalf(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var inputs []string
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
inputs = append(inputs, messages[0].Content)
|
||
|
|
if len(inputs) == 1 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
if len(inputs) == 2 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[{"name":"Alpha","slug":"entity/alpha"}],"concepts":[]}`}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"Beta","slug":"concept/beta"}]}`}, nil
|
||
|
|
}}
|
||
|
|
content := "Alpha " + strings.Repeat("a", 1194) + "\n" + strings.Repeat("b", 1200)
|
||
|
|
result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": content})
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(inputs) != 3 || inputs[1]+inputs[2] != content || len(result.Entities) != 1 || len(result.Concepts) != 1 {
|
||
|
|
t.Fatalf("partition recovery lost data: calls=%d entities=%d concepts=%d", len(inputs), len(result.Entities), len(result.Concepts))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionInvalidJSONIsBoundedAndDoesNotReturnPartialOutput(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var calls atomic.Int32
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`}, nil
|
||
|
|
}}
|
||
|
|
result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
|
||
|
|
if err == nil || calls.Load() != 4 || len(result.Entities) != 0 {
|
||
|
|
t.Fatalf("invalid output must fail after first failing half, calls=%d err=%v", calls.Load(), err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionLargeTruncatedHalfGetsOneAdditionalSplit(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var leaves []string
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
content := messages[0].Content
|
||
|
|
if len([]rune(content)) > 8000 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
leaves = append(leaves, content)
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
source := "Valid" + strings.Repeat("甲", 7995) + "Valid" + strings.Repeat("乙", 7995) + "Valid" + strings.Repeat("丙", 7995) + "Valid" + strings.Repeat("丁", 7995)
|
||
|
|
result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": source})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 7, calls)
|
||
|
|
require.Equal(t, source, strings.Join(leaves, ""), "every source character must reach a successful leaf")
|
||
|
|
require.Len(t, leaves, 4)
|
||
|
|
require.Len(t, result.Entities, 1, "identical fresh candidates coalesce only after all leaves succeed")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionAdditionalSplitIsBoundedAndCannotPublishPartialResult(t *testing.T) {
|
||
|
|
for _, lastFailure := range []string{"length", "invalid"} {
|
||
|
|
t.Run(lastFailure, func(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 3 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[{"name":"x","slug":"entity/valid"}],"concepts":[]}`}, nil
|
||
|
|
}
|
||
|
|
if calls == 4 && lastFailure == "invalid" {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
|
||
|
|
}}
|
||
|
|
result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 32000)})
|
||
|
|
require.Error(t, err)
|
||
|
|
expected := 4
|
||
|
|
if lastFailure == "invalid" {
|
||
|
|
expected++
|
||
|
|
}
|
||
|
|
require.Equal(t, expected, calls)
|
||
|
|
require.Empty(t, result.Entities, "a successful first leaf must not publish when its sibling fails")
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionNestedProviderFailureStopsFurtherPartitions(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(77))
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls < 3 {
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
return nil, errors.New("API request failed with status 429: private provider detail")
|
||
|
|
}}
|
||
|
|
key := wikiProviderCircuitKey(77, model.id)
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
|
||
|
|
result, err := s.generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 32000)})
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Equal(t, 3, calls)
|
||
|
|
require.Empty(t, result.Entities)
|
||
|
|
require.NotContains(t, err.Error(), "private provider detail")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionInvalidCandidatesUseOneWholeSourceCorrection(t *testing.T) {
|
||
|
|
for _, invalid := range []string{
|
||
|
|
`{"entities":[{"name":" ","slug":"entity/private"}],"concepts":[]}`,
|
||
|
|
`{"entities":[{"name":"Private","slug":"concept/private"}],"concepts":[]}`,
|
||
|
|
`{"entities":[],"concepts":[{"name":"Private","slug":"concept/ "}]}`,
|
||
|
|
} {
|
||
|
|
t.Run(invalid, func(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
var inputs []string
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
inputs = append(inputs, messages[0].Content)
|
||
|
|
if len(inputs) == 1 {
|
||
|
|
return &types.ChatResponse{Content: invalid}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`}, nil
|
||
|
|
}}
|
||
|
|
source := "Valid " + strings.Repeat("a", 1194) + "\n" + strings.Repeat("b", 1200)
|
||
|
|
result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": source})
|
||
|
|
if err != nil || len(inputs) != 2 || !strings.HasPrefix(inputs[1], source+"\n\n"+wikiCandidateRepairInstruction) || len(result.Entities) != 1 {
|
||
|
|
t.Fatalf("invalid candidate repair lost source or exceeded budget: calls=%d err=%v", len(inputs), err)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionInvalidCandidateInSecondHalfCannotPublishFirstHalf(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
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":"x","slug":"entity/valid"}],"concepts":[]}`}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"private-source-name","slug":"entity/private-source-slug"}]}`}, nil
|
||
|
|
}}
|
||
|
|
result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
|
||
|
|
if err == nil || calls != 4 || len(result.Entities) != 0 || len(result.Concepts) != 0 || !strings.Contains(err.Error(), "partition 2/2") || !strings.Contains(err.Error(), `"field":"concepts.slug","index":0`) || strings.Contains(err.Error(), "private-source") {
|
||
|
|
t.Fatalf("invalid candidate must preserve scope and report only field path: calls=%d err=%v", calls, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiExtractionDoesNotSplitOrImmediatelyRetryRateLimit(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(77))
|
||
|
|
var calls atomic.Int32
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
return nil, errors.New("API request failed with status 429: private text")
|
||
|
|
}}
|
||
|
|
key := wikiProviderCircuitKey(77, model.id)
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
|
||
|
|
_, err := s.generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
|
||
|
|
if err == nil || calls.Load() != 1 || strings.Contains(err.Error(), "private text") {
|
||
|
|
t.Fatalf("rate limit retry storm or raw provider leak: calls=%d err=%v", calls.Load(), err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiRejectsTruncatedEvenSyntacticallyValidOutput(t *testing.T) {
|
||
|
|
for _, reason := range []string{"length", "max_tokens", "max_output_tokens", "incomplete"} {
|
||
|
|
for _, content := range []string{`{"entities":[],"concepts":[]}`, "", " "} {
|
||
|
|
err := validateWikiResponse(&types.ChatResponse{Content: content, FinishReason: reason})
|
||
|
|
var outputErr *wikiOutputError
|
||
|
|
if !errors.As(err, &outputErr) || !outputErr.incomplete {
|
||
|
|
t.Fatalf("lost incomplete response reason %s: %v", reason, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if err := validateWikiResponse(&types.ChatResponse{Content: " ", FinishReason: "stop"}); err == nil {
|
||
|
|
t.Fatal("accepted empty output")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationFailureCannotBecomePartialSuccess(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
if strings.Contains(messages[0].Content, "BROKEN") {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":`}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
|
||
|
|
}}
|
||
|
|
chunks := []*types.Chunk{
|
||
|
|
{ID: "first", Content: strings.Repeat("a", 12000), ChunkIndex: 0},
|
||
|
|
{ID: "second", Content: "BROKEN", ChunkIndex: 1},
|
||
|
|
}
|
||
|
|
citations, slugs, _, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err == nil || citations != nil || slugs != nil {
|
||
|
|
t.Fatalf("partial classification returned as success: %v", err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationValidBatchesTranslateHandlesAndIgnoreNilChunks(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
|
||
|
|
}}
|
||
|
|
chunks := []*types.Chunk{nil, {ID: "first", Content: strings.Repeat("a", 12000), ChunkIndex: 0}, {ID: "second", Content: "b", ChunkIndex: 1}}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
|
||
|
|
if err != nil || completed != 2 || strings.Join(citations["entity/alpha"], ",") != "first,second" {
|
||
|
|
t.Fatalf("bad citation coverage: completed=%d result=%v err=%v", completed, citations, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationRejectsMissingSchemaAndUnknownSourceHandles(t *testing.T) {
|
||
|
|
for _, output := range []string{`null`, `{}`, `{"citations":{"entity/alpha":["c999"]}}`, `{"citations":{},"new_slugs":[{"slug":"concept/new","name":"New","source_chunks":["c999"]}]}`} {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
return &types.ChatResponse{Content: output}, nil
|
||
|
|
}}
|
||
|
|
citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Source text"}}, "English", &WikiBatchContext{})
|
||
|
|
if err == nil || citations != nil || completed != 0 {
|
||
|
|
t.Fatalf("accepted citations without valid source coverage: completed=%d err=%v", completed, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|