WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,234 @@
+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])
+			}
+		})
+	}
+}