From 9d0cce77cec635bee4f4ef3599ba0bfdd24adea0 Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Fri, 22 May 2026 13:58:06 +0530 Subject: [PATCH] feat(embeddings): rate-limit cloud embedding requests to the backend's hard 60/min cap (#2461) Co-authored-by: sanil-23 Co-authored-by: Claude --- src/openhuman/config/schema/load.rs | 17 + src/openhuman/config/schema/storage_memory.rs | 18 + src/openhuman/embeddings/mod.rs | 1 + src/openhuman/embeddings/openai.rs | 8 + src/openhuman/embeddings/rate_limit.rs | 387 ++++++++++++++++++ 5 files changed, 431 insertions(+) create mode 100644 src/openhuman/embeddings/rate_limit.rs diff --git a/src/openhuman/config/schema/load.rs b/src/openhuman/config/schema/load.rs index 937b3d5e9..11dc8e39e 100644 --- a/src/openhuman/config/schema/load.rs +++ b/src/openhuman/config/schema/load.rs @@ -1284,6 +1284,14 @@ impl Config { } set_runtime_proxy_config(self.proxy.clone()); + + // Push the embedding request budget into its process-global limiter so + // every cloud embed (via the shared `OpenAiEmbedding` chokepoint) is + // throttled to the configured rate. Kept here, with the proxy commit, + // so the pure overlay stays side-effect-free for tests. + crate::openhuman::embeddings::rate_limit::set_embedding_rate_limit( + self.memory.embedding_rate_limit_per_min, + ); } /// Pure-ish env overlay: applies overrides read from `env` to `self`. @@ -1727,6 +1735,15 @@ impl Config { self.memory_tree.embedding_strict = strict; } } + // Cloud embedding request budget (requests/min) on `memory.*`. `0` + // disables throttling. A blank or non-numeric value leaves the + // configured/default budget untouched. Committed to the process-global + // limiter in `apply_env_overrides`. + if let Some(val) = env.get("OPENHUMAN_MEMORY_EMBED_RATE_LIMIT") { + if let Ok(per_min) = val.trim().parse::() { + self.memory.embedding_rate_limit_per_min = per_min; + } + } // LLM entity extractor overrides — set endpoint + model to route // ingest scoring through Ollama NER (Phase 2 follow-up). Empty diff --git a/src/openhuman/config/schema/storage_memory.rs b/src/openhuman/config/schema/storage_memory.rs index 1f22c7637..22430734a 100644 --- a/src/openhuman/config/schema/storage_memory.rs +++ b/src/openhuman/config/schema/storage_memory.rs @@ -47,6 +47,14 @@ pub struct MemoryConfig { pub embedding_model: String, #[serde(default = "default_embedding_dims")] pub embedding_dimensions: usize, + /// Outbound embedding-request budget for cloud providers, in requests per + /// minute. Cloud backends (OpenHuman/Voyage, OpenAI, remote `custom:` + /// endpoints) cap requests per account; the client throttles to stay under + /// that quota rather than tripping 429s. `0` disables throttling. Loopback + /// endpoints are always exempt. Env override: + /// `OPENHUMAN_MEMORY_EMBED_RATE_LIMIT`. + #[serde(default = "default_embedding_rate_limit_per_min")] + pub embedding_rate_limit_per_min: u32, #[serde(default = "default_min_relevance_score")] pub min_relevance_score: f64, #[serde(default)] @@ -97,6 +105,11 @@ fn default_embedding_dims() -> usize { // Keep this in sync with `embeddings::cloud::DEFAULT_CLOUD_EMBEDDING_DIMENSIONS`. 1024 } +fn default_embedding_rate_limit_per_min() -> u32 { + // Cloud embedding backends cap requests at ~60/min per account. Keep in + // sync with `embeddings::rate_limit::DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN`. + 60 +} fn default_min_relevance_score() -> f64 { 0.4 } @@ -109,6 +122,7 @@ impl Default for MemoryConfig { embedding_provider: default_embedding_provider(), embedding_model: default_embedding_model(), embedding_dimensions: default_embedding_dims(), + embedding_rate_limit_per_min: default_embedding_rate_limit_per_min(), min_relevance_score: default_min_relevance_score(), sqlite_open_timeout_secs: None, agentmemory_url: None, @@ -131,6 +145,10 @@ impl std::fmt::Debug for MemoryConfig { .field("embedding_provider", &self.embedding_provider) .field("embedding_model", &self.embedding_model) .field("embedding_dimensions", &self.embedding_dimensions) + .field( + "embedding_rate_limit_per_min", + &self.embedding_rate_limit_per_min, + ) .field("min_relevance_score", &self.min_relevance_score) .field("sqlite_open_timeout_secs", &self.sqlite_open_timeout_secs) .field("agentmemory_url", &self.agentmemory_url) diff --git a/src/openhuman/embeddings/mod.rs b/src/openhuman/embeddings/mod.rs index e9c7130e8..90fbde0be 100644 --- a/src/openhuman/embeddings/mod.rs +++ b/src/openhuman/embeddings/mod.rs @@ -17,6 +17,7 @@ pub mod noop; pub mod ollama; pub mod openai; mod provider_trait; +pub mod rate_limit; pub mod store; pub use cloud::{ diff --git a/src/openhuman/embeddings/openai.rs b/src/openhuman/embeddings/openai.rs index 369ca25b7..64a98d7e7 100644 --- a/src/openhuman/embeddings/openai.rs +++ b/src/openhuman/embeddings/openai.rs @@ -94,6 +94,14 @@ impl EmbeddingProvider for OpenAiEmbedding { return Ok(Vec::new()); } + // Proactively gate the outbound request against the per-endpoint rate + // budget so cloud backends (OpenHuman/Voyage, OpenAI, custom remote + // endpoints) stay under their account quota instead of tripping 429s. + // This is the single chokepoint every cloud embed funnels through — + // the `cloud` provider delegates here, and `openai`/`custom:` use it + // directly. Loopback endpoints are exempt (see `rate_limit`). + super::rate_limit::acquire_embedding_slot(&self.base_url).await; + let url = self.embeddings_url(); tracing::debug!( diff --git a/src/openhuman/embeddings/rate_limit.rs b/src/openhuman/embeddings/rate_limit.rs new file mode 100644 index 000000000..c609829aa --- /dev/null +++ b/src/openhuman/embeddings/rate_limit.rs @@ -0,0 +1,387 @@ +//! Client-side request-rate limiting for cloud embedding backends. +//! +//! Cloud embedding backends (the OpenHuman backend / Voyage, OpenAI, and any +//! OpenAI-compatible `custom:` endpoint) cap requests at a fixed rate per +//! account — ~60/min by default. Every [`super::openai::OpenAiEmbedding::embed`] +//! call is exactly one HTTP POST, and memory-tree ingest fans out one call per +//! chunk across several job workers, so without throttling we routinely trip +//! the backend limiter and absorb `429`s (which `openai.rs` currently +//! downgrades to a warning breadcrumb). This module spends the budget +//! *proactively* so requests stay under the quota instead of reacting to 429s. +//! +//! ## Why a process-global registry keyed by endpoint +//! +//! The quota is account-wide, not per-instance — and `OpenAiEmbedding` +//! instances are ephemeral (the cloud provider builds a fresh one on every +//! `embed` call, and the embedder is constructed from several independent +//! sites: the memory factory, `default_embedding_provider`, and the +//! memory-tree score path). A per-instance limiter would therefore reset +//! constantly and enforce nothing. Instead the buckets live in a +//! process-global registry keyed by the resolved base URL, so all ephemeral +//! instances pointing at the same backend share one budget while distinct +//! backends get independent budgets. +//! +//! Loopback endpoints (`localhost` / `127.0.0.1` / `::1`) are exempt: a local +//! LocalAI- or Ollama-compatible `custom:` server is not the remote quota this +//! guards, and capping it to 60/min would needlessly throttle local +//! throughput. + +use std::collections::HashMap; +use std::net::IpAddr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Mutex, OnceLock, PoisonError}; +use std::time::Duration; + +use tokio::time::Instant; + +/// Default outbound embedding request budget, in requests per minute. +/// +/// Cloud embedding backends cap requests at ~60/min per account. Used when the +/// operator hasn't overridden `memory.embedding_rate_limit_per_min`. +/// +/// Keep in sync with `default_embedding_rate_limit_per_min` in +/// `config::schema::storage_memory`. +pub const DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN: u32 = 60; + +/// Process-global configured budget (requests/min). `0` disables throttling. +static CONFIGURED_LIMIT: AtomicU32 = AtomicU32::new(DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN); + +/// Process-global registry of per-endpoint token buckets, keyed by base URL. +static BUCKETS: OnceLock>>> = OnceLock::new(); + +/// Override the process-global embedding request budget. `0` disables +/// throttling entirely. +/// +/// Wired from config load (`config::schema::load::apply_env_overrides`) so the +/// live budget tracks `memory.embedding_rate_limit_per_min`. When the rate +/// changes, existing buckets are dropped so the new rate takes effect on the +/// next request — mirroring how `proxy::set_runtime_proxy_config` clears its +/// client cache on reconfigure. +pub fn set_embedding_rate_limit(per_minute: u32) { + let prev = CONFIGURED_LIMIT.swap(per_minute, Ordering::Relaxed); + // Only drop buckets when the rate actually changes. Clearing on every call + // (e.g. repeated config reloads with an unchanged value) would keep handing + // out a fresh burst token and erode the hard-cap pacing guarantee. + if prev != per_minute { + if let Some(registry) = BUCKETS.get() { + registry + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clear(); + } + } + tracing::debug!( + target: "embeddings::rate_limit", + "[embeddings] rate limit set to {per_minute}/min ({})", + if per_minute == 0 { "disabled" } else { "enabled" } + ); +} + +/// The current process-global embedding request budget (requests/min). +#[must_use] +pub fn embedding_rate_limit() -> u32 { + CONFIGURED_LIMIT.load(Ordering::Relaxed) +} + +/// Gate one outbound embedding HTTP request for `base_url`. +/// +/// Blocks cooperatively until the per-endpoint token bucket has a slot, then +/// consumes it. No-ops when throttling is disabled (`limit == 0`) or when +/// `base_url` is a loopback host. +pub async fn acquire_embedding_slot(base_url: &str) { + acquire_with_limit(base_url, embedding_rate_limit()).await; +} + +/// Inner of [`acquire_embedding_slot`] with the budget passed explicitly, so +/// tests can exercise the gating logic without mutating the process-global +/// limit (which would race across parallel tests). +async fn acquire_with_limit(base_url: &str, limit: u32) { + if limit == 0 || is_loopback_url(base_url) { + return; + } + bucket_for(base_url, limit).acquire().await; +} + +/// Get-or-create the bucket for `base_url`. The first caller's `per_minute` +/// fixes the rate for that endpoint until [`set_embedding_rate_limit`] clears +/// the registry; in practice the limit is configured once at startup before +/// any embedding request, so the rate is uniform across the process. +fn bucket_for(base_url: &str, per_minute: u32) -> Arc { + let registry = BUCKETS.get_or_init(|| Mutex::new(HashMap::new())); + let mut map = registry.lock().unwrap_or_else(PoisonError::into_inner); + map.entry(base_url.to_string()) + .or_insert_with(|| { + tracing::debug!( + target: "embeddings::rate_limit", + "[embeddings] new rate limiter endpoint={base_url} limit={per_minute}/min" + ); + Arc::new(TokenBucket::per_minute(per_minute)) + }) + .clone() +} + +/// True when `base_url`'s host is loopback (`localhost`, `127.0.0.0/8`, `::1`). +/// Unparseable URLs are treated as non-loopback so a malformed remote endpoint +/// still gets throttled rather than silently bypassing the limiter. +fn is_loopback_url(base_url: &str) -> bool { + let Ok(url) = reqwest::Url::parse(base_url) else { + return false; + }; + match url.host_str() { + Some(host) => { + host.eq_ignore_ascii_case("localhost") + || host + .trim_start_matches('[') + .trim_end_matches(']') + .parse::() + .map(|ip| ip.is_loopback()) + .unwrap_or(false) + } + None => false, + } +} + +/// Minimum-interval token bucket sized for a **hard** requests-per-minute cap. +/// +/// Capacity is intentionally a single token, not `per_minute`. The backend +/// enforces 60/min as a hard limit, and a token bucket admits up to +/// `capacity + refill × window` requests in any window — so a `per_minute`-sized +/// burst could momentarily reach ~`2 × per_minute` in the first rolling minute +/// and trip the cap. With capacity 1 the bucket paces requests at one per +/// `60 / per_minute` seconds — a steady `per_minute` per minute with no burst. +/// An idle bucket refills that one token, so an occasional lone request (e.g. an +/// interactive retrieval query embed) still goes out immediately; only +/// back-to-back requests are spaced. +/// [`Self::acquire`] consumes one token, sleeping until one accrues if empty. +struct TokenBucket { + state: tokio::sync::Mutex, + capacity: f64, + refill_per_sec: f64, +} + +struct BucketState { + tokens: f64, + last_refill: Instant, +} + +/// Burst allowance, in tokens. One token = no burst beyond a single request, +/// which is what keeps us strictly under a hard per-minute cap. +const BURST_TOKENS: f64 = 1.0; + +impl TokenBucket { + fn per_minute(per_minute: u32) -> Self { + // `max(1)` is purely defensive against divide-by-zero; the `limit == 0` + // path in `acquire_embedding_slot` never constructs a bucket. + let refill_per_sec = f64::from(per_minute.max(1)) / 60.0; + Self { + state: tokio::sync::Mutex::new(BucketState { + tokens: BURST_TOKENS, + last_refill: Instant::now(), + }), + capacity: BURST_TOKENS, + refill_per_sec, + } + } + + async fn acquire(&self) { + loop { + // Compute refill + the wait-until-next-token *while holding the + // lock*, but drop the guard before sleeping so concurrent callers + // aren't blocked on the mutex during the sleep. + let wait = { + let mut state = self.state.lock().await; + let now = Instant::now(); + let elapsed = now.duration_since(state.last_refill).as_secs_f64(); + state.last_refill = now; + match refill_and_take( + &mut state.tokens, + self.capacity, + self.refill_per_sec, + elapsed, + ) { + None => return, + Some(wait) => wait, + } + }; + tracing::debug!( + target: "embeddings::rate_limit", + "[embeddings] throttling embed request: waiting {:.0}ms for a slot", + wait.as_secs_f64() * 1000.0 + ); + tokio::time::sleep(wait).await; + } + } +} + +/// Refill the bucket for `elapsed_secs`, then try to consume one token. +/// Pure in its inputs (no clock) so the rate math is unit-testable directly. +/// Returns `None` when a token was consumed, or `Some(wait)` — the time until +/// the next whole token accrues — when the bucket is dry. +fn refill_and_take( + tokens: &mut f64, + capacity: f64, + refill_per_sec: f64, + elapsed_secs: f64, +) -> Option { + *tokens = (*tokens + elapsed_secs * refill_per_sec).min(capacity); + if *tokens >= 1.0 { + *tokens -= 1.0; + None + } else { + let deficit = 1.0 - *tokens; + Some(Duration::from_secs_f64(deficit / refill_per_sec)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn loopback_hosts_are_exempt() { + assert!(is_loopback_url("http://localhost:11434")); + assert!(is_loopback_url("http://LOCALHOST/v1")); + assert!(is_loopback_url("http://127.0.0.1:1234/v1")); + assert!(is_loopback_url("http://127.5.6.7")); + assert!(is_loopback_url("http://[::1]:8080")); + } + + #[test] + fn remote_and_malformed_hosts_are_throttled() { + assert!(!is_loopback_url("https://api.openai.com")); + assert!(!is_loopback_url("https://api.openhuman.example/openai/v1")); + assert!(!is_loopback_url("https://10.0.0.5/v1")); // private but not loopback + assert!(!is_loopback_url("not a url")); // malformed → throttled, not bypassed + } + + // ── Bucket math (pure, no clock) ───────────────────────── + + #[test] + fn take_consumes_when_token_available() { + let mut tokens = 5.0; + assert!( + refill_and_take(&mut tokens, 60.0, 1.0, 0.0).is_none(), + "a token is available → consume, no wait" + ); + assert!( + (tokens - 4.0).abs() < 1e-9, + "one token consumed, got {tokens}" + ); + } + + #[test] + fn take_waits_a_full_period_when_empty() { + let mut tokens = 0.0; + // Empty bucket at 1 token/sec → wait ~1s, consuming nothing yet. + let wait = refill_and_take(&mut tokens, 60.0, 1.0, 0.0).expect("must wait"); + assert!((wait.as_secs_f64() - 1.0).abs() < 1e-6, "got {wait:?}"); + assert!(tokens.abs() < 1e-9, "no token consumed while waiting"); + } + + #[test] + fn partial_refill_shortens_the_wait() { + let mut tokens = 0.0; + // 0.25s at 1/sec accrues 0.25 tokens → still <1 → wait the remaining 0.75s. + let wait = refill_and_take(&mut tokens, 60.0, 1.0, 0.25).expect("must wait"); + assert!((wait.as_secs_f64() - 0.75).abs() < 1e-6, "got {wait:?}"); + } + + #[test] + fn refill_is_capped_at_capacity() { + let mut tokens = 50.0; + // A huge idle gap must not let the bucket overflow capacity. + assert!(refill_and_take(&mut tokens, 60.0, 1.0, 10_000.0).is_none()); + assert!( + (tokens - 59.0).abs() < 1e-9, + "capped at 60 then consumed 1 → 59, got {tokens}" + ); + } + + // ── Gating glue (explicit limit, no global mutation) ───── + + #[tokio::test] + async fn disabled_limit_never_blocks() { + for _ in 0..1000 { + acquire_with_limit("https://api.example.com/v1", 0).await; + } + } + + #[tokio::test] + async fn loopback_bypasses_even_at_tight_limit() { + // limit=1 would throttle hard if applied; loopback must bypass it. + for _ in 0..50 { + acquire_with_limit("http://127.0.0.1:11434/v1", 1).await; + } + } + + #[tokio::test] + async fn first_remote_call_does_not_block() { + // Capacity is one token, so a fresh (or idle-refilled) bucket lets + // exactly one request through immediately; pacing of subsequent + // back-to-back requests is covered by `acquire_traverses_wait_branch_*`. + // Unique URL → a bucket isolated from other tests. + let url = "https://burst-test.example/v1"; + let start = Instant::now(); + acquire_with_limit(url, 60).await; + assert!( + start.elapsed() < Duration::from_millis(500), + "first call on a fresh bucket must not block, elapsed {:?}", + start.elapsed() + ); + } + + #[tokio::test] + async fn back_to_back_acquires_are_paced_no_burst() { + // The hard-cap guarantee: capacity is one token, so the first acquire + // passes instantly but the second must wait for refill. 600/min → 10 + // tokens/sec → ~100ms spacing (kept short so the test is fast). A local + // bucket avoids the global registry (and its clear-on-reconfigure). + let bucket = TokenBucket::per_minute(600); + bucket.acquire().await; // consumes the single burst token + let start = Instant::now(); + bucket.acquire().await; // must be paced + assert!( + start.elapsed() >= Duration::from_millis(50), + "second back-to-back acquire must be paced (no burst), elapsed {:?}", + start.elapsed() + ); + } + + #[tokio::test] + async fn acquire_traverses_wait_branch_when_dry() { + // Drive the real throttle path (refill → still dry → sleep → re-check → + // consume) without a slow test: pre-seed the bucket just under one + // token so the wait is ~0.1s of real time. 60/min == 1 token/sec. + let bucket = TokenBucket::per_minute(60); + { + let mut state = bucket.state.lock().await; + state.tokens = 0.9; + state.last_refill = Instant::now(); + } + let start = Instant::now(); + bucket.acquire().await; // must sleep ~100ms for the final 0.1 token + assert!( + start.elapsed() >= Duration::from_millis(50), + "dry bucket must wait for a token, elapsed {:?}", + start.elapsed() + ); + } + + #[tokio::test] + async fn public_entrypoint_reads_global_and_delegates() { + // Exercises the read-global + delegate path; loopback returns at once + // regardless of the configured rate. + acquire_embedding_slot("http://localhost:11434/v1").await; + } + + // ── Global limit setter/getter ─────────────────────────── + + #[test] + fn set_and_read_round_trip() { + set_embedding_rate_limit(0); + assert_eq!(embedding_rate_limit(), 0); + set_embedding_rate_limit(120); + assert_eq!(embedding_rate_limit(), 120); + set_embedding_rate_limit(DEFAULT_EMBEDDING_RATE_LIMIT_PER_MIN); // restore global + } +}