Files
promptkit/internal/generation/coordinator.go

415 lines
12 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"
"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
"gitea.maximumdirect.net/eric/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)
}
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, err, 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
}
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
outcome, dispatched, capacityWait, generationDuration, startedAt := c.execute(ctx, req.Target.BackendID, prepared, captureEvidence)
if dispatched {
attempts++
}
if captureEvidence {
attempt := AttemptEvidence{
StartedAt: startedAt.UTC(),
Dispatched: dispatched,
ResponseReceived: outcome.Evidence.ResponseReceived,
HTTP: cloneHTTPEvidence(outcome.Evidence),
CapacityWait: capacityWait,
RequestDuration: generationDuration,
Failure: ClassifyFailure(outcome.Err),
ResponseDecoded: outcome.ResponseDecoded,
GenerationCompleted: outcome.GenerationCompleted,
}
if outcome.StructuredOutputRequested != nil {
requested := *outcome.StructuredOutputRequested
attempt.StructuredOutputRequested = &requested
}
if dispatched {
usage.Add(outcome.Usage, outcome.UsagePresence)
var attemptUsage observation.UsageAccumulator
attemptUsage.Add(outcome.Usage, outcome.UsagePresence)
attempt.Usage = attemptUsage.Observed()
}
evidence.Attempts = append(evidence.Attempts, attempt)
}
if recorder != nil {
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
Dispatched: dispatched, CapacityWait: capacityWait, GenerationDuration: generationDuration,
Outcome: attemptOutcome(outcome.Err), Failure: ClassifyFailure(outcome.Err),
StatusCode: outcome.StatusCode, Usage: outcome.Usage, UsagePresence: outcome.UsagePresence})
}
if outcome.Err == nil {
if err := ctx.Err(); err != nil {
if captureEvidence {
return finish(outcome.Response, err)
}
return finish(nil, err)
}
return finish(outcome.Response, nil)
}
if !dispatched || !eligibleStatus(outcome) {
if captureEvidence {
return finish(outcome.Response, outcome.Err)
}
return finish(nil, outcome.Err)
}
if retriesUsed == req.TransientRetries {
if retriesUsed == 0 {
return finish(nil, outcome.Err)
}
return finish(nil, &ExhaustedError{attempts: attempts, cause: outcome.Err})
}
if err := ctx.Err(); err != nil {
return finish(nil, err)
}
retryOrdinal := retriesUsed + 1
delay := retryDelay(c.backoff(retryOrdinal), 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)
}
func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration, captureEvidence bool) (outcome llm.AttemptOutcome, dispatched bool, capacityWait, generationDuration time.Duration, startedAt time.Time) {
start := c.now()
startedAt = start
release, err := c.manager.Acquire(ctx, backendID)
capacityWait = c.capacityWait(backendID, start)
if err != nil {
return llm.AttemptOutcome{Err: err}, false, capacityWait, 0, startedAt
}
defer release()
if err := ctx.Err(); err != nil {
return llm.AttemptOutcome{Err: err}, false, capacityWait, 0, startedAt
}
dispatchStart := c.now()
if captureEvidence {
outcome = prepared.ExecuteWithEvidence(ctx)
dispatched = outcome.Evidence.DispatchInvoked
} else {
outcome = prepared.Execute(ctx)
dispatched = true
}
return outcome, dispatched, capacityWait, elapsed(dispatchStart, c.now()), startedAt
}
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()
}
}