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

Add bounded structured-output repair

Browse Source
This commit is contained in:
Eric Rakestraw
2026-05-04 21:35:31 -05:00
parent 328703df36
commit ea05945457
3 changed files with 431 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
internal/usecase/repairer.go Normal file
View File

@@ -0,0 +1,71 @@
package usecase
import (
"context"
"errors"
"fmt"
"strings"
"gitea.maximumdirect.net/eric/scriptorium/internal/domain"
"gitea.maximumdirect.net/eric/scriptorium/internal/llm"
)
type OutputRepairer interface {
Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error)
}
type RepairRequest struct {
PreviousOutput string
ValidationErrors []string
Target domain.ModelTarget
Attempt int
MaxAttempts int
Mode domain.ValidationMode
}
type defaultOutputRepairer struct {
llm llm.Client
}
func NewDefaultOutputRepairer(llmClient llm.Client) OutputRepairer {
return &defaultOutputRepairer{llm: llmClient}
}
func (r *defaultOutputRepairer) Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error) {
if r.llm == nil {
return nil, errors.New("llm client is required for repair")
}
errs := "(none provided)"
if len(req.ValidationErrors) > 0 {
errs = strings.Join(req.ValidationErrors, "\n")
}
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,
),
},
}}
resp, err := r.llm.Generate(ctx, domain.GenerateRequest{Prompt: prompt, Target: req.Target})
if err != nil {
return nil, err
}
if resp == nil {
return nil, errors.New("repair llm returned nil response")
}
return resp, nil
}

87
internal/usecase/runner.go
View File

