refactor(redis-store): own the router cooldown and usage encodings

Cooldown converts to and from its CooldownCacheValue dict and router_usage::Kind
derives its tpm and rpm key segment with strum, so the bridge carries those
values through PyValue and FromStr instead of spelling the fields itself.
PyValue moves down from the auth cache store to serve both.

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",

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

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,

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<RouterCooldowns>);
 
+/// 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<Self> {
-        let fields = value.cast::<PyDict>()?;
-        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::<PyNumber>()?.0,
-            cooldown_time: item("cooldown_time")?.extract::<PyNumber>()?.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<Self::Output> {
-        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)
     }
 }
 

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<RouterUsage>);
 
-/// `(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<Self> {
         let (deployment_id, model, minute, kind): (String, Option<String>, 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,

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<Self> {
+        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<Self::Output> {
+        to_python(py, self.0)
+    }
+}
+
+fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult<Value> {
+    if depth > MAX_DEPTH {
+        return Err(PyValueError::new_err("value is nested too deeply"));
+    }
+    let nested = |items: Bound<'_, PyAny>| -> PyResult<Vec<Value>> {
+        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::<PyBool>() {
+        return Ok(Value::Bool(flag.is_true()));
+    }
+    if value.is_instance_of::<PyInt>() {
+        return Ok(Value::Int(match value.extract::<i64>() {
+            Ok(small) => small.into(),
+            Err(_) => value
+                .str()?
+                .to_str()?
+                .parse::<BigInt>()
+                .map_err(|error| PyValueError::new_err(error.to_string()))?,
+        }));
+    }
+    if let Ok(number) = value.cast::<PyFloat>() {
+        return Ok(Value::Float(number.value()));
+    }
+    if let Ok(text) = value.cast::<PyString>() {
+        return Ok(Value::Str(text.to_str()?.to_owned()));
+    }
+    if let Ok(bytes) = value.cast::<PyBytes>() {
+        return Ok(Value::Bytes(bytes.as_bytes().to_vec()));
+    }
+    if let Ok(complex) = value.cast::<PyComplex>() {
+        return Ok(Value::Complex {
+            re: complex.real(),
+            im: complex.imag(),
+        });
+    }
+    if let Ok(entries) = value.cast::<PyDict>() {
+        return entries
+            .iter()
+            .map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?)))
+            .collect::<PyResult<_>>()
+            .map(Value::Dict);
+    }
+    if value.is_instance_of::<PyList>() {
+        return nested(value.clone()).map(Value::List);
+    }
+    if value.is_instance_of::<PyTuple>() {
+        return nested(value.clone()).map(Value::Tuple);
+    }
+    if value.is_instance_of::<PySet>() || value.is_instance_of::<PyFrozenSet>() {
+        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<Bound<'_, PyAny>> {
+    let all = |items: Vec<Value>| -> PyResult<Vec<Bound<'_, PyAny>>> {
+        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::<PyInt>()
+                .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()
+        }
+    })
+}

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

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<Self, DecodeError> {
-        let Value::Dict(fields) = python_value(raw)? else {
+    pub fn from_value(value: &Value) -> Result<Self, DecodeError> {
+        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, DecodeError> {
+        Self::from_value(&python_value(raw)?)
+    }
 }
 
 fn key(model_id: &str) -> String {

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}:{}",

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::<Kind>().unwrap(), kind);
+}
+
+#[rstest]
+#[case::unknown("tokens")]
+#[case::wrong_case("TPM")]
+fn other_text_is_not_a_counter_kind(#[case] segment: &str) {
+    assert!(segment.parse::<Kind>().is_err());
+}
+
 #[rstest]
 #[tokio::test]
 async fn reads_counters_python_incremented_under_mains_keys(