Add prepared execution lifecycle to the runner
This commit is contained in:
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)
|
||||
}
|
||||
Reference in New Issue
Block a user