434 lines
13 KiB
Go
434 lines
13 KiB
Go
// Package generation coordinates model attempts, backend permits, and bounded
|
|
// built-in HTTP retries for one generation request.
|
|
package generation
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"math/rand/v2"
|
|
"time"
|
|
|
|
"go.maximumdirect.net/promptkit/internal/capacity"
|
|
"go.maximumdirect.net/promptkit/internal/domain"
|
|
"go.maximumdirect.net/promptkit/internal/llm"
|
|
"go.maximumdirect.net/promptkit/internal/observation"
|
|
)
|
|
|
|
const (
|
|
initialBackoff = 500 * time.Millisecond
|
|
maximumBackoff = 8 * time.Second
|
|
)
|
|
|
|
// Coordinator implements the internal model client using the existing backend
|
|
// capacity manager. Exactly one of builtIn and injected is selected.
|
|
type Coordinator struct {
|
|
manager *capacity.Manager
|
|
builtIn *llm.OpenAICompatibleClient
|
|
injected llm.Client
|
|
now func() time.Time
|
|
jitter func(time.Duration) time.Duration
|
|
wait func(context.Context, time.Duration) error
|
|
}
|
|
|
|
// NewBuiltIn selects the built-in HTTP retry path.
|
|
func NewBuiltIn(manager *capacity.Manager, client *llm.OpenAICompatibleClient) *Coordinator {
|
|
return &Coordinator{manager: manager, builtIn: client, now: time.Now, jitter: randomJitter, wait: waitContext}
|
|
}
|
|
|
|
// NewInjected selects one opaque client dispatch with no automatic retry.
|
|
func NewInjected(manager *capacity.Manager, client llm.Client) *Coordinator {
|
|
return &Coordinator{manager: manager, injected: client, now: time.Now, jitter: randomJitter, wait: waitContext}
|
|
}
|
|
|
|
// ExhaustedError retains the final eligible status failure and the number of
|
|
// attempts actually dispatched.
|
|
type ExhaustedError struct {
|
|
attempts uint64
|
|
cause error
|
|
}
|
|
|
|
func (e *ExhaustedError) Error() string {
|
|
if e == nil {
|
|
return "transient model retries exhausted"
|
|
}
|
|
return fmt.Sprintf("transient model retries exhausted after %d attempts", e.attempts)
|
|
}
|
|
|
|
func (e *ExhaustedError) Unwrap() error {
|
|
if e == nil {
|
|
return nil
|
|
}
|
|
return e.cause
|
|
}
|
|
|
|
func (e *ExhaustedError) Attempts() uint64 {
|
|
if e == nil {
|
|
return 0
|
|
}
|
|
return e.attempts
|
|
}
|
|
|
|
// Generate executes either one opaque client call or a frozen built-in HTTP
|
|
// sequence with one permit per actual dispatch.
|
|
func (c *Coordinator) Generate(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
|
response, err, _ := c.generate(ctx, req, false)
|
|
return response, err
|
|
}
|
|
|
|
// GenerateWithEvidence runs the built-in retry path and returns its factual
|
|
// attempt evidence alongside any response retained before cancellation.
|
|
func (c *Coordinator) GenerateWithEvidence(
|
|
ctx context.Context,
|
|
req domain.GenerateRequest,
|
|
) (*domain.GenerateResponse, error, GenerationEvidence) {
|
|
return c.generate(ctx, req, true)
|
|
}
|
|
|
|
// PreparationError identifies a built-in request preparation failure.
|
|
type PreparationError struct{ Cause error }
|
|
|
|
func (e *PreparationError) Error() string { return e.Cause.Error() }
|
|
func (e *PreparationError) Unwrap() error { return e.Cause }
|
|
|
|
func (c *Coordinator) generate(
|
|
ctx context.Context,
|
|
req domain.GenerateRequest,
|
|
captureEvidence bool,
|
|
) (*domain.GenerateResponse, error, GenerationEvidence) {
|
|
if c.injected != nil {
|
|
return c.generateInjected(ctx, req, captureEvidence)
|
|
}
|
|
prepared, err := c.builtIn.PrepareGeneration(req)
|
|
if err != nil {
|
|
return nil, &PreparationError{Cause: err}, GenerationEvidence{PreparationFailed: true}
|
|
}
|
|
|
|
evidence := GenerationEvidence{}
|
|
if captureEvidence {
|
|
evidence.InitialEndpoint = prepared.InitialEndpoint()
|
|
}
|
|
var usage observation.UsageAccumulator
|
|
finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
|
|
if captureEvidence {
|
|
evidence.Usage = usage.Observed()
|
|
}
|
|
return response, err, evidence
|
|
}
|
|
var attempts, retriesUsed uint64
|
|
recorder := observation.FromContext(ctx)
|
|
generation := req.Generation
|
|
if generation == 0 {
|
|
generation = 1
|
|
}
|
|
for {
|
|
if err := ctx.Err(); err != nil {
|
|
return finish(nil, err)
|
|
}
|
|
ordinal := attempts + 1
|
|
result := c.execute(ctx, req.Target.BackendID, prepared, captureEvidence)
|
|
if result.Dispatched {
|
|
attempts++
|
|
}
|
|
if captureEvidence {
|
|
attempt := AttemptEvidence{
|
|
StartedAt: result.StartedAt.UTC(),
|
|
Dispatched: result.Dispatched,
|
|
ResponseReceived: result.Outcome.Evidence.ResponseReceived,
|
|
HTTP: cloneHTTPEvidence(result.Outcome.Evidence),
|
|
CapacityWait: result.CapacityWait,
|
|
RequestDuration: result.RequestDuration,
|
|
Failure: ClassifyFailure(result.Outcome.Err),
|
|
ResponseDecoded: result.Outcome.ResponseDecoded,
|
|
GenerationCompleted: result.Outcome.GenerationCompleted,
|
|
}
|
|
if result.Outcome.StructuredOutputRequested != nil {
|
|
requested := *result.Outcome.StructuredOutputRequested
|
|
attempt.StructuredOutputRequested = &requested
|
|
}
|
|
if result.Dispatched {
|
|
usage.Add(result.Outcome.Usage, result.Outcome.UsagePresence)
|
|
var attemptUsage observation.UsageAccumulator
|
|
attemptUsage.Add(result.Outcome.Usage, result.Outcome.UsagePresence)
|
|
attempt.Usage = attemptUsage.Observed()
|
|
}
|
|
evidence.Attempts = append(evidence.Attempts, attempt)
|
|
}
|
|
if recorder != nil {
|
|
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
|
|
Dispatched: result.Dispatched, CapacityWait: result.CapacityWait, GenerationDuration: result.RequestDuration,
|
|
Outcome: attemptOutcome(result.Outcome.Err), Failure: ClassifyFailure(result.Outcome.Err),
|
|
StatusCode: result.Outcome.StatusCode, Usage: result.Outcome.Usage, UsagePresence: result.Outcome.UsagePresence})
|
|
}
|
|
if result.Outcome.Err == nil {
|
|
if err := ctx.Err(); err != nil {
|
|
if captureEvidence {
|
|
return finish(result.Outcome.Response, err)
|
|
}
|
|
return finish(nil, err)
|
|
}
|
|
return finish(result.Outcome.Response, nil)
|
|
}
|
|
if !result.Dispatched || !eligibleStatus(result.Outcome) {
|
|
if captureEvidence {
|
|
return finish(result.Outcome.Response, result.Outcome.Err)
|
|
}
|
|
return finish(nil, result.Outcome.Err)
|
|
}
|
|
if retriesUsed == req.TransientRetries {
|
|
if retriesUsed == 0 {
|
|
return finish(nil, result.Outcome.Err)
|
|
}
|
|
return finish(nil, &ExhaustedError{attempts: attempts, cause: result.Outcome.Err})
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return finish(nil, err)
|
|
}
|
|
retryOrdinal := retriesUsed + 1
|
|
delay := retryDelay(c.backoff(retryOrdinal), result.Outcome.RetryNotBefore, c.now())
|
|
waitStart := c.now()
|
|
waitErr := c.wait(ctx, delay)
|
|
waitDuration := elapsed(waitStart, c.now())
|
|
if captureEvidence {
|
|
evidence.RetryDelays = append(evidence.RetryDelays, waitDuration)
|
|
}
|
|
if recorder != nil {
|
|
recorder.Delay(observation.DelayData{Generation: generation, Attempt: attempts,
|
|
Elapsed: waitDuration, Failure: ClassifyFailure(waitErr)})
|
|
}
|
|
if waitErr != nil {
|
|
return finish(nil, waitErr)
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return finish(nil, err)
|
|
}
|
|
retriesUsed++
|
|
}
|
|
}
|
|
|
|
func (c *Coordinator) generateInjected(
|
|
ctx context.Context,
|
|
req domain.GenerateRequest,
|
|
captureEvidence bool,
|
|
) (*domain.GenerateResponse, error, GenerationEvidence) {
|
|
evidence := GenerationEvidence{}
|
|
var usage observation.UsageAccumulator
|
|
finish := func(response *domain.GenerateResponse, err error) (*domain.GenerateResponse, error, GenerationEvidence) {
|
|
if captureEvidence {
|
|
evidence.Usage = usage.Observed()
|
|
}
|
|
return response, err, evidence
|
|
}
|
|
recorder := observation.FromContext(ctx)
|
|
generation := req.Generation
|
|
if generation == 0 {
|
|
generation = 1
|
|
}
|
|
start := c.now()
|
|
release, err := c.manager.Acquire(ctx, req.Target.BackendID)
|
|
capacityWait := c.capacityWait(req.Target.BackendID, start)
|
|
if err != nil {
|
|
if captureEvidence {
|
|
evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
|
|
StartedAt: start.UTC(),
|
|
CapacityWait: capacityWait,
|
|
Failure: ClassifyFailure(err),
|
|
})
|
|
}
|
|
if recorder != nil {
|
|
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
|
|
}
|
|
return finish(nil, err)
|
|
}
|
|
defer release()
|
|
if err := ctx.Err(); err != nil {
|
|
release()
|
|
if captureEvidence {
|
|
evidence.Attempts = append(evidence.Attempts, AttemptEvidence{
|
|
StartedAt: start.UTC(),
|
|
CapacityWait: capacityWait,
|
|
Failure: ClassifyFailure(err),
|
|
})
|
|
}
|
|
if recorder != nil {
|
|
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: ClassifyFailure(err)})
|
|
}
|
|
return finish(nil, err)
|
|
}
|
|
dispatchStart := c.now()
|
|
response, err := c.injected.Generate(ctx, req)
|
|
duration := elapsed(dispatchStart, c.now())
|
|
if err == nil && response == nil {
|
|
err = llm.ErrMalformedResponse
|
|
}
|
|
if captureEvidence {
|
|
attempt := AttemptEvidence{
|
|
StartedAt: start.UTC(),
|
|
Dispatched: true,
|
|
ResponseReceived: response != nil,
|
|
RequestDuration: duration,
|
|
Failure: ClassifyFailure(err),
|
|
GenerationCompleted: err == nil && response != nil,
|
|
}
|
|
requested := req.StructuredOutput != nil
|
|
attempt.StructuredOutputRequested = &requested
|
|
var reportedUsage domain.TokenUsage
|
|
var reportedPresence *domain.TokenUsagePresence
|
|
if err == nil && response != nil {
|
|
reportedUsage = response.Usage
|
|
reportedPresence = response.UsagePresence
|
|
}
|
|
usage.Add(reportedUsage, reportedPresence)
|
|
var attemptUsage observation.UsageAccumulator
|
|
attemptUsage.Add(reportedUsage, reportedPresence)
|
|
attempt.Usage = attemptUsage.Observed()
|
|
evidence.Attempts = append(evidence.Attempts, attempt)
|
|
}
|
|
release()
|
|
if recorder != nil {
|
|
data := observation.AttemptData{Generation: generation, Attempt: 1, Dispatched: true, CapacityWait: capacityWait, GenerationDuration: duration,
|
|
Outcome: attemptOutcome(err), Failure: ClassifyFailure(err)}
|
|
if err == nil {
|
|
data.Usage = response.Usage
|
|
data.UsagePresence = response.UsagePresence
|
|
}
|
|
recorder.Attempt(data)
|
|
}
|
|
if err != nil {
|
|
return finish(nil, err)
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
if captureEvidence {
|
|
return finish(response, err)
|
|
}
|
|
return finish(nil, err)
|
|
}
|
|
return finish(response, nil)
|
|
}
|
|
|
|
type attemptResult struct {
|
|
Outcome llm.AttemptOutcome
|
|
Dispatched bool
|
|
CapacityWait time.Duration
|
|
RequestDuration time.Duration
|
|
StartedAt time.Time
|
|
}
|
|
|
|
func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration, captureEvidence bool) attemptResult {
|
|
result := attemptResult{StartedAt: c.now()}
|
|
release, err := c.manager.Acquire(ctx, backendID)
|
|
result.CapacityWait = c.capacityWait(backendID, result.StartedAt)
|
|
if err != nil {
|
|
result.Outcome.Err = err
|
|
return result
|
|
}
|
|
defer release()
|
|
if err := ctx.Err(); err != nil {
|
|
result.Outcome.Err = err
|
|
return result
|
|
}
|
|
dispatchStart := c.now()
|
|
if captureEvidence {
|
|
result.Outcome = prepared.ExecuteWithEvidence(ctx)
|
|
result.Dispatched = result.Outcome.Evidence.DispatchInvoked
|
|
} else {
|
|
result.Outcome = prepared.Execute(ctx)
|
|
result.Dispatched = true
|
|
}
|
|
result.RequestDuration = elapsed(dispatchStart, c.now())
|
|
return result
|
|
}
|
|
|
|
func (c *Coordinator) capacityWait(backendID string, start time.Time) time.Duration {
|
|
if !c.manager.Limited(backendID) {
|
|
return 0
|
|
}
|
|
return elapsed(start, c.now())
|
|
}
|
|
|
|
func elapsed(start, end time.Time) time.Duration {
|
|
if d := end.Sub(start); d > 0 {
|
|
return d
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func attemptOutcome(err error) domain.ExecutionOutcome {
|
|
if err == nil {
|
|
return domain.ExecutionOutcomeSuccess
|
|
}
|
|
return domain.ExecutionOutcomeError
|
|
}
|
|
|
|
// ClassifyFailure maps an execution error to the shared content-free category.
|
|
func ClassifyFailure(err error) domain.ExecutionFailureKind {
|
|
if err == nil {
|
|
return domain.ExecutionFailureNone
|
|
}
|
|
if errors.Is(err, context.Canceled) {
|
|
return domain.ExecutionFailureCanceled
|
|
}
|
|
if errors.Is(err, context.DeadlineExceeded) {
|
|
return domain.ExecutionFailureDeadline
|
|
}
|
|
if errors.Is(err, llm.ErrUnexpectedStatus) {
|
|
return domain.ExecutionFailureProviderStatus
|
|
}
|
|
if errors.Is(err, llm.ErrTransport) {
|
|
return domain.ExecutionFailureTransport
|
|
}
|
|
if errors.Is(err, llm.ErrMalformedResponse) {
|
|
return domain.ExecutionFailureMalformedResponse
|
|
}
|
|
if errors.Is(err, llm.ErrInvalidRequest) {
|
|
return domain.ExecutionFailureInvalidRequest
|
|
}
|
|
return domain.ExecutionFailureGeneration
|
|
}
|
|
|
|
func eligibleStatus(outcome llm.AttemptOutcome) bool {
|
|
if !errors.Is(outcome.Err, llm.ErrUnexpectedStatus) {
|
|
return false
|
|
}
|
|
switch outcome.StatusCode {
|
|
case 429, 500, 502, 503, 504:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (c *Coordinator) backoff(ordinal uint64) time.Duration {
|
|
base := initialBackoff
|
|
for step := uint64(1); step < ordinal && base < maximumBackoff; step++ {
|
|
base *= 2
|
|
}
|
|
half := base / 2
|
|
return half + c.jitter(half)
|
|
}
|
|
|
|
func retryDelay(local time.Duration, notBefore, now time.Time) time.Duration {
|
|
if providerDelay := notBefore.Sub(now); providerDelay > local {
|
|
return providerDelay
|
|
}
|
|
return local
|
|
}
|
|
|
|
func randomJitter(max time.Duration) time.Duration {
|
|
return time.Duration(rand.Int64N(int64(max) + 1))
|
|
}
|
|
|
|
func waitContext(ctx context.Context, delay time.Duration) error {
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
timer := time.NewTimer(delay)
|
|
defer timer.Stop()
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-timer.C:
|
|
return ctx.Err()
|
|
}
|
|
}
|