@@ -33,6 +33,7 @@ type Runner struct {
renderer prompt.Renderer
llm llm.Client
validator validate.Validator
repairer OutputRepairer
}
func NewRunner(
@@ -41,6 +42,17 @@ func NewRunner(
renderer prompt.Renderer,
llmClient llm.Client,
validator validate.Validator,
) *Runner {
return NewRunnerWithRepairer(profiles, artifacts, renderer, llmClient, validator, nil)
}
func NewRunnerWithRepairer(
profiles profile.Repository,
artifacts artifact.Reader,
renderer prompt.Renderer,
llmClient llm.Client,
validator validate.Validator,
repairer OutputRepairer,
) *Runner {
return &Runner{
profiles: profiles,
@@ -48,6 +60,7 @@ func NewRunner(
renderer: renderer,
llm: llmClient,
validator: validator,
repairer: repairer,
}
}
@@ -96,18 +109,38 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
}
outputArtifact := buildOutputArtifact(genResp.Content, effectiveContract.Format)
validationResult := domain.ValidationResult{
Status: domain.ValidationSkipped,
Mode: effectiveContract.ValidationMode,
SchemaPath: effectiveContract.SchemaPath,
RepairAttempts: effectiveContract.RepairAttempts,
IsValid: true,
validationResult, err := r.validateOutput(ctx, &outputArtifact, effectiveContract, 0)
if err != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
}
if r.validator != nil && effectiveContract.ValidationMode != domain.ValidationNone {
validationResult, err = r.validator.Validate(ctx, &outputArtifact, effectiveContract)
if err != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
if r.shouldAttemptRepair(effectiveContract, validationResult) {
attemptsUsed := 0
for attemptsUsed < effectiveContract.RepairAttempts && validationResult.Status == domain.ValidationFailed {
attemptsUsed++
repairResp, repairErr := r.repairer.Repair(ctx, RepairRequest{
PreviousOutput: genResp.Content,
ValidationErrors: validationResult.Errors,
Target: effectiveModel,
Attempt: attemptsUsed,
MaxAttempts: effectiveContract.RepairAttempts,
Mode: effectiveContract.ValidationMode,
})
if repairErr != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, repairErr)
}
if repairResp == nil {
return nil, fmt.Errorf("%w: repairer returned nil response", ErrValidation)
}
genResp = repairResp
outputArtifact = buildOutputArtifact(genResp.Content, effectiveContract.Format)
validationResult, err = r.validateOutput(ctx, &outputArtifact, effectiveContract, attemptsUsed)
if err != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
}
}
}
@@ -129,6 +162,38 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
}, nil
}
func (r *Runner) validateOutput(ctx context.Context, artifact *domain.Artifact, contract domain.OutputContract, attemptsUsed int) (domain.ValidationResult, error) {
if r.validator == nil || contract.ValidationMode == domain.ValidationNone {
return domain.ValidationResult{
Status: domain.ValidationSkipped,
Mode: contract.ValidationMode,
SchemaPath: contract.SchemaPath,
RepairAttempts: attemptsUsed,
IsValid: true,
}, nil
}
res, err := r.validator.Validate(ctx, artifact, contract)
if err != nil {
return domain.ValidationResult{}, err
}
res.RepairAttempts = attemptsUsed
return res, nil
}
func (r *Runner) shouldAttemptRepair(contract domain.OutputContract, validationResult domain.ValidationResult) bool {
if r.repairer == nil {
return false
}
if contract.RepairAttempts <= 0 {
return false
}
if validationResult.Status != domain.ValidationFailed {
return false
}
return contract.ValidationMode == domain.ValidationJSON || contract.ValidationMode == domain.ValidationJSONSchema
}
func mergeModelTarget(base domain.ModelTarget, override *domain.ModelTarget) domain.ModelTarget {
if override == nil {
return base

284
internal/usecase/runner_test.go
View File

@@ -87,6 +87,29 @@ func (f *fakeValidator) Validate(ctx context.Context, artifact *domain.Artifact,
return f.result, nil
}
type fakeRepairer struct {
responses []*domain.GenerateResponse
err error
calls int
lastReq RepairRequest
}
func (f *fakeRepairer) Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error) {
f.calls++
f.lastReq = req
if f.err != nil {
return nil, f.err
}
if len(f.responses) == 0 {
return nil, errors.New("no repair response configured")
}
idx := f.calls - 1
if idx >= len(f.responses) {
idx = len(f.responses) - 1
}
return f.responses[idx], nil
}
func TestRunnerRunSuccessful(t *testing.T) {
repo := &fakeProfileRepo{
profile: &domain.PromptProfile{
@@ -377,6 +400,267 @@ func TestRunnerRunValidationFailureWithRealValidatorPreservesRawOutput(t *testin
}
}
func TestRunnerRunNoRepairWhenDisabled(t *testing.T) {
repairer := &fakeRepairer{
responses: []*domain.GenerateResponse{{Content: `{"ok":true}`}},
}
runner := NewRunnerWithRepairer(
&fakeProfileRepo{profile: &domain.PromptProfile{
ID: "p-json",
Version: "1",
OutputFormat: domain.FormatJSON,
ModelDefaults: domain.ModelTarget{
Endpoint: "ep",
Model: "m",
},
Validation: domain.OutputContract{
ValidationMode: domain.ValidationJSON,
Format: domain.FormatJSON,
RepairAttempts: 0,
},
}},
&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
&fakeRenderer{rendered: &domain.RenderedPrompt{}},
&fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":`}},
validate.NewStandardValidator(t.TempDir()),
repairer,
)
res, err := runner.Run(context.Background(), domain.RunRequest{
ProfileID: "p-json",
Inputs: map[string]domain.ArtifactRef{
"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
},
})
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
if repairer.calls != 0 {
t.Fatalf("expected no repair calls, got %d", repairer.calls)
}
if res.Validation.Status != domain.ValidationFailed {
t.Fatalf("expected failed validation, got %q", res.Validation.Status)
}
if res.RawOutput != `{"broken":` {
t.Fatalf("expected original output preserved, got %q", res.RawOutput)
}
if res.Validation.RepairAttempts != 0 {
t.Fatalf("expected repair attempts 0, got %d", res.Validation.RepairAttempts)
}
}
func TestRunnerRunSuccessfulRepairAfterInvalidJSON(t *testing.T) {
repairer := &fakeRepairer{
responses: []*domain.GenerateResponse{{Content: `{"ok":true}`, Usage: domain.TokenUsage{TotalTokens: 5}}},
}
runner := NewRunnerWithRepairer(
&fakeProfileRepo{profile: &domain.PromptProfile{
ID: "p-json",
Version: "1",
OutputFormat: domain.FormatJSON,
ModelDefaults: domain.ModelTarget{
Endpoint: "ep",
Model: "m",
},
Validation: domain.OutputContract{
ValidationMode: domain.ValidationJSON,
Format: domain.FormatJSON,
RepairAttempts: 1,
},
}},
&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
&fakeRenderer{rendered: &domain.RenderedPrompt{}},
&fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":`, Usage: domain.TokenUsage{TotalTokens: 3}}},
validate.NewStandardValidator(t.TempDir()),
repairer,
)
res, err := runner.Run(context.Background(), domain.RunRequest{
ProfileID: "p-json",
Inputs: map[string]domain.ArtifactRef{
"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
},
})
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
if repairer.calls != 1 {
t.Fatalf("expected one repair call, got %d", repairer.calls)
}
if res.Validation.Status != domain.ValidationPassed {
t.Fatalf("expected passed validation, got %q", res.Validation.Status)
}
if res.Validation.RepairAttempts != 1 {
t.Fatalf("expected repair attempts 1, got %d", res.Validation.RepairAttempts)
}
if res.RawOutput != `{"ok":true}` {
t.Fatalf("expected repaired output, got %q", res.RawOutput)
}
}
func TestRunnerRunSuccessfulRepairAfterSchemaFailure(t *testing.T) {
tmp := t.TempDir()
if err := os.WriteFile(filepath.Join(tmp, "schema.json"), []byte(`{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"required": ["name"],
"properties": {
"name": {"type": "string"}
}
}`), 0644); err != nil {
t.Fatal(err)
}
repairer := &fakeRepairer{
responses: []*domain.GenerateResponse{{Content: `{"name":"eris"}`}},
}
runner := NewRunnerWithRepairer(
&fakeProfileRepo{profile: &domain.PromptProfile{
ID: "p-json",
Version: "1",
OutputFormat: domain.FormatJSON,
ModelDefaults: domain.ModelTarget{
Endpoint: "ep",
Model: "m",
},
Validation: domain.OutputContract{
ValidationMode: domain.ValidationJSONSchema,
SchemaPath: "schema.json",
Format: domain.FormatJSON,
RepairAttempts: 1,
},
}},
&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
&fakeRenderer{rendered: &domain.RenderedPrompt{}},
&fakeLLM{resp: &domain.GenerateResponse{Content: `{"count":1}`}},
validate.NewStandardValidator(tmp),
repairer,
)
res, err := runner.Run(context.Background(), domain.RunRequest{
ProfileID: "p-json",
Inputs: map[string]domain.ArtifactRef{
"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
},
})
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
if res.Validation.Status != domain.ValidationPassed {
t.Fatalf("expected passed validation, got %q", res.Validation.Status)
}
if res.Validation.RepairAttempts != 1 {
t.Fatalf("expected repair attempts 1, got %d", res.Validation.RepairAttempts)
}
if res.RawOutput != `{"name":"eris"}` {
t.Fatalf("expected repaired output, got %q", res.RawOutput)
}
}
func TestRunnerRunFailedRepairPreservesRawOutputAndErrors(t *testing.T) {
repairer := &fakeRepairer{
responses: []*domain.GenerateResponse{
{Content: `{"repair1":`},
{Content: `{"repair2":`},
},
}
runner := NewRunnerWithRepairer(
&fakeProfileRepo{profile: &domain.PromptProfile{
ID: "p-json",
Version: "1",
OutputFormat: domain.FormatJSON,
ModelDefaults: domain.ModelTarget{
Endpoint: "ep",
Model: "m",
},
Validation: domain.OutputContract{
ValidationMode: domain.ValidationJSON,
Format: domain.FormatJSON,
RepairAttempts: 2,
},
}},
&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
&fakeRenderer{rendered: &domain.RenderedPrompt{}},
&fakeLLM{resp: &domain.GenerateResponse{Content: `{"initial":`}},
validate.NewStandardValidator(t.TempDir()),
repairer,
)
res, err := runner.Run(context.Background(), domain.RunRequest{
ProfileID: "p-json",
Inputs: map[string]domain.ArtifactRef{
"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
},
})
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
if res.Validation.Status != domain.ValidationFailed {
t.Fatalf("expected failed validation, got %q", res.Validation.Status)
}
if len(res.Validation.Errors) == 0 {
t.Fatal("expected validation errors after failed repair")
}
if res.Validation.RepairAttempts != 2 {
t.Fatalf("expected repair attempts 2, got %d", res.Validation.RepairAttempts)
}
if res.RawOutput != `{"repair2":` {
t.Fatalf("expected final repaired output preserved, got %q", res.RawOutput)
}
}
func TestRunnerRunRepairAttemptsBounded(t *testing.T) {
repairer := &fakeRepairer{
responses: []*domain.GenerateResponse{
{Content: `{"repair1":`},
{Content: `{"repair2":`},
{Content: `{"repair3":`},
},
}
runner := NewRunnerWithRepairer(
&fakeProfileRepo{profile: &domain.PromptProfile{
ID: "p-json",
Version: "1",
OutputFormat: domain.FormatJSON,
ModelDefaults: domain.ModelTarget{
Endpoint: "ep",
Model: "m",
},
Validation: domain.OutputContract{
ValidationMode: domain.ValidationJSON,
Format: domain.FormatJSON,
RepairAttempts: 1,
},
}},
&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
&fakeRenderer{rendered: &domain.RenderedPrompt{}},
&fakeLLM{resp: &domain.GenerateResponse{Content: `{"initial":`}},
validate.NewStandardValidator(t.TempDir()),
repairer,
)
res, err := runner.Run(context.Background(), domain.RunRequest{
ProfileID: "p-json",
Inputs: map[string]domain.ArtifactRef{
"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
},
})
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
if repairer.calls != 1 {
t.Fatalf("expected repair calls bounded to 1, got %d", repairer.calls)
}
if res.Validation.RepairAttempts != 1 {
t.Fatalf("expected repair attempts 1, got %d", res.Validation.RepairAttempts)
}
}
func minimalProfile() *domain.PromptProfile {
return &domain.PromptProfile{
ID: "p",