235 lines
12 KiB
Go
235 lines
12 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"path/filepath"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
type wikiPageRepairChat struct {
|
||
|
|
templateCaptureChatModel
|
||
|
|
id string
|
||
|
|
call func(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *wikiPageRepairChat) GetModelID() string { return m.id }
|
||
|
|
func (m *wikiPageRepairChat) Chat(ctx context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
return m.call(ctx, messages, opts)
|
||
|
|
}
|
||
|
|
|
||
|
|
func pageRepairInput() map[string]string {
|
||
|
|
return map[string]string{
|
||
|
|
"HasAdditions": "1", "HasRetractions": "1", "PageSlug": "concept/alpha", "PageTitle": "Alpha", "PageType": "concept",
|
||
|
|
"ExistingContent": "Existing fact is not valid unless condition A holds.",
|
||
|
|
"SharedSourceContexts": "Source context and uncertainty.", "NewContent": "Source A states 42, except under condition B.",
|
||
|
|
"DeletedContent": "Obsolete assertion.", "RemainingSourcesContent": "Preserved source B.",
|
||
|
|
"AvailableSlugs": "ref-1 (Alpha)", "Language": "English", "CustomInstructions": "Keep attribution.", "InstructionScope": "wiki_content",
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageOutputRepairRegeneratesWholePageWithIdenticalScope(t *testing.T) {
|
||
|
|
svc := &wikiIngestService{}
|
||
|
|
input := pageRepairInput()
|
||
|
|
before := make(map[string]string, len(input))
|
||
|
|
for key, value := range input {
|
||
|
|
before[key] = value
|
||
|
|
}
|
||
|
|
var requests [][]chat.Message
|
||
|
|
var prefixes []string
|
||
|
|
model := &wikiPageRepairChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
requests = append(requests, append([]chat.Message(nil), messages...))
|
||
|
|
purpose, prefix := types.LLMCallMetadataFromContext(ctx)
|
||
|
|
require.Equal(t, "wiki_page_modify", purpose)
|
||
|
|
prefixes = append(prefixes, prefix)
|
||
|
|
require.Equal(t, wikiLLMMaxTokens, opts.MaxTokens)
|
||
|
|
if len(requests) == 1 {
|
||
|
|
return &types.ChatResponse{Content: "PARTIAL_OUTPUT_MUST_NOT_BE_REUSED", FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Complete revision\n# Alpha\nNot valid unless condition A holds; source A states 42 except under condition B.", FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
got, err := svc.generateWikiPageModification(context.Background(), model, input)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Contains(t, got, "Complete revision")
|
||
|
|
require.Len(t, requests, 2)
|
||
|
|
require.Equal(t, before, input)
|
||
|
|
require.Equal(t, requests[0][0], requests[1][0], "base system prompt must remain unchanged")
|
||
|
|
require.Equal(t, requests[0][1].Content+"\n\n"+wikiPageRepairInstruction, requests[1][1].Content)
|
||
|
|
require.Equal(t, prefixes[0], prefixes[1], "page source prefix remains reusable")
|
||
|
|
require.NotContains(t, requests[1][1].Content, "PARTIAL_OUTPUT_MUST_NOT_BE_REUSED")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageOutputRepairSuccessNeedsNoSecondCall(t *testing.T) {
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Complete\nComplete body", FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
got, err := (&wikiIngestService{}).generateWikiPageModification(context.Background(), model, pageRepairInput())
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NotEmpty(t, got)
|
||
|
|
require.Equal(t, 1, calls)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageOutputRepairStopsAfterOneFailedRegeneration(t *testing.T) {
|
||
|
|
for _, response := range []*types.ChatResponse{
|
||
|
|
{Content: "still partial", FinishReason: "length"},
|
||
|
|
{Content: "still partial", FinishReason: "max_tokens"},
|
||
|
|
{Content: "still partial", FinishReason: "max_output_tokens"},
|
||
|
|
{Content: "still partial", FinishReason: types.FinishReasonIncomplete},
|
||
|
|
{Content: " ", FinishReason: "stop"},
|
||
|
|
{Content: "WIKI_PAGE_REPAIR_UNSAFE", FinishReason: "stop"},
|
||
|
|
{Content: "SUMMARY: Summary without body", FinishReason: "stop"},
|
||
|
|
{Content: "SUMMARY: \nBody without required summary", FinishReason: "stop"},
|
||
|
|
{Content: "Missing required summary line", FinishReason: "stop"},
|
||
|
|
{Content: "I cannot complete the requested page.", FinishReason: "stop"},
|
||
|
|
{Content: strings.Repeat("x", 12001), FinishReason: "stop"},
|
||
|
|
} {
|
||
|
|
t.Run(fmt.Sprintf("%s-%d", response.FinishReason, len(response.Content)), func(t *testing.T) {
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
return response, nil
|
||
|
|
}}
|
||
|
|
got, err := (&wikiIngestService{}).generateWikiPageModification(context.Background(), model, pageRepairInput())
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Empty(t, got)
|
||
|
|
require.Equal(t, 2, calls)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageSummaryOnlyDoesNotOverwritePublishedPage(t *testing.T) {
|
||
|
|
for _, repaired := range []string{"SUMMARY: Only a summary", "Missing required SUMMARY line", "I cannot complete the requested page.", "SUMMARY: \nMissing summary"} {
|
||
|
|
t.Run(repaired, func(t *testing.T) {
|
||
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
||
|
|
original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept,
|
||
|
|
Status: types.WikiPageStatusPublished, Content: "Original supported content", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}}
|
||
|
|
before := *original
|
||
|
|
pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
|
||
|
|
svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}}
|
||
|
|
batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
|
||
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "unchanged-fingerprint"})
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Incomplete first page", FinishReason: "stop"}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: repaired, FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
update := SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: original.Slug, Type: types.WikiPageTypeConcept, DocTitle: "Source", SourceChunks: []string{"source-chunk"}, Item: extractedItem{Name: "Alpha", Slug: original.Slug, Details: "New fact"}}
|
||
|
|
changed, _, failed, err := svc.reduceSlugUpdates(context.Background(), model, "kb", original.Slug, []SlugUpdate{update}, 7, batch, nil)
|
||
|
|
require.Error(t, err)
|
||
|
|
require.False(t, changed)
|
||
|
|
require.True(t, failed)
|
||
|
|
require.Equal(t, 2, calls)
|
||
|
|
require.Zero(t, pages.writes)
|
||
|
|
require.Equal(t, before, *pages.pages[original.Slug])
|
||
|
|
var checkpoints int64
|
||
|
|
require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&checkpoints).Error)
|
||
|
|
require.Zero(t, checkpoints)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageOutputRepairDoesNotMultiplyProviderFailures(t *testing.T) {
|
||
|
|
for _, afterOutputFailure := range []bool{false, true} {
|
||
|
|
for _, failure := range []string{"API request failed with status 401", "API request failed with status 402", "API request failed with status 403", "API request failed with status 429", "connection closed", "timeout"} {
|
||
|
|
if failure == "timeout" && !afterOutputFailure {
|
||
|
|
continue // The first request retains its existing transient transport policy.
|
||
|
|
}
|
||
|
|
t.Run(fmt.Sprintf("repair-%t-%s", afterOutputFailure, failure), func(t *testing.T) {
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(700))
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if afterOutputFailure && calls == 1 {
|
||
|
|
return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
|
||
|
|
}
|
||
|
|
return nil, errors.New(failure)
|
||
|
|
}}
|
||
|
|
key := wikiProviderCircuitKey(700, model.id)
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
|
||
|
|
got, err := (&wikiIngestService{}).generateWikiPageModification(ctx, model, pageRepairInput())
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Empty(t, got)
|
||
|
|
expected := 1
|
||
|
|
if afterOutputFailure {
|
||
|
|
expected++
|
||
|
|
}
|
||
|
|
require.Equal(t, expected, calls)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageOutputRepairHonorsCancellation(t *testing.T) {
|
||
|
|
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++
|
||
|
|
cancel()
|
||
|
|
return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
|
||
|
|
}}
|
||
|
|
got, err := (&wikiIngestService{}).generateWikiPageModification(ctx, model, pageRepairInput())
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Empty(t, got)
|
||
|
|
require.Equal(t, 1, calls)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiPageOutputRepairFailurePreservesPublishedPageAndCheckpoint(t *testing.T) {
|
||
|
|
for _, repairsSuccessfully := range []bool{false, true} {
|
||
|
|
t.Run(fmt.Sprint(repairsSuccessfully), func(t *testing.T) {
|
||
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
||
|
|
original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept,
|
||
|
|
Status: types.WikiPageStatusPublished, Content: "Original supported content", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}}
|
||
|
|
before := *original
|
||
|
|
pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
|
||
|
|
svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}}
|
||
|
|
batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
|
||
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "unchanged-fingerprint"})
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
if repairsSuccessfully && calls == 2 {
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Complete revised summary\nOriginal supported content. New fact has a condition.", FinishReason: "stop"}, nil
|
||
|
|
}
|
||
|
|
return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
|
||
|
|
}}
|
||
|
|
update := SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: original.Slug, Type: types.WikiPageTypeConcept, DocTitle: "Source", SourceChunks: []string{"source-chunk"}, Item: extractedItem{Name: "Alpha", Slug: original.Slug, Details: "New fact has a condition."}}
|
||
|
|
changed, _, failed, err := svc.reduceSlugUpdates(context.Background(), model, "kb", original.Slug, []SlugUpdate{update}, 7, batch, nil)
|
||
|
|
require.Equal(t, 2, calls)
|
||
|
|
var checkpoints int64
|
||
|
|
require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&checkpoints).Error)
|
||
|
|
if repairsSuccessfully {
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, changed)
|
||
|
|
require.False(t, failed)
|
||
|
|
require.Equal(t, 1, pages.writes)
|
||
|
|
require.Equal(t, int64(1), checkpoints)
|
||
|
|
require.Contains(t, pages.pages[original.Slug].Content, "Original supported content")
|
||
|
|
require.Contains(t, pages.pages[original.Slug].SourceRefs, "old-doc")
|
||
|
|
require.Contains(t, pages.pages[original.Slug].SourceRefs, "doc")
|
||
|
|
} else {
|
||
|
|
require.Error(t, err)
|
||
|
|
require.False(t, changed)
|
||
|
|
require.True(t, failed)
|
||
|
|
require.Zero(t, pages.writes)
|
||
|
|
require.Zero(t, checkpoints)
|
||
|
|
require.Equal(t, before, *pages.pages[original.Slug])
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|