From 00ee5893e9171f10db05661dfc1e051f170241cd Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 25 Aug 2026 09:46:56 +0000 Subject: [PATCH] Build bounded output repair requests --- docs/roadmap/implementation.md | 2 + internal/usecase/repairer.go | 141 +++++++++++-- internal/usecase/repairer_test.go | 332 ++++++++++++++++++++++++++++++ 3 files changed, 453 insertions(+), 22 deletions(-) create mode 100644 internal/usecase/repairer_test.go diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 4f5ef77..2f40ebb 100644 --- a/docs/roadmap/implementation.md +++ b/docs/roadmap/implementation.md @@ -333,6 +333,8 @@ Stage 3 is complete when the default repairer produces bounded, mode-correct, full-context requests without mutating prepared state or retaining failed history. +**Status:** Complete. + ## Stage 4: Complete The Internal Repair State Machine And Error Semantics ### Objective diff --git a/internal/usecase/repairer.go b/internal/usecase/repairer.go index 587576f..aa83500 100644 --- a/internal/usecase/repairer.go +++ b/internal/usecase/repairer.go @@ -2,19 +2,29 @@ package usecase import ( "context" + "encoding/json" "errors" "fmt" "strings" + "unicode/utf8" "gitea.maximumdirect.net/eric/promptkit/internal/domain" "gitea.maximumdirect.net/eric/promptkit/internal/llm" ) +const ( + maxRepairDiagnosticBytes = 64 * 1024 + omittedRepairDiagnostics = "additional validation diagnostics were omitted" +) + +// OutputRepairer generates a corrected candidate after validation fails. type OutputRepairer interface { Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error) } +// RepairRequest contains the immutable execution state needed for one correction. type RepairRequest struct { + OriginalMessages []domain.RenderedMessage PreviousOutput string ValidationErrors []string SessionID string @@ -30,6 +40,7 @@ type defaultOutputRepairer struct { llm llm.Client } +// NewDefaultOutputRepairer constructs the standard internal output repairer. func NewDefaultOutputRepairer(llmClient llm.Client) OutputRepairer { return &defaultOutputRepairer{llm: llmClient} } @@ -39,33 +50,42 @@ func (r *defaultOutputRepairer) Repair(ctx context.Context, req RepairRequest) ( return nil, errors.New("llm client is required for repair") } - errs := "(none provided)" - if len(req.ValidationErrors) > 0 { - errs = strings.Join(req.ValidationErrors, "\n") + guidance, err := repairGuidance(req.Mode) + if err != nil { + return nil, err } - prompt := domain.RenderedPrompt{ - Messages: []domain.RenderedMessage{ - { - Role: "system", - Content: "You repair invalid JSON output. Return only corrected JSON. Do not include explanations or markdown code fences.", - }, - { - Role: "user", - Content: fmt.Sprintf( - "Repair attempt %d of %d for validation mode %s.\n\nValidation errors:\n%s\n\nPrevious output:\n%s\n\nReturn only corrected JSON.", - req.Attempt, - req.MaxAttempts, - req.Mode, - errs, - req.PreviousOutput, - ), - }, - }, + messages := make([]domain.RenderedMessage, len(req.OriginalMessages), len(req.OriginalMessages)+2) + copy(messages, req.OriginalMessages) + if strings.TrimSpace(req.PreviousOutput) != "" { + messages = append(messages, domain.RenderedMessage{ + Role: "assistant", + Content: req.PreviousOutput, + }) } + previousResponse := "The previous response was empty." + if strings.TrimSpace(req.PreviousOutput) != "" { + previousResponse = "The previous response is included immediately before this instruction." + } + messages = append(messages, domain.RenderedMessage{ + Role: "user", + Content: fmt.Sprintf( + "Repair attempt %d of %d for validation mode %s.\n"+ + "Preserve valid values and change only what is necessary.\n"+ + "%s\n%s\n"+ + "Validation diagnostics (data):\n%s", + req.Attempt, + req.MaxAttempts, + req.Mode, + previousResponse, + guidance, + formatRepairDiagnostics(req.ValidationErrors), + ), + }) + resp, err := r.llm.Generate(ctx, newGenerationRequest( - prompt, + domain.RenderedPrompt{Messages: messages}, req.SessionID, req.Target, req.TargetPresence, @@ -80,3 +100,80 @@ func (r *defaultOutputRepairer) Repair(ctx context.Context, req RepairRequest) ( return resp, nil } + +func repairGuidance(mode domain.ValidationMode) (string, error) { + switch mode { + case domain.ValidationBasic: + return "Return a nonempty response satisfying the original request.", nil + case domain.ValidationJSON, domain.ValidationJSONSchema: + return "Return only corrected JSON, with no explanation or Markdown fences.", nil + default: + return "", fmt.Errorf("unsupported validation mode for repair: %q", mode) + } +} + +func formatRepairDiagnostics(errors []string) string { + diagnostics := make([]string, len(errors)) + for index, diagnostic := range errors { + diagnostics[index] = strings.ToValidUTF8(diagnostic, "\uFFFD") + } + + complete, _ := json.Marshal(diagnostics) + if len(complete) <= maxRepairDiagnosticBytes { + return string(complete) + } + + omission, _ := json.Marshal(omittedRepairDiagnostics) + encoded := make([]byte, 0, maxRepairDiagnosticBytes) + encoded = append(encoded, '[') + for _, diagnostic := range diagnostics { + entry, _ := json.Marshal(diagnostic) + separator := 0 + if len(encoded) > 1 { + separator = 1 + } + if len(encoded)+separator+len(entry)+1+len(omission)+1 <= maxRepairDiagnosticBytes { + if separator != 0 { + encoded = append(encoded, ',') + } + encoded = append(encoded, entry...) + continue + } + + available := maxRepairDiagnosticBytes - len(encoded) - separator - 1 - len(omission) - 1 + if separator != 0 { + encoded = append(encoded, ',') + } + encoded = append(encoded, truncateDiagnosticJSONValue(diagnostic, available)...) + break + } + if len(encoded) > 1 { + encoded = append(encoded, ',') + } + encoded = append(encoded, omission...) + encoded = append(encoded, ']') + return string(encoded) +} + +func truncateDiagnosticJSONValue(value string, maxBytes int) []byte { + if maxBytes < len(`""`) { + return []byte(`""`) + } + + low, high := 0, len(value) + best := []byte(`""`) + for low <= high { + mid := low + (high-low)/2 + for mid > 0 && mid < len(value) && !utf8.RuneStart(value[mid]) { + mid-- + } + candidate, _ := json.Marshal(value[:mid]) + if len(candidate) <= maxBytes { + best = candidate + low = mid + 1 + continue + } + high = mid - 1 + } + return best +} diff --git a/internal/usecase/repairer_test.go b/internal/usecase/repairer_test.go new file mode 100644 index 0000000..f575bf3 --- /dev/null +++ b/internal/usecase/repairer_test.go @@ -0,0 +1,332 @@ +package usecase + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + "testing" + "unicode/utf8" + + "gitea.maximumdirect.net/eric/promptkit/internal/domain" +) + +type recordingRepairClient struct { + requests []domain.GenerateRequest + response *domain.GenerateResponse + err error +} + +func (c *recordingRepairClient) Generate(_ context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) { + c.requests = append(c.requests, req) + return c.response, c.err +} + +func TestDefaultOutputRepairerBuildsFullContextRequest(t *testing.T) { + client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}} + repairer := NewDefaultOutputRepairer(client) + original := []domain.RenderedMessage{ + {Role: "system", Content: "Follow the task.", CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"}}, + {Role: "user", Content: "Summarize the report."}, + } + before := append([]domain.RenderedMessage(nil), original...) + previous := strings.Repeat("candidate ", 12_000) + target := domain.ExecutionTarget{BackendID: "backend", Endpoint: "https://provider.example/v1", Model: "model"} + presence := domain.ExecutionTargetPresence{Temperature: true, TopP: true} + structured := &domain.StructuredOutputSpec{} + + response, err := repairer.Repair(context.Background(), RepairRequest{ + OriginalMessages: original, + PreviousOutput: previous, + ValidationErrors: []string{"invalid JSON"}, + SessionID: "session", + Target: target, + TargetPresence: presence, + StructuredOutput: structured, + Attempt: 1, + MaxAttempts: 3, + Mode: domain.ValidationJSONSchema, + }) + if err != nil { + t.Fatalf("repair: %v", err) + } + if response == nil || response.Content != "corrected" { + t.Fatalf("response = %+v", response) + } + if !reflect.DeepEqual(original, before) { + t.Fatalf("original messages changed: got %#v, want %#v", original, before) + } + if len(client.requests) != 1 { + t.Fatalf("generation requests = %d, want 1", len(client.requests)) + } + + request := client.requests[0] + if request.Prompt.SessionID != "session" || !reflect.DeepEqual(request.Target, target) || request.TargetPresence != presence || request.StructuredOutput != structured { + t.Fatalf("generation request fields = %+v", request) + } + if len(request.Prompt.Messages) != len(original)+2 { + t.Fatalf("message count = %d, want %d", len(request.Prompt.Messages), len(original)+2) + } + if !reflect.DeepEqual(request.Prompt.Messages[:len(original)], original) { + t.Fatalf("original messages = %#v, want %#v", request.Prompt.Messages[:len(original)], original) + } + assistant := request.Prompt.Messages[len(original)] + if assistant.Role != "assistant" || assistant.Content != previous { + t.Fatalf("assistant candidate = %+v", assistant) + } + correction := request.Prompt.Messages[len(original)+1] + if correction.Role != "user" || !strings.Contains(correction.Content, "Repair attempt 1 of 3") || + !strings.Contains(correction.Content, "Preserve valid values") || + !strings.Contains(correction.Content, "Return only corrected JSON") { + t.Fatalf("correction message = %q", correction.Content) + } + if diagnostics := repairDiagnosticsFromMessage(t, correction.Content); !reflect.DeepEqual(diagnostics, []string{"invalid JSON"}) { + t.Fatalf("diagnostics = %#v", diagnostics) + } +} + +func TestDefaultOutputRepairerDoesNotAccumulateCandidates(t *testing.T) { + client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}} + repairer := NewDefaultOutputRepairer(client) + backing := make([]domain.RenderedMessage, 1, 4) + backing[0] = domain.RenderedMessage{Role: "user", Content: "Original task"} + before := append([]domain.RenderedMessage(nil), backing...) + + for _, candidate := range []string{"first invalid", "second invalid"} { + _, err := repairer.Repair(context.Background(), RepairRequest{ + OriginalMessages: backing, + PreviousOutput: candidate, + Attempt: 1, + MaxAttempts: 3, + Mode: domain.ValidationJSON, + }) + if err != nil { + t.Fatalf("repair %q: %v", candidate, err) + } + } + if !reflect.DeepEqual(backing, before) { + t.Fatalf("caller messages changed: got %#v, want %#v", backing, before) + } + if len(client.requests) != 2 { + t.Fatalf("generation requests = %d, want 2", len(client.requests)) + } + for index, request := range client.requests { + messages := request.Prompt.Messages + if len(messages) != 3 || messages[0] != backing[0] || messages[1].Role != "assistant" || messages[1].Content != []string{"first invalid", "second invalid"}[index] { + t.Fatalf("request %d messages = %#v", index, messages) + } + } +} + +func TestDefaultOutputRepairerOmitsEmptyCandidateMessage(t *testing.T) { + for _, tc := range []struct { + name string + candidate string + }{ + {name: "empty", candidate: ""}, + {name: "whitespace", candidate: " \n\t "}, + } { + t.Run(tc.name, func(t *testing.T) { + client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}} + repairer := NewDefaultOutputRepairer(client) + _, err := repairer.Repair(context.Background(), RepairRequest{ + OriginalMessages: []domain.RenderedMessage{{Role: "user", Content: "Original task"}}, + PreviousOutput: tc.candidate, + Attempt: 1, + MaxAttempts: 1, + Mode: domain.ValidationBasic, + }) + if err != nil { + t.Fatalf("repair: %v", err) + } + messages := client.requests[0].Prompt.Messages + if len(messages) != 2 || messages[1].Role != "user" || !strings.Contains(messages[1].Content, "previous response was empty") { + t.Fatalf("messages = %#v", messages) + } + if strings.Contains(messages[1].Content, tc.candidate) && tc.candidate != "" { + t.Fatalf("correction message repeated whitespace candidate: %q", messages[1].Content) + } + }) + } +} + +func TestDefaultOutputRepairerUsesModeSpecificGuidance(t *testing.T) { + tests := []struct { + mode domain.ValidationMode + want string + }{ + {mode: domain.ValidationBasic, want: "Return a nonempty response"}, + {mode: domain.ValidationJSON, want: "Return only corrected JSON"}, + {mode: domain.ValidationJSONSchema, want: "Return only corrected JSON"}, + } + + for _, tc := range tests { + t.Run(string(tc.mode), func(t *testing.T) { + client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}} + _, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: tc.mode, Attempt: 1, MaxAttempts: 1}) + if err != nil { + t.Fatalf("repair: %v", err) + } + if !strings.Contains(client.requests[0].Prompt.Messages[0].Content, tc.want) { + t.Fatalf("correction message = %q", client.requests[0].Prompt.Messages[0].Content) + } + }) + } + + client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "unexpected"}} + _, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: domain.ValidationNone}) + if err == nil || !strings.Contains(err.Error(), "unsupported validation mode") { + t.Fatalf("unsupported mode error = %v", err) + } + if len(client.requests) != 0 { + t.Fatalf("generation requests = %d, want 0", len(client.requests)) + } +} + +func TestFormatRepairDiagnosticsBoundsAndPreservesData(t *testing.T) { + tests := []struct { + name string + errors []string + wantOmission bool + check func(*testing.T, []string) + }{ + { + name: "below limit", + errors: []string{"first", "second"}, + check: func(t *testing.T, got []string) { + t.Helper() + if !reflect.DeepEqual(got, []string{"first", "second"}) { + t.Fatalf("diagnostics = %#v", got) + } + }, + }, + { + name: "at limit", + errors: []string{strings.Repeat("x", maxRepairDiagnosticBytes-4)}, + check: func(t *testing.T, got []string) { + t.Helper() + if len(got) != 1 || len(got[0]) != maxRepairDiagnosticBytes-4 { + t.Fatalf("diagnostics lengths = %#v", got) + } + }, + }, + { + name: "multibyte truncation preserves prior entries", + errors: []string{"first", strings.Repeat("界", maxRepairDiagnosticBytes)}, + wantOmission: true, + check: func(t *testing.T, got []string) { + t.Helper() + if len(got) != 3 || got[0] != "first" || !strings.HasPrefix(strings.Repeat("界", maxRepairDiagnosticBytes), got[1]) { + t.Fatalf("diagnostics = %#v", got) + } + }, + }, + { + name: "many diagnostics", + errors: manyRepairDiagnostics(), + wantOmission: true, + check: func(t *testing.T, got []string) { + t.Helper() + if len(got) < 2 || got[0] != manyRepairDiagnostics()[0] { + t.Fatalf("diagnostics = %#v", got) + } + }, + }, + { + name: "invalid UTF-8", + errors: []string{"broken\xffinput"}, + check: func(t *testing.T, got []string) { + t.Helper() + if !reflect.DeepEqual(got, []string{"broken\uFFFDinput"}) { + t.Fatalf("diagnostics = %#v", got) + } + }, + }, + { + name: "one huge diagnostic", + errors: []string{strings.Repeat("x", maxRepairDiagnosticBytes*2)}, + wantOmission: true, + check: func(t *testing.T, got []string) { + t.Helper() + if len(got) != 2 || !strings.HasPrefix(strings.Repeat("x", maxRepairDiagnosticBytes*2), got[0]) { + t.Fatalf("diagnostics = %#v", got) + } + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + before := append([]string(nil), tc.errors...) + encoded := formatRepairDiagnostics(tc.errors) + if len(encoded) > maxRepairDiagnosticBytes || !utf8.ValidString(encoded) { + t.Fatalf("encoded diagnostic length/UTF-8 = (%d, %v)", len(encoded), utf8.ValidString(encoded)) + } + var got []string + if err := json.Unmarshal([]byte(encoded), &got); err != nil { + t.Fatalf("decode diagnostics: %v; encoded=%q", err, encoded) + } + if !reflect.DeepEqual(tc.errors, before) { + t.Fatalf("input diagnostics changed: got %#v, want %#v", tc.errors, before) + } + if hasOmission := len(got) > 0 && got[len(got)-1] == omittedRepairDiagnostics; hasOmission != tc.wantOmission { + t.Fatalf("omission = %v, want %v; diagnostics=%#v", hasOmission, tc.wantOmission, got) + } + tc.check(t, got) + }) + } +} + +func manyRepairDiagnostics() []string { + diagnostics := make([]string, 1_000) + for index := range diagnostics { + diagnostics[index] = fmt.Sprintf("diagnostic %04d %s", index, strings.Repeat("x", 128)) + } + return diagnostics +} + +func TestDefaultOutputRepairerPropagatesGenerationFailures(t *testing.T) { + expected := errors.New("generation failed") + for _, tc := range []struct { + name string + client *recordingRepairClient + want error + }{ + {name: "nil client", want: nil}, + {name: "generation error", client: &recordingRepairClient{err: expected}, want: expected}, + {name: "nil response", client: &recordingRepairClient{}, want: nil}, + } { + t.Run(tc.name, func(t *testing.T) { + var repairer OutputRepairer + if tc.client == nil { + repairer = NewDefaultOutputRepairer(nil) + } else { + repairer = NewDefaultOutputRepairer(tc.client) + } + response, err := repairer.Repair(context.Background(), RepairRequest{Mode: domain.ValidationJSON}) + if response != nil || err == nil { + t.Fatalf("response/error = (%+v, %v)", response, err) + } + if tc.want != nil && !errors.Is(err, tc.want) { + t.Fatalf("error = %v, want %v", err, tc.want) + } + }) + } +} + +func repairDiagnosticsFromMessage(t *testing.T, message string) []string { + t.Helper() + const marker = "Validation diagnostics (data):\n" + index := strings.Index(message, marker) + if index < 0 { + t.Fatalf("missing diagnostics marker in %q", message) + } + encoded := message[index+len(marker):] + var diagnostics []string + if err := json.Unmarshal([]byte(encoded), &diagnostics); err != nil { + t.Fatalf("decode diagnostics: %v", err) + } + return diagnostics +}