Preserve repair settings and cumulative usage
This commit is contained in:
@@ -228,6 +228,30 @@ type fakeRepairer struct {
|
||||
reqs []RepairRequest
|
||||
}
|
||||
|
||||
type sequenceLLM struct {
|
||||
responses []*domain.GenerateResponse
|
||||
requests []domain.GenerateRequest
|
||||
}
|
||||
|
||||
func (c *sequenceLLM) Generate(_ context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
||||
c.requests = append(c.requests, req)
|
||||
index := len(c.requests) - 1
|
||||
if index >= len(c.responses) {
|
||||
return nil, errors.New("no generation response configured")
|
||||
}
|
||||
return c.responses[index], nil
|
||||
}
|
||||
|
||||
type recordingRepairer struct {
|
||||
next OutputRepairer
|
||||
reqs []RepairRequest
|
||||
}
|
||||
|
||||
func (r *recordingRepairer) Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error) {
|
||||
r.reqs = append(r.reqs, req)
|
||||
return r.next.Repair(ctx, req)
|
||||
}
|
||||
|
||||
type fakeRunAdmitter struct {
|
||||
backendIDs []string
|
||||
err error
|
||||
@@ -2070,50 +2094,255 @@ func TestRunnerRunValidationStillWorks(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunStructuredRepairRemainsBoundedAndUsesEffectiveModelSettings(t *testing.T) {
|
||||
repairer := &fakeRepairer{responses: []*domain.GenerateResponse{{Content: `{"broken":`}, {Content: `{"still":`}}}
|
||||
llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: `{"initial":`}}
|
||||
func TestRunnerRepairStateMachine(t *testing.T) {
|
||||
failed := func(mode domain.ValidationMode, diagnostic string) domain.ValidationResult {
|
||||
return domain.ValidationResult{
|
||||
Status: domain.ValidationFailed,
|
||||
Mode: mode,
|
||||
Errors: []string{diagnostic},
|
||||
IsValid: false,
|
||||
}
|
||||
}
|
||||
passed := func(mode domain.ValidationMode) domain.ValidationResult {
|
||||
return domain.ValidationResult{
|
||||
Status: domain.ValidationPassed,
|
||||
Mode: mode,
|
||||
IsValid: true,
|
||||
}
|
||||
}
|
||||
responses := func(count int) []*domain.GenerateResponse {
|
||||
values := make([]*domain.GenerateResponse, count)
|
||||
for index := range values {
|
||||
unit := index + 1
|
||||
values[index] = &domain.GenerateResponse{
|
||||
Content: fmt.Sprintf(`{"candidate":%d}`, index),
|
||||
Usage: domain.TokenUsage{
|
||||
PromptTokens: unit,
|
||||
CompletionTokens: unit * 10,
|
||||
TotalTokens: unit * 100,
|
||||
CachedTokens: unit * 1000,
|
||||
CacheWriteTokens: unit * 10000,
|
||||
},
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
zeroOverrides := &domain.ExecutionTargetOverride{
|
||||
Temperature: float64Ptr(0),
|
||||
MaxTokens: intPtr(0),
|
||||
TopP: float64Ptr(0),
|
||||
TimeoutSeconds: intPtr(0),
|
||||
}
|
||||
allPresent := domain.ExecutionTargetPresence{
|
||||
Temperature: true,
|
||||
MaxTokens: true,
|
||||
TopP: true,
|
||||
TimeoutSeconds: true,
|
||||
}
|
||||
|
||||
runner := NewRunnerWithRepairer(
|
||||
&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{
|
||||
"exec": {ID: "exec", BackendID: "custom", Model: "profile-model", TimeoutSeconds: 55},
|
||||
}}, fakeBackendResolver{backends: map[string]domain.Backend{
|
||||
"custom": {ID: "custom", Endpoint: "http://backend/v1"},
|
||||
}},
|
||||
tests := []struct {
|
||||
name string
|
||||
mode domain.ValidationMode
|
||||
budget int
|
||||
validationResults []domain.ValidationResult
|
||||
responses []*domain.GenerateResponse
|
||||
execution *domain.ExecutionTargetOverride
|
||||
wantPresence domain.ExecutionTargetPresence
|
||||
wantRepairs int
|
||||
wantStatus domain.ValidationStatus
|
||||
structured bool
|
||||
}{
|
||||
{
|
||||
name: "initial success does not repair",
|
||||
mode: domain.ValidationJSON,
|
||||
budget: 3,
|
||||
validationResults: []domain.ValidationResult{passed(domain.ValidationJSON)},
|
||||
responses: responses(1),
|
||||
wantStatus: domain.ValidationPassed,
|
||||
},
|
||||
{
|
||||
name: "basic failure is ineligible despite budget",
|
||||
mode: domain.ValidationBasic,
|
||||
budget: 3,
|
||||
validationResults: []domain.ValidationResult{failed(domain.ValidationBasic, "empty output")},
|
||||
responses: responses(1),
|
||||
wantStatus: domain.ValidationFailed,
|
||||
},
|
||||
{
|
||||
name: "inherited numeric values remain absent",
|
||||
mode: domain.ValidationJSON,
|
||||
budget: 1,
|
||||
validationResults: []domain.ValidationResult{
|
||||
failed(domain.ValidationJSON, "initial syntax"),
|
||||
passed(domain.ValidationJSON),
|
||||
},
|
||||
responses: responses(2),
|
||||
wantRepairs: 1,
|
||||
wantStatus: domain.ValidationPassed,
|
||||
},
|
||||
{
|
||||
name: "explicit numeric zeros remain present",
|
||||
mode: domain.ValidationJSONSchema,
|
||||
budget: 1,
|
||||
execution: zeroOverrides,
|
||||
structured: true,
|
||||
validationResults: []domain.ValidationResult{
|
||||
failed(domain.ValidationJSONSchema, "initial schema mismatch"),
|
||||
passed(domain.ValidationJSONSchema),
|
||||
},
|
||||
responses: responses(2),
|
||||
wantPresence: allPresent,
|
||||
wantRepairs: 1,
|
||||
wantStatus: domain.ValidationPassed,
|
||||
},
|
||||
{
|
||||
name: "successful repair stops below larger budget",
|
||||
mode: domain.ValidationJSON,
|
||||
budget: 4,
|
||||
validationResults: []domain.ValidationResult{
|
||||
failed(domain.ValidationJSON, "candidate zero"),
|
||||
failed(domain.ValidationJSON, "candidate one"),
|
||||
passed(domain.ValidationJSON),
|
||||
},
|
||||
responses: responses(3),
|
||||
wantRepairs: 2,
|
||||
wantStatus: domain.ValidationPassed,
|
||||
},
|
||||
{
|
||||
name: "failed repairs exhaust exact larger budget",
|
||||
mode: domain.ValidationJSON,
|
||||
budget: 3,
|
||||
validationResults: []domain.ValidationResult{
|
||||
failed(domain.ValidationJSON, "candidate zero"),
|
||||
failed(domain.ValidationJSON, "candidate one"),
|
||||
failed(domain.ValidationJSON, "candidate two"),
|
||||
failed(domain.ValidationJSON, "candidate three"),
|
||||
},
|
||||
responses: responses(4),
|
||||
wantRepairs: 3,
|
||||
wantStatus: domain.ValidationFailed,
|
||||
},
|
||||
}
|
||||
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
llmClient,
|
||||
validate.NewStandardValidator("."),
|
||||
repairer, nil)
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
definition := promptDef(domain.FormatJSON, tc.mode, tc.budget)
|
||||
plan := &recordingPreparedValidation{results: tc.validationResults}
|
||||
if tc.structured {
|
||||
definition.Validation.SchemaPath = "schema.json"
|
||||
plan.schemaDocument = map[string]any{"type": "object"}
|
||||
}
|
||||
validator := &recordingValidationPreparer{plan: plan}
|
||||
client := &sequenceLLM{responses: tc.responses}
|
||||
repairer := &recordingRepairer{next: NewDefaultOutputRepairer(client)}
|
||||
runner := NewRunnerWithRepairer(
|
||||
&fakePromptRepo{def: definition},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{
|
||||
"exec": {ID: "exec", BackendID: "custom", Model: "profile-model"},
|
||||
}},
|
||||
fakeBackendResolver{backends: map[string]domain.Backend{
|
||||
"custom": {ID: "custom", Endpoint: "http://backend.example/v1"},
|
||||
}},
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
client,
|
||||
validator,
|
||||
repairer,
|
||||
nil,
|
||||
)
|
||||
|
||||
res, err := runner.Run(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
Execution: &domain.ExecutionTargetOverride{Endpoint: "http://override/v1", Model: "override-model", TimeoutSeconds: intPtr(22)},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if repairer.calls != 1 || res.Validation.RepairAttempts != 1 {
|
||||
t.Fatalf("expected one bounded repair, calls=%d attempts=%d", repairer.calls, res.Validation.RepairAttempts)
|
||||
}
|
||||
if len(repairer.reqs) != 1 {
|
||||
t.Fatalf("expected one repair request, got %d", len(repairer.reqs))
|
||||
}
|
||||
if repairer.reqs[0].Target.Endpoint != "http://override/v1" || repairer.reqs[0].Target.Model != "override-model" {
|
||||
t.Fatalf("expected repair to use effective target, got %+v", repairer.reqs[0].Target)
|
||||
}
|
||||
if repairer.reqs[0].Target.TimeoutSeconds != 22 {
|
||||
t.Fatalf("expected repair to use effective timeout, got %d", repairer.reqs[0].Target.TimeoutSeconds)
|
||||
}
|
||||
if llmClient.lastReq.Target.BackendID != "custom" ||
|
||||
repairer.reqs[0].Target.BackendID != "custom" ||
|
||||
res.SelectedBackendID != "custom" {
|
||||
t.Fatalf("expected backend identity in generation, repair, and result: generate=%q repair=%q result=%q",
|
||||
llmClient.lastReq.Target.BackendID, repairer.reqs[0].Target.BackendID, res.SelectedBackendID)
|
||||
result, err := runner.Run(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
SessionID: " repair-session ",
|
||||
APIKey: "direct-secret",
|
||||
Inputs: singleInputRef(),
|
||||
Execution: tc.execution,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
if len(client.requests) != tc.wantRepairs+1 || len(repairer.reqs) != tc.wantRepairs {
|
||||
t.Fatalf(
|
||||
"generation/repair calls = (%d, %d), want (%d, %d)",
|
||||
len(client.requests), len(repairer.reqs), tc.wantRepairs+1, tc.wantRepairs,
|
||||
)
|
||||
}
|
||||
if len(plan.artifacts) != tc.wantRepairs+1 {
|
||||
t.Fatalf("validation calls = %d, want %d", len(plan.artifacts), tc.wantRepairs+1)
|
||||
}
|
||||
|
||||
initialRequest := client.requests[0]
|
||||
if initialRequest.TargetPresence != tc.wantPresence {
|
||||
t.Fatalf("initial target presence = %+v, want %+v", initialRequest.TargetPresence, tc.wantPresence)
|
||||
}
|
||||
if initialRequest.Target.APIKey != "direct-secret" || initialRequest.Target.BackendID != "custom" ||
|
||||
initialRequest.Prompt.SessionID != "repair-session" {
|
||||
t.Fatalf("initial common request fields = %+v", initialRequest)
|
||||
}
|
||||
if initialRequest.Target.Endpoint != "http://backend.example/v1" ||
|
||||
initialRequest.Target.Model != "profile-model" ||
|
||||
initialRequest.Target.Temperature != 0 || initialRequest.Target.MaxTokens != 0 || initialRequest.Target.TopP != 0 {
|
||||
t.Fatalf("initial effective target = %+v", initialRequest.Target)
|
||||
}
|
||||
if tc.structured != (initialRequest.StructuredOutput != nil) {
|
||||
t.Fatalf("initial structured output = %+v, want present %v", initialRequest.StructuredOutput, tc.structured)
|
||||
}
|
||||
if tc.structured && (initialRequest.StructuredOutput.JSONSchema == nil ||
|
||||
initialRequest.StructuredOutput.JSONSchema.Name != "p_1") {
|
||||
t.Fatalf("initial JSON Schema metadata = %+v, want derived schema name p_1", initialRequest.StructuredOutput)
|
||||
}
|
||||
|
||||
for index, req := range repairer.reqs {
|
||||
if req.Attempt != index+1 || req.MaxAttempts != tc.budget || req.Mode != tc.mode {
|
||||
t.Fatalf("repair request %d progression = %+v", index, req)
|
||||
}
|
||||
if req.PreviousOutput != tc.responses[index].Content ||
|
||||
!reflect.DeepEqual(req.ValidationErrors, tc.validationResults[index].Errors) {
|
||||
t.Fatalf("repair request %d prior state = %+v", index, req)
|
||||
}
|
||||
if req.TargetPresence != tc.wantPresence || !reflect.DeepEqual(req.Target, initialRequest.Target) ||
|
||||
req.SessionID != initialRequest.Prompt.SessionID ||
|
||||
!reflect.DeepEqual(req.StructuredOutput, initialRequest.StructuredOutput) {
|
||||
t.Fatalf("repair request %d common fields drifted: %+v", index, req)
|
||||
}
|
||||
|
||||
generated := client.requests[index+1]
|
||||
if generated.TargetPresence != initialRequest.TargetPresence ||
|
||||
!reflect.DeepEqual(generated.Target, initialRequest.Target) ||
|
||||
generated.Prompt.SessionID != initialRequest.Prompt.SessionID ||
|
||||
!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)
|
||||
}
|
||||
}
|
||||
|
||||
lastResponse := tc.responses[tc.wantRepairs]
|
||||
if result.RawOutput != lastResponse.Content || string(result.Artifact.Body) != lastResponse.Content {
|
||||
t.Fatalf("final output = (%q, %q), want %q", result.RawOutput, result.Artifact.Body, lastResponse.Content)
|
||||
}
|
||||
if result.Validation.Status != tc.wantStatus || result.Validation.RepairAttempts != tc.wantRepairs {
|
||||
t.Fatalf("final validation = %+v, want status %q and %d repairs", result.Validation, tc.wantStatus, tc.wantRepairs)
|
||||
}
|
||||
if result.SelectedBackendID != "custom" || result.SessionID != "repair-session" ||
|
||||
result.EffectiveModelParams.APIKey != "" {
|
||||
t.Fatalf("result execution metadata = %+v", result)
|
||||
}
|
||||
|
||||
var wantUsage domain.TokenUsage
|
||||
for _, response := range tc.responses[:tc.wantRepairs+1] {
|
||||
wantUsage.PromptTokens += response.Usage.PromptTokens
|
||||
wantUsage.CompletionTokens += response.Usage.CompletionTokens
|
||||
wantUsage.TotalTokens += response.Usage.TotalTokens
|
||||
wantUsage.CachedTokens += response.Usage.CachedTokens
|
||||
wantUsage.CacheWriteTokens += response.Usage.CacheWriteTokens
|
||||
}
|
||||
if result.Usage != wantUsage {
|
||||
t.Fatalf("cumulative usage = %+v, want %+v", result.Usage, wantUsage)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2208,99 +2437,6 @@ func TestRunnerSchedulesInitialAndRepairGenerationThroughOneBackendPool(t *testi
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunRepairCarriesEffectiveSessionID(t *testing.T) {
|
||||
llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":`}}
|
||||
runner := NewRunnerWithRepairer(
|
||||
&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{
|
||||
"exec": {ID: "exec", Endpoint: "http://example.test/v1", Model: "model"},
|
||||
}},
|
||||
nil,
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
llmClient,
|
||||
validate.NewStandardValidator("."),
|
||||
NewDefaultOutputRepairer(llmClient), nil)
|
||||
|
||||
result, err := runner.Run(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
SessionID: " repair-session ",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if llmClient.calls != 2 {
|
||||
t.Fatalf("expected initial generation and one repair, got %d calls", llmClient.calls)
|
||||
}
|
||||
if llmClient.lastReq.Prompt.SessionID != "repair-session" {
|
||||
t.Fatalf("expected repair generation to retain effective session, got %q", llmClient.lastReq.Prompt.SessionID)
|
||||
}
|
||||
if result.SessionID != "repair-session" {
|
||||
t.Fatalf("expected result to retain effective session, got %q", result.SessionID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunJSONSchemaRepairCarriesStructuredOutputSpec(t *testing.T) {
|
||||
def := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 1)
|
||||
def.Validation.SchemaPath = "events.schema.json"
|
||||
|
||||
validator := &fakeValidator{
|
||||
result: domain.ValidationResult{
|
||||
Status: domain.ValidationFailed,
|
||||
Mode: domain.ValidationJSONSchema,
|
||||
Errors: []string{"schema mismatch"},
|
||||
IsValid: false,
|
||||
},
|
||||
schemaDoc: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"events": map[string]any{"type": "array"},
|
||||
},
|
||||
},
|
||||
}
|
||||
repairer := &fakeRepairer{
|
||||
responses: []*domain.GenerateResponse{
|
||||
{Content: `{"events":[]}`},
|
||||
},
|
||||
}
|
||||
llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: `{"events":[1]}`}}
|
||||
runner := NewRunnerWithRepairer(
|
||||
&fakePromptRepo{def: def},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}}, nil,
|
||||
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
llmClient,
|
||||
validator,
|
||||
repairer, nil)
|
||||
|
||||
_, err := runner.Run(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got %v", err)
|
||||
}
|
||||
if llmClient.lastReq.StructuredOutput == nil || llmClient.lastReq.StructuredOutput.JSONSchema == nil {
|
||||
t.Fatalf("expected initial llm request to include structured output, got %+v", llmClient.lastReq.StructuredOutput)
|
||||
}
|
||||
if len(repairer.reqs) != 1 {
|
||||
t.Fatalf("expected one repair request, got %d", len(repairer.reqs))
|
||||
}
|
||||
if repairer.reqs[0].StructuredOutput == nil || repairer.reqs[0].StructuredOutput.JSONSchema == nil {
|
||||
t.Fatalf("expected repair request structured output, got %+v", repairer.reqs[0].StructuredOutput)
|
||||
}
|
||||
if repairer.reqs[0].StructuredOutput.JSONSchema.Name != "p_1" {
|
||||
t.Fatalf("expected derived schema name p_1, got %q", repairer.reqs[0].StructuredOutput.JSONSchema.Name)
|
||||
}
|
||||
if validator.schemaLoads != 1 || validator.validateCalls != 2 {
|
||||
t.Fatalf("schema preparation/validation calls = (%d, %d), want (1, 2)", validator.schemaLoads, validator.validateCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecutionProfileToTargetPopulatesAllFieldsAndCopiesExtraParams(t *testing.T) {
|
||||
src := &domain.ExecutionProfile{
|
||||
ID: "exec",
|
||||
|
||||
Reference in New Issue
Block a user