Baton.RateLimiter behaviour (Baton v0.27.4)

Copy Markdown View Source

A pluggable pre-call gate for token throughput.

Oban caps the number of concurrent LLM jobs, but never the tokens they send per minute — so a burst of steps (e.g. compare mode fanning one flow out across many models) can trip a provider's input-tokens-per-minute (ITPM) or output-tokens-per-minute (OTPM) ceiling even while the queue is within its concurrency limit. This behaviour is the seam where a host inserts a token accountant.

Baton.LLMStep calls acquire/3 immediately before every model call and reconcile/3 immediately after every call it gated — with the real usage on success, and with zero actuals when the call failed outright (a 429, a timeout, a transport error), so the reservation is credited back in full rather than leaking out of a tight bucket:

  • acquire/3 is handed the model, an estimate of the input tokens the call will spend, and the client opts. It either reserves that budget and returns {:ok, opts} — with opts possibly rewritten (an implementation may set :rate_account, and in a fallback design may rewrite :provider/:model) — or returns {:snooze, seconds} to park the job attempt-free (no retry consumed, the Oban slot is freed) until budget is available. The snooze reuses Baton.LLMStep's existing 429 snooze contract. Output-token budget is not part of acquire/3 — the call's own opts[:max_tokens] is already the caller's requested output ceiling, so an implementation that tracks OTPM can reserve against it directly without Baton computing anything new.
  • reconcile/3 is handed the account that was reserved, an estimate map, and an actual map — each %{input: non_neg_integer(), output: non_neg_integer()} — so the implementation can correct both the input and output estimates' drift against the real spend. It is bookkeeping — Baton.LLMStep swallows its errors so it can never fail a paid job.

The default implementation (Baton.RateLimiter.Noop) reserves nothing and is wired unless a host configures its own:

config :baton, rate_limiter: MyApp.LLM.RateLimiter

Estimates vs. actuals

acquire/3's est_tokens is a pre-call input estimate (Baton derives it from the prompt character count — see Baton.LLMStep), because the true input-token count is only known once the provider replies. reconcile/3 closes that loop for both dimensions: estimate.input is est_tokens from the paired acquire/3 call, and estimate.output is the call's opts[:max_tokens] (0 when unset) — the reservation an OTPM-tracking implementation would have made against its own output bucket, using data it already had at acquire/3 time. actual.input and actual.output come from the response's reported usage. An implementation credits (or, on an under-estimate, further debits) each dimension's own bucket independently — the two are unrelated budgets, not a combined total.

Summary

Types

Opaque account identifier an implementation associates with a reservation. acquire/3 may stash it in the returned opts under :rate_account; Baton.LLMStep threads that value back into reconcile/3.

Token counts split by direction — the shape of both reconcile/3 arguments after the account.

Callbacks

Reserve budget for an upcoming call, or ask the caller to snooze.

Gate a batch submission, or ask the caller to snooze.

Reconcile a gated call's estimated input/output spend against its actual input/output tokens. Always returns :ok; treated as best-effort by Baton.LLMStep.

Types

account()

@type account() :: term()

Opaque account identifier an implementation associates with a reservation. acquire/3 may stash it in the returned opts under :rate_account; Baton.LLMStep threads that value back into reconcile/3.

token_counts()

@type token_counts() :: %{input: non_neg_integer(), output: non_neg_integer()}

Token counts split by direction — the shape of both reconcile/3 arguments after the account.

Callbacks

acquire(model, est_tokens, opts)

@callback acquire(
  model :: String.t() | nil,
  est_tokens :: non_neg_integer(),
  opts :: keyword()
) ::
  {:ok, keyword()} | {:snooze, pos_integer()}

Reserve budget for an upcoming call, or ask the caller to snooze.

{:ok, opts} proceeds with the (possibly rewritten) opts. {:snooze, secs} parks the job attempt-free for secs seconds and re-runs the step later.

acquire_batch(opts)

(optional)
@callback acquire_batch(opts :: keyword()) :: {:ok, keyword()} | {:snooze, pos_integer()}

Gate a batch submission, or ask the caller to snooze.

Optional — when a limiter doesn't export it, Baton.LLMStep proceeds as if it returned {:ok, opts}.

Batch traffic is metered differently from live traffic: providers draw it from a separate pool, so there is nothing to reserve against the ITPM/OTPM budgets acquire/3 protects and no reconcile/3 counterpart here. What can still be exceeded is the rate of submissions (and any cap on batches in flight), which is what this callback is for. Polling is ungated — it is cheap and asking a limiter about it would only add latency to a step that is already waiting hours.

reconcile(account, estimate, actual)

@callback reconcile(account(), estimate :: token_counts(), actual :: token_counts()) ::
  :ok

Reconcile a gated call's estimated input/output spend against its actual input/output tokens. Always returns :ok; treated as best-effort by Baton.LLMStep.

A call that failed without usable usage (rate-limited, timed out, transport error) reconciles with %{input: 0, output: 0} — an implementation doing estimate - actual bucket arithmetic credits the full reservation back with no special-casing.