diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 469b994559d..87d70d0799c 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -41,9 +41,7 @@ mod _native { #[pymodule_export] use crate::redis::client::native_redis_breaker_states; #[pymodule_export] - use crate::redis::store::auth_cache::{ - NativeAuthCache, NativeAuthCacheError, NativeDeclaredAuthReads, NativeDeclaredAuthWrite, - }; + use crate::redis::store::auth_cache::NativeAuthCache; #[pymodule_export] use crate::redis::store::request::NativeRequestRedis; #[pymodule_export] @@ -130,9 +128,6 @@ mod tests { "NativeRedisClient", "NativeRequestRedis", "NativeAuthCache", - "NativeAuthCacheError", - "NativeDeclaredAuthReads", - "NativeDeclaredAuthWrite", "NativePending", "NativeRouterCooldowns", "NativeRouterUsage", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs index ac65d97020e..4bc800b4a23 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -1,39 +1,24 @@ -use std::{ - sync::{Arc, Mutex, PoisonError}, - time::Duration, -}; +use std::{sync::Arc, time::Duration}; use litellm_python_compat::Value; -use litellm_redis_store::auth_cache::{AuthCache, DeclaredGets, DeclaredWrite, Error, Settings}; -use pyo3::{exceptions::PyRuntimeError, prelude::*}; +use litellm_redis_store::auth_cache::{AuthCache, DeclaredWrite, Error, Settings}; +use pyo3::prelude::*; -use super::{PyNumber, call, request::NativeRequestRedis, store_error, value::PyValue}; -use crate::execution::{run_async_value, run_sync_value}; +use super::{ + NativePending, PyNumber, kind_error, redis_async, redis_sync, request::NativeRequestRedis, + store_error, ttl, value::PyValue, +}; use crate::redis::client::NativeRedisClient; -pyo3::create_exception!( - _native, - NativeAuthCacheError, - pyo3::exceptions::PyException, - "An auth cache operation failed. Args are (kind, message), kind one of \"breaker_open\", \"timeout\", \"transport\", \"rejected\", \"decode\", \"encode\" or \"memory\"." -); - fn auth_cache_error(error: Error) -> PyErr { - use litellm_redis_client::Error as RedisError; - let kind = match &error { - Error::Redis(RedisError::BreakerOpen) => "breaker_open", - Error::Redis(RedisError::Transport(error)) if error.is_timeout() => "timeout", - Error::Redis(RedisError::Transport(_)) => "transport", - Error::Redis(RedisError::Server(_)) => "rejected", - Error::Redis(RedisError::Decode(_)) | Error::Decode(_) => "decode", - Error::Encode(_) => "encode", - Error::Memory(_) => "memory", - }; - NativeAuthCacheError::new_err((kind, litellm_error::Report::new(&error).to_string())) + kind_error(error.kind(), &error) } -fn ttl(seconds: Option) -> Option { - seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0))) +/// A `set_pre_call` or `delete_pre_call`, settled once and raising when Redis refused it. +fn declared_write(declared: DeclaredWrite) -> NativePending { + NativePending::new("async_set_cache", auth_cache_error, async move { + declared.resolve().await + }) } fn values(values: Vec>) -> Vec> { @@ -72,19 +57,15 @@ impl NativeAuthCache { fn get<'py>(&self, py: Python<'py>, key: String) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "async_get_cache")?; - run_async_value(py, async move { - let value = scope.scope(store.get(&key)).await.map_err(store_error)?; - Ok(value.map(PyValue)) + redis_async(py, "async_get_cache", store_error, async move { + store.get(&key).await.map(|value| value.map(PyValue)) }) } fn get_sync(&self, py: Python<'_>, key: String) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "get_cache")?; - run_sync_value(py, async move { - let value = scope.scope(store.get(&key)).await.map_err(store_error)?; - Ok(value.map(PyValue)) + redis_sync(py, "get_cache", store_error, async move { + store.get(&key).await.map(|value| value.map(PyValue)) }) } @@ -97,12 +78,8 @@ impl NativeAuthCache { ttl_seconds: Option, ) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "async_set_cache")?; - run_async_value(py, async move { - scope - .scope(store.set(&key, value.0, ttl(ttl_seconds))) - .await - .map_err(store_error) + redis_async(py, "async_set_cache", store_error, async move { + store.set(&key, value.0, ttl(ttl_seconds)).await }) } @@ -115,12 +92,8 @@ impl NativeAuthCache { ttl_seconds: Option, ) -> PyResult<()> { let store = self.0.clone(); - let scope = call(py, "set_cache")?; - run_sync_value(py, async move { - scope - .scope(store.set_str(&key, value.0, ttl(ttl_seconds))) - .await - .map_err(store_error) + redis_sync(py, "set_cache", store_error, async move { + store.set_str(&key, value.0, ttl(ttl_seconds)).await }) } @@ -132,24 +105,24 @@ impl NativeAuthCache { key: String, value: PyValue, ttl_seconds: Option, - ) -> PyResult> { + ) -> PyResult> { Ok(self .0 .set_pre_call(&request.0, &key, value.0, ttl(ttl_seconds)) .map_err(store_error)? - .map(NativeDeclaredAuthWrite::new)) + .map(declared_write)) } fn delete_pre_call( &self, request: &NativeRequestRedis, key: String, - ) -> PyResult> { + ) -> PyResult> { Ok(self .0 .delete_pre_call(&request.0, &key) .map_err(store_error)? - .map(NativeDeclaredAuthWrite::new)) + .map(declared_write)) } #[pyo3(signature = (entries, ttl_seconds = None))] @@ -164,12 +137,8 @@ impl NativeAuthCache { .into_iter() .map(|(key, value)| (key, value.0)) .collect(); - let scope = call(py, "async_set_cache_pipeline")?; - run_async_value(py, async move { - scope - .scope(store.set_many(entries, ttl(ttl_seconds))) - .await - .map_err(store_error) + redis_async(py, "async_set_cache_pipeline", store_error, async move { + store.set_many(entries, ttl(ttl_seconds)).await }) } @@ -186,34 +155,25 @@ impl NativeAuthCache { .into_iter() .map(|(key, value, seconds)| (key, value.0, ttl(seconds))) .collect(); - let scope = call(py, "async_set_cache_pipeline_with_ttls")?; - run_async_value(py, async move { - scope - .scope(store.write_back(request.as_ref(), entries)) - .await - .map_err(auth_cache_error) - }) + redis_async( + py, + "async_set_cache_pipeline_with_ttls", + auth_cache_error, + async move { store.write_back(request.as_ref(), entries).await }, + ) } fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "async_delete_cache")?; - run_async_value(py, async move { - scope - .scope(store.delete(&keys)) - .await - .map_err(auth_cache_error) + redis_async(py, "async_delete_cache", auth_cache_error, async move { + store.delete(&keys).await }) } fn delete_sync(&self, py: Python<'_>, keys: Vec) -> PyResult<()> { let store = self.0.clone(); - let scope = call(py, "delete_cache")?; - run_sync_value(py, async move { - scope - .scope(store.delete(&keys)) - .await - .map_err(auth_cache_error) + redis_sync(py, "delete_cache", auth_cache_error, async move { + store.delete(&keys).await }) } @@ -223,23 +183,23 @@ impl NativeAuthCache { &self, request: &NativeRequestRedis, keys: Vec, - ) -> Option { - self.0 - .declare_get_many(&request.0, &keys) - .map(|declared| NativeDeclaredAuthReads(Mutex::new(Some(declared)))) + ) -> Option { + self.0.declare_get_many(&request.0, &keys).map(|declared| { + NativePending::new("async_batch_get_cache", auth_cache_error, async move { + declared.resolve().await.map(values) + }) + }) } /// One round trip for keys missing in memory; resolves to `None` without Redis. fn get_many<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "async_batch_get_cache")?; - run_async_value(py, async move { - scope - .scope(store.get_many(&keys)) + redis_async(py, "async_batch_get_cache", auth_cache_error, async move { + store + .get_many(&keys) .await .transpose() .map(|read| read.map(values)) - .map_err(auth_cache_error) }) } @@ -250,13 +210,12 @@ impl NativeAuthCache { amount: PyNumber, ) -> PyResult> { let store = self.0.clone(); - let scope = call(py, "async_increment_if_present")?; - run_async_value(py, async move { - scope - .scope(store.increment_if_present(&key, amount.0)) - .await - .map_err(auth_cache_error) - }) + redis_async( + py, + "async_increment_if_present", + auth_cache_error, + async move { store.increment_if_present(&key, amount.0).await }, + ) } fn get_local(&self, key: &str) -> PyResult> { @@ -283,56 +242,3 @@ impl NativeAuthCache { self.0.len_local().map_err(store_error) } } - -/// A `declare_get_many` on a request's pre-call batch. `resolve` answers it once. -#[pyclass(frozen, name = "NativeDeclaredAuthReads")] -pub(crate) struct NativeDeclaredAuthReads(Mutex>); - -#[pymethods] -impl NativeDeclaredAuthReads { - fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { - let declared = self - .0 - .lock() - .unwrap_or_else(PoisonError::into_inner) - .take() - .ok_or_else(|| PyRuntimeError::new_err("this auth cache read was already resolved"))?; - let scope = call(py, "async_batch_get_cache")?; - run_async_value(py, async move { - scope - .scope(declared.resolve()) - .await - .map(values) - .map_err(auth_cache_error) - }) - } -} - -/// A `set_pre_call` or `delete_pre_call`. `resolve` settles it once, raising when Redis refused it. -#[pyclass(frozen, name = "NativeDeclaredAuthWrite")] -pub(crate) struct NativeDeclaredAuthWrite(Mutex>); - -impl NativeDeclaredAuthWrite { - fn new(declared: DeclaredWrite) -> Self { - Self(Mutex::new(Some(declared))) - } -} - -#[pymethods] -impl NativeDeclaredAuthWrite { - fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { - let declared = self - .0 - .lock() - .unwrap_or_else(PoisonError::into_inner) - .take() - .ok_or_else(|| PyRuntimeError::new_err("this auth cache write was already resolved"))?; - let scope = call(py, "async_set_cache")?; - run_async_value(py, async move { - scope - .scope(declared.resolve()) - .await - .map_err(auth_cache_error) - }) - } -} diff --git a/litellm-rust/crates/redis-store/src/auth_cache.rs b/litellm-rust/crates/redis-store/src/auth_cache.rs index ae333e539a7..07e8348dcf7 100644 --- a/litellm-rust/crates/redis-store/src/auth_cache.rs +++ b/litellm-rust/crates/redis-store/src/auth_cache.rs @@ -6,7 +6,7 @@ use litellm_python_compat::{Value, json, repr::to_str}; use litellm_redis_client::{Declared, RedisBatch, RedisClient}; use crate::tier::{Clock, system_clock}; -use crate::{DecodeError, Number, RequestRedis, codec::python_value}; +use crate::{DecodeError, ErrorKind, Number, RequestRedis, codec::python_value}; /// `InMemoryCache`'s TTL for a write that names none and has no `default_in_memory_ttl`. pub const MEMORY_TTL: Duration = Duration::from_secs(600); @@ -50,6 +50,17 @@ pub enum Error { Encode(#[from] litellm_python_compat::Error), } +impl Error { + pub fn kind(&self) -> ErrorKind { + match self { + Self::Redis(error) => error.into(), + Self::Decode(_) => ErrorKind::Decode, + Self::Encode(_) => ErrorKind::Encode, + Self::Memory(_) => ErrorKind::Memory, + } + } +} + /// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every /// other object under a prefixed key. pub fn is_key_object(key: &str) -> bool { diff --git a/litellm-rust/crates/redis-store/src/error.rs b/litellm-rust/crates/redis-store/src/error.rs index ff6c30f6859..89dd61433ad 100644 --- a/litellm-rust/crates/redis-store/src/error.rs +++ b/litellm-rust/crates/redis-store/src/error.rs @@ -12,6 +12,8 @@ pub enum ErrorKind { Rejected, #[strum(serialize = "decode")] Decode, + #[strum(serialize = "encode")] + Encode, #[strum(serialize = "memory")] Memory, } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index fb24f97b4b1..55998083af5 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -15,7 +15,6 @@ class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... class NativeRedisStoreError(Exception): ... -class NativeAuthCacheError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... @@ -93,14 +92,6 @@ class NativeRouterUsage: def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... -@final -class NativeDeclaredAuthReads: - def resolve(self) -> Future[list[object]]: ... - -@final -class NativeDeclaredAuthWrite: - def resolve(self) -> Future[None]: ... - @final class NativeAuthCache: def __new__( @@ -120,15 +111,17 @@ class NativeAuthCache: def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... def set_pre_call( self, request: NativeRequestRedis, key: str, value: object, ttl_seconds: float | None = None - ) -> NativeDeclaredAuthWrite | None: ... - def delete_pre_call(self, request: NativeRequestRedis, key: str) -> NativeDeclaredAuthWrite | None: ... + ) -> NativePending[None] | None: ... + def delete_pre_call(self, request: NativeRequestRedis, key: str) -> NativePending[None] | None: ... def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... def write_back( self, request: NativeRequestRedis | None, entries: Sequence[tuple[str, object, float | None]] ) -> Future[None]: ... def delete(self, keys: Sequence[str]) -> Future[None]: ... def delete_sync(self, keys: Sequence[str]) -> None: ... - def declare_get_many(self, request: NativeRequestRedis, keys: Sequence[str]) -> NativeDeclaredAuthReads | None: ... + def declare_get_many( + self, request: NativeRequestRedis, keys: Sequence[str] + ) -> NativePending[list[object]] | None: ... def get_many(self, keys: Sequence[str]) -> Future[list[object] | None]: ... def increment_if_present(self, key: str, amount: float) -> Future[float | None]: ... def get_local(self, key: str) -> object: ... @@ -349,12 +342,9 @@ __all__ = [ "ForkedAfterNativeRuntimeStarted", "HuggingFaceEncoding", "NativeAuthCache", - "NativeAuthCacheError", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", - "NativeDeclaredAuthReads", - "NativeDeclaredAuthWrite", "NativeDiagnosticProcessor", "NativePending", "NativeRedisClient",