311 lines
13 KiB
Go
311 lines
13 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"path/filepath"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
"unicode/utf8"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func wikiCatalogTestCandidates() []extractedItem {
|
||
|
|
return []extractedItem{
|
||
|
|
{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{strings.Repeat("甲", 13000)}},
|
||
|
|
{Name: "Beta", Slug: "entity/beta", Aliases: []string{strings.Repeat("乙", 13000)}},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiCatalogTestChunks(n int) []*types.Chunk {
|
||
|
|
var chunks []*types.Chunk
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
chunks = append(chunks, &types.Chunk{ID: fmt.Sprintf("source-%02d", i), ChunkIndex: i,
|
||
|
|
Content: fmt.Sprintf("SOURCE_%02d\n", i) + strings.Repeat("原", 7000)})
|
||
|
|
}
|
||
|
|
return chunks
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateCatalogWholeEntriesStableAndInvalidated(t *testing.T) {
|
||
|
|
entities := wikiCatalogTestCandidates()
|
||
|
|
entities[0].Details = "Audit detail is retained but not rendered"
|
||
|
|
entities[0].Description = "Audit description is retained but not rendered"
|
||
|
|
entities[0].SourceChunks = []string{"original-source"}
|
||
|
|
concepts := []extractedItem{{Name: "Topic", Slug: "concept/topic", Aliases: []string{"one\ntwo", "three"}}}
|
||
|
|
shards, err := buildWikiCandidateShards(entities, concepts, wikiReadCandidateRunes)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.GreaterOrEqual(t, len(shards), 2)
|
||
|
|
var got []wikiCandidateEntry
|
||
|
|
for _, shard := range shards {
|
||
|
|
require.Equal(t, shard.XML, renderWikiCandidateEntries(shard.Entries))
|
||
|
|
require.LessOrEqual(t, utf8.RuneCountInString(shard.XML), wikiReadCandidateRunes)
|
||
|
|
got = append(got, shard.Entries...)
|
||
|
|
}
|
||
|
|
require.ElementsMatch(t, []wikiCandidateEntry{{"entity", entities[0]}, {"entity", entities[1]}, {"concept", concepts[0]}}, got)
|
||
|
|
reordered, err := buildWikiCandidateShards([]extractedItem{entities[1], entities[0]}, concepts, wikiReadCandidateRunes)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, shards, reordered)
|
||
|
|
duplicates, err := buildWikiCandidateShards([]extractedItem{entities[0], entities[0]}, nil, wikiReadCandidateRunes)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, duplicates, 2, "even duplicate input entries are not silently discarded by planning")
|
||
|
|
require.NotEqual(t, duplicates[0].Hash, duplicates[1].Hash, "identical shards retain distinct occurrence checkpoints")
|
||
|
|
chunks := wikiCatalogTestChunks(3)
|
||
|
|
plan, err := buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
|
||
|
|
require.NoError(t, err)
|
||
|
|
reorderedPlan, err := buildWikiCitationPlan([]extractedItem{entities[1], entities[0]}, concepts, []*types.Chunk{chunks[2], chunks[0], chunks[1]}, "Chinese")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, plan, reorderedPlan)
|
||
|
|
changed, err := buildWikiCitationPlan(entities, concepts, chunks, "English")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NotEqual(t, plan.Hash, changed.Hash)
|
||
|
|
entities[0].Details += " changed audit field"
|
||
|
|
changed, err = buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NotEqual(t, plan.Hash, changed.Hash)
|
||
|
|
chunks[2].Content += "tail condition"
|
||
|
|
changedSource, err := buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NotEqual(t, changed.Hash, changedSource.Hash)
|
||
|
|
_, err = buildWikiCandidateShards([]extractedItem{{Name: "Oversize", Slug: "entity/oversize", Aliases: []string{strings.Repeat("名", wikiReadCandidateRunes)}}}, nil, wikiReadCandidateRunes)
|
||
|
|
require.ErrorContains(t, err, "single-entry budget")
|
||
|
|
_, err = buildWikiCandidateShards([]extractedItem{{Slug: "entity/no-name"}}, nil, wikiReadCandidateRunes)
|
||
|
|
require.ErrorContains(t, err, "empty identity")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateCatalogSourcesPreserveExactRangesAndWrappers(t *testing.T) {
|
||
|
|
chunks := []*types.Chunk{
|
||
|
|
{ID: "second", ChunkIndex: 1, Content: "\nsecond\n\n"},
|
||
|
|
{ID: "first", ChunkIndex: 0, Content: strings.Repeat("原文😀\n", 9100) + " trailing qualification \n"},
|
||
|
|
}
|
||
|
|
sources, err := buildWikiCitationSources(chunks, maxRunesPerCitationBatch)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.GreaterOrEqual(t, len(sources), 3)
|
||
|
|
joined := make(map[string]string)
|
||
|
|
ends := make(map[string]int)
|
||
|
|
for _, source := range sources {
|
||
|
|
require.LessOrEqual(t, utf8.RuneCountInString(renderChunksXML(source.Batch)), maxRunesPerCitationBatch)
|
||
|
|
for i, span := range source.Ranges {
|
||
|
|
text := source.Batch.chunks[i].Content
|
||
|
|
require.Equal(t, ends[span.ChunkID], span.StartRune)
|
||
|
|
require.Equal(t, span.EndRune-span.StartRune, utf8.RuneCountInString(text))
|
||
|
|
require.Equal(t, wikiCheckpointHash(text), span.TextHash)
|
||
|
|
joined[span.ChunkID] += text
|
||
|
|
ends[span.ChunkID] = span.EndRune
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for _, chunk := range chunks {
|
||
|
|
require.Equal(t, chunk.Content, joined[chunk.ID])
|
||
|
|
}
|
||
|
|
require.NotEqual(t, sources[0].Hash, sources[1].Hash, "source ranges remain distinct even for repeating content")
|
||
|
|
_, err = buildWikiCitationSources([]*types.Chunk{chunks[0], chunks[0]}, maxRunesPerCitationBatch)
|
||
|
|
require.ErrorContains(t, err, "duplicate")
|
||
|
|
_, err = buildWikiCitationSources(chunks, 1)
|
||
|
|
require.ErrorContains(t, err, "wrapper")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateCatalogAllPairsResumeAndLateEvidence(t *testing.T) {
|
||
|
|
path := filepath.Join(t.TempDir(), "db")
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "pairs"}
|
||
|
|
chunks := wikiCatalogTestChunks(9)
|
||
|
|
plan, err := buildWikiCitationPlan(wikiCatalogTestCandidates(), nil, chunks, "Chinese")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, plan.Shards, 2)
|
||
|
|
require.Len(t, plan.Sources, 9)
|
||
|
|
var calls, active, peak atomic.Int32
|
||
|
|
var mu sync.Mutex
|
||
|
|
seen := map[string]int{}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
concurrent := active.Add(1)
|
||
|
|
defer active.Add(-1)
|
||
|
|
for old := peak.Load(); concurrent > old; old = peak.Load() {
|
||
|
|
if peak.CompareAndSwap(old, concurrent) {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
prompt := messages[0].Content
|
||
|
|
if utf8.RuneCountInString(prompt) > wikiCitationRequestRunes {
|
||
|
|
return nil, errors.New("unbounded request")
|
||
|
|
}
|
||
|
|
slug := "entity/beta"
|
||
|
|
if strings.Contains(prompt, `name: "Alpha"`) {
|
||
|
|
slug = "entity/alpha"
|
||
|
|
}
|
||
|
|
for i := range chunks {
|
||
|
|
if strings.Contains(prompt, fmt.Sprintf("SOURCE_%02d", i)) {
|
||
|
|
mu.Lock()
|
||
|
|
seen[fmt.Sprintf("%s/%d", slug, i)]++
|
||
|
|
mu.Unlock()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if slug == "entity/alpha" && !strings.Contains(prompt, "SOURCE_08") {
|
||
|
|
return &types.ChatResponse{Content: `{"citations":{}}`}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: fmt.Sprintf(`{"citations":{%q:["c000"]}}`, slug)}, nil
|
||
|
|
}}
|
||
|
|
got, _, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
|
||
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
||
|
|
require.Nil(t, got)
|
||
|
|
require.Equal(t, 16, completed)
|
||
|
|
require.Equal(t, int32(16), calls.Load())
|
||
|
|
var progress wikiCitationProgress
|
||
|
|
_, err = svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.False(t, progress.StageComplete)
|
||
|
|
require.Equal(t, 18, progress.TotalPairs)
|
||
|
|
sqlDB, err := svc.checkpointDB.DB()
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, sqlDB.Close())
|
||
|
|
svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
|
||
|
|
got, _, completed, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 18, completed)
|
||
|
|
require.Equal(t, []string{chunks[8].ID}, got["entity/alpha"], "an early shard must inspect the last source batch")
|
||
|
|
require.Len(t, got["entity/beta"], 9)
|
||
|
|
require.Equal(t, int32(18), calls.Load())
|
||
|
|
require.LessOrEqual(t, peak.Load(), int32(maxCitationBatchConcurrency))
|
||
|
|
for key, count := range seen {
|
||
|
|
require.Equal(t, 1, count, key)
|
||
|
|
}
|
||
|
|
require.Len(t, seen, 18)
|
||
|
|
_, err = svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, progress.StageComplete)
|
||
|
|
_, _, _, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, int32(18), calls.Load())
|
||
|
|
_, _, _, err = (&wikiIngestService{}).classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", nil)
|
||
|
|
require.ErrorContains(t, err, "需要持久断点")
|
||
|
|
require.Equal(t, int32(18), calls.Load())
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateCatalogFailedPairKeepsSuccessfulSiblings(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "failure"}
|
||
|
|
chunks := wikiCatalogTestChunks(3)
|
||
|
|
plan, err := buildWikiCitationPlan(wikiCatalogTestCandidates(), nil, chunks, "English")
|
||
|
|
require.NoError(t, err)
|
||
|
|
var failed atomic.Bool
|
||
|
|
var mu sync.Mutex
|
||
|
|
calls := map[string]int{}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
prompt := messages[0].Content
|
||
|
|
slug := "entity/beta"
|
||
|
|
if strings.Contains(prompt, `name: "Alpha"`) {
|
||
|
|
slug = "entity/alpha"
|
||
|
|
}
|
||
|
|
index := 0
|
||
|
|
for i := range chunks {
|
||
|
|
if strings.Contains(prompt, fmt.Sprintf("SOURCE_%02d", i)) {
|
||
|
|
index = i
|
||
|
|
}
|
||
|
|
}
|
||
|
|
key := fmt.Sprintf("%s/%d", slug, index)
|
||
|
|
mu.Lock()
|
||
|
|
calls[key]++
|
||
|
|
mu.Unlock()
|
||
|
|
if slug == "entity/alpha" && index == 1 && failed.CompareAndSwap(false, true) {
|
||
|
|
deadline := time.Now().Add(3 * time.Second)
|
||
|
|
for time.Now().Before(deadline) {
|
||
|
|
var count int64
|
||
|
|
if err := svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "citation_pair").Count(&count).Error; err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if count > 0 {
|
||
|
|
return nil, errors.New("invalid request")
|
||
|
|
}
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
return nil, ctx.Err()
|
||
|
|
case <-time.After(time.Millisecond):
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil, errors.New("sibling checkpoint did not persist")
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: fmt.Sprintf(`{"citations":{%q:["c000"]}}`, slug)}, nil
|
||
|
|
}}
|
||
|
|
got, _, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Nil(t, got)
|
||
|
|
require.GreaterOrEqual(t, completed, 1)
|
||
|
|
require.Less(t, completed, 6)
|
||
|
|
var progress wikiCitationProgress
|
||
|
|
_, loadErr := svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
|
||
|
|
require.NoError(t, loadErr)
|
||
|
|
require.False(t, progress.StageComplete)
|
||
|
|
got, _, completed, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 6, completed)
|
||
|
|
require.Len(t, got["entity/alpha"], 3)
|
||
|
|
for key, count := range calls {
|
||
|
|
want := 1
|
||
|
|
if key == "entity/alpha/1" {
|
||
|
|
want = 2
|
||
|
|
}
|
||
|
|
require.Equal(t, want, count, key)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateCatalogPairIsolationAndEvidenceUnion(t *testing.T) {
|
||
|
|
for _, corrupt := range []bool{false, true} {
|
||
|
|
t.Run(fmt.Sprintf("corrupt-%t", corrupt), func(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "isolation"}
|
||
|
|
entities := wikiCatalogTestCandidates()
|
||
|
|
chunks := []*types.Chunk{{ID: "first", Content: "Alpha Beta first", ChunkIndex: 0}, {ID: "second", Content: "Alpha Beta second", ChunkIndex: 1}}
|
||
|
|
plan, err := buildWikiCitationPlan(entities, nil, chunks, "English")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, plan.Shards, 2)
|
||
|
|
require.Len(t, plan.Sources, 1)
|
||
|
|
for i, shard := range plan.Shards {
|
||
|
|
result := citationBatchResult{Citations: map[string][]string{"entity/alpha": {"first"}}}
|
||
|
|
if i == 1 && !corrupt {
|
||
|
|
result.NewSlugs = []newSlugFromCitation{{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"second"}}}
|
||
|
|
result.Citations["entity/beta"] = []string{"second"}
|
||
|
|
}
|
||
|
|
key := wikiCheckpointHash([]any{wikiCandidateCatalogVersion, plan.Hash, shard.Hash, plan.Sources[0].Hash, "English"})
|
||
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation_pair", key, result))
|
||
|
|
}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
return nil, errors.New("cached validation must not invoke model repair")
|
||
|
|
}}
|
||
|
|
citations, declarations, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
|
||
|
|
if corrupt {
|
||
|
|
require.ErrorContains(t, err, "undeclared slug")
|
||
|
|
require.Nil(t, citations)
|
||
|
|
require.Less(t, completed, 2)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 2, completed)
|
||
|
|
got, _, _ := mergeCitationsIntoItems(entities, nil, citations, declarations)
|
||
|
|
require.Equal(t, []string{"first", "second"}, findBySlug(got, "entity/alpha").SourceChunks)
|
||
|
|
require.Equal(t, []string{"second"}, findBySlug(got, "entity/beta").SourceChunks)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCandidateCatalogCompleteRequestGuard(t *testing.T) {
|
||
|
|
var calls atomic.Int32
|
||
|
|
inner := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
return &types.ChatResponse{}, nil
|
||
|
|
}}
|
||
|
|
model := wikiCitationBoundedChat{inner}
|
||
|
|
_, err := model.Chat(context.Background(), []chat.Message{{Content: strings.Repeat("x", wikiCitationRequestRunes-1)}, {Content: "repair instruction"}}, nil)
|
||
|
|
require.ErrorContains(t, err, "complete-request budget")
|
||
|
|
require.Zero(t, calls.Load())
|
||
|
|
require.Equal(t, inner.GetModelID(), model.GetModelID())
|
||
|
|
_, err = model.Chat(context.Background(), []chat.Message{{Content: "bounded request"}}, nil)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, int32(1), calls.Load())
|
||
|
|
}
|