From ae6f1a98657a0ae6f600384ba066d84e9fdaee39 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 25 Aug 2026 09:55:19 +0000 Subject: [PATCH] Enable bounded output repair in the engine --- docs/roadmap/implementation.md | 2 + engine.go | 23 ++++---- engine_test.go | 14 ++++- generation_error_contract_test.go | 38 +++++++++++++ prepared_execution_contract_test.go | 63 +++++++++++++++++++-- public_contract_test.go | 85 ++++++++++++++++++++++++++++- types.go | 19 +++---- 7 files changed, 213 insertions(+), 31 deletions(-) diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 9775d31..72d1167 100644 --- a/docs/roadmap/implementation.md +++ b/docs/roadmap/implementation.md @@ -515,6 +515,8 @@ rather than touching it mechanically. Stage 5 is complete when ordinary and prepared public workflows repair through either client path with unchanged API shape and correct public result and error semantics. +**Status:** Complete. + ## Stage 6: Publish Canonical Documentation And Run Maintainer Validation ### Objective diff --git a/engine.go b/engine.go index 4673e37..8eef840 100644 --- a/engine.go +++ b/engine.go @@ -430,7 +430,7 @@ func NewEngine(cfg Config, opts ...Option) (*Engine, error) { } return &Engine{ - runner: usecase.NewRunner( + runner: usecase.NewRunnerWithRepairer( promptDefs, profiles, backendRegistry, @@ -438,6 +438,7 @@ func NewEngine(cfg Config, opts ...Option) (*Engine, error) { prompt.NewGoRenderer(), llmClient, validator, + usecase.NewDefaultOutputRepairer(llmClient), capacityManager, ), }, nil @@ -639,10 +640,12 @@ func (e *Engine) PrepareExecution(ctx context.Context, req RunRequest) (*Prepare // generated output. // // A content-validation failure is a successful run whose -// RunResult.Validation has Status ValidationFailed. An inability to perform -// validation returns an error matching ErrValidation and no partial result. -// The public Engine does not perform output repair, so validation is -// single-pass even when OutputContract.RepairAttempts is positive. +// RunResult.Validation has Status ValidationFailed. When its output contract +// has a positive repair budget, a failed eligible validation can make bounded +// additional model calls and stops at the first valid candidate. Exhaustion +// returns the final failed validation result with cumulative usage and actual +// repair attempts. An inability to generate or validate returns an error and +// no partial result. // // Run can return every error category documented by [Engine.Prepare], plus // ErrCapacityExceeded and ErrLLMGenerate. An engine admission rejection is @@ -686,17 +689,17 @@ func (e *Engine) Run(ctx context.Context, req RunRequest) (*RunResult, error) { // // The supplied context governs this execution attempt independently of the // preparation context. It covers credential revalidation, admission, -// generation, validation, and any internal repair. Result timing begins after -// the claim and excludes preparation and consumer-held delay. +// generation, validation, and any bounded output repair. Result timing begins +// after the claim and excludes preparation and consumer-held delay. // // RunPrepared can return ErrInvalidRequest, ErrAPIKeyEnvMissing, // ErrCapacityExceeded, ErrLLMGenerate, or ErrValidation as applicable while // preserving documented collaborator and context identities. An engine // admission rejection is discoverable as [CapacityError] and still matches // ErrCapacityExceeded. A built-in OpenAI-compatible non-2xx response is -// discoverable as [GenerationError]. A completed content-validation rejection -// is returned in RunResult, not as an operational error. An operational error -// returns no partial RunResult. +// discoverable as [GenerationError]. A completed content-validation rejection, +// including repair exhaustion, is returned in RunResult, not as an operational +// error. An operational error returns no partial RunResult. func (e *Engine) RunPrepared(ctx context.Context, prepared *PreparedExecution) (*RunResult, error) { if e == nil || e.runner == nil { return nil, fmt.Errorf("%w: engine is nil", ErrInvalidConfig) diff --git a/engine_test.go b/engine_test.go index db5f462..a60ef34 100644 --- a/engine_test.go +++ b/engine_test.go @@ -3486,9 +3486,10 @@ extra_params: } type fakeLLMClient struct { - response *promptkit.GenerateResponse - err error - requests []promptkit.GenerateRequest + response *promptkit.GenerateResponse + responses []*promptkit.GenerateResponse + err error + requests []promptkit.GenerateRequest } type recordingArtifactReader struct { @@ -3664,5 +3665,12 @@ func (f *fakeLLMClient) Generate(_ context.Context, req promptkit.GenerateReques if f.err != nil { return nil, f.err } + if len(f.responses) > 0 { + index := len(f.requests) - 1 + if index >= len(f.responses) { + return nil, fmt.Errorf("no response configured for generation %d", index+1) + } + return f.responses[index], nil + } return f.response, nil } diff --git a/generation_error_contract_test.go b/generation_error_contract_test.go index b6beaac..a6e77b8 100644 --- a/generation_error_contract_test.go +++ b/generation_error_contract_test.go @@ -39,6 +39,44 @@ func TestBuiltInGenerationError(t *testing.T) { assertGenerationError(t, err, http.StatusServiceUnavailable, "", "", "") } +func TestBuiltInRepairGenerationError(t *testing.T) { + const ( + codeMarker = "repair-code-marker" + typeMarker = "repair-type-marker" + messageMarker = "repair-message-marker" + ) + calls := 0 + config := contractConfig(frameworkSchemaDir) + config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + calls++ + if calls == 1 { + body := `{"choices":[{"message":{"content":"not-json"}}]}` + return &http.Response{StatusCode: http.StatusOK, ContentLength: int64(len(body)), Body: io.NopCloser(strings.NewReader(body))}, nil + } + body := `{"error":{"code":"` + codeMarker + `","type":"` + typeMarker + `","message":"` + messageMarker + `"}}` + return &http.Response{StatusCode: http.StatusUnprocessableEntity, ContentLength: int64(len(body)), Body: io.NopCloser(strings.NewReader(body))}, nil + })} + engine, err := promptkit.NewEngine(config) + if err != nil { + t.Fatalf("NewEngine: %v", err) + } + req := generationErrorRunRequest() + req.Validation = &promptkit.OutputContract{ + Format: promptkit.FormatJSON, + ValidationMode: promptkit.ValidationJSON, + RepairAttempts: 1, + } + + result, err := engine.Run(context.Background(), req) + if result != nil { + t.Fatalf("Run result = %#v, want nil", result) + } + if calls != 2 { + t.Fatalf("provider calls = %d, want 2", calls) + } + assertGenerationError(t, err, http.StatusUnprocessableEntity, codeMarker, typeMarker, messageMarker) +} + func assertGenerationError(t *testing.T, err error, statusCode int, code, providerType, message string) { t.Helper() diff --git a/prepared_execution_contract_test.go b/prepared_execution_contract_test.go index 646aee1..87a3300 100644 --- a/prepared_execution_contract_test.go +++ b/prepared_execution_contract_test.go @@ -256,6 +256,50 @@ func TestPreparedExecutionLifecycleAndEngineBinding(t *testing.T) { } } +func TestPreparedExecutionRepairsEmptyBasicOutput(t *testing.T) { + client := &preparedRecordingClient{responses: []*promptkit.GenerateResponse{ + { + Content: "", + Usage: promptkit.TokenUsage{PromptTokens: 3, CompletionTokens: 5, TotalTokens: 8}, + }, + { + Content: "Corrected summary.", + Usage: promptkit.TokenUsage{PromptTokens: 7, CompletionTokens: 11, TotalTokens: 18}, + }, + }} + engine := newPreparedContractEngine(t, client, "Summarize the source.") + prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{ + PromptID: "prepared", + Validation: &promptkit.OutputContract{ + Format: promptkit.FormatMarkdown, + ValidationMode: promptkit.ValidationBasic, + RepairAttempts: 1, + }, + }) + if err != nil { + t.Fatalf("prepare execution: %v", err) + } + details := prepared.Details() + + result, err := engine.RunPrepared(context.Background(), prepared) + if err != nil { + t.Fatalf("run prepared: %v", err) + } + if result.RawOutput != "Corrected summary." || result.Validation.Status != promptkit.ValidationPassed || + result.Validation.RepairAttempts != 1 || result.Usage != (promptkit.TokenUsage{PromptTokens: 10, CompletionTokens: 16, TotalTokens: 26}) { + t.Fatalf("repaired result = %+v", result) + } + requests := client.snapshot() + if len(requests) != 2 || len(requests[1].Prompt.Messages) != len(details.Messages)+1 || + !reflect.DeepEqual(requests[1].Prompt.Messages[:len(details.Messages)], details.Messages) || + requests[1].Prompt.Messages[len(requests[1].Prompt.Messages)-1].Role != "user" { + t.Fatalf("prepared repair requests = %#v", requests) + } + if _, err := engine.RunPrepared(context.Background(), prepared); !errors.Is(err, promptkit.ErrInvalidRequest) { + t.Fatalf("second RunPrepared error = %v, want ErrInvalidRequest", err) + } +} + func TestPreparedExecutionConcurrentClaimAllowsOneGeneration(t *testing.T) { release := make(chan struct{}) client := &preparedRecordingClient{ @@ -671,12 +715,13 @@ func (r *mutablePreparedArtifactReader) callCount() int { } type preparedRecordingClient struct { - mu sync.Mutex - response *promptkit.GenerateResponse - err error - requests []promptkit.GenerateRequest - started chan struct{} - release <-chan struct{} + mu sync.Mutex + response *promptkit.GenerateResponse + responses []*promptkit.GenerateResponse + err error + requests []promptkit.GenerateRequest + started chan struct{} + release <-chan struct{} } func (c *preparedRecordingClient) Generate( @@ -700,6 +745,12 @@ func (c *preparedRecordingClient) Generate( if c.err != nil { return nil, c.err } + if len(c.responses) > 0 { + if index := len(c.requests) - 1; index < len(c.responses) { + return c.responses[index], nil + } + return nil, fmt.Errorf("no response configured for generation %d", len(c.requests)) + } return c.response, nil } diff --git a/public_contract_test.go b/public_contract_test.go index 56e366c..96499df 100644 --- a/public_contract_test.go +++ b/public_contract_test.go @@ -844,7 +844,7 @@ func TestBackendExtraParamsAreDeeplyCopiedAtConstructionAndLookup(t *testing.T) } } -func TestEngineValidationIsSinglePass(t *testing.T) { +func TestEngineValidationWithZeroRepairBudgetIsSinglePass(t *testing.T) { client := &fakeLLMClient{ response: &promptkit.GenerateResponse{Content: "not-json"}, } @@ -859,7 +859,7 @@ func TestEngineValidationIsSinglePass(t *testing.T) { Validation: &promptkit.OutputContract{ Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, - RepairAttempts: 3, + RepairAttempts: 0, }, }) if err != nil { @@ -874,6 +874,87 @@ func TestEngineValidationIsSinglePass(t *testing.T) { } } +func TestEngineRunRepairsJSONSchemaOutput(t *testing.T) { + client := &fakeLLMClient{responses: []*promptkit.GenerateResponse{ + { + Content: "{}", + Usage: promptkit.TokenUsage{PromptTokens: 2, CompletionTokens: 3, TotalTokens: 5, CachedTokens: 7, CacheWriteTokens: 11}, + }, + { + Content: `{"events":[{"title":"Repaired event"}]}`, + Usage: promptkit.TokenUsage{PromptTokens: 13, CompletionTokens: 17, TotalTokens: 19, CachedTokens: 23, CacheWriteTokens: 29}, + }, + }} + engine := newContractEngineWithOptions(t, frameworkSchemaDir, promptkit.WithLLMClient(client)) + + result, err := engine.Run(context.Background(), promptkit.RunRequest{ + PromptID: frameworkStructuredEventsPromptID, + SessionID: " repair-session ", + Inputs: map[string]promptkit.ArtifactRef{ + "transcript": promptkit.Inline("Rin opens the gate."), + "glossary": promptkit.Inline("gate: A guarded passage."), + }, + Validation: &promptkit.OutputContract{ + Format: promptkit.FormatJSON, + ValidationMode: promptkit.ValidationJSONSchema, + SchemaPath: "structured_events.schema.json", + RepairAttempts: 1, + }, + }) + if err != nil { + t.Fatalf("run: %v", err) + } + if result.RawOutput != client.responses[1].Content || result.Validation.Status != promptkit.ValidationPassed || + result.Validation.RepairAttempts != 1 { + t.Fatalf("repaired result = %+v", result) + } + wantUsage := promptkit.TokenUsage{PromptTokens: 15, CompletionTokens: 20, TotalTokens: 24, CachedTokens: 30, CacheWriteTokens: 40} + if result.Usage != wantUsage { + t.Fatalf("usage = %+v, want %+v", result.Usage, wantUsage) + } + if len(client.requests) != 2 { + t.Fatalf("generation calls = %d, want 2", len(client.requests)) + } + initial, repaired := client.requests[0], client.requests[1] + if initial.Prompt.SessionID != "repair-session" || repaired.Prompt.SessionID != initial.Prompt.SessionID || + !reflect.DeepEqual(repaired.Target, initial.Target) || repaired.TargetPresence != initial.TargetPresence || + !reflect.DeepEqual(repaired.StructuredOutput, initial.StructuredOutput) { + t.Fatalf("generation request state drifted: initial=%+v repaired=%+v", initial, repaired) + } + if initial.StructuredOutput == nil || initial.StructuredOutput.JSONSchema == nil { + t.Fatalf("expected structured output on initial request: %+v", initial) + } +} + +func TestEngineRunReturnsFinalResultAfterRepairExhaustion(t *testing.T) { + client := &fakeLLMClient{responses: []*promptkit.GenerateResponse{ + {Content: "not-json", Usage: promptkit.TokenUsage{PromptTokens: 2, CompletionTokens: 3, TotalTokens: 5}}, + {Content: "still-not-json", Usage: promptkit.TokenUsage{PromptTokens: 7, CompletionTokens: 11, TotalTokens: 18}}, + }} + engine := newContractEngineWithOptions(t, frameworkSchemaDir, promptkit.WithLLMClient(client)) + + result, err := engine.Run(context.Background(), promptkit.RunRequest{ + PromptID: frameworkMarkdownSummaryPromptID, + Inputs: map[string]promptkit.ArtifactRef{ + "transcript": promptkit.Inline("Rin opens the gate."), + "glossary": promptkit.Inline("gate: A guarded passage."), + }, + Validation: &promptkit.OutputContract{ + Format: promptkit.FormatJSON, + ValidationMode: promptkit.ValidationJSON, + RepairAttempts: 1, + }, + }) + if err != nil || result == nil { + t.Fatalf("run = (%+v, %v), want exhausted result", result, err) + } + if result.RawOutput != "still-not-json" || result.Validation.Status != promptkit.ValidationFailed || + result.Validation.RepairAttempts != 1 || len(result.Validation.Errors) == 0 || + result.Usage != (promptkit.TokenUsage{PromptTokens: 9, CompletionTokens: 14, TotalTokens: 23}) { + t.Fatalf("exhausted result = %+v", result) + } +} + func TestRepeatedOptionsUseLastValueInEachCategory(t *testing.T) { profile := promptkit.Profile{ID: "profile", Endpoint: "http://example.test/v1", Model: "model"} diff --git a/types.go b/types.go index 4c7b74c..732bc06 100644 --- a/types.go +++ b/types.go @@ -565,8 +565,8 @@ type ExecutionTargetPresence struct { // JSON representation. // // A non-nil RunRequest.Validation replaces the complete prompt contract. It -// does not merge fields. The public Engine validates generated output once and -// does not install an output repairer. +// does not merge fields. The public Engine performs bounded correction after a +// failed eligible validation when RepairAttempts is positive. type OutputContract struct { // Format selects generated artifact metadata. An empty value in a non-nil // request replacement defaults to FormatText. @@ -577,9 +577,9 @@ type OutputContract struct { // SchemaPath is required when ValidationMode is ValidationJSONSchema and is // ignored by other modes. SchemaPath string `json:"schema_path"` - // RepairAttempts is a non-negative requested repair limit. Zero requests no - // repairs. The public Engine performs no repairs even when this value is - // positive, so its runs report zero attempts used. + // RepairAttempts is an additional generation-call budget from zero through + // three. Zero is single-pass. A positive value is valid only with basic, + // json, or json_schema validation. RepairAttempts int `json:"repair_attempts"` } @@ -595,8 +595,8 @@ type ValidationResult struct { Errors []string `json:"errors,omitempty"` // SchemaPath is the effective schema path for JSON Schema validation. SchemaPath string `json:"schema_path,omitempty"` - // RepairAttempts is the number of repairs actually attempted. It is always - // zero for the public Engine. + // RepairAttempts is the number of corrective generation calls actually + // started for this result. RepairAttempts int `json:"repair_attempts"` // IsValid is true for ValidationPassed and ValidationSkipped and false for // ValidationFailed. @@ -714,9 +714,8 @@ type GenerateRequest struct { // GenerateResponse is returned by an injected LLM client and has a stable JSON // representation. type GenerateResponse struct { - // Content is the generated output. It must be non-empty when using the - // built-in client; injected clients may return empty content for Promptkit - // validation to classify. + // Content is the generated output. It may be explicitly empty; Promptkit + // applies the effective output contract to classify it. Content string `json:"content"` // Usage is the client's token accounting. Usage TokenUsage `json:"usage"`