Build bounded output repair requests
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
332
internal/usecase/repairer_test.go
Normal 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
|
||||
}
|
||||
Reference in New Issue
Block a user