llm.NewRetrying wraps a Provider and retries ErrRateLimited, ErrTimeout and 5xx ErrProvider at most twice. A rate limit waits RetryAfter, capped at 20 s; other retries use jittered exponential backoff. ErrAuth, 4xx and other errors are never retried, and no retry is made whose wait would not end before the caller's context deadline. Generate gets a 60 s per-attempt timeout when Request.Timeout is 0. Maestro task: 1a882923-fdb4-4576-ab82-9529e9974b26 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
108 lines
3.2 KiB
Go
108 lines
3.2 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"math/rand/v2"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
// DefaultTimeout bounds each Generate attempt whose Request sets no
|
|
// Timeout.
|
|
DefaultTimeout = 60 * time.Second
|
|
// MaxRetries is how many times a failed Generate is retried, so a call
|
|
// makes at most MaxRetries+1 attempts.
|
|
MaxRetries = 2
|
|
// MaxRetryAfter caps the wait a rate-limited vendor may ask for.
|
|
MaxRetryAfter = 20 * time.Second
|
|
// backoffBase is the first retry's backoff before jitter; each later
|
|
// retry doubles it.
|
|
backoffBase = time.Second
|
|
)
|
|
|
|
// Retrying is a Provider that retries transient Generate failures of the
|
|
// Provider it wraps. It retries ErrRateLimited, ErrTimeout and a 5xx
|
|
// ErrProvider at most MaxRetries times, and nothing else. A rate-limited
|
|
// retry waits for the vendor's RetryAfter, capped at MaxRetryAfter; every
|
|
// other retry, and a rate limit that names no wait, waits a jittered
|
|
// exponential backoff. A retry whose wait would not end before the
|
|
// caller's context deadline is not made: the last error is returned at
|
|
// once. ListModels and Verify pass straight through.
|
|
type Retrying struct {
|
|
Provider
|
|
// now is the clock; tests replace it.
|
|
now func() time.Time
|
|
// sleep waits d or until ctx is done; tests replace it.
|
|
sleep func(ctx context.Context, d time.Duration) error
|
|
// jitter returns a random duration in [0, n) for n > 0; tests replace
|
|
// it.
|
|
jitter func(n time.Duration) time.Duration
|
|
}
|
|
|
|
// NewRetrying returns inner wrapped with retries and the default Generate
|
|
// timeout.
|
|
func NewRetrying(inner Provider) *Retrying {
|
|
return &Retrying{
|
|
Provider: inner,
|
|
now: time.Now,
|
|
sleep: sleep,
|
|
jitter: rand.N[time.Duration],
|
|
}
|
|
}
|
|
|
|
// Generate sends req to the wrapped Provider, retrying transient failures.
|
|
// A req.Timeout of 0 becomes DefaultTimeout; either bounds each attempt,
|
|
// not the whole call, which ctx bounds.
|
|
func (r *Retrying) Generate(ctx context.Context, req Request) (Response, error) {
|
|
if req.Timeout == 0 {
|
|
req.Timeout = DefaultTimeout
|
|
}
|
|
for retry := 0; ; retry++ {
|
|
resp, err := r.Provider.Generate(ctx, req)
|
|
if err == nil || retry == MaxRetries || ctx.Err() != nil {
|
|
return resp, err
|
|
}
|
|
wait, ok := r.wait(err, retry)
|
|
if !ok {
|
|
return resp, err
|
|
}
|
|
if deadline, has := ctx.Deadline(); has && !r.now().Add(wait).Before(deadline) {
|
|
return resp, err
|
|
}
|
|
if serr := r.sleep(ctx, wait); serr != nil {
|
|
return resp, err
|
|
}
|
|
}
|
|
}
|
|
|
|
// wait returns how long to wait before retrying after err, the retry-th
|
|
// retry counting from 0, and false when err is not retried.
|
|
func (r *Retrying) wait(err error, retry int) (time.Duration, bool) {
|
|
var rl *RateLimitedError
|
|
if errors.As(err, &rl) && rl.RetryAfter > 0 {
|
|
return min(rl.RetryAfter, MaxRetryAfter), true
|
|
}
|
|
var pe *ProviderError
|
|
switch {
|
|
case errors.Is(err, ErrRateLimited), errors.Is(err, ErrTimeout):
|
|
case errors.As(err, &pe) && pe.Status() >= 500 && pe.Status() <= 599:
|
|
default:
|
|
return 0, false
|
|
}
|
|
half := (backoffBase << retry) / 2
|
|
return half + r.jitter(half), true
|
|
}
|
|
|
|
// sleep waits d or until ctx is done, returning ctx's error in that case.
|
|
func sleep(ctx context.Context, d time.Duration) error {
|
|
t := time.NewTimer(d)
|
|
defer t.Stop()
|
|
select {
|
|
case <-t.C:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|