//// Rate-limit rotation for providers that hold several accounts or keys. A
//// request that hits a limit is recorded against its account and handed to
//// the next one with room; when none has room, a brief limit is waited out
//// and a lasting one fails fast. A 429 arrives before any output, so a retry
//// repeats nothing the model or the user saw.
////
//// A provider supplies a `Pool`: how to find the session's account, how to
//// record a limit against one, and how to stream on one. `albedo_accounts`
//// holds the ordering and limit bookkeeping those usually share.
import albedo/harness/extension
import albedo/openai_api/types
import gleam/bit_array
import gleam/crypto
import gleam/dynamic/decode
import gleam/json
import gleam/list
import gleam/option.{type Option, None, Some}
import gleam/result
import gleam/string
/// What a 429 said about the account that received it.
pub type Limit {
/// Seconds to minutes, e.g. a burst or tokens-per-minute cap: worth waiting
/// for when every account has hit it.
Brief
/// Hours, e.g. a spent subscription quota: a sibling may have room, but
/// waiting is pointless.
Lasting
}
/// A limit once recorded: its kind, and whether another account still has room.
pub type Marked {
Marked(limit: Limit, room: Bool)
}
/// A Marked from a provider's own flags: whether the limit lasts, and whether
/// a sibling still has room.
fn marked(lasting: Bool, room: Bool) -> Marked {
case lasting {
True -> Marked(Lasting, room)
False -> Marked(Brief, room)
}
}
/// A rate or usage limit parsed from an account provider's 429 response.
pub type Limited {
Limited(account: String, until: String, next: String, lasting: Bool)
}
/// Parses a provider's limit JSON describing when an account resets and its sibling.
pub fn decode_limited(encoded: String) -> Result(Limited, Nil) {
let decoder = {
use account <- decode.optional_field("account", "", decode.string)
use until <- decode.field("until", decode.string)
use next <- decode.field("next", decode.string)
use lasting <- decode.optional_field("lasting", True, decode.bool)
decode.success(Limited(account, until, next, lasting))
}
json.parse(encoded, decoder) |> result.replace_error(Nil)
}
/// Marks a successful limit decode against its lasting duration and sibling availability.
pub fn mark_limit(result: Result(Limited, a)) -> Option(Marked) {
result
|> result.map(fn(limit) { marked(limit.lasting, limit.next != "") })
|> option.from_result
}
/// Formats a user-facing rate-limit or quota-exhausted explanation.
pub fn limit_message(head: String, next: String, none_left: String) -> String {
case next {
"" -> head <> "; " <> none_left
_ -> next_turn_message(head, next)
}
}
/// Explains that a request limit was hit and names the sibling account used next.
pub fn next_turn_message(head: String, next: String) -> String {
head
<> "; the next turn will use "
<> next
<> ". Send your message again to continue"
}
/// Extracts an error message from an API error response body, or falls back to the body.
pub fn error_message(body: String) -> String {
json.parse(body, decode.at(["error", "message"], decode.string))
|> result.unwrap(body)
}
/// Guards upstream resolution against mismatched provider names and protocol requirements.
pub fn require_provider(
context: extension.ModelContext,
provider: String,
label: String,
protocol: types.Protocol,
build: fn() -> Result(extension.Upstream, String),
) -> Option(Result(extension.Upstream, String)) {
case context.provider == provider, context.protocol == protocol {
False, _ -> None
True, True -> Some(build())
True, False ->
Some(Error(
label
<> " provider requires the "
<> types.protocol_name(protocol)
<> " protocol",
))
}
}
/// Tests whether two OpenAI client configurations share the same API key.
pub fn same_client(a: types.Client, b: types.Client) -> Bool {
a.api_key == b.api_key
}
/// An upstream for an OpenAI-compatible client rotating through pool.
/// `label` names an account without naming its credential.
pub fn client_upstream(
pool: Pool(types.Client),
first: types.Client,
explain: fn(types.Client, types.Error) -> Option(String),
label: fn(types.Client) -> String,
) -> extension.Upstream {
upstream(first.base_url, first.protocol, pool, first, explain, label)
}
/// A non-secret account label for a key-authenticated client: the first
/// eight hex characters of the key's SHA-256, enough to tell siblings apart.
pub fn key_label(key: String) -> String {
key
|> bit_array.from_string
|> crypto.hash(crypto.Sha256, _)
|> bit_array.base16_encode
|> string.slice(0, 8)
|> string.lowercase
}
pub type Pool(account) {
Pool(
/// The account the session should use now. Limited accounts come last,
/// so after a `mark` this names a sibling when one has room.
current: fn() -> Result(account, String),
/// Records the limit a 429 `body` reports against `account`. None when the
/// body is not a limit this provider recognises; it is then treated as
/// brief and never moves the session.
mark: fn(account, String) -> Option(Marked),
same: fn(account, account) -> Bool,
stream: fn(account, types.Request, fn(types.Event) -> types.Control) ->
Result(types.Turn, types.Error),
)
}
/// How far one request may go looking for room: account hops left, waits
/// taken so far, and the sleep it uses (tests replace it).
pub type Budget {
Budget(hops: Int, waits: Int, pause: fn(Int) -> Nil)
}
/// Accounts one request may move through after limits, beyond the first.
fn budget() -> Budget {
Budget(8, 0, sleep)
}
/// Backoff when every account is briefly limited: about a minute in all.
const waits = [4000, 8000, 15_000, 30_000]
/// An upstream that rotates through `pool` starting from `first`. `explain`
/// is given the account the failing attempt went to, not necessarily `first`;
/// `label` names the account that served, without naming its credential.
pub fn upstream(
endpoint: String,
protocol: types.Protocol,
pool: Pool(account),
first: account,
explain: fn(account, types.Error) -> Option(String),
label: fn(account) -> String,
) -> extension.Upstream {
let slot = new_slot()
extension.Upstream(
endpoint,
protocol,
fn(request, on_event) {
served(slot, first)
stream(pool, first, request, on_event, budget(), fn(account) {
served(slot, account)
})
},
fn(error) { explain(last_served(slot, first), error) },
fn() {
case string.trim(label(last_served(slot, first))) {
"" -> None
account -> Some(account)
}
},
fn(_) { [] },
types.any_images,
)
}
/// Streams on `account`, moving to siblings or waiting as limits allow.
/// `on_account` hears each account an attempt is sent to after the first.
pub fn stream(
pool: Pool(account),
account: account,
request: types.Request,
on_event: fn(types.Event) -> types.Control,
budget: Budget,
on_account: fn(account) -> Nil,
) -> Result(types.Turn, types.Error) {
case pool.stream(account, request, on_event) {
Error(types.HttpError(429, body)) as failed -> {
let marked = pool.mark(account, body)
let next = case budget.hops > 0, marked {
True, Some(Marked(_, True)) ->
case pool.current() {
Ok(next) ->
case pool.same(next, account) {
True -> None
False -> Some(next)
}
Error(_) -> None
}
_, _ -> None
}
case next, list.drop(waits, budget.waits), marked {
Some(next), _, _ -> {
on_account(next)
stream(
pool,
next,
request,
on_event,
Budget(..budget, hops: budget.hops - 1),
on_account,
)
}
None, _, Some(Marked(Lasting, _)) -> failed
None, [delay, ..], _ ->
case pool.current() {
Ok(retry) -> {
budget.pause(delay)
on_account(retry)
stream(
pool,
retry,
request,
on_event,
Budget(..budget, waits: budget.waits + 1),
on_account,
)
}
Error(_) -> failed
}
None, [], _ -> failed
}
}
other -> other
}
}
type Slot
@external(erlang, "erlang", "make_ref")
fn new_slot() -> Slot
@external(erlang, "albedo_accounts", "served")
fn served(slot: Slot, account: account) -> Nil
@external(erlang, "albedo_accounts", "last_served")
fn last_served(slot: Slot, first: account) -> account
@external(erlang, "albedo_retry", "sleep")
fn sleep(milliseconds: Int) -> Nil