diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d4cc4db147d..3b67588b506 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4377,6 +4377,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-python-compat", "litellm-redis-client", "litellm-redis-connect", "litellm-redis-store", @@ -4493,6 +4494,7 @@ dependencies = [ "litellm-secret-value", "rstest", "serde_json", + "strum", "thiserror 2.0.19", "tokio", "tracing", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index bd62f16fed9..b574e9c23c6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -34,6 +34,7 @@ litellm-redis-client.workspace = true litellm-redis-connect.workspace = true litellm-redis-store.workspace = true litellm-redis-resolve.workspace = true +litellm-python-compat.workspace = true litellm-redis-types.workspace = true tracing.workspace = true litellm-cache-response.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index e609316efe9..f4cd89b342d 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -2,6 +2,7 @@ mod policy; pub(crate) mod request; pub(crate) mod router_cooldowns; pub(crate) mod router_usage; +mod value; use std::{ pin::Pin, diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs index 69f005dcb3e..044c64dcf98 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs @@ -2,13 +2,14 @@ use std::sync::Arc; use litellm_redis_store::router_cooldowns::{Cooldown, RouterCooldowns}; use pyo3::{ + exceptions::PyValueError, prelude::*, types::{PyDict, PyTuple}, }; use super::{ NativePending, PyNumber, policy::router_tier_policies, redis_async, redis_sync, - request::NativeRequestRedis, store_error, ttl, + request::NativeRequestRedis, store_error, ttl, value::PyValue, }; use crate::redis::client::NativeRedisClient; @@ -17,35 +18,27 @@ use crate::redis::client::NativeRedisClient; #[pyclass(frozen, name = "NativeRouterCooldowns")] pub(crate) struct NativeRouterCooldowns(Arc); +/// A `CooldownCacheValue` dict, whose shape `Cooldown` owns. pub(crate) struct PyCooldown(Cooldown); impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown { type Error = PyErr; fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { - let fields = value.cast::()?; - let item = |name: &str| fields.as_any().get_item(name); - Ok(Self(Cooldown { - exception_received: item("exception_received")?.extract()?, - status_code: item("status_code")?.extract()?, - timestamp: item("timestamp")?.extract::()?.0, - cooldown_time: item("cooldown_time")?.extract::()?.0, - })) + let PyValue(value) = value.extract()?; + Cooldown::from_value(&value) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) } } impl<'py> IntoPyObject<'py> for PyCooldown { - type Target = PyDict; - type Output = Bound<'py, PyDict>; + type Target = PyAny; + type Output = Bound<'py, PyAny>; type Error = PyErr; fn into_pyobject(self, py: Python<'py>) -> PyResult { - let fields = PyDict::new(py); - fields.set_item("exception_received", self.0.exception_received)?; - fields.set_item("status_code", self.0.status_code)?; - fields.set_item("timestamp", PyNumber(self.0.timestamp))?; - fields.set_item("cooldown_time", PyNumber(self.0.cooldown_time))?; - Ok(fields) + PyValue(self.0.to_value()).into_pyobject(py) } } diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs index fe2945e7497..26597a16731 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -1,6 +1,6 @@ use std::{sync::Arc, time::Duration}; -use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use litellm_redis_store::router_usage::{Counter, RouterUsage}; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; use super::{PyNumber, policy::router_tier_policies, redis_async, store_error}; @@ -10,7 +10,7 @@ use crate::redis::client::NativeRedisClient; #[pyclass(frozen, name = "NativeRouterUsage")] pub(crate) struct NativeRouterUsage(Arc); -/// `(deployment_id, model, minute, "tpm" | "rpm")`. +/// `(deployment_id, model, minute, kind)`, the kind spelled as its key segment. pub(crate) struct PyCounter(Counter); impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { @@ -19,15 +19,9 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { let (deployment_id, model, minute, kind): (String, Option, String, String) = value.extract()?; - let kind = match kind.as_str() { - "tpm" => Kind::Tpm, - "rpm" => Kind::Rpm, - _ => { - return Err(PyValueError::new_err(format!( - "unknown usage counter kind {kind:?}" - ))); - } - }; + let kind = kind + .parse() + .map_err(|_| PyValueError::new_err(format!("unknown usage counter kind {kind:?}")))?; Ok(Self(Counter { deployment_id, model, diff --git a/litellm-rust/crates/python-bridge/src/redis/store/value.rs b/litellm-rust/crates/python-bridge/src/redis/store/value.rs new file mode 100644 index 00000000000..c86c9e9bb20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/value.rs @@ -0,0 +1,125 @@ +use litellm_python_compat::{BigInt, MAX_DEPTH, Value}; +use pyo3::{ + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{ + PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString, + PyTuple, + }, +}; + +/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge. +/// Anything else, a live object included, is refused with `TypeError`. +pub(crate) struct PyValue(pub(crate) Value); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyValue { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + to_value(&value, 0).map(Self) + } +} + +impl<'py> IntoPyObject<'py> for PyValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + to_python(py, self.0) + } +} + +fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult { + if depth > MAX_DEPTH { + return Err(PyValueError::new_err("value is nested too deeply")); + } + let nested = |items: Bound<'_, PyAny>| -> PyResult> { + items + .try_iter()? + .map(|item| to_value(&item?, depth + 1)) + .collect() + }; + if value.is_none() { + return Ok(Value::None); + } + if let Ok(flag) = value.cast::() { + return Ok(Value::Bool(flag.is_true())); + } + if value.is_instance_of::() { + return Ok(Value::Int(match value.extract::() { + Ok(small) => small.into(), + Err(_) => value + .str()? + .to_str()? + .parse::() + .map_err(|error| PyValueError::new_err(error.to_string()))?, + })); + } + if let Ok(number) = value.cast::() { + return Ok(Value::Float(number.value())); + } + if let Ok(text) = value.cast::() { + return Ok(Value::Str(text.to_str()?.to_owned())); + } + if let Ok(bytes) = value.cast::() { + return Ok(Value::Bytes(bytes.as_bytes().to_vec())); + } + if let Ok(complex) = value.cast::() { + return Ok(Value::Complex { + re: complex.real(), + im: complex.imag(), + }); + } + if let Ok(entries) = value.cast::() { + return entries + .iter() + .map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?))) + .collect::>() + .map(Value::Dict); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::List); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::Tuple); + } + if value.is_instance_of::() || value.is_instance_of::() { + return nested(value.clone()).map(Value::Set); + } + Err(PyTypeError::new_err(format!( + "{} is not a plain data value", + value.get_type().name()? + ))) +} + +fn to_python(py: Python<'_>, value: Value) -> PyResult> { + let all = |items: Vec| -> PyResult>> { + items.into_iter().map(|item| to_python(py, item)).collect() + }; + Ok(match value { + Value::None => py.None().into_bound(py), + Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(), + Value::Int(number) => match i64::try_from(&number) { + Ok(small) => small.into_pyobject(py)?.into_any(), + Err(_) => py + .get_type::() + .call1((number.to_string(),))? + .into_any(), + }, + Value::Float(number) => number.into_pyobject(py)?.into_any(), + Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(), + Value::Str(text) => text.into_pyobject(py)?.into_any(), + Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(), + Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(), + Value::List(items) => PyList::new(py, all(items)?)?.into_any(), + Value::Set(items) => PySet::new(py, all(items)?)?.into_any(), + Value::Dict(entries) => { + let dict = PyDict::new(py); + for (key, item) in entries { + dict.set_item(to_python(py, key)?, to_python(py, item)?)?; + } + dict.into_any() + } + }) +} diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml index b5d3c0cf88c..7b4466ffa4a 100644 --- a/litellm-rust/crates/redis-store/Cargo.toml +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -12,6 +12,7 @@ litellm-cache-memory.workspace = true litellm-python-compat.workspace = true litellm-redis-client.workspace = true serde_json.workspace = true +strum.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } tracing.workspace = true diff --git a/litellm-rust/crates/redis-store/src/router_cooldowns.rs b/litellm-rust/crates/redis-store/src/router_cooldowns.rs index 52643eabd0b..fe194ce2c33 100644 --- a/litellm-rust/crates/redis-store/src/router_cooldowns.rs +++ b/litellm-rust/crates/redis-store/src/router_cooldowns.rs @@ -36,10 +36,10 @@ impl Cooldown { self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64() } - /// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`. - pub fn encode(&self) -> String { + /// The `CooldownCacheValue` dict the router reads and writes. + pub fn to_value(&self) -> Value { let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value); - repr(&Value::Dict(vec![ + Value::Dict(vec![ field( "exception_received", Value::Str(self.exception_received.clone()), @@ -47,12 +47,11 @@ impl Cooldown { field("status_code", Value::Str(self.status_code.clone())), field("timestamp", self.timestamp.to_value()), field("cooldown_time", self.cooldown_time.to_value()), - ])) + ]) } - /// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does. - pub fn decode(raw: &[u8]) -> Result { - let Value::Dict(fields) = python_value(raw)? else { + pub fn from_value(value: &Value) -> Result { + let Value::Dict(fields) = value else { return Err(DecodeError::Shape("a cooldown is a dict")); }; let field = |name: &str| { @@ -77,6 +76,16 @@ impl Cooldown { cooldown_time: number("cooldown_time")?, }) } + + /// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`. + pub fn encode(&self) -> String { + repr(&self.to_value()) + } + + /// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does. + pub fn decode(raw: &[u8]) -> Result { + Self::from_value(&python_value(raw)?) + } } fn key(model_id: &str) -> String { diff --git a/litellm-rust/crates/redis-store/src/router_usage.rs b/litellm-rust/crates/redis-store/src/router_usage.rs index 68353aa5050..d3ce0fdfea9 100644 --- a/litellm-rust/crates/redis-store/src/router_usage.rs +++ b/litellm-rust/crates/redis-store/src/router_usage.rs @@ -22,9 +22,12 @@ const MEMORY_ENTRIES: usize = 200; /// `RoutingArgs.ttl`, the window of a per-minute counter. pub const COUNTER_TTL: Duration = Duration::from_secs(60); -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +/// A counter's key segment, which is also how the Python router names it. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)] pub enum Kind { + #[strum(serialize = "tpm")] Tpm, + #[strum(serialize = "rpm")] Rpm, } @@ -39,10 +42,7 @@ pub struct Counter { impl Counter { /// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`. fn key(&self) -> String { - let kind = match self.kind { - Kind::Tpm => "tpm", - Kind::Rpm => "rpm", - }; + let kind: &str = self.kind.into(); let model = self.model.as_deref().unwrap_or("None"); format!( "global_router:{}:{model}:{kind}:{}", diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs index 75cbfdf2ebb..ed27351c807 100644 --- a/litellm-rust/crates/redis-store/tests/router_usage.rs +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -20,6 +20,21 @@ fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] { }) } +#[rstest] +#[case::tpm(Kind::Tpm, "tpm")] +#[case::rpm(Kind::Rpm, "rpm")] +fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) { + assert_eq!(<&str>::from(kind), segment); + assert_eq!(segment.parse::().unwrap(), kind); +} + +#[rstest] +#[case::unknown("tokens")] +#[case::wrong_case("TPM")] +fn other_text_is_not_a_counter_kind(#[case] segment: &str) { + assert!(segment.parse::().is_err()); +} + #[rstest] #[tokio::test] async fn reads_counters_python_incremented_under_mains_keys(