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,189 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/agent"
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWikiCandidateRepairClassifiesStrictOutputWithoutModelValues(t *testing.T) {
+	for _, fixture := range []struct{ content, code, field string }{
+		{`{"entities":[`, "invalid_json_syntax", ""},
+		{`{"entities":"PRIVATE_VALUE","concepts":[]}`, "json_type_mismatch", ""},
+		{`null`, "missing_entities_array", "entities"},
+		{`{"concepts":[]}`, "missing_entities_array", "entities"},
+		{`{"entities":[],"concepts":null}`, "missing_concepts_array", "concepts"},
+		{`{"entities":[{"name":" ","slug":"entity/PRIVATE_VALUE"}],"concepts":[]}`, "candidate_name_empty", "entities.name"},
+		{`{"entities":[],"concepts":[{"name":"PRIVATE_VALUE","slug":"concept/ "}]}`, "candidate_slug_invalid", "concepts.slug"},
+		{`{"entities":[{"name":"PRIVATE_VALUE","slug":"concept/private"}],"concepts":[]}`, "candidate_slug_invalid", "entities.slug"},
+		{`{"entities":[{"name":"PRIVATE_VALUE","slug":"entity/private","aliases":17}],"concepts":[]}`, "json_type_mismatch", ""},
+	} {
+		t.Run(fixture.code+fixture.field, func(t *testing.T) {
+			result, validation := parseWikiCandidateOutput(fixture.content)
+			require.NotNil(t, validation)
+			require.Empty(t, result.Entities)
+			require.Empty(t, result.Concepts)
+			require.Equal(t, fixture.code, validation.Code)
+			require.Equal(t, fixture.field, validation.Field)
+			require.NotContains(t, validation.String(), "PRIVATE_VALUE")
+			require.NotContains(t, wikiCandidateCorrectionInstruction(*validation), "PRIVATE_VALUE")
+			if fixture.code == "invalid_json_syntax" || fixture.code == "json_type_mismatch" {
+				require.Positive(t, validation.JSONOffset)
+			}
+		})
+	}
+}
+
+func TestWikiCandidateRepairUsesExactSourceSharedContextAndOneFreshRequest(t *testing.T) {
+	guard := &candidateRepairTestGuard{}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	ctx = requestbudget.WithGuard(context.WithValue(ctx, types.TenantIDContextKey, uint64(71)), guard)
+	deadline, _ := ctx.Deadline()
+	source := "Valid\n" + strings.Repeat("original-source\n", 500)
+	data := map[string]string{"Content": source, "Language": "English", "PreviousSlugs": "entity/known"}
+	wantData := make(map[string]string)
+	for key, value := range data {
+		wantData[key] = value
+	}
+	calls := 0
+	var first string
+	model := &wikiScriptedChat{id: t.Name(), call: func(callCtx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
+		calls++
+		require.Same(t, guard, requestbudget.FromContext(callCtx))
+		actualDeadline, ok := callCtx.Deadline()
+		require.True(t, ok)
+		require.Equal(t, deadline, actualDeadline)
+		require.NoError(t, callCtx.Err())
+		require.Len(t, messages, 1)
+		if calls == 1 {
+			first = messages[0].Content
+			require.Contains(t, first, source)
+			return &types.ChatResponse{Content: `{"entities":"PRIVATE_REJECTED_OUTPUT","concepts":[]}`, FinishReason: "stop"}, nil
+		}
+		require.Equal(t, 2, calls)
+		require.True(t, strings.HasPrefix(messages[0].Content, first+"\n\n"+wikiCandidateRepairInstruction))
+		require.NotContains(t, messages[0].Content, "PRIVATE_REJECTED_OUTPUT")
+		require.Contains(t, messages[0].Content, `"code":"json_type_mismatch"`)
+		return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
+	}}
+	result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, agent.WikiCandidateSlugPrompt, data)
+	require.NoError(t, err)
+	require.Equal(t, 2, calls)
+	require.Len(t, result.Entities, 1)
+	require.Equal(t, wantData, data)
+	require.GreaterOrEqual(t, guard.checks, 4)
+}
+
+type candidateRepairTestGuard struct {
+	checks  int
+	stopped bool
+}
+
+func (g *candidateRepairTestGuard) Key() string { return "candidate-repair-offline" }
+func (g *candidateRepairTestGuard) Check(context.Context) error {
+	g.checks++
+	if g.stopped {
+		return requestbudget.ErrStopped
+	}
+	return nil
+}
+func (g *candidateRepairTestGuard) Reserve(ctx context.Context, _, _ string) error {
+	return g.Check(ctx)
+}
+
+func TestWikiCandidateRepairCannotBypassCancellationOrTransportBudget(t *testing.T) {
+	for _, stop := range []string{"cancel", "budget"} {
+		t.Run(stop, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			guard := &candidateRepairTestGuard{}
+			ctx = requestbudget.WithGuard(ctx, guard)
+			calls := 0
+			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+				calls++
+				if stop == "cancel" {
+					cancel()
+				} else {
+					guard.stopped = true
+				}
+				return &types.ChatResponse{Content: `{"entities":[],"concepts":null}`, FinishReason: "stop"}, nil
+			}}
+			result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
+			require.Error(t, err)
+			require.Equal(t, 1, calls)
+			require.Empty(t, result.Entities)
+			if stop == "cancel" {
+				require.ErrorIs(t, err, context.Canceled)
+			} else {
+				require.ErrorIs(t, err, requestbudget.ErrStopped)
+			}
+		})
+	}
+}
+
+func TestWikiCandidateRepairProviderFailuresNeverStartAnotherCorrection(t *testing.T) {
+	for _, when := range []string{"initial", "correction"} {
+		for _, code := range []string{"401", "403", "429", "503"} {
+			t.Run(when+code, func(t *testing.T) {
+				ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(72))
+				calls := 0
+				model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
+					calls++
+					if when == "correction" && calls == 1 {
+						return &types.ChatResponse{Content: `{"concepts":[]}`, FinishReason: "stop"}, nil
+					}
+					return nil, errors.New("API request failed with status " + code + ": PRIVATE_PROVIDER_DETAIL")
+				}}
+				key := wikiProviderCircuitKey(72, model.id)
+				t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
+				result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
+				require.Error(t, err)
+				want := 1
+				if when == "correction" {
+					want = 2
+				} else if code == "503" {
+					want = wikiLLMMaxAttempts
+				}
+				require.Equal(t, want, calls)
+				require.Empty(t, result.Entities)
+			})
+		}
+	}
+}
+
+func TestWikiCandidateRepairPartialFailureCannotAdvanceWindowCheckpoint(t *testing.T) {
+	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
+	scope := &wikiCheckpointScope{7, "kb", "doc", "candidate-correction"}
+	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":"WINDOW_000","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
+		}
+		return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"PRIVATE_VALUE","slug":"entity/private"}]}`, FinishReason: "stop"}, nil
+	}}
+	entities, concepts, items, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", wikiFullReadTestSource(2), "English", nil, &WikiBatchContext{}, scope)
+	require.ErrorContains(t, err, "partition 2/2")
+	require.ErrorContains(t, err, "candidate_slug_invalid")
+	require.NotContains(t, err.Error(), "PRIVATE_VALUE")
+	require.Equal(t, 4, calls)
+	require.Nil(t, entities)
+	require.Nil(t, concepts)
+	require.Nil(t, items)
+	var count int64
+	require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage IN ?", []string{"read_candidate", "read_progress", "candidates", "map", "page"}).Count(&count).Error)
+	require.Zero(t, count, "partial successful leaf must not become a completed window or progress")
+}