diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index e41f25113db..79ab9aa5b40 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -43,7 +43,7 @@ mod _native { #[pymodule_export] use crate::redis::store::auth_cache::NativeAuthCache; #[pymodule_export] - use crate::redis::store::rate_limits::{NativeRateLimitPending, NativeRateLimits}; + use crate::redis::store::rate_limits::NativeRateLimits; #[pymodule_export] use crate::redis::store::request::NativeRequestRedis; #[pymodule_export] @@ -131,7 +131,6 @@ mod tests { "NativeRequestRedis", "NativeAuthCache", "NativeRateLimits", - "NativeRateLimitPending", "NativePending", "NativeRouterCooldowns", "NativeRouterUsage", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs index 540d984c6e6..c0d9452e591 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -1,8 +1,4 @@ -use std::{ - pin::Pin, - sync::{Arc, Mutex, PoisonError}, - time::Duration, -}; +use std::{sync::Arc, time::Duration}; use litellm_redis_store::rate_limits::{ AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Error, Gauge, @@ -15,7 +11,7 @@ use pyo3::{ types::{PyDict, PyString}, }; -use super::{PyNumber, call, request::NativeRequestRedis, store_error}; +use super::{NativePending, PyNumber, redis_async, request::NativeRequestRedis, store_error, ttl}; use crate::execution::run_async_value; use crate::redis::client::NativeRedisClient; @@ -167,116 +163,74 @@ impl<'py> IntoPyObject<'py> for PyLocal { } } -enum Reply { - WindowCounts(Vec), - Values(Vec>), - Atomic(AtomicReply), - Guarded(GuardedReply), - Acquire(AcquireReply), - Counts(Vec), - Totals(Vec), +type OverLimit = Option<(usize, i64, i64)>; + +fn window_counts(counts: Vec) -> Vec<(String, i64)> { + counts + .into_iter() + .map( + |WindowCount { + window_start, + count, + }| (window_start, count), + ) + .collect() } -impl<'py> IntoPyObject<'py> for Reply { - type Target = PyAny; - type Output = Bound<'py, PyAny>; - type Error = PyErr; - - fn into_pyobject(self, py: Python<'py>) -> PyResult { - match self { - Self::WindowCounts(counts) => counts +fn atomic(reply: AtomicReply) -> (OverLimit, Vec<(i64, i64)>) { + match reply { + AtomicReply::Applied(applied) => ( + None, + applied .into_iter() .map( - |WindowCount { - window_start, + |Applied { count, - }| (window_start, count), + window_start, + }| (count, window_start), ) - .collect::>() - .into_pyobject(py), - Self::Values(values) => values.into_pyobject(py), - Self::Atomic(AtomicReply::Applied(applied)) => ( - None::<(usize, i64, i64)>, - applied - .into_iter() - .map( - |Applied { - count, - window_start, - }| (count, window_start), - ) - .collect::>(), - ) - .into_pyobject(py) - .map(Bound::into_any), - Self::Atomic(AtomicReply::OverLimit { - index, - current, - limit, - }) => (Some((index, current, limit)), Vec::<(i64, i64)>::new()) - .into_pyobject(py) - .map(Bound::into_any), - Self::Guarded(GuardedReply { applied, count }) => { - (applied, count).into_pyobject(py).map(Bound::into_any) - } - Self::Acquire(AcquireReply::Acquired(counts)) => (None::<(usize, i64, i64)>, counts) - .into_pyobject(py) - .map(Bound::into_any), - Self::Acquire(AcquireReply::OverLimit { - index, - in_flight, - limit, - }) => (Some((index, in_flight, limit)), Vec::::new()) - .into_pyobject(py) - .map(Bound::into_any), - Self::Counts(counts) => counts.into_pyobject(py), - Self::Totals(totals) => totals.into_pyobject(py), - } + .collect(), + ), + AtomicReply::OverLimit { + index, + current, + limit, + } => (Some((index, current, limit)), Vec::new()), } } -type Resolve = Pin> + Send>>; +fn guarded(GuardedReply { applied, count }: GuardedReply) -> (bool, i64) { + (applied, count) +} + +fn acquire(reply: AcquireReply) -> (OverLimit, Vec) { + match reply { + AcquireReply::Acquired(counts) => (None, counts), + AcquireReply::OverLimit { + index, + in_flight, + limit, + } => (Some((index, in_flight, limit)), Vec::new()), + } +} /// One script call declared on a batch: the request's pre-call batch, or a batch of its own when -/// no request was passed. `resolve` flushes that batch if it has not gone out, once. -#[pyclass(frozen, name = "NativeRateLimitPending")] -pub(crate) struct NativeRateLimitPending { - resolve: Mutex>, +/// no request was passed. Resolving it flushes that batch if it has not gone out. +fn pending( call_type: &'static str, -} - -impl NativeRateLimitPending { - fn new( - pending: Result, Error>, - reply: fn(T) -> Reply, - call_type: &'static str, - ) -> PyResult { - let pending = pending.map_err(rate_limit_error)?; - Ok(Self { - resolve: Mutex::new(Some(Box::pin( - async move { pending.resolve().await.map(reply) }, - ))), - call_type, - }) - } -} - -#[pymethods] -impl NativeRateLimitPending { - fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { - let resolve = self - .resolve - .lock() - .unwrap_or_else(PoisonError::into_inner) - .take() - .ok_or_else(|| { - PyRuntimeError::new_err("this rate-limit script was already resolved") - })?; - let scope = call(py, self.call_type)?; - run_async_value(py, async move { - scope.scope(resolve).await.map_err(rate_limit_error) - }) - } + pending: Result, Error>, + reply: fn(T) -> R, +) -> PyResult +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + let pending = pending.map_err(rate_limit_error)?; + Ok(NativePending::new( + call_type, + rate_limit_error, + async move { pending.resolve().await.map(reply) }, + )) } fn rate_limit_error(error: Error) -> PyErr { @@ -313,7 +267,7 @@ impl NativeRateLimits { backfill_ttl_seconds: Option, redis_ttl_seconds: u64, ) -> Self { - let backfill_ttl = backfill_ttl_seconds.map(|ttl| Duration::from_secs_f64(ttl.max(0.0))); + let backfill_ttl = ttl(backfill_ttl_seconds); Self(Arc::new(RateLimits::new( client.map(|client| client.0.clone()), backfill_ttl, @@ -344,11 +298,7 @@ impl NativeRateLimits { #[pyo3(signature = (entry, value, ttl_seconds))] fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option) -> PyResult<()> { self.0 - .set_local( - &entry.0, - value.0, - ttl_seconds.map(|ttl| Duration::from_secs_f64(ttl.max(0.0))), - ) + .set_local(&entry.0, value.0, ttl(ttl_seconds)) .map_err(store_error) } @@ -360,14 +310,12 @@ impl NativeRateLimits { backfill_ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); - let backfill_ttl = backfill_ttl_seconds.map(|ttl| Duration::from_secs_f64(ttl.max(0.0))); - let scope = call(py, "async_get_cache")?; - run_async_value(py, async move { - scope - .scope(store.get(&entry.0, backfill_ttl)) + let backfill_ttl = ttl(backfill_ttl_seconds); + redis_async(py, "async_get_cache", store_error, async move { + store + .get(&entry.0, backfill_ttl) .await .map(|local| local.map(PyLocal)) - .map_err(store_error) }) } @@ -378,13 +326,13 @@ impl NativeRateLimits { now: i64, window_size: i64, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); - NativeRateLimitPending::new( + pending( + "batch_rate_limiter_script", self.0 .sliding_window(request.as_ref(), &counters(counters_), now, window_size), - Reply::WindowCounts, - "batch_rate_limiter_script", + window_counts, ) } @@ -393,12 +341,12 @@ impl NativeRateLimits { &self, counters_: Vec, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); - NativeRateLimitPending::new( - self.0.read_counters(request.as_ref(), &counters(counters_)), - Reply::Values, + pending( "batch_counter_read_script", + self.0.read_counters(request.as_ref(), &counters(counters_)), + std::convert::identity, ) } @@ -407,16 +355,16 @@ impl NativeRateLimits { &self, counters_: Vec, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); let counters_ = counters_ .into_iter() .map(|counter| counter.0) .collect::>(); - NativeRateLimitPending::new( - self.0.check_and_increment(request.as_ref(), &counters_), - Reply::Atomic, + pending( "check_and_increment_by_n_script", + self.0.check_and_increment(request.as_ref(), &counters_), + atomic, ) } @@ -428,7 +376,7 @@ impl NativeRateLimits { amount: PyNumber, ttl_seconds: Option, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); let increment = GuardedIncrement { counter: counter.0, @@ -436,10 +384,10 @@ impl NativeRateLimits { amount: amount.0, ttl: ttl_seconds.map(Duration::from_secs), }; - NativeRateLimitPending::new( - self.0.guarded_increment(request.as_ref(), &increment), - Reply::Guarded, + pending( "window_guarded_token_increment_script", + self.0.guarded_increment(request.as_ref(), &increment), + guarded, ) } @@ -450,7 +398,7 @@ impl NativeRateLimits { slot_ttl_seconds: u64, slot_id: &str, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); let gauges = gauges .into_iter() @@ -459,15 +407,15 @@ impl NativeRateLimits { limit, }) .collect::>(); - NativeRateLimitPending::new( + pending( + "parallel_acquire_script", self.0.parallel_acquire( request.as_ref(), &gauges, Duration::from_secs(slot_ttl_seconds), slot_id, ), - Reply::Acquire, - "parallel_acquire_script", + acquire, ) } @@ -477,13 +425,13 @@ impl NativeRateLimits { gauges: Vec, slot_id: &str, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); - NativeRateLimitPending::new( + pending( + "parallel_release_script", self.0 .parallel_release(request.as_ref(), &scopes(gauges), slot_id), - Reply::Counts, - "parallel_release_script", + std::convert::identity, ) } @@ -493,16 +441,16 @@ impl NativeRateLimits { gauges: Vec, slot_ttl_seconds: u64, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); - NativeRateLimitPending::new( + pending( + "parallel_count_script", self.0.parallel_count( request.as_ref(), &scopes(gauges), Duration::from_secs(slot_ttl_seconds), ), - Reply::Counts, - "parallel_count_script", + std::convert::identity, ) } @@ -511,13 +459,13 @@ impl NativeRateLimits { &self, increments_: Vec, request: Option>, - ) -> PyResult { + ) -> PyResult { let request = request_of(request); - NativeRateLimitPending::new( + pending( + "token_increment_script", self.0 .token_increment(request.as_ref(), &increments(increments_)), - Reply::Totals, - "token_increment_script", + std::convert::identity, ) } @@ -560,13 +508,8 @@ impl NativeRateLimits { increment: PyIncrement, ) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "async_increment")?; - run_async_value(py, async move { - scope - .scope(store.increment(&increment.0)) - .await - .map(PyNumber) - .map_err(store_error) + redis_async(py, "async_increment", store_error, async move { + store.increment(&increment.0).await.map(PyNumber) }) } @@ -577,12 +520,8 @@ impl NativeRateLimits { ) -> PyResult> { let store = self.0.clone(); let increments_ = increments(increments_); - let scope = call(py, "async_increment_pipeline")?; - run_async_value(py, async move { - scope - .scope(store.increment_many(&increments_)) - .await - .map_err(store_error) + redis_async(py, "async_increment_pipeline", store_error, async move { + store.increment_many(&increments_).await }) } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index bff07f6142a..8f7b62f62da 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -19,6 +19,7 @@ class NativeRedisStoreError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... + @final class NativeClickHouseSpendConfig: def __new__(cls, database: str, url: str, retention_days: int) -> NativeClickHouseSpendConfig: ... @@ -100,10 +101,6 @@ _RateLimitLocal = int | float | str | dict[str, float] _RateLimitIncrement = tuple[_RateLimitCounter, int | float, int | None] _RateLimitAtomicCounter = tuple[_RateLimitCounter, int, int, int, int] -@final -class NativeRateLimitPending: - def resolve(self) -> Future[object]: ... - @final class NativeRateLimits: def __new__( @@ -127,13 +124,13 @@ class NativeRateLimits: window_size: int, *, request: NativeRequestRedis | None = None, - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def read_counters( self, counters_: Sequence[_RateLimitCounter], *, request: NativeRequestRedis | None = None - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def check_and_increment( self, counters_: Sequence[_RateLimitAtomicCounter], *, request: NativeRequestRedis | None = None - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def guarded_increment( self, counter: _RateLimitCounter, @@ -142,7 +139,7 @@ class NativeRateLimits: ttl_seconds: int | None, *, request: NativeRequestRedis | None = None, - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def parallel_acquire( self, gauges: Sequence[tuple[_RateLimitScope, int]], @@ -150,16 +147,16 @@ class NativeRateLimits: slot_id: str, *, request: NativeRequestRedis | None = None, - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def parallel_release( self, gauges: Sequence[_RateLimitScope], slot_id: str, *, request: NativeRequestRedis | None = None - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def parallel_count( self, gauges: Sequence[_RateLimitScope], slot_ttl_seconds: int, *, request: NativeRequestRedis | None = None - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def token_increment( self, increments_: Sequence[_RateLimitIncrement], *, request: NativeRequestRedis | None = None - ) -> NativeRateLimitPending: ... + ) -> NativePending[object]: ... def token_increment_post_call( self, request: NativeRequestRedis, increments_: Sequence[_RateLimitIncrement] ) -> Future[None]: ... @@ -295,6 +292,7 @@ def messages( def amessages( call: NativeCall, ) -> Execution: ... + @final class ResponsesWebSocketConnection: def __new__(cls, _uninstantiable: Never, /) -> Never: ...