diff --git a/internal/usecase/repairer.go b/internal/usecase/repairer.go new file mode 100644 index 0000000..ac23098 --- /dev/null +++ b/internal/usecase/repairer.go @@ -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 +} diff --git a/internal/usecase/runner.go b/internal/usecase/runner.go index b5a15c3..6d82b70 100644 --- a/internal/usecase/runner.go +++ b/internal/usecase/runner.go @@ -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 diff --git a/internal/usecase/runner_test.go b/internal/usecase/runner_test.go index 6c2a040..929188b 100644 --- a/internal/usecase/runner_test.go +++ b/internal/usecase/runner_test.go @@ -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",