diff --git a/README.md b/README.md index cc57888..b53a05c 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,10 @@ implementations — one adapter serves all three providers. Order matters. The outermost wrapper sees the call first; the innermost hits the network. Put `cache.New` below `budget.New` so cache hits don't burn budget. Put `retry.New` on top so retries run regardless -of which inner layer tripped. +of which inner layer tripped. Put `limit.Client` (or `limit.Judge`) +just inside `retry`, so each attempt takes a slot and the 429 that +triggers a retry has already narrowed the limit before the retry +sleeps. ![request flow](./docs/request-flow.svg) @@ -65,6 +68,7 @@ underneath. Roadmap lives in [ROADMAP.md](./ROADMAP.md). - Anthropic, OpenAI, Gemini — all backed by `langchaingo/llms`, in the `provider/` subpackages - A single internal adapter; add a provider = add a ~30-line file - `retry.New` middleware for exp-backoff on transient errors +- `limit.New` — an adaptive per-provider concurrency limiter (halves on a 429 or transport failure, honours `Retry-After`, widens on sustained success); wraps `Client` and `Judge`, and reports every change through `OnChange` - Cost tracking via a static `pricing` table + `Usage.CostUSD` on every Response - Capability flags (`MaxContext`, `SupportsJSONMode`, `SupportsStreaming`, `SupportsTools`, `SupportsVision`) with a `capabilities.Capable` interface - `router.New` with per-provider fallback (`router.OnRateLimit`, `router.OnTransient`, or a custom `router.FallbackPolicy`) diff --git a/middleware/limit/limit.go b/middleware/limit/limit.go new file mode 100644 index 0000000..97aa2ae --- /dev/null +++ b/middleware/limit/limit.go @@ -0,0 +1,298 @@ +// Package limit is an adaptive concurrency limiter for a provider. +// +// A fixed concurrency number is a guess about a limit the provider does +// not publish and changes without notice. Guess high and the provider +// answers 429 — which, unless every caller handles it, is a request +// silently lost; guess low and the wall clock is spent waiting on a +// semaphore that could have been wider. Measured on 2026-09-18: a chat +// provider at concurrency 6 rate-limited 13 of 21 documents into empty +// results that reported success; a System One provider under moderate +// load failed 3 of 40 batches after four minutes of retries. +// +// The limiter moves the number instead of fixing it — additive increase, +// multiplicative decrease, the way TCP finds a link's capacity: +// +// - a slot is taken before every call and released with the outcome; +// - a rate-limit or transport failure halves the limit at once and, if +// the provider said how long to wait, pauses new calls until then; +// - a window of consecutive successes widens the limit by one. +// +// It sits INSIDE retry — retry.New(cfg)(limit.Client(l)(provider)) — so +// each attempt takes a slot and the failure that triggers a retry has +// already narrowed the limit before the retry sleeps. Every change is +// reported through OnChange, because a throttled run is a fact the +// operator must be able to see. +package limit + +import ( + "context" + "sync" + "time" + + "github.com/hallelx2/llmgate" + "github.com/hallelx2/llmgate/capabilities" +) + +// Config tunes a Limiter. Zero values select the defaults. +type Config struct { + // Initial is the starting limit. Default 4. + Initial int + // Min is the floor the limit never drops below. Default 1. + Min int + // Max is the ceiling the limit never rises above. Default 64. + Max int + // SuccessWindow is how many consecutive successes widen the limit by + // one. Default 20 — wide enough that a lucky run does not overshoot. + SuccessWindow int + // Decrease is applied on failure: limit = max(Min, limit*Decrease). + // Default 0.5. + Decrease float64 + // OnChange, when set, is called synchronously with every limit change. + OnChange func(Event) + // Now overrides the clock, for tests. + Now func() time.Time +} + +// Event is one change of the limit. +type Event struct { + // Cause is "rate_limited", "transient", or "success". + Cause string + // From and To are the limit before and after. + From, To int + // PausedFor is how long new calls are held, when the provider said + // so; zero otherwise. + PausedFor time.Duration + // Err is the error that caused a decrease; nil on an increase. + Err error +} + +// Stats is a snapshot for a dashboard. +type Stats struct { + Limit int + InFlight int + Waiting int + Successes int // consecutive successes toward the next increase + // Decreases and Increases count changes since construction. + Decreases, Increases int + // PausedUntil is when a provider-requested pause ends; zero if none. + PausedUntil time.Time +} + +// Limiter is an adaptive semaphore. Share one per provider across every +// caller that talks to it; the point is that all of them see one limit. +type Limiter struct { + cfg Config + + mu sync.Mutex + cond *sync.Cond + limit int + inflight int + waiting int + successes int + paused time.Time + decs int + incs int +} + +// New builds a Limiter. +func New(cfg Config) *Limiter { + if cfg.Initial <= 0 { + cfg.Initial = 4 + } + if cfg.Min <= 0 { + cfg.Min = 1 + } + if cfg.Max <= 0 { + cfg.Max = 64 + } + if cfg.Max < cfg.Min { + cfg.Max = cfg.Min + } + if cfg.Initial < cfg.Min { + cfg.Initial = cfg.Min + } + if cfg.Initial > cfg.Max { + cfg.Initial = cfg.Max + } + if cfg.SuccessWindow <= 0 { + cfg.SuccessWindow = 20 + } + if cfg.Decrease <= 0 || cfg.Decrease >= 1 { + cfg.Decrease = 0.5 + } + if cfg.Now == nil { + cfg.Now = time.Now + } + l := &Limiter{cfg: cfg, limit: cfg.Initial} + l.cond = sync.NewCond(&l.mu) + return l +} + +// Acquire blocks until a slot is free and no provider-requested pause is +// in force, or ctx ends. The returned release must be called exactly once +// with the call's outcome; nil means success. +func (l *Limiter) Acquire(ctx context.Context) (release func(err error), err error) { + l.mu.Lock() + l.waiting++ + // A cancelled context must wake this waiter; a paused limiter must + // wake it when the pause ends. Both are broadcasts on the cond. + stop := context.AfterFunc(ctx, func() { + l.mu.Lock() + l.cond.Broadcast() + l.mu.Unlock() + }) + for { + if ctx.Err() != nil { + l.waiting-- + l.mu.Unlock() + stop() + return nil, ctx.Err() + } + now := l.cfg.Now() + if l.inflight < l.limit && !now.Before(l.paused) { + break + } + if now.Before(l.paused) { + // Wake when the pause ends, whatever else happens. + time.AfterFunc(l.paused.Sub(now), func() { + l.mu.Lock() + l.cond.Broadcast() + l.mu.Unlock() + }) + } + l.cond.Wait() + } + l.waiting-- + l.inflight++ + l.mu.Unlock() + stop() + + var once sync.Once + return func(err error) { + once.Do(func() { l.release(err) }) + }, nil +} + +func (l *Limiter) release(err error) { + l.mu.Lock() + defer l.mu.Unlock() + l.inflight-- + switch { + case err == nil: + l.successes++ + if l.successes >= l.cfg.SuccessWindow && l.limit < l.cfg.Max { + l.successes = 0 + from := l.limit + l.limit++ + l.incs++ + l.emit(Event{Cause: "success", From: from, To: l.limit}) + } + case llmgate.IsRateLimited(err), llmgate.Classify(err) == llmgate.ErrClassTransient: + cause := "transient" + if llmgate.IsRateLimited(err) { + cause = "rate_limited" + } + l.successes = 0 + from := l.limit + to := int(float64(l.limit) * l.cfg.Decrease) + if to < l.cfg.Min { + to = l.cfg.Min + } + var pause time.Duration + if d, ok := llmgate.RetryAfter(err); ok { + pause = d + until := l.cfg.Now().Add(d) + if until.After(l.paused) { + l.paused = until + } + } + if to != from { + l.limit = to + l.decs++ + } + if to != from || pause > 0 { + l.emit(Event{Cause: cause, From: from, To: to, PausedFor: pause, Err: err}) + } + default: + // Auth, bad request, context length, cancellation: not about + // capacity. The limit stands; the success streak does not. + l.successes = 0 + } + l.cond.Broadcast() +} + +func (l *Limiter) emit(ev Event) { + if l.cfg.OnChange != nil { + l.cfg.OnChange(ev) + } +} + +// Limit is the current limit. +func (l *Limiter) Limit() int { + l.mu.Lock() + defer l.mu.Unlock() + return l.limit +} + +// Stats is a snapshot. +func (l *Limiter) Stats() Stats { + l.mu.Lock() + defer l.mu.Unlock() + s := Stats{Limit: l.limit, InFlight: l.inflight, Waiting: l.waiting, Successes: l.successes, Decreases: l.decs, Increases: l.incs} + if l.cfg.Now().Before(l.paused) { + s.PausedUntil = l.paused + } + return s +} + +// Client wraps a chat client so every Complete takes a slot. +func Client(l *Limiter) llmgate.Middleware { + return func(inner llmgate.Client) llmgate.Client { + return &limitedClient{inner: inner, l: l} + } +} + +type limitedClient struct { + inner llmgate.Client + l *Limiter +} + +func (c *limitedClient) Complete(ctx context.Context, req llmgate.Request) (*llmgate.Response, error) { + release, err := c.l.Acquire(ctx) + if err != nil { + return nil, err + } + resp, err := c.inner.Complete(ctx, req) + release(err) + return resp, err +} + +// CountTokens is local and takes no slot. +func (c *limitedClient) CountTokens(ctx context.Context, text string) (int, error) { + return c.inner.CountTokens(ctx, text) +} + +// Capabilities delegates so Capable is not lost through wrapping. +func (c *limitedClient) Capabilities() capabilities.Capabilities { return capabilities.Of(c.inner) } + +// Judge wraps a Judge so every batch takes a slot. +func Judge(l *Limiter) llmgate.JudgeMiddleware { + return func(inner llmgate.Judge) llmgate.Judge { + return &limitedJudge{inner: inner, l: l} + } +} + +type limitedJudge struct { + inner llmgate.Judge + l *Limiter +} + +func (j *limitedJudge) Judge(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + release, err := j.l.Acquire(ctx) + if err != nil { + return nil, err + } + res, err := j.inner.Judge(ctx, req) + release(err) + return res, err +} diff --git a/middleware/limit/limit_test.go b/middleware/limit/limit_test.go new file mode 100644 index 0000000..7b370a5 --- /dev/null +++ b/middleware/limit/limit_test.go @@ -0,0 +1,216 @@ +package limit + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/hallelx2/llmgate" + "github.com/hallelx2/llmgate/middleware/retry" +) + +// capacityClient is a provider that serves at most cap calls at once and +// answers 429 to the rest — the shape every real provider has and none +// documents. +type capacityClient struct { + cap int + inflight int32 + peak int32 + rejected int32 + served int32 + delay time.Duration +} + +func (c *capacityClient) Complete(ctx context.Context, _ llmgate.Request) (*llmgate.Response, error) { + n := atomic.AddInt32(&c.inflight, 1) + defer atomic.AddInt32(&c.inflight, -1) + for { + p := atomic.LoadInt32(&c.peak) + if n <= p || atomic.CompareAndSwapInt32(&c.peak, p, n) { + break + } + } + if int(n) > c.cap { + atomic.AddInt32(&c.rejected, 1) + return nil, llmgate.NewLLMError(llmgate.ProviderOpenAI, 429, "too many requests", nil) + } + time.Sleep(c.delay) + atomic.AddInt32(&c.served, 1) + return &llmgate.Response{Content: "ok"}, nil +} + +func (c *capacityClient) CountTokens(context.Context, string) (int, error) { return 0, nil } + +// The property the package exists for: against a provider that admits N +// at once, the limit converges to about N, and — inside retry — no +// request is lost. +func TestLimiterConvergesOnTheProvidersCapacityAndLosesNothing(t *testing.T) { + prov := &capacityClient{cap: 4, delay: 5 * time.Millisecond} + var events []Event + var emu sync.Mutex + l := New(Config{Initial: 16, Max: 32, SuccessWindow: 5, OnChange: func(e Event) { + emu.Lock() + events = append(events, e) + emu.Unlock() + }}) + client := retry.New(retry.Config{MaxRetries: 8, BaseDelay: time.Millisecond, MaxDelay: 5 * time.Millisecond})(Client(l)(prov)) + + const N = 200 + var wg sync.WaitGroup + var failed int32 + for i := 0; i < N; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if _, err := client.Complete(context.Background(), llmgate.Request{}); err != nil { + atomic.AddInt32(&failed, 1) + } + }() + } + wg.Wait() + + if failed != 0 { + t.Errorf("%d of %d requests failed past retry; the limiter should have kept them under capacity", failed, N) + } + if got := prov.served; got != N { + t.Errorf("served %d want %d", got, N) + } + lim := l.Limit() + if lim < 2 || lim > 8 { + t.Errorf("limit %d after the run; want it settled near the provider's capacity of 4", lim) + } + emu.Lock() + defer emu.Unlock() + var dec, inc int + for _, e := range events { + if e.Cause == "rate_limited" { + dec++ + } else if e.Cause == "success" { + inc++ + } + } + if dec == 0 { + t.Errorf("starting at 16 against capacity 4, at least one decrease was expected") + } + if inc == 0 { + t.Errorf("after settling, successes should have widened the limit at least once") + } +} + +func TestDecreaseHalvesAndHonoursRetryAfter(t *testing.T) { + now := time.Unix(1000, 0) + l := New(Config{Initial: 8, Now: func() time.Time { return now }}) + rel, _ := l.Acquire(context.Background()) + rl := &llmgate.LLMError{Class: llmgate.ErrClassRateLimited, StatusCode: 429, RetryAfterDur: 3 * time.Second} + rel(rl) + if l.Limit() != 4 { + t.Errorf("limit after a 429: %d want 4", l.Limit()) + } + st := l.Stats() + if st.PausedUntil.IsZero() || !st.PausedUntil.Equal(now.Add(3*time.Second)) { + t.Errorf("pause not recorded from Retry-After: %+v", st) + } + // While paused, Acquire must not hand out a slot. + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + if _, err := l.Acquire(ctx); err == nil { + t.Errorf("acquired a slot during a provider-requested pause") + } + // Pause over: slots flow again. + now = now.Add(4 * time.Second) + rel2, err := l.Acquire(context.Background()) + if err != nil { + t.Fatalf("after the pause: %v", err) + } + rel2(nil) +} + +func TestSuccessWindowWidensAndFloorCeilingHold(t *testing.T) { + l := New(Config{Initial: 2, Min: 1, Max: 3, SuccessWindow: 3}) + for i := 0; i < 3; i++ { + rel, _ := l.Acquire(context.Background()) + rel(nil) + } + if l.Limit() != 3 { + t.Errorf("after a window of successes: %d want 3", l.Limit()) + } + for i := 0; i < 6; i++ { + rel, _ := l.Acquire(context.Background()) + rel(nil) + } + if l.Limit() != 3 { + t.Errorf("ceiling: %d want 3", l.Limit()) + } + transient := &llmgate.LLMError{Class: llmgate.ErrClassTransient, StatusCode: 503} + for i := 0; i < 5; i++ { + rel, _ := l.Acquire(context.Background()) + rel(transient) + } + if l.Limit() != 1 { + t.Errorf("floor: %d want 1", l.Limit()) + } +} + +func TestNonCapacityErrorsDoNotMoveTheLimit(t *testing.T) { + l := New(Config{Initial: 4}) + rel, _ := l.Acquire(context.Background()) + rel(&llmgate.LLMError{Class: llmgate.ErrClassAuth, StatusCode: 401}) + rel2, _ := l.Acquire(context.Background()) + rel2(errors.New("bad request")) + if l.Limit() != 4 { + t.Errorf("an auth or bad-request error is not about capacity: limit %d", l.Limit()) + } +} + +func TestAcquireHonoursContextWhileWaiting(t *testing.T) { + l := New(Config{Initial: 1}) + rel, _ := l.Acquire(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if _, err := l.Acquire(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("waiting past ctx: %v", err) + } + rel(nil) + if l.Stats().Waiting != 0 { + t.Errorf("a cancelled waiter is still counted: %+v", l.Stats()) + } +} + +func TestJudgeWrapperTakesASlotPerBatch(t *testing.T) { + l := New(Config{Initial: 1}) + var inflight, peak int32 + mock := &llmgate.MockJudge{Respond: func(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + n := atomic.AddInt32(&inflight, 1) + defer atomic.AddInt32(&inflight, -1) + for { + p := atomic.LoadInt32(&peak) + if n <= p || atomic.CompareAndSwapInt32(&peak, p, n) { + break + } + } + time.Sleep(2 * time.Millisecond) + return &llmgate.Judgment{Model: "mock", Answers: map[string]llmgate.Answer{}}, nil + }} + j := Judge(l)(mock) + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, err := j.Judge(context.Background(), llmgate.JudgeRequest{ + State: map[string]any{"x": "y"}, + Questions: map[string]llmgate.Question{"q": llmgate.Noul{Instructions: "Is `x` y?", Criteria: &llmgate.NoulCriteria{True: "yes", False: "no"}}}, + }) + if err != nil { + t.Errorf("judge: %v", err) + } + }() + } + wg.Wait() + if peak != 1 { + t.Errorf("limit 1 let %d batches run at once", peak) + } +}