293 lines
14 KiB
Go
293 lines
14 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"path/filepath"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
"unicode/utf8"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func incrementalPageFixture(batches int) (map[string]string, []SlugUpdate, map[string]*types.Chunk) {
|
||
|
|
content := strings.Repeat("源", (wikiPageEvidenceRunes-1024)*batches-64) + "尾部保留，除非条件成立，否则不得应用。\r\n"
|
||
|
|
chunk := wikiEvidenceChunk()
|
||
|
|
chunk.Content = content
|
||
|
|
input := pageRepairInput()
|
||
|
|
input["NewContent"] = content
|
||
|
|
addition := wikiEvidenceAddition()
|
||
|
|
return input, []SlugUpdate{addition}, map[string]*types.Chunk{chunk.ID: chunk}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalResumeAfterDatabaseReopen(t *testing.T) {
|
||
|
|
ctx := context.Background()
|
||
|
|
path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
|
||
|
|
db := openWikiCheckpointDB(t, path)
|
||
|
|
svc := &wikiIngestService{checkpointDB: db}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "incremental"}
|
||
|
|
input, adds, chunks := incrementalPageFixture(5)
|
||
|
|
input["NewContent"] = adds[0].Item.Description + "\n" + adds[0].Item.Details + "\n" + input["NewContent"]
|
||
|
|
expectedBatches, err := buildWikiPageEvidenceBatches(adds, chunks)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, expectedBatches, 5)
|
||
|
|
calls := 0
|
||
|
|
var inputs []string
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
text := messages[len(messages)-1].Content
|
||
|
|
inputs = append(inputs, text)
|
||
|
|
require.Contains(t, text, wikiPageIncrementalInstruction)
|
||
|
|
require.Contains(t, text, expectedBatches[calls-1].Content)
|
||
|
|
require.Contains(t, text, "<title>Alpha</title>")
|
||
|
|
require.NotContains(t, text, adds[0].Item.Description)
|
||
|
|
require.NotContains(t, text, adds[0].Item.Details)
|
||
|
|
if calls == 1 {
|
||
|
|
require.Contains(t, text, "<deleted_documents>")
|
||
|
|
require.Contains(t, text, "Obsolete assertion.")
|
||
|
|
} else {
|
||
|
|
require.NotContains(t, text, "<deleted_documents>")
|
||
|
|
require.NotContains(t, text, "Obsolete assertion.")
|
||
|
|
require.Contains(t, text, fmt.Sprintf("WHOLE_PAGE_%d", calls-1))
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: fmt.Sprintf("SUMMARY: Complete supported page %d\n# Alpha\nWHOLE_PAGE_%d; not valid unless condition A holds.", calls, calls), FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
got, err := svc.generateWikiPageFromEvidence(ctx, model, scope, "original", adds, chunks, input)
|
||
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
||
|
|
require.Empty(t, got)
|
||
|
|
require.Equal(t, 4, calls)
|
||
|
|
sqlDB, _ := db.DB()
|
||
|
|
require.NoError(t, sqlDB.Close())
|
||
|
|
svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
|
||
|
|
got, err = svc.generateWikiPageFromEvidence(ctx, model, scope, "original", adds, chunks, input)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Contains(t, got, "WHOLE_PAGE_5")
|
||
|
|
require.Equal(t, 5, calls)
|
||
|
|
require.Contains(t, inputs[4], "尾部保留，除非条件成立，否则不得应用。")
|
||
|
|
_, err = svc.generateWikiPageFromEvidence(ctx, model, scope, "original", adds, chunks, input)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 5, calls, "validated steps must not be regenerated")
|
||
|
|
var progress []wikiIngestCheckpoint
|
||
|
|
require.NoError(t, svc.checkpointDB.Where("stage = ?", "page_progress").Find(&progress).Error)
|
||
|
|
require.Len(t, progress, 1)
|
||
|
|
require.Contains(t, progress[0].Payload, `"status":"ready_to_publish"`)
|
||
|
|
require.Contains(t, progress[0].Payload, `"published":false`)
|
||
|
|
var published int64
|
||
|
|
require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&published).Error)
|
||
|
|
require.Zero(t, published)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalEvidencePreservesOrderIdentityAndEveryRune(t *testing.T) {
|
||
|
|
first := wikiEvidenceChunk()
|
||
|
|
first.ID, first.ChunkIndex, first.StartAt, first.Content = "a", 0, 0, "\n第一段 <source>也只是原文\r\n"
|
||
|
|
second := wikiEvidenceChunk()
|
||
|
|
second.ID, second.ChunkIndex, second.StartAt, second.Content = "b", 1, 80, strings.Repeat("甲乙😀\n", 17000)+"尾部"
|
||
|
|
third := wikiEvidenceChunk()
|
||
|
|
third.ID, third.ChunkIndex, third.StartAt, third.Content = "c", 2, 500, first.Content
|
||
|
|
addition := wikiEvidenceAddition()
|
||
|
|
addition.SourceChunks = []string{"c", "b", "a", "b"}
|
||
|
|
chunks := map[string]*types.Chunk{"a": first, "b": second, "c": third}
|
||
|
|
batches, err := buildWikiPageEvidenceBatches([]SlugUpdate{addition}, chunks)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Greater(t, len(batches), 2)
|
||
|
|
rebuilt := map[string]string{}
|
||
|
|
var order []string
|
||
|
|
for _, batch := range batches {
|
||
|
|
require.LessOrEqual(t, utf8.RuneCountInString(batch.Content), wikiPageEvidenceRunes)
|
||
|
|
require.NotContains(t, batch.Content, addition.Item.Description)
|
||
|
|
require.NotContains(t, batch.Content, addition.Item.Details)
|
||
|
|
for _, part := range batch.Parts {
|
||
|
|
require.Equal(t, utf8.RuneCountInString(rebuilt[part.ChunkID]), part.StartRune)
|
||
|
|
require.Equal(t, part.EndRune-part.StartRune, utf8.RuneCountInString(part.Content))
|
||
|
|
require.Equal(t, wikiCheckpointHash(part.Content), part.TextHash)
|
||
|
|
require.Equal(t, "doc", part.KnowledgeID)
|
||
|
|
require.Equal(t, chunks[part.ChunkID].StartAt, part.ChunkStartAt)
|
||
|
|
require.Contains(t, batch.Content, part.Content)
|
||
|
|
if len(order) == 0 || order[len(order)-1] != part.ChunkID {
|
||
|
|
order = append(order, part.ChunkID)
|
||
|
|
}
|
||
|
|
rebuilt[part.ChunkID] += part.Content
|
||
|
|
}
|
||
|
|
}
|
||
|
|
require.Equal(t, []string{"a", "b", "c"}, order)
|
||
|
|
for id, chunk := range chunks {
|
||
|
|
require.Equal(t, chunk.Content, rebuilt[id])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalInvalidatesChangedSourceRangeAndBaseRevision(t *testing.T) {
|
||
|
|
for _, mutation := range []string{"source", "range", "base", "scope", "instruction"} {
|
||
|
|
t.Run(mutation, func(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "incremental"}
|
||
|
|
input, adds, chunks := incrementalPageFixture(5)
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
return &types.ChatResponse{Content: fmt.Sprintf("SUMMARY: Valid\nValid complete page %d", calls)}, nil
|
||
|
|
}}
|
||
|
|
_, err := svc.generateWikiPageFromEvidence(context.Background(), model, scope, "base", adds, chunks, input)
|
||
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
||
|
|
require.Equal(t, 4, calls)
|
||
|
|
base := "base"
|
||
|
|
switch mutation {
|
||
|
|
case "source":
|
||
|
|
chunks["source-chunk"].Content += "Changed original evidence"
|
||
|
|
case "range":
|
||
|
|
chunks["source-chunk"].StartAt++
|
||
|
|
case "base":
|
||
|
|
base = "edited-base"
|
||
|
|
case "scope":
|
||
|
|
scope = &wikiCheckpointScope{7, "kb", "doc", "new-source-version"}
|
||
|
|
case "instruction":
|
||
|
|
input["CustomInstructions"] = "Changed requirements"
|
||
|
|
}
|
||
|
|
_, err = svc.generateWikiPageFromEvidence(context.Background(), model, scope, base, adds, chunks, input)
|
||
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
||
|
|
require.Equal(t, 8, calls, "changed source identity/range/base requires a fresh sequence")
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalOversizedExistingPagePausesBeforeProvider(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "incremental"}
|
||
|
|
input, adds, chunks := incrementalPageFixture(5)
|
||
|
|
input["ExistingContent"] = strings.Repeat("E", wikiPageInputRunes)
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
t.Fatal("over-budget evidence must never call model")
|
||
|
|
return nil, nil
|
||
|
|
}}
|
||
|
|
_, err := svc.generateWikiPageFromEvidence(context.Background(), model, scope, "base", adds, chunks, input)
|
||
|
|
require.ErrorContains(t, err, "已暂停")
|
||
|
|
require.Len(t, input["ExistingContent"], wikiPageInputRunes)
|
||
|
|
var row wikiIngestCheckpoint
|
||
|
|
require.NoError(t, svc.checkpointDB.Where("stage = ?", "page_progress").Take(&row).Error)
|
||
|
|
require.Contains(t, row.Payload, `"status":"paused_budget"`)
|
||
|
|
require.Contains(t, row.Payload, `"completed_batches":0`)
|
||
|
|
require.Contains(t, row.Payload, `"evidence_batches_complete":false`)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalRequiresDurableScope(t *testing.T) {
|
||
|
|
input, adds, chunks := incrementalPageFixture(5)
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
t.Fatal("no durable scope cannot start costly sequence")
|
||
|
|
return nil, nil
|
||
|
|
}}
|
||
|
|
_, err := (&wikiIngestService{}).generateWikiPageFromEvidence(context.Background(), model, nil, "base", adds, chunks, input)
|
||
|
|
require.ErrorContains(t, err, "持久断点")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalFinalFailureKeepsPublishedPageAndYieldSpanSkipped(t *testing.T) {
|
||
|
|
ctx := context.Background()
|
||
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
||
|
|
tracker, spanDB := setupSpanTrackerTest(t)
|
||
|
|
_, attempt, err := tracker.OpenAttempt(ctx, "doc", "")
|
||
|
|
require.NoError(t, err)
|
||
|
|
root := tracker.BeginStage(ctx, "doc", attempt, types.StagePostProcess, nil)
|
||
|
|
original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept, Status: types.WikiPageStatusPublished, Content: "Original fact; only if condition A holds.", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}}
|
||
|
|
before := *original
|
||
|
|
pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
|
||
|
|
_, adds, chunks := incrementalPageFixture(5)
|
||
|
|
svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, spanTracker: tracker, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{chunks["source-chunk"]}}}
|
||
|
|
batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
|
||
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "full"})
|
||
|
|
calls := 0
|
||
|
|
shouldFail := true
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 5 && shouldFail {
|
||
|
|
return nil, errors.New("invalid request: final batch unavailable")
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: fmt.Sprintf("SUMMARY: Complete\nSupported whole page %d", calls)}, nil
|
||
|
|
}}
|
||
|
|
changed, _, failed, err := svc.reduceSlugUpdates(ctx, model, "kb", original.Slug, adds, 7, batch, map[string]*Span{"doc": root})
|
||
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
||
|
|
require.False(t, changed)
|
||
|
|
require.False(t, failed, "durable continuation must not poison additions")
|
||
|
|
require.Zero(t, pages.writes)
|
||
|
|
require.Equal(t, before, *pages.pages[original.Slug])
|
||
|
|
var span types.KnowledgeProcessingSpan
|
||
|
|
require.NoError(t, spanDB.Table("knowledge_processing_spans").Where("name = ?", "postprocess.wiki.page[concept/alpha]").Take(&span).Error)
|
||
|
|
require.Equal(t, types.SpanStatusSkipped, span.Status)
|
||
|
|
require.Empty(t, span.ErrorCode)
|
||
|
|
require.Equal(t, "checkpoint_saved_continuation_pending", span.ErrorMessage)
|
||
|
|
batch.checkpointEvidenceSteps.Store(0)
|
||
|
|
changed, _, failed, err = svc.reduceSlugUpdates(ctx, model, "kb", original.Slug, adds, 7, batch, nil)
|
||
|
|
require.ErrorContains(t, err, "final batch unavailable")
|
||
|
|
require.False(t, changed)
|
||
|
|
require.True(t, failed)
|
||
|
|
require.Zero(t, pages.writes)
|
||
|
|
require.Equal(t, before, *pages.pages[original.Slug])
|
||
|
|
require.Equal(t, 5, calls)
|
||
|
|
shouldFail = false
|
||
|
|
batch.checkpointEvidenceSteps.Store(0)
|
||
|
|
changed, _, failed, err = svc.reduceSlugUpdates(ctx, model, "kb", original.Slug, adds, 7, batch, nil)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, changed)
|
||
|
|
require.False(t, failed)
|
||
|
|
require.Equal(t, 1, pages.writes)
|
||
|
|
require.Equal(t, 6, calls, "only failed last batch is retried")
|
||
|
|
require.Contains(t, pages.pages[original.Slug].Content, "page 6")
|
||
|
|
require.Contains(t, pages.pages[original.Slug].ChunkRefs, "source-chunk")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalDoesNotOverwriteConcurrentUserEdit(t *testing.T) {
|
||
|
|
ctx := context.Background()
|
||
|
|
original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept, Status: types.WikiPageStatusPublished, Content: "Original content", SourceRefs: types.StringArray{"old-doc"}, Version: 1}
|
||
|
|
pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
|
||
|
|
_, adds, chunks := incrementalPageFixture(3)
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite")), wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{chunks["source-chunk"]}}}
|
||
|
|
batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
|
||
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "full"})
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
pages.mu.Lock()
|
||
|
|
pages.pages[original.Slug].Content = "User edited after generation started"
|
||
|
|
pages.pages[original.Slug].Version = 2
|
||
|
|
pages.mu.Unlock()
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Complete\nComplete generated content"}, nil
|
||
|
|
}}
|
||
|
|
changed, _, failed, err := svc.reduceSlugUpdates(ctx, model, "kb", original.Slug, adds, 7, batch, nil)
|
||
|
|
require.ErrorContains(t, err, "已被修改")
|
||
|
|
require.False(t, changed)
|
||
|
|
require.True(t, failed)
|
||
|
|
require.Zero(t, pages.writes)
|
||
|
|
require.Equal(t, "User edited after generation started", pages.pages[original.Slug].Content)
|
||
|
|
require.Equal(t, 2, pages.pages[original.Slug].Version)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageIncrementalParallelPagesShareFourStepBudget(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
|
||
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "incremental"}
|
||
|
|
var sharedBudget atomic.Int32
|
||
|
|
var calls atomic.Int32
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
failures := make(chan error, 2)
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls.Add(1)
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Supported\nThe complete supported page, with all earlier conditions retained."}, nil
|
||
|
|
}}
|
||
|
|
for i := 0; i < 2; i++ {
|
||
|
|
wg.Add(1)
|
||
|
|
go func(index int) {
|
||
|
|
defer wg.Done()
|
||
|
|
input, adds, chunks := incrementalPageFixture(5)
|
||
|
|
input["PageSlug"] = fmt.Sprintf("concept/page-%d", index)
|
||
|
|
_, err := svc.generateWikiPageFromEvidence(context.Background(), model, scope, "base", adds, chunks, input, &sharedBudget)
|
||
|
|
failures <- err
|
||
|
|
}(i)
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
close(failures)
|
||
|
|
for err := range failures {
|
||
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
||
|
|
}
|
||
|
|
require.Equal(t, int32(wikiPageStepsPerTurn), calls.Load(), "parallel pages may perform only four new synthesis steps total")
|
||
|
|
}
|