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

Run bounded output repair attempts

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-25 09:51:28 +00:00
parent 00ee5893e9
commit ee99dc9478
4 changed files with 109 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
internal/usecase/prepared_execution_test.go
View File

@@ -3,10 +3,12 @@ package usecase
import (
"context"
"errors"
"fmt"
"reflect"
"testing"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
"gitea.maximumdirect.net/eric/promptkit/internal/validate"
)
@@ -385,19 +387,20 @@ func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *test
admitter := &fakeRunAdmitter{}
reader := defaultArtifactReader()
renderer := defaultRenderer()
client := &sequenceLLM{responses: []*domain.GenerateResponse{{
Content: `{"broken":true}`,
Usage: domain.TokenUsage{
PromptTokens: 13, CompletionTokens: 17, TotalTokens: 19,
CachedTokens: 23, CacheWriteTokens: 29,
},
}}}
runner := NewRunnerWithRepairer(
&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
nil,
reader,
renderer,
&fakeLLM{resp: &domain.GenerateResponse{
Content: `{"broken":true}`,
Usage: domain.TokenUsage{
PromptTokens: 13, CompletionTokens: 17, TotalTokens: 19,
CachedTokens: 23, CacheWriteTokens: 29,
},
}},
client,
validator,
repairer,
admitter,
@@ -421,6 +424,10 @@ func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *test
if validator.directValidateCalls != 0 || repairer.calls != 1 {
t.Fatalf("validation/repair calls=(direct=%d repair=%d), want (0, 1)", validator.directValidateCalls, repairer.calls)
}
if len(client.requests) != 1 || len(repairer.reqs) != 1 ||
!reflect.DeepEqual(repairer.reqs[0].OriginalMessages, client.requests[0].Prompt.Messages) {
t.Fatalf("initial and repair messages = (%#v, %#v)", client.requests, repairer.reqs)
}
if result.Validation.Status != domain.ValidationPassed || result.Validation.RepairAttempts != 1 {
t.Fatalf("unexpected repaired validation result: %+v", result.Validation)
}
@@ -443,6 +450,7 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
generationFailure := errors.New("generation failed")
validationFailure := errors.New("validation failed")
repairFailure := errors.New("repair failed")
repairInvalidRequest := fmt.Errorf("repair request: %w", llm.ErrInvalidRequest)
tests := []struct {
name string
@@ -450,6 +458,7 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
validation *recordingPreparedValidation
repairer *fakeRepairer
wantError error
wantSource error
}{
{
name: "generation failure",
@@ -474,8 +483,51 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
IsValid: false,
}},
},
repairer: &fakeRepairer{err: repairFailure},
wantError: ErrValidation,
repairer: &fakeRepairer{err: repairFailure},
wantError: ErrLLMGenerate,
wantSource: repairFailure,
},
{
name: "repair invalid request",
validation: &recordingPreparedValidation{
results: []domain.ValidationResult{{
Status: domain.ValidationFailed,
Mode: domain.ValidationJSON,
Errors: []string{"invalid"},
IsValid: false,
}},
},
repairer: &fakeRepairer{err: repairInvalidRequest},
wantError: ErrInvalidRequest,
wantSource: repairInvalidRequest,
},
{
name: "repair cancellation",
validation: &recordingPreparedValidation{
results: []domain.ValidationResult{{
Status: domain.ValidationFailed,
Mode: domain.ValidationJSON,
Errors: []string{"invalid"},
IsValid: false,
}},
},
repairer: &fakeRepairer{err: context.Canceled},
wantError: ErrLLMGenerate,
wantSource: context.Canceled,
},
{
name: "repair deadline",
validation: &recordingPreparedValidation{
results: []domain.ValidationResult{{
Status: domain.ValidationFailed,
Mode: domain.ValidationJSON,
Errors: []string{"invalid"},
IsValid: false,
}},
},
repairer: &fakeRepairer{err: context.DeadlineExceeded},
wantError: ErrLLMGenerate,
wantSource: context.DeadlineExceeded,
},
}
@@ -514,6 +566,9 @@ func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
if result != nil || !errors.Is(err, test.wantError) {
t.Fatalf("run prepared=(%+v, %v), want %v", result, err, test.wantError)
}
if test.wantSource != nil && !errors.Is(err, test.wantSource) {
t.Fatalf("run prepared error = %v, want source %v", err, test.wantSource)
}
if len(admitter.backendIDs) != 1 || admitter.releaseCalls != 1 {
t.Fatalf(
"admission calls=%#v releases=%d, want one each",

27
internal/usecase/runner.go
View File

@@ -184,10 +184,10 @@ func (r *Runner) executePreparedRun(
prepared.StructuredOutput,
))
if err != nil {
if errors.Is(err, llm.ErrInvalidRequest) {
return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
}
return nil, fmt.Errorf("%w: %w", ErrLLMGenerate, err)
return nil, wrapGenerationError(err)
}
if genResp == nil {
return nil, fmt.Errorf("%w: model returned nil response", ErrLLMGenerate)
}
usage := genResp.Usage
@@ -203,6 +203,7 @@ func (r *Runner) executePreparedRun(
attemptsUsed++
repairResp, repairErr := r.repairer.Repair(ctx, RepairRequest{
OriginalMessages: prepared.Messages,
PreviousOutput: genResp.Content,
ValidationErrors: validationResult.Errors,
SessionID: prepared.SessionID,
@@ -214,10 +215,10 @@ func (r *Runner) executePreparedRun(
Mode: prepared.OutputContract.ValidationMode,
})
if repairErr != nil {
return nil, fmt.Errorf("%w: %w", ErrValidation, repairErr)
return nil, wrapGenerationError(repairErr)
}
if repairResp == nil {
return nil, fmt.Errorf("%w: repairer returned nil response", ErrValidation)
return nil, fmt.Errorf("%w: repairer returned nil response", ErrLLMGenerate)
}
genResp = repairResp
@@ -257,6 +258,13 @@ func (r *Runner) executePreparedRun(
}, nil
}
func wrapGenerationError(err error) error {
if errors.Is(err, llm.ErrInvalidRequest) {
return fmt.Errorf("%w: %w", ErrInvalidRequest, err)
}
return fmt.Errorf("%w: %w", ErrLLMGenerate, err)
}
func addTokenUsage(total, next domain.TokenUsage) domain.TokenUsage {
return domain.TokenUsage{
PromptTokens: total.PromptTokens + next.PromptTokens,
@@ -524,7 +532,12 @@ func (r *Runner) shouldAttemptRepair(contract domain.OutputContract, validationR
if validationResult.Status != domain.ValidationFailed {
return false
}
return contract.ValidationMode == domain.ValidationJSON || contract.ValidationMode == domain.ValidationJSONSchema
switch contract.ValidationMode {
case domain.ValidationBasic, domain.ValidationJSON, domain.ValidationJSONSchema:
return true
default:
return false
}
}
func mergeExecutionTarget(base domain.ExecutionTarget, override domain.ExecutionTarget) domain.ExecutionTarget {

29
internal/usecase/runner_test.go
View File

@@ -2143,6 +2143,8 @@ func TestRunnerRepairStateMachine(t *testing.T) {
TopP: true,
TimeoutSeconds: true,
}
emptyThenValid := responses(2)
emptyThenValid[0].Content = " \t "
tests := []struct {
name string
@@ -2165,12 +2167,13 @@ func TestRunnerRepairStateMachine(t *testing.T) {
wantStatus: domain.ValidationPassed,
},
{
name: "basic failure is ineligible despite budget",
name: "empty basic output repairs successfully",
mode: domain.ValidationBasic,
budget: 3,
validationResults: []domain.ValidationResult{failed(domain.ValidationBasic, "empty output")},
responses: responses(1),
wantStatus: domain.ValidationFailed,
validationResults: []domain.ValidationResult{failed(domain.ValidationBasic, "empty output"), passed(domain.ValidationBasic)},
responses: emptyThenValid,
wantRepairs: 1,
wantStatus: domain.ValidationPassed,
},
{
name: "inherited numeric values remain absent",
@@ -2305,6 +2308,9 @@ func TestRunnerRepairStateMachine(t *testing.T) {
!reflect.DeepEqual(req.ValidationErrors, tc.validationResults[index].Errors) {
t.Fatalf("repair request %d prior state = %+v", index, req)
}
if !reflect.DeepEqual(req.OriginalMessages, initialRequest.Prompt.Messages) {
t.Fatalf("repair request %d original messages drifted: %#v", index, req.OriginalMessages)
}
if req.TargetPresence != tc.wantPresence || !reflect.DeepEqual(req.Target, initialRequest.Target) ||
req.SessionID != initialRequest.Prompt.SessionID ||
!reflect.DeepEqual(req.StructuredOutput, initialRequest.StructuredOutput) {
@@ -2318,8 +2324,19 @@ func TestRunnerRepairStateMachine(t *testing.T) {
!reflect.DeepEqual(generated.StructuredOutput, initialRequest.StructuredOutput) {
t.Fatalf("repair generation request %d common fields drifted: %+v", index, generated)
}
if reflect.DeepEqual(generated.Prompt.Messages, initialRequest.Prompt.Messages) {
t.Fatalf("repair generation request %d reused the initial prompt", index)
expectedMessages := len(initialRequest.Prompt.Messages) + 1
if strings.TrimSpace(tc.responses[index].Content) != "" {
expectedMessages++
}
if len(generated.Prompt.Messages) != expectedMessages ||
generated.Prompt.Messages[len(generated.Prompt.Messages)-1].Role != "user" {
t.Fatalf("repair generation request %d messages = %#v", index, generated.Prompt.Messages)
}
if strings.TrimSpace(tc.responses[index].Content) != "" {
assistant := generated.Prompt.Messages[len(generated.Prompt.Messages)-2]
if assistant.Role != "assistant" || assistant.Content != tc.responses[index].Content {
t.Fatalf("repair generation request %d candidate = %+v", index, assistant)
}
}
}