eric/promptkit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Build bounded output repair requests

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-25 09:46:56 +00:00
parent 64d1cffd89
commit 00ee5893e9
3 changed files with 453 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
internal/usecase/repairer.go
View File

@@ -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
}

332
internal/usecase/repairer_test.go Normal file
View File

@@ -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
}