Add prepared execution lifecycle to the runner
This commit is contained in:
@@ -18,13 +18,19 @@ type visit struct {
|
||||
ptr uintptr
|
||||
}
|
||||
|
||||
// Copy validates and deeply copies a JSON-compatible value while preserving
|
||||
// compatible concrete map, slice, array, scalar, and number types.
|
||||
func Copy(src any) (any, error) {
|
||||
return copyValue(reflect.ValueOf(src), "value", make(map[visit]struct{}), true)
|
||||
}
|
||||
|
||||
// CopyMap validates and deeply copies an extra-parameter map while preserving
|
||||
// compatible concrete map, slice, array, scalar, and number types.
|
||||
func CopyMap(src map[string]any) (map[string]any, error) {
|
||||
if src == nil {
|
||||
return nil, nil
|
||||
}
|
||||
copied, err := copyValue(reflect.ValueOf(src), "extra_params", make(map[visit]struct{}))
|
||||
copied, err := copyValue(reflect.ValueOf(src), "extra_params", make(map[visit]struct{}), false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -35,7 +41,12 @@ func CopyMap(src map[string]any) (map[string]any, error) {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func copyValue(value reflect.Value, path string, seen map[visit]struct{}) (any, error) {
|
||||
func copyValue(
|
||||
value reflect.Value,
|
||||
path string,
|
||||
seen map[visit]struct{},
|
||||
allowEmptyMapKeys bool,
|
||||
) (any, error) {
|
||||
if !value.IsValid() {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -43,7 +54,7 @@ func copyValue(value reflect.Value, path string, seen map[visit]struct{}) (any,
|
||||
if value.IsNil() {
|
||||
return nil, nil
|
||||
}
|
||||
return copyValue(value.Elem(), path, seen)
|
||||
return copyValue(value.Elem(), path, seen, allowEmptyMapKeys)
|
||||
}
|
||||
if !value.CanInterface() {
|
||||
return nil, fmt.Errorf("%s: value cannot be copied", path)
|
||||
@@ -88,22 +99,27 @@ func copyValue(value reflect.Value, path string, seen map[visit]struct{}) (any,
|
||||
}
|
||||
seen[current] = struct{}{}
|
||||
defer delete(seen, current)
|
||||
return copyValue(value.Elem(), path, seen)
|
||||
return copyValue(value.Elem(), path, seen, allowEmptyMapKeys)
|
||||
case reflect.Map:
|
||||
return copyMapValue(value, path, seen)
|
||||
return copyMapValue(value, path, seen, allowEmptyMapKeys)
|
||||
case reflect.Slice:
|
||||
if value.IsNil() {
|
||||
return nil, nil
|
||||
}
|
||||
return copySequenceValue(value, path, seen)
|
||||
return copySequenceValue(value, path, seen, allowEmptyMapKeys)
|
||||
case reflect.Array:
|
||||
return copySequenceValue(value, path, seen)
|
||||
return copySequenceValue(value, path, seen, allowEmptyMapKeys)
|
||||
default:
|
||||
return nil, fmt.Errorf("%s: unsupported JSON value type %s", path, value.Type())
|
||||
}
|
||||
}
|
||||
|
||||
func copyMapValue(value reflect.Value, path string, seen map[visit]struct{}) (any, error) {
|
||||
func copyMapValue(
|
||||
value reflect.Value,
|
||||
path string,
|
||||
seen map[visit]struct{},
|
||||
allowEmptyMapKeys bool,
|
||||
) (any, error) {
|
||||
if value.IsNil() {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -133,10 +149,10 @@ func copyMapValue(value reflect.Value, path string, seen map[visit]struct{}) (an
|
||||
elementType := value.Type().Elem()
|
||||
for _, key := range keys {
|
||||
name := key.String()
|
||||
if name == "" {
|
||||
if name == "" && !allowEmptyMapKeys {
|
||||
return nil, fmt.Errorf("%s: map key must not be empty", path)
|
||||
}
|
||||
copied, err := copyValue(value.MapIndex(key), path+"."+name, seen)
|
||||
copied, err := copyValue(value.MapIndex(key), path+"."+name, seen, allowEmptyMapKeys)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -171,7 +187,12 @@ func copyMapValue(value reflect.Value, path string, seen map[visit]struct{}) (an
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func copySequenceValue(value reflect.Value, path string, seen map[visit]struct{}) (any, error) {
|
||||
func copySequenceValue(
|
||||
value reflect.Value,
|
||||
path string,
|
||||
seen map[visit]struct{},
|
||||
allowEmptyMapKeys bool,
|
||||
) (any, error) {
|
||||
var current visit
|
||||
if value.Kind() == reflect.Slice {
|
||||
current = visit{typ: value.Type(), ptr: value.Pointer()}
|
||||
@@ -186,7 +207,12 @@ func copySequenceValue(value reflect.Value, path string, seen map[visit]struct{}
|
||||
preserveType := true
|
||||
elementType := value.Type().Elem()
|
||||
for i := 0; i < value.Len(); i++ {
|
||||
copied, err := copyValue(value.Index(i), fmt.Sprintf("%s[%d]", path, i), seen)
|
||||
copied, err := copyValue(
|
||||
value.Index(i),
|
||||
fmt.Sprintf("%s[%d]", path, i),
|
||||
seen,
|
||||
allowEmptyMapKeys,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -44,6 +44,21 @@ func TestCopyMapPreservesTypesAndIsolatesMutations(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyAllowsEmptyObjectKeysAndIsolatesMutations(t *testing.T) {
|
||||
nested := map[string]any{"": []any{"original"}}
|
||||
|
||||
copiedValue, err := jsonvalue.Copy(nested)
|
||||
if err != nil {
|
||||
t.Fatalf("copy value: %v", err)
|
||||
}
|
||||
nested[""].([]any)[0] = "changed"
|
||||
|
||||
copied := copiedValue.(map[string]any)
|
||||
if got := copied[""].([]any)[0]; got != "original" {
|
||||
t.Fatalf("copied value was not isolated: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyMapRejectsInvalidValues(t *testing.T) {
|
||||
cyclicMap := map[string]any{}
|
||||
cyclicMap["self"] = cyclicMap
|
||||
|
||||
300
internal/usecase/prepared_execution.go
Normal file
300
internal/usecase/prepared_execution.go
Normal file
@@ -0,0 +1,300 @@
|
||||
package usecase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/jsonvalue"
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/validate"
|
||||
)
|
||||
|
||||
type preparedExecutionState uint8
|
||||
|
||||
const (
|
||||
preparedExecutionReady preparedExecutionState = iota
|
||||
preparedExecutionClaimed
|
||||
preparedExecutionDiscarded
|
||||
)
|
||||
|
||||
// PreparedExecution owns one frozen, single-use runner execution.
|
||||
type PreparedExecution struct {
|
||||
owner *Runner
|
||||
mu sync.Mutex
|
||||
state preparedExecutionState
|
||||
details *domain.PreparedRun
|
||||
payload *preparedExecutionPayload
|
||||
}
|
||||
|
||||
type preparedExecutionPayload struct {
|
||||
prepared *domain.PreparedRun
|
||||
validation validate.PreparedValidation
|
||||
}
|
||||
|
||||
// PrepareExecution completes preparation without generation or admission and
|
||||
// returns a runner-bound, single-use execution.
|
||||
func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*PreparedExecution, error) {
|
||||
state, err := r.resolvePreparation(ctx, req, time.Now().UTC())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
validationPlan, err := r.prepareValidation(ctx, state.effectiveContract)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
structuredOutput, err := r.structuredOutputFromValidationPlan(
|
||||
state.definition,
|
||||
state.effectiveContract,
|
||||
validationPlan,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prepared, err := r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
executionSnapshot, err := clonePreparedRun(prepared)
|
||||
if err != nil {
|
||||
prepared.EffectiveModelParams.APIKey = ""
|
||||
return nil, fmt.Errorf("%w: failed to copy prepared execution: %v", ErrInvalidRequest, err)
|
||||
}
|
||||
executionSnapshot.StructuredOutput = prepared.StructuredOutput
|
||||
prepared.EffectiveModelParams.APIKey = ""
|
||||
|
||||
details, err := clonePreparedRun(executionSnapshot)
|
||||
if err != nil {
|
||||
executionSnapshot.EffectiveModelParams.APIKey = ""
|
||||
return nil, fmt.Errorf("%w: failed to copy prepared execution details: %v", ErrInvalidRequest, err)
|
||||
}
|
||||
details.EffectiveModelParams.APIKey = ""
|
||||
|
||||
return &PreparedExecution{
|
||||
owner: r,
|
||||
state: preparedExecutionReady,
|
||||
details: details,
|
||||
payload: &preparedExecutionPayload{
|
||||
prepared: executionSnapshot,
|
||||
validation: validationPlan,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *Runner) prepareValidation(
|
||||
ctx context.Context,
|
||||
contract domain.OutputContract,
|
||||
) (validate.PreparedValidation, error) {
|
||||
if r.validator == nil {
|
||||
return noOpPreparedValidation{contract: contract}, nil
|
||||
}
|
||||
preparer, ok := r.validator.(validate.ValidationPreparer)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%w: validator does not support prepared validation", ErrValidation)
|
||||
}
|
||||
plan, err := preparer.PrepareValidation(ctx, contract)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
|
||||
}
|
||||
if plan == nil {
|
||||
return nil, fmt.Errorf("%w: validator returned nil prepared validation", ErrValidation)
|
||||
}
|
||||
return plan, nil
|
||||
}
|
||||
|
||||
func (r *Runner) structuredOutputFromValidationPlan(
|
||||
def *domain.PromptDefinition,
|
||||
contract domain.OutputContract,
|
||||
plan validate.PreparedValidation,
|
||||
) (*domain.StructuredOutputSpec, error) {
|
||||
if contract.ValidationMode != domain.ValidationJSONSchema {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
schemaDocument := plan.SchemaDocument()
|
||||
if schemaDocument == nil {
|
||||
if r.validator == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("%w: prepared json_schema validation has no schema document", ErrValidation)
|
||||
}
|
||||
return structuredOutputSpec(def, schemaDocument), nil
|
||||
}
|
||||
|
||||
// Details returns a fresh credential-redacted copy of the prepared run.
|
||||
func (p *PreparedExecution) Details() *domain.PreparedRun {
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
p.mu.Lock()
|
||||
detailsSnapshot := p.details
|
||||
p.mu.Unlock()
|
||||
|
||||
details, err := clonePreparedRun(detailsSnapshot)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return details
|
||||
}
|
||||
|
||||
// Discard invalidates an unclaimed execution and drops its private payload.
|
||||
func (p *PreparedExecution) Discard() {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
if p.state != preparedExecutionReady {
|
||||
p.mu.Unlock()
|
||||
return
|
||||
}
|
||||
p.state = preparedExecutionDiscarded
|
||||
payload := p.payload
|
||||
p.payload = nil
|
||||
p.mu.Unlock()
|
||||
|
||||
payload.clear()
|
||||
}
|
||||
|
||||
// RunPrepared claims and executes one prepared execution owned by this runner.
|
||||
func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (*domain.RunResult, error) {
|
||||
payload, err := prepared.claim(r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer payload.clear()
|
||||
|
||||
runID, err := newRunID()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create run id: %w", err)
|
||||
}
|
||||
start := time.Now().UTC()
|
||||
|
||||
target := payload.prepared.EffectiveModelParams
|
||||
if err := validateAPIKey(target.APIKeyEnv, target.APIKey, target.APIKeyRequired); err != nil {
|
||||
return nil, fmt.Errorf("%w: %w", ErrInvalidRequest, err)
|
||||
}
|
||||
|
||||
release, err := r.admitRun(ctx, target.BackendID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer release()
|
||||
|
||||
return r.executePreparedRun(ctx, payload.prepared, runID, start, func(
|
||||
ctx context.Context,
|
||||
artifact *domain.Artifact,
|
||||
attemptsUsed int,
|
||||
) (domain.ValidationResult, error) {
|
||||
result, validationErr := payload.validation.Validate(ctx, artifact)
|
||||
if validationErr != nil {
|
||||
return domain.ValidationResult{}, validationErr
|
||||
}
|
||||
result.RepairAttempts = attemptsUsed
|
||||
return result, nil
|
||||
})
|
||||
}
|
||||
|
||||
func (p *PreparedExecution) claim(owner *Runner) (*preparedExecutionPayload, error) {
|
||||
if p == nil || owner == nil || p.owner != owner {
|
||||
return nil, fmt.Errorf("%w: prepared execution does not belong to this runner", ErrInvalidRequest)
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.state != preparedExecutionReady || p.payload == nil {
|
||||
return nil, fmt.Errorf("%w: prepared execution is not ready", ErrInvalidRequest)
|
||||
}
|
||||
p.state = preparedExecutionClaimed
|
||||
payload := p.payload
|
||||
p.payload = nil
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func (p *preparedExecutionPayload) clear() {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
if p.prepared != nil {
|
||||
p.prepared.EffectiveModelParams.APIKey = ""
|
||||
}
|
||||
p.prepared = nil
|
||||
p.validation = nil
|
||||
}
|
||||
|
||||
type noOpPreparedValidation struct {
|
||||
contract domain.OutputContract
|
||||
}
|
||||
|
||||
func (p noOpPreparedValidation) Validate(
|
||||
ctx context.Context,
|
||||
_ *domain.Artifact,
|
||||
) (domain.ValidationResult, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return domain.ValidationResult{}, err
|
||||
}
|
||||
return domain.ValidationResult{
|
||||
Status: domain.ValidationSkipped,
|
||||
Mode: p.contract.ValidationMode,
|
||||
SchemaPath: p.contract.SchemaPath,
|
||||
RepairAttempts: p.contract.RepairAttempts,
|
||||
IsValid: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (noOpPreparedValidation) SchemaDocument() any {
|
||||
return nil
|
||||
}
|
||||
|
||||
func clonePreparedRun(source *domain.PreparedRun) (*domain.PreparedRun, error) {
|
||||
if source == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
copied := *source
|
||||
extraParams, err := jsonvalue.CopyMap(source.EffectiveModelParams.ExtraParams)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
copied.EffectiveModelParams.ExtraParams = extraParams
|
||||
|
||||
if source.InputHashes != nil {
|
||||
copied.InputHashes = make(map[string]string, len(source.InputHashes))
|
||||
for name, hash := range source.InputHashes {
|
||||
copied.InputHashes[name] = hash
|
||||
}
|
||||
}
|
||||
if source.Messages != nil {
|
||||
copied.Messages = make([]domain.RenderedMessage, len(source.Messages))
|
||||
for i, message := range source.Messages {
|
||||
copied.Messages[i] = message
|
||||
if message.CacheControl != nil {
|
||||
cacheControl := *message.CacheControl
|
||||
copied.Messages[i].CacheControl = &cacheControl
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if source.StructuredOutput != nil {
|
||||
structuredOutput := *source.StructuredOutput
|
||||
copied.StructuredOutput = &structuredOutput
|
||||
if source.StructuredOutput.JSONSchema != nil {
|
||||
jsonSchema := *source.StructuredOutput.JSONSchema
|
||||
copied.StructuredOutput.JSONSchema = &jsonSchema
|
||||
schema, copyErr := cloneJSONValue(source.StructuredOutput.JSONSchema.Schema)
|
||||
if copyErr != nil {
|
||||
return nil, copyErr
|
||||
}
|
||||
copied.StructuredOutput.JSONSchema.Schema = schema
|
||||
}
|
||||
}
|
||||
return &copied, nil
|
||||
}
|
||||
|
||||
func cloneJSONValue(source any) (any, error) {
|
||||
return jsonvalue.Copy(source)
|
||||
}
|
||||
468
internal/usecase/prepared_execution_test.go
Normal file
468
internal/usecase/prepared_execution_test.go
Normal file
@@ -0,0 +1,468 @@
|
||||
package usecase
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
||||
"gitea.maximumdirect.net/eric/promptkit/internal/validate"
|
||||
)
|
||||
|
||||
type recordingPreparedValidation struct {
|
||||
contract domain.OutputContract
|
||||
schemaDocument any
|
||||
results []domain.ValidationResult
|
||||
errs []error
|
||||
artifacts []string
|
||||
}
|
||||
|
||||
func (p *recordingPreparedValidation) Validate(
|
||||
_ context.Context,
|
||||
artifact *domain.Artifact,
|
||||
) (domain.ValidationResult, error) {
|
||||
p.artifacts = append(p.artifacts, string(artifact.Body))
|
||||
index := len(p.artifacts) - 1
|
||||
if index < len(p.errs) && p.errs[index] != nil {
|
||||
return domain.ValidationResult{}, p.errs[index]
|
||||
}
|
||||
if len(p.results) == 0 {
|
||||
return domain.ValidationResult{
|
||||
Status: domain.ValidationPassed,
|
||||
Mode: p.contract.ValidationMode,
|
||||
IsValid: true,
|
||||
}, nil
|
||||
}
|
||||
if index >= len(p.results) {
|
||||
index = len(p.results) - 1
|
||||
}
|
||||
return p.results[index], nil
|
||||
}
|
||||
|
||||
func (p *recordingPreparedValidation) SchemaDocument() any {
|
||||
return p.schemaDocument
|
||||
}
|
||||
|
||||
type recordingValidationPreparer struct {
|
||||
plan *recordingPreparedValidation
|
||||
prepareErr error
|
||||
prepareCalls int
|
||||
directValidateCalls int
|
||||
}
|
||||
|
||||
func (v *recordingValidationPreparer) Validate(
|
||||
context.Context,
|
||||
*domain.Artifact,
|
||||
domain.OutputContract,
|
||||
) (domain.ValidationResult, error) {
|
||||
v.directValidateCalls++
|
||||
return domain.ValidationResult{}, errors.New("live validation must not be used")
|
||||
}
|
||||
|
||||
func (v *recordingValidationPreparer) PrepareValidation(
|
||||
_ context.Context,
|
||||
contract domain.OutputContract,
|
||||
) (validate.PreparedValidation, error) {
|
||||
v.prepareCalls++
|
||||
if v.prepareErr != nil {
|
||||
return nil, v.prepareErr
|
||||
}
|
||||
v.plan.contract = contract
|
||||
return v.plan, nil
|
||||
}
|
||||
|
||||
func TestRunnerPrepareExecutionCompletesWithoutAdmissionOrGeneration(t *testing.T) {
|
||||
schemaDocument := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"value": map[string]any{"type": "string"},
|
||||
"": map[string]any{"type": "boolean"},
|
||||
},
|
||||
}
|
||||
def := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 0)
|
||||
def.Validation.SchemaPath = "schema.json"
|
||||
reader := defaultArtifactReader()
|
||||
renderer := &fakeRenderer{rendered: &domain.RenderedPrompt{
|
||||
SessionID: "prepared-session",
|
||||
Messages: []domain.RenderedMessage{{
|
||||
Role: "user",
|
||||
Content: "original message",
|
||||
}},
|
||||
}}
|
||||
llmClient := &fakeLLM{forbid: true}
|
||||
validator := &recordingValidationPreparer{
|
||||
plan: &recordingPreparedValidation{schemaDocument: schemaDocument},
|
||||
}
|
||||
admitter := &fakeRunAdmitter{}
|
||||
profile := defaultExecutionProfile()
|
||||
profile.ExtraParams = map[string]any{
|
||||
"metadata": map[string]any{"source": "original"},
|
||||
}
|
||||
runner := NewRunner(
|
||||
&fakePromptRepo{def: def},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
|
||||
nil,
|
||||
reader,
|
||||
renderer,
|
||||
llmClient,
|
||||
validator,
|
||||
admitter,
|
||||
)
|
||||
|
||||
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
APIKey: "direct-test-key",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("prepare execution: %v", err)
|
||||
}
|
||||
defer prepared.Discard()
|
||||
|
||||
if validator.prepareCalls != 1 || validator.directValidateCalls != 0 {
|
||||
t.Fatalf(
|
||||
"validation calls=(prepare=%d direct=%d), want (1, 0)",
|
||||
validator.prepareCalls,
|
||||
validator.directValidateCalls,
|
||||
)
|
||||
}
|
||||
if reader.calls != 1 || renderer.calls != 1 {
|
||||
t.Fatalf("completion calls=(artifact=%d render=%d), want (1, 1)", reader.calls, renderer.calls)
|
||||
}
|
||||
if len(admitter.backendIDs) != 0 || llmClient.calls != 0 {
|
||||
t.Fatalf("prepare invoked execution collaborators: admission=%v generation=%d", admitter.backendIDs, llmClient.calls)
|
||||
}
|
||||
|
||||
first := prepared.Details()
|
||||
if first == nil {
|
||||
t.Fatal("prepared details are nil")
|
||||
}
|
||||
if first.EffectiveModelParams.APIKey != "" {
|
||||
t.Fatal("prepared details retained the direct API key")
|
||||
}
|
||||
if first.StructuredOutput == nil ||
|
||||
first.StructuredOutput.JSONSchema == nil ||
|
||||
!reflect.DeepEqual(first.StructuredOutput.JSONSchema.Schema, schemaDocument) {
|
||||
t.Fatalf("prepared details have unexpected structured output: %#v", first.StructuredOutput)
|
||||
}
|
||||
|
||||
first.Messages[0].Content = "caller mutation"
|
||||
first.InputHashes["input"] = "caller mutation"
|
||||
first.EffectiveModelParams.ExtraParams["metadata"].(map[string]any)["source"] = "caller mutation"
|
||||
first.StructuredOutput.JSONSchema.Schema.(map[string]any)["type"] = "string"
|
||||
renderer.rendered.Messages[0].Content = "source mutation"
|
||||
|
||||
second := prepared.Details()
|
||||
if second.Messages[0].Content != "original message" ||
|
||||
second.InputHashes["input"] == "caller mutation" ||
|
||||
second.EffectiveModelParams.ExtraParams["metadata"].(map[string]any)["source"] != "original" ||
|
||||
second.StructuredOutput.JSONSchema.Schema.(map[string]any)["type"] != "object" {
|
||||
t.Fatalf("details did not preserve an independent snapshot: %#v", second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunPreparedRechecksEnvironmentCredentialBeforeAdmission(t *testing.T) {
|
||||
const environmentName = "PROMPTKIT_PREPARED_EXECUTION_TEST_KEY"
|
||||
t.Setenv(environmentName, "available-during-preparation")
|
||||
|
||||
profile := defaultExecutionProfile()
|
||||
profile.APIKeyEnv = environmentName
|
||||
validator := &recordingValidationPreparer{plan: &recordingPreparedValidation{}}
|
||||
admitter := &fakeRunAdmitter{}
|
||||
llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: "unexpected"}}
|
||||
runner := NewRunner(
|
||||
&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}},
|
||||
nil,
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
llmClient,
|
||||
validator,
|
||||
admitter,
|
||||
)
|
||||
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("prepare execution: %v", err)
|
||||
}
|
||||
if err := os.Unsetenv(environmentName); err != nil {
|
||||
t.Fatalf("unset credential environment: %v", err)
|
||||
}
|
||||
|
||||
result, err := runner.RunPrepared(context.Background(), prepared)
|
||||
if result != nil {
|
||||
t.Fatalf("credential failure returned partial result: %+v", result)
|
||||
}
|
||||
if !errors.Is(err, ErrInvalidRequest) || !errors.Is(err, ErrAPIKeyEnvMissing) {
|
||||
t.Fatalf("credential error identities are missing: %v", err)
|
||||
}
|
||||
if len(admitter.backendIDs) != 0 || llmClient.calls != 0 {
|
||||
t.Fatalf("credential failure reached admission or generation: admission=%v generation=%d", admitter.backendIDs, llmClient.calls)
|
||||
}
|
||||
if _, err := runner.RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
|
||||
t.Fatalf("credential failure did not consume execution: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunPreparedUsesFrozenValidationForInitialAndRepairOutputs(t *testing.T) {
|
||||
validator := &recordingValidationPreparer{
|
||||
plan: &recordingPreparedValidation{
|
||||
results: []domain.ValidationResult{
|
||||
{
|
||||
Status: domain.ValidationFailed,
|
||||
Mode: domain.ValidationJSON,
|
||||
Errors: []string{"invalid"},
|
||||
IsValid: false,
|
||||
},
|
||||
{
|
||||
Status: domain.ValidationPassed,
|
||||
Mode: domain.ValidationJSON,
|
||||
IsValid: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
repairer := &fakeRepairer{
|
||||
responses: []*domain.GenerateResponse{{Content: `{"repaired":true}`}},
|
||||
}
|
||||
admitter := &fakeRunAdmitter{}
|
||||
reader := defaultArtifactReader()
|
||||
renderer := defaultRenderer()
|
||||
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}`}},
|
||||
validator,
|
||||
repairer,
|
||||
admitter,
|
||||
)
|
||||
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("prepare execution: %v", err)
|
||||
}
|
||||
|
||||
result, err := runner.RunPrepared(context.Background(), prepared)
|
||||
if err != nil {
|
||||
t.Fatalf("run prepared: %v", err)
|
||||
}
|
||||
if !reflect.DeepEqual(validator.plan.artifacts, []string{`{"broken":true}`, `{"repaired":true}`}) {
|
||||
t.Fatalf("prepared validation artifacts=%#v", validator.plan.artifacts)
|
||||
}
|
||||
if validator.directValidateCalls != 0 || repairer.calls != 1 {
|
||||
t.Fatalf("validation/repair calls=(direct=%d repair=%d), want (0, 1)", validator.directValidateCalls, repairer.calls)
|
||||
}
|
||||
if result.Validation.Status != domain.ValidationPassed || result.Validation.RepairAttempts != 1 {
|
||||
t.Fatalf("unexpected repaired validation result: %+v", result.Validation)
|
||||
}
|
||||
if admitter.releaseCalls != 1 {
|
||||
t.Fatalf("admission releases=%d, want 1", admitter.releaseCalls)
|
||||
}
|
||||
if reader.calls != 1 || renderer.calls != 1 {
|
||||
t.Fatalf("execution reopened preparation sources: artifact=%d render=%d", reader.calls, renderer.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRunPreparedReleasesAdmissionAcrossExecutionErrors(t *testing.T) {
|
||||
generationFailure := errors.New("generation failed")
|
||||
validationFailure := errors.New("validation failed")
|
||||
repairFailure := errors.New("repair failed")
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
generationErr error
|
||||
validation *recordingPreparedValidation
|
||||
repairer *fakeRepairer
|
||||
wantError error
|
||||
}{
|
||||
{
|
||||
name: "generation failure",
|
||||
generationErr: generationFailure,
|
||||
validation: &recordingPreparedValidation{},
|
||||
wantError: ErrLLMGenerate,
|
||||
},
|
||||
{
|
||||
name: "validation failure",
|
||||
validation: &recordingPreparedValidation{
|
||||
errs: []error{validationFailure},
|
||||
},
|
||||
wantError: ErrValidation,
|
||||
},
|
||||
{
|
||||
name: "repair failure",
|
||||
validation: &recordingPreparedValidation{
|
||||
results: []domain.ValidationResult{{
|
||||
Status: domain.ValidationFailed,
|
||||
Mode: domain.ValidationJSON,
|
||||
Errors: []string{"invalid"},
|
||||
IsValid: false,
|
||||
}},
|
||||
},
|
||||
repairer: &fakeRepairer{err: repairFailure},
|
||||
wantError: ErrValidation,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
def := promptDef(domain.FormatJSON, domain.ValidationJSON, 1)
|
||||
if test.repairer == nil {
|
||||
def.Validation.RepairAttempts = 0
|
||||
}
|
||||
validator := &recordingValidationPreparer{plan: test.validation}
|
||||
admitter := &fakeRunAdmitter{}
|
||||
runner := NewRunnerWithRepairer(
|
||||
&fakePromptRepo{def: def},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
|
||||
nil,
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
&fakeLLM{
|
||||
resp: &domain.GenerateResponse{Content: `{"value":true}`},
|
||||
err: test.generationErr,
|
||||
},
|
||||
validator,
|
||||
test.repairer,
|
||||
admitter,
|
||||
)
|
||||
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("prepare execution: %v", err)
|
||||
}
|
||||
|
||||
result, err := runner.RunPrepared(context.Background(), prepared)
|
||||
if result != nil || !errors.Is(err, test.wantError) {
|
||||
t.Fatalf("run prepared=(%+v, %v), want %v", result, err, test.wantError)
|
||||
}
|
||||
if len(admitter.backendIDs) != 1 || admitter.releaseCalls != 1 {
|
||||
t.Fatalf(
|
||||
"admission calls=%#v releases=%d, want one each",
|
||||
admitter.backendIDs,
|
||||
admitter.releaseCalls,
|
||||
)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerPreparedExecutionOwnershipUseAndDiscard(t *testing.T) {
|
||||
newRunner := func() *Runner {
|
||||
return NewRunner(
|
||||
&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationNone, 0)},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
|
||||
nil,
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
&fakeLLM{resp: &domain.GenerateResponse{Content: "ok"}},
|
||||
&recordingValidationPreparer{plan: &recordingPreparedValidation{}},
|
||||
nil,
|
||||
)
|
||||
}
|
||||
request := domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
}
|
||||
|
||||
owner := newRunner()
|
||||
prepared, err := owner.PrepareExecution(context.Background(), request)
|
||||
if err != nil {
|
||||
t.Fatalf("prepare execution: %v", err)
|
||||
}
|
||||
if _, err := newRunner().RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
|
||||
t.Fatalf("foreign runner error=%v, want ErrInvalidRequest", err)
|
||||
}
|
||||
if _, err := owner.RunPrepared(context.Background(), prepared); err != nil {
|
||||
t.Fatalf("owner run prepared: %v", err)
|
||||
}
|
||||
if _, err := owner.RunPrepared(context.Background(), prepared); !errors.Is(err, ErrInvalidRequest) {
|
||||
t.Fatalf("second owner run error=%v, want ErrInvalidRequest", err)
|
||||
}
|
||||
if prepared.Details() == nil {
|
||||
t.Fatal("details unavailable after execution")
|
||||
}
|
||||
|
||||
discarded, err := owner.PrepareExecution(context.Background(), request)
|
||||
if err != nil {
|
||||
t.Fatalf("prepare discarded execution: %v", err)
|
||||
}
|
||||
discarded.Discard()
|
||||
discarded.Discard()
|
||||
if _, err := owner.RunPrepared(context.Background(), discarded); !errors.Is(err, ErrInvalidRequest) {
|
||||
t.Fatalf("discarded execution error=%v, want ErrInvalidRequest", err)
|
||||
}
|
||||
if discarded.Details() == nil {
|
||||
t.Fatal("details unavailable after discard")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerPreparedExecutionWithoutValidatorSkipsValidation(t *testing.T) {
|
||||
runner := NewRunner(
|
||||
&fakePromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 0)},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
|
||||
nil,
|
||||
defaultArtifactReader(),
|
||||
defaultRenderer(),
|
||||
&fakeLLM{resp: &domain.GenerateResponse{Content: `{}`}},
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("prepare execution: %v", err)
|
||||
}
|
||||
|
||||
result, err := runner.RunPrepared(context.Background(), prepared)
|
||||
if err != nil {
|
||||
t.Fatalf("run prepared: %v", err)
|
||||
}
|
||||
if result.Validation.Status != domain.ValidationSkipped || !result.Validation.IsValid {
|
||||
t.Fatalf("unexpected no-validator result: %+v", result.Validation)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerPrepareExecutionRequiresValidationPreparer(t *testing.T) {
|
||||
reader := defaultArtifactReader()
|
||||
runner := NewRunner(
|
||||
&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationBasic, 0)},
|
||||
&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
|
||||
nil,
|
||||
reader,
|
||||
defaultRenderer(),
|
||||
&fakeLLM{forbid: true},
|
||||
&fakeValidator{},
|
||||
nil,
|
||||
)
|
||||
|
||||
prepared, err := runner.PrepareExecution(context.Background(), domain.RunRequest{
|
||||
PromptID: "p",
|
||||
ProfileID: "exec",
|
||||
Inputs: singleInputRef(),
|
||||
})
|
||||
if prepared != nil || !errors.Is(err, ErrValidation) {
|
||||
t.Fatalf("prepare execution=(%+v, %v), want ErrValidation", prepared, err)
|
||||
}
|
||||
if reader.calls != 0 {
|
||||
t.Fatalf("unsupported validator allowed completion, artifact calls=%d", reader.calls)
|
||||
}
|
||||
}
|
||||
@@ -131,26 +131,39 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if r.admitter != nil {
|
||||
release, admitErr := r.admitter.Admit(ctx, state.effectiveModel.BackendID)
|
||||
if admitErr != nil {
|
||||
if errors.Is(admitErr, capacity.ErrCapacityExceeded) {
|
||||
return nil, fmt.Errorf(
|
||||
"backend %q admission: %w",
|
||||
state.effectiveModel.BackendID,
|
||||
admitErr,
|
||||
)
|
||||
}
|
||||
return nil, admitErr
|
||||
}
|
||||
defer release()
|
||||
release, err := r.admitRun(ctx, state.effectiveModel.BackendID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer release()
|
||||
|
||||
prepared, err := r.completePreparation(ctx, req, state)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return r.executePreparedRun(ctx, prepared, runID, start, func(
|
||||
ctx context.Context,
|
||||
artifact *domain.Artifact,
|
||||
attemptsUsed int,
|
||||
) (domain.ValidationResult, error) {
|
||||
return r.validateOutput(ctx, artifact, prepared.OutputContract, attemptsUsed)
|
||||
})
|
||||
}
|
||||
|
||||
type preparedValidationFunc func(
|
||||
context.Context,
|
||||
*domain.Artifact,
|
||||
int,
|
||||
) (domain.ValidationResult, error)
|
||||
|
||||
func (r *Runner) executePreparedRun(
|
||||
ctx context.Context,
|
||||
prepared *domain.PreparedRun,
|
||||
runID string,
|
||||
start time.Time,
|
||||
validateArtifact preparedValidationFunc,
|
||||
) (*domain.RunResult, error) {
|
||||
genResp, err := r.llm.Generate(ctx, domain.GenerateRequest{
|
||||
Prompt: domain.RenderedPrompt{SessionID: prepared.SessionID, Messages: prepared.Messages},
|
||||
Target: prepared.EffectiveModelParams,
|
||||
@@ -165,7 +178,7 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
|
||||
}
|
||||
|
||||
outputArtifact := buildOutputArtifact(genResp.Content, prepared.OutputContract.Format)
|
||||
validationResult, err := r.validateOutput(ctx, &outputArtifact, prepared.OutputContract, 0)
|
||||
validationResult, err := validateArtifact(ctx, &outputArtifact, 0)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
|
||||
}
|
||||
@@ -195,7 +208,7 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
|
||||
genResp = repairResp
|
||||
outputArtifact = buildOutputArtifact(genResp.Content, prepared.OutputContract.Format)
|
||||
|
||||
validationResult, err = r.validateOutput(ctx, &outputArtifact, prepared.OutputContract, attemptsUsed)
|
||||
validationResult, err = validateArtifact(ctx, &outputArtifact, attemptsUsed)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %w", ErrValidation, err)
|
||||
}
|
||||
@@ -324,7 +337,15 @@ func (r *Runner) completePreparation(
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.completePreparationWithStructuredOutput(ctx, req, state, structuredOutput)
|
||||
}
|
||||
|
||||
func (r *Runner) completePreparationWithStructuredOutput(
|
||||
ctx context.Context,
|
||||
req domain.RunRequest,
|
||||
state *preparationState,
|
||||
structuredOutput *domain.StructuredOutputSpec,
|
||||
) (*domain.PreparedRun, error) {
|
||||
resolvedInputs := make(map[string]*domain.Artifact, len(req.Inputs))
|
||||
inputHashes := make(map[string]string, len(req.Inputs))
|
||||
for name, ref := range req.Inputs {
|
||||
@@ -374,6 +395,20 @@ func (r *Runner) completePreparation(
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *Runner) admitRun(ctx context.Context, backendID string) (func(), error) {
|
||||
if r.admitter == nil {
|
||||
return func() {}, nil
|
||||
}
|
||||
release, err := r.admitter.Admit(ctx, backendID)
|
||||
if err != nil {
|
||||
if errors.Is(err, capacity.ErrCapacityExceeded) {
|
||||
return nil, fmt.Errorf("backend %q admission: %w", backendID, err)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return release, nil
|
||||
}
|
||||
|
||||
func (r *Runner) resolveStructuredOutput(ctx context.Context, def *domain.PromptDefinition, contract domain.OutputContract) (*domain.StructuredOutputSpec, error) {
|
||||
if contract.ValidationMode != domain.ValidationJSONSchema {
|
||||
return nil, nil
|
||||
@@ -389,14 +424,18 @@ func (r *Runner) resolveStructuredOutput(ctx context.Context, def *domain.Prompt
|
||||
return nil, fmt.Errorf("%w: failed to load json schema for structured output: %v", ErrValidation, err)
|
||||
}
|
||||
|
||||
return structuredOutputSpec(def, schemaDoc), nil
|
||||
}
|
||||
|
||||
func structuredOutputSpec(def *domain.PromptDefinition, schemaDocument any) *domain.StructuredOutputSpec {
|
||||
return &domain.StructuredOutputSpec{
|
||||
Type: domain.StructuredOutputJSONSchema,
|
||||
JSONSchema: &domain.StructuredOutputJSONSpec{
|
||||
Name: deriveStructuredSchemaName(def.ID, def.Version),
|
||||
Strict: true,
|
||||
Schema: schemaDoc,
|
||||
Schema: schemaDocument,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func deriveStructuredSchemaName(promptID string, promptVersion string) string {
|
||||
|
||||
Reference in New Issue
Block a user