diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 1a4eb51af08..b13bce2ed67 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2611,6 +2611,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-gcs" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-cache", + "percent-encoding", + "reqwest 0.12.28", + "serde_json", + "tokio", + "wiremock", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2815,6 +2830,7 @@ dependencies = [ "litellm-cache", "litellm-cache-azure-blob", "litellm-cache-disk", + "litellm-cache-gcs", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index fabec9bcd6c..3ab1e592c6a 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -32,6 +32,7 @@ litellm-cache = { path = "crates/cache" } litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } @@ -70,6 +71,7 @@ base64 = "0.22" moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" +percent-encoding = "2.3" webpki-roots = "1" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 8aeddae9efc..534d85acdb0 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -128,6 +128,14 @@ impl VertexAuth { } } + pub async fn access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + self.load_provider(config, env_lookup).await?.token().await + } + pub async fn validate_environment( &self, headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml new file mode 100644 index 00000000000..4ec60bcfa3b --- /dev/null +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-cache-gcs" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-cache.workspace = true +percent-encoding.workspace = true +reqwest.workspace = true +tokio.workspace = true + +[dev-dependencies] +serde_json.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs new file mode 100644 index 00000000000..65282ac99d5 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -0,0 +1,260 @@ +use std::{future::Future, sync::Arc, time::Duration}; + +use futures_util::future::try_join_all; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, + FlushCache, +}; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; +use reqwest::Client; + +use crate::{GcpTokenSource, TokenSource}; + +pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; + +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); + +pub fn key_prefix(gcs_path: Option<&str>) -> String { + match gcs_path { + Some(path) if !path.is_empty() => format!("{}/", path.trim_end_matches('/')), + _ => String::new(), + } +} + +#[derive(Clone, Debug)] +pub struct GcsConfig { + pub bucket_name: String, + pub gcs_path: Option, + pub path_service_account: Option, + pub endpoint: String, +} + +impl GcsConfig { + pub fn new(bucket_name: impl Into) -> Self { + Self { + bucket_name: bucket_name.into(), + gcs_path: None, + path_service_account: None, + endpoint: DEFAULT_ENDPOINT.to_string(), + } + } +} + +pub struct GcsCache { + config: GcsConfig, + key_prefix: String, + client: Client, + token: Arc, + codec: S, +} + +impl GcsCache { + pub fn new(config: GcsConfig, codec: S) -> Result { + let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + Self::with_token_source(config, codec, token) + } + + pub fn with_token_source( + config: GcsConfig, + codec: S, + token: Arc, + ) -> Result { + let client = Client::builder().build().map_err(|_| Error::Unavailable)?; + let key_prefix = key_prefix(config.gcs_path.as_deref()); + Ok(Self { + config, + key_prefix, + client, + token, + codec, + }) + } + + pub fn bucket_name(&self) -> &str { + &self.config.bucket_name + } + + pub fn key_prefix(&self) -> &str { + &self.key_prefix + } + + pub fn path_service_account(&self) -> Option<&str> { + self.config.path_service_account.as_deref() + } + + pub fn object_name(&self, key: &str) -> String { + format!("{}{}", self.key_prefix, key) + } + + fn encoded_object_name(&self, key: &str) -> String { + percent_encode(self.object_name(key).as_bytes(), OBJECT_NAME_ENCODE_SET).to_string() + } + + fn endpoint(&self, path: &str) -> String { + format!("{}{}", self.config.endpoint.trim_end_matches('/'), path) + } + + async fn async_set(&self, key: &str, value: S::Value) -> Result<(), Error> { + let token = self.token.bearer_token().await?; + let payload = self.codec.encode(&value)?; + let url = self.endpoint(&format!( + "/upload/storage/v1/b/{}/o?uploadType=media&name={}", + self.config.bucket_name, + self.encoded_object_name(key) + )); + let response = self + .client + .post(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(payload) + .send() + .await + .map_err(|_| Error::Unavailable)?; + if !response.status().is_success() { + return Err(Error::Unavailable); + } + Ok(()) + } + + async fn async_get(&self, key: &str) -> Result, Error> { + let token = self.token.bearer_token().await?; + let url = self.endpoint(&format!( + "/storage/v1/b/{}/o/{}?alt=media", + self.config.bucket_name, + self.encoded_object_name(key) + )); + let response = self + .client + .get(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .send() + .await + .map_err(|_| Error::Unavailable)?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if !response.status().is_success() { + return Err(Error::Unavailable); + } + let body = response.bytes().await.map_err(|_| Error::Unavailable)?; + self.codec + .decode(&body) + .map(Some) + .map_err(|_| Error::InvalidEntry) + } + + fn run_sync(future: F) -> Result + where + F: Future> + Send, + T: Send, + { + let run = || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|_| Error::Unavailable) + .and_then(|runtime| runtime.block_on(future)) + }; + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread { + return tokio::task::block_in_place(run); + } + return std::thread::scope(|scope| { + scope + .spawn(run) + .join() + .map_err(|_| Error::Unavailable) + .and_then(|result| result) + }); + } + run() + } +} + +impl BaseCache for GcsCache { + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache(&self, key: &str, value: Self::Value, _: &Self::Context) -> Result<(), Error> { + Self::run_sync(self.async_set(key, value)) + } + + fn get_cache(&self, key: &str, _: &Self::Context) -> Result, Error> { + Self::run_sync(self.async_get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + _: Self::Context, + ) -> Result<(), Error> { + self.async_set(key, value).await + } + + async fn async_get_cache( + &self, + key: &str, + _: &Self::Context, + ) -> Result, Error> { + self.async_get(key).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, Self::Value)>, + context: Self::Context, + ) -> Result<(), Error> { + try_join_all(entries.into_iter().map(|(key, value)| { + let context = context.clone(); + async move { self.async_set_cache(&key, value, context).await } + })) + .await + .map(|_| ()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for GcsCache { + async fn async_batch_get_cache( + &self, + keys: Vec, + context: Self::Context, + ) -> Result>, Error> { + try_join_all(keys.into_iter().map(|key| { + let context = context.clone(); + async move { + match self.async_get_cache(&key, &context).await { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + } + } + })) + .await + } +} + +impl FlushCache for GcsCache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs new file mode 100644 index 00000000000..cbb61cf0685 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod token; + +pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; +pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs new file mode 100644 index 00000000000..adb601c276c --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -0,0 +1,44 @@ +use std::{future::Future, pin::Pin}; + +use litellm_auth_gcp::{VertexAuth, VertexConfig}; +use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_cache::Error; + +pub trait TokenSource: Send + Sync + 'static { + fn bearer_token(&self) -> Pin> + Send + '_>>; +} + +pub struct GcpTokenSource { + auth: VertexAuth, + config: VertexConfig, +} + +impl GcpTokenSource { + pub fn new(path_service_account: Option) -> Self { + let credentials = path_service_account + .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); + Self { + auth: VertexAuth::default(), + config: VertexConfig::new(credentials, None, None), + } + } +} + +impl TokenSource for GcpTokenSource { + fn bearer_token(&self) -> Pin> + Send + '_>> { + Box::pin(async move { + self.auth + .access_token(&self.config, &|name| std::env::var(name).ok()) + .await + .map_err(|_| Error::Unavailable) + }) + } +} + +pub struct StaticTokenSource(pub String); + +impl TokenSource for StaticTokenSource { + fn bearer_token(&self) -> Pin> + Send + '_>> { + Box::pin(async move { Ok(self.0.clone()) }) + } +} diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs new file mode 100644 index 00000000000..45eecf01cec --- /dev/null +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -0,0 +1,324 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache, + JsonCodec, +}; +use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix}; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_bytes, header, method, path, query_param}, +}; + +fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig { + GcsConfig { + bucket_name: "bucket".into(), + gcs_path: gcs_path.map(str::to_string), + path_service_account: None, + endpoint: server.uri(), + } +} + +fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache> { + GcsCache::with_token_source( + config(server, gcs_path), + JsonCodec::new(), + Arc::new(StaticTokenSource("tok".into())), + ) + .unwrap() +} + +#[tokio::test] +async fn set_writes_encoded_object_and_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .and(header("authorization", "Bearer tok")) + .and(header("content-type", "application/json")) + .and(body_bytes(br#"{"value":"entry"}"#)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + cache(&server, Some("cache/")) + .set_cache( + "team:a b/c", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].url.query(), + Some("uploadType=media&name=cache%2Fteam%3Aa%20b%2Fc") + ); +} + +#[tokio::test] +async fn get_maps_statuses_and_decode_failures() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/hit")) + .and(query_param("alt", "media")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/missing")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/server-error")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/invalid")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .mount(&server) + .await; + + let cache = cache(&server, None); + assert_eq!( + cache + .get_cache("hit", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); + assert_eq!( + cache + .get_cache("missing", &ExactCacheContext::default()) + .unwrap(), + None + ); + assert_eq!( + cache + .get_cache("server-error", &ExactCacheContext::default()) + .unwrap_err(), + Error::Unavailable + ); + assert_eq!( + cache + .get_cache("invalid", &ExactCacheContext::default()) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn key_prefix_normalizes_paths() { + assert_eq!(key_prefix(None), ""); + assert_eq!(key_prefix(Some("a/b/")), "a/b/"); + assert_eq!(key_prefix(Some("a/b")), "a/b/"); + assert_eq!(key_prefix(Some("")), ""); +} + +#[tokio::test] +async fn object_names_use_python_quote_encoding() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + let cache = cache(&server, Some("p/")); + cache + .set_cache( + "a~b-c_d.e/f g%h", + json!({"value": "punctuation"}), + &ExactCacheContext::default(), + ) + .unwrap(); + cache + .set_cache( + "ключ", + json!({"value": "utf8"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let queries: Vec<_> = requests + .iter() + .filter_map(|request| request.url.query()) + .collect(); + assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")); + assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")); +} + +#[tokio::test] +async fn ignores_ttl_and_writes_pipeline_concurrently() { + let server = MockServer::start().await; + for key in ["one", "two", "three"] { + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .and(query_param("name", key)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + } + let cache = cache(&server, None); + assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))), + None + ); + cache + .async_set_cache_pipeline( + vec![ + ("one".into(), json!({"key": "one"})), + ("two".into(), json!({"key": "two"})), + ("three".into(), json!({"key": "three"})), + ], + ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn async_batch_get_preserves_hits_misses_and_invalid_entries() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/hit")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/missing")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/invalid")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .mount(&server) + .await; + + assert_eq!( + cache(&server, None) + .async_batch_get_cache( + vec!["hit".into(), "missing".into(), "invalid".into()], + ExactCacheContext::default(), + ) + .await + .unwrap(), + vec![ + BatchEntry::Hit(json!({"value": "entry"})), + BatchEntry::Miss, + BatchEntry::Invalid, + ] + ); +} + +#[tokio::test] +async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() { + let server = MockServer::start().await; + let cache = cache(&server, None); + assert_eq!(cache.flush_cache(), Ok(())); + assert_eq!(cache.disconnect().await, Ok(())); + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); +} + +#[test] +fn sync_operations_work_without_an_active_runtime() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let server = runtime.block_on(MockServer::start()); + runtime.block_on( + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server), + ); + runtime.block_on( + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server), + ); + let cache = cache(&server, None); + cache + .set_cache( + "key", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn sync_operations_work_inside_a_multi_thread_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + let cache = cache(&server, None); + cache + .set_cache( + "key", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); +} + +struct FailingTokenSource; + +impl TokenSource for FailingTokenSource { + fn bearer_token( + &self, + ) -> std::pin::Pin> + Send + '_>> + { + Box::pin(async { Err(Error::Unavailable) }) + } +} + +#[tokio::test] +async fn token_source_failure_skips_http() { + let server = MockServer::start().await; + let cache = GcsCache::with_token_source( + config(&server, None), + JsonCodec::::new(), + Arc::new(FailingTokenSource), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap_err(), + Error::Unavailable + ); + assert_eq!(server.received_requests().await.unwrap().len(), 0); +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..1a381d0afd8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("operation is not supported by this cache")] + UnsupportedOperation, } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index f5d08850b39..4a9ea633ac5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,6 +24,7 @@ litellm-cache.workspace = true litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-gcs.workspace = true litellm-cache-disk.workspace = true litellm-cache-response.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 927017e2e60..6c180708335 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -79,6 +79,13 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[derive(Debug, PartialEq)] +pub(super) struct GcsCacheConfig { + pub(super) bucket_name: String, + pub(super) key_prefix: String, + pub(super) path_service_account: Option, +} + pub(super) struct AzureBlobCacheConfig { pub(super) account_url: String, pub(super) container: String, @@ -98,6 +105,7 @@ const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + Gcs(GcsCacheConfig), Disk(DiskCacheConfig), AzureBlob(AzureBlobCacheConfig), } @@ -114,6 +122,7 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + GcsBucket, DiskStore, } @@ -125,6 +134,7 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::GcsBucket => "native GCS cache requires a configured bucket name", Self::DiskStore => "native disk cache requires the built-in diskcache store", } } @@ -168,6 +178,13 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::Gcs) => match project_gcs(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Gcs(backend), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some(CacheType::Disk) => match project_disk(&backend)? { Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { policy, @@ -185,8 +202,7 @@ impl NativeCacheConfig { CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 - | CacheType::QdrantSemantic - | CacheType::Gcs, + | CacheType::QdrantSemantic, ) | None => Ok(CacheConfigProjection::Unsupported( UnsupportedCacheConfig::Backend, @@ -198,8 +214,9 @@ impl NativeCacheConfig { let default_ttl = match &self.backend { CacheBackendConfig::Memory(config) => Some(config.default_ttl), CacheBackendConfig::Redis(config) => Some(config.default_ttl), - CacheBackendConfig::Disk(_) => None, - CacheBackendConfig::AzureBlob(_) => None, + CacheBackendConfig::Disk(_) + | CacheBackendConfig::AzureBlob(_) + | CacheBackendConfig::Gcs(_) => None, }; if service.default_ttl() != default_ttl { return Some("facade and native backend default TTLs must match"); @@ -226,6 +243,31 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Gcs(config) + if service + .gcs_backend() + .is_none_or(|backend| backend.bucket_name() != config.bucket_name) => + { + Some("facade and native backend buckets must match") + } + CacheBackendConfig::Gcs(config) + if service + .gcs_backend() + .is_none_or(|backend| backend.key_prefix() != config.key_prefix) => + { + Some("facade and native backend key prefixes must match") + } + CacheBackendConfig::Gcs(config) + if service.gcs_backend().is_none_or(|backend| { + backend.path_service_account() != config.path_service_account.as_deref() + }) => + { + Some("facade and native backend credentials must match") + } + CacheBackendConfig::Gcs(_) => None, CacheBackendConfig::Disk(_) if service.kind() != "disk" => { Some("facade and native backend types must match") } @@ -277,6 +319,23 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { }) } +#[inline(never)] +fn project_gcs( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let bucket_name = match backend.getattr("bucket_name")?.extract::>() { + Ok(Some(bucket_name)) if !bucket_name.is_empty() => bucket_name, + _ => return Ok(Err(UnsupportedCacheConfig::GcsBucket)), + }; + Ok(Ok(GcsCacheConfig { + bucket_name, + key_prefix: backend.getattr("key_prefix")?.extract::()?, + path_service_account: backend + .getattr("path_service_account")? + .extract::>()?, + })) +} + #[inline(never)] fn project_disk( backend: &Bound<'_, PyAny>, @@ -680,7 +739,7 @@ mod tests { use super::{ CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement, - DiskCacheConfig, NativeCacheConfig, RedisProtocol, + DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, }; use crate::cache::native::NativeResponseCache; @@ -755,6 +814,71 @@ mod tests { }); } + #[test] + fn projects_gcs_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(bucket_name='bucket', key_prefix='cache/', path_service_account='credentials.json')\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("GCS cache should be supported"); + }; + let CacheBackendConfig::Gcs(gcs) = config.backend else { + panic!("expected GCS configuration"); + }; + assert_eq!( + gcs, + GcsCacheConfig { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: Some("credentials.json".into()), + } + ); + let matching = NativeResponseCache::gcs( + litellm_cache_gcs::GcsConfig { + bucket_name: "bucket".into(), + gcs_path: Some("cache/".into()), + path_service_account: Some("credentials.json".into()), + endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(), + }, + Some("token".into()), + ) + .unwrap(); + let matching_config = NativeCacheConfig { + policy: config.policy, + backend: CacheBackendConfig::Gcs(gcs), + }; + assert_eq!(matching_config.service_mismatch(&matching), None); + }); + } + + #[test] + fn rejects_gcs_without_a_bucket_name() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("GCS cache without a bucket should be unsupported"); + }; + assert!(matches!(&reason, UnsupportedCacheConfig::GcsBucket)); + assert_eq!( + reason.message(), + "native GCS cache requires a configured bucket name" + ); + }); + } + #[test] fn projects_resolved_redis_tls_configuration() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index ee6423daf6e..0fa4511b1c1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -320,6 +320,7 @@ impl FacadeGuard { "RedisClusterCache", "redis", ), + ("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"), ("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"), ("azure-blob", _) => ( "litellm.caching.azure_blob_cache", @@ -369,6 +370,9 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "bucket_name", + "key_prefix", + "path_service_account", ], )?, disk_store: (kind == "disk") diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 6e12e12040b..64ac3696e80 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -2,6 +2,8 @@ use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_host_python::{release_gil, run_sync_value}; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; + use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; #[pyclass(frozen, name = "_CacheTestHandle")] @@ -64,6 +66,31 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))] + fn gcs( + py: Python<'_>, + bucket_name: String, + gcs_path: Option, + path_service_account: Option, + endpoint: Option, + token: Option, + ) -> PyResult { + let config = GcsConfig { + bucket_name, + gcs_path, + path_service_account, + endpoint: endpoint.unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), + }; + let service = release_gil(py, move || NativeResponseCache::gcs(config, token)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] #[pyo3(signature = (directory))] fn disk(py: Python<'_>, directory: String) -> PyResult { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..ac4494150d9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -10,7 +10,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +21,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 80789cc279a..b10f9e2a443 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -3,6 +3,7 @@ use std::{path::Path, sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_disk::DiskCache; +use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_response::{ @@ -17,6 +18,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + Gcs(Arc>>), Disk(Arc>>), AzureBlob(Arc>>), } @@ -54,6 +56,18 @@ impl NativeResponseCache { Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) } + pub fn gcs(config: GcsConfig, token: Option) -> Result { + let backend = match token { + Some(token) => GcsCache::with_token_source( + config, + ResponseCacheCodec, + Arc::new(StaticTokenSource(token)), + )?, + None => GcsCache::new(config, ResponseCacheCodec)?, + }; + Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend))))) + } + pub async fn azure_blob(account_url: &str, container: &str) -> Result { let backend = AzureBlobCache::connect( account_url, @@ -73,7 +87,7 @@ impl NativeResponseCache { cache.backend().account_url(), cache.backend().container_name(), )), - Self::Memory(_) | Self::Redis { .. } | Self::Disk(_) => None, + Self::Memory(_) | Self::Redis { .. } | Self::Disk(_) | Self::Gcs(_) => None, } } } @@ -83,6 +97,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::Gcs(_) => "gcs", Self::Disk(_) => "disk", Self::AzureBlob(_) => "azure-blob", } @@ -92,6 +107,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::Gcs(cache) => cache.default_ttl(), Self::Disk(cache) => cache.default_ttl(), Self::AzureBlob(cache) => cache.default_ttl(), } @@ -101,12 +117,13 @@ impl NativeResponseCache { match self { Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), + Self::Gcs(_) => None, } } pub fn topology(&self) -> Option<&RedisTopology> { match self { - Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None, + Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), } } @@ -114,14 +131,14 @@ impl NativeResponseCache { pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) => None, + Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) => None, + Self::Redis { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, } } @@ -138,7 +155,7 @@ impl NativeResponseCache { pub fn directory(&self) -> Option<&Path> { match self { Self::Disk(cache) => Some(cache.backend().directory()), - Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) => None, + Self::Memory(_) | Self::Redis { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None, } } @@ -150,6 +167,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(request, now), Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Gcs(cache) => cache.lookup(request, now), Self::Disk(cache) => cache.lookup(request, now), Self::AzureBlob(cache) => cache.lookup(request, now), } @@ -164,6 +182,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(request, response, now), Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Gcs(cache) => cache.store(request, response, now), Self::Disk(cache) => cache.store(request, response, now), Self::AzureBlob(cache) => cache.store(request, response, now), } @@ -177,6 +196,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup_batch(requests, now), Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Gcs(cache) => cache.lookup_batch(requests, now), Self::Disk(cache) => cache.lookup_batch(requests, now), Self::AzureBlob(cache) => cache.lookup_batch(requests, now), } @@ -190,6 +210,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Gcs(cache) => cache.async_lookup(request, now).await, Self::Disk(cache) => cache.async_lookup(request, now).await, Self::AzureBlob(cache) => cache.async_lookup(request, now).await, } @@ -211,6 +232,7 @@ impl NativeResponseCache { cache, buffer: Some(buffer), } => buffer.async_store(cache, request, response, now).await, + Self::Gcs(cache) => cache.async_store(request, response, now).await, Self::Disk(cache) => cache.async_store(request, response, now).await, Self::AzureBlob(cache) => cache.async_store(request, response, now).await, } @@ -224,6 +246,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await, Self::Disk(cache) => cache.async_lookup_batch(requests, now).await, Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await, } @@ -237,6 +260,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_store_batch(entries, now).await, Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Gcs(cache) => cache.async_store_batch(entries, now).await, Self::Disk(cache) => cache.async_store_batch(entries, now).await, Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await, } @@ -251,6 +275,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::Gcs(cache) => cache.async_flush().await, Self::Disk(cache) => cache.async_flush().await, Self::AzureBlob(cache) => cache.async_flush().await, } @@ -260,8 +285,16 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::Gcs(cache) => cache.test_connection().await, Self::Disk(cache) => cache.test_connection().await, Self::AzureBlob(cache) => cache.test_connection().await, } } + + pub fn gcs_backend(&self) -> Option<&GcsCache> { + match self { + Self::Gcs(cache) => Some(cache.backend()), + _ => None, + } + } } diff --git a/tests/test_litellm_rust/support/fake_gcs.py b/tests/test_litellm_rust/support/fake_gcs.py new file mode 100644 index 00000000000..67eb61798b9 --- /dev/null +++ b/tests/test_litellm_rust/support/fake_gcs.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from functools import partial +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from socket import socket +from types import MappingProxyType +from typing import Final, cast +from urllib.parse import unquote, urlsplit + + +@dataclass(frozen=True, slots=True) +class RecordedRequest: + method: str + path: str + query: str + headers: Mapping[str, str] + body: bytes + + +class _FakeGcsHandler(BaseHTTPRequestHandler): + def __init__( + self, + request: socket | tuple[bytes, socket], + client_address: tuple[str, int], + server: ThreadingHTTPServer, + *, + fake: FakeGcs, + ) -> None: + self._fake: Final = fake + super().__init__(request, client_address, server) + + def _handle(self) -> None: + parsed: Final = urlsplit(self.path) + content_length: Final = int(self.headers.get("Content-Length", "0")) + body: Final = self.rfile.read(content_length) if content_length else b"" + headers: Final = MappingProxyType( + {name.title(): value for name, value in self.headers.items()} + ) + self._fake.record( + RecordedRequest( + method=self.command, + path=parsed.path, + query=parsed.query, + headers=headers, + body=body, + ) + ) + if self.headers.get("Authorization") != f"Bearer {self._fake.token}": + self._send_json(401, {"error": "unauthorized"}) + return + + upload_prefix: Final = "/upload/storage/v1/b/" + download_prefix: Final = "/storage/v1/b/" + if parsed.path.startswith(upload_prefix) and parsed.path.endswith("/o"): + self._upload(parsed.path[len(upload_prefix) : -2], parsed.query, body) + return + if parsed.path.startswith(download_prefix): + self._download(parsed.path[len(download_prefix) :], parsed.query) + return + self._send_json(404, {"error": "not found"}) + + def _upload(self, path: str, query: str, body: bytes) -> None: + values: Final = { + unquote(pair.partition("=")[0]): unquote(pair.partition("=")[2]) + for pair in query.split("&") + if pair + } + if not path or values.get("uploadType") != "media" or "name" not in values: + self._send_json(404, {"error": "not found"}) + return + self._fake.put_object(path, values["name"], body) + self._send_json(200, {"name": values["name"], "bucket": path}) + + def _download(self, path: str, query: str) -> None: + bucket, separator, encoded_name = path.partition("/o/") + if not separator or query != "alt=media": + self._send_json(404, {"error": "not found"}) + return + name: Final = unquote(encoded_name) + if name.endswith("/server-error") or name == "server-error": + self._send_json(500, {"error": "server error"}) + return + body: Final = self._fake.get_object(bucket, name) + if body is None: + self._send_json(404, {"error": "not found"}) + return + self._send(200, body, "application/octet-stream") + + def _send_json(self, status: int, value: object) -> None: + payload: Final = json.dumps(value).encode() + self._send(status, payload, "application/json") + + def _send(self, status: int, body: bytes, content_type: str) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + do_GET = _handle + do_POST = _handle + + +class FakeGcs: + def __init__(self) -> None: + self._objects: dict[tuple[str, str], bytes] = {} # mutable-ok: fake object store + self._requests: list[RecordedRequest] = [] # mutable-ok: recorded request history + self._server = ThreadingHTTPServer( + ("127.0.0.1", 0), + partial(_FakeGcsHandler, fake=self), + ) + self._worker = threading.Thread(target=self._server.serve_forever, daemon=True) + self._worker.start() + self.token: Final = "test-token" + + @property + def url(self) -> str: + address: Final = cast(tuple[str, int], self._server.server_address) + host, port = address + return f"http://{host}:{port}" + + @property + def objects(self) -> Mapping[tuple[str, str], bytes]: + return MappingProxyType(self._objects) + + @property + def requests(self) -> tuple[RecordedRequest, ...]: + return tuple(self._requests) + + def put(self, bucket: str, name: str, body: bytes) -> None: + self.put_object(bucket, name, body) + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._worker.join(timeout=5) + + def record(self, request: RecordedRequest) -> None: + self._requests.append(request) + + def put_object(self, bucket: str, name: str, body: bytes) -> None: + self._objects[(bucket, name)] = body + + def get_object(self, bucket: str, name: str) -> bytes | None: + return self._objects.get((bucket, name)) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 8cddc868073..ba95b21d530 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -22,11 +22,13 @@ from azure.storage.blob import ContainerClient import litellm from litellm.caching.azure_blob_cache import AzureBlobCache from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.gcs_cache import GCSCache from litellm.caching.disk_cache import DiskCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound pytestmark: Final = pytest.mark.requires_rust_extension @@ -34,6 +36,7 @@ pytestmark: Final = pytest.mark.requires_rust_extension class CacheLookup(Protocol): def get_cache(self, **kwargs: object) -> object: ... + def flush_cache(self) -> object: ... def request(key: str = "key") -> dict[str, object]: @@ -53,6 +56,15 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def fake_gcs() -> Generator[FakeGcs]: + server: Final = FakeGcs() + try: + yield server + finally: + server.close() + + @pytest.fixture def azure_blob_facade() -> Generator[Cache]: account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") @@ -630,6 +642,222 @@ async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: } +async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + fake_gcs.put( + "bucket", + "cache/sync", + json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), + ) + fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) + fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("missing")) is None + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = fake_gcs.objects[("bucket", "cache/native")] + stored_value: Final = cast(dict[str, object], json.loads(stored)) + assert stored_value["response"] == response + assert isinstance(stored_value["timestamp"], float) + upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") + assert upload.path == "/upload/storage/v1/b/bucket/o" + assert upload.query == "uploadType=media&name=cache%2Fnative" + assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" + assert upload.headers["Content-Type"] == "application/json" + upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" + assert "ttl" not in upload_text.lower() + assert "expiry" not in upload_text.lower() + download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) + assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" + assert download.query == "alt=media" + + binding.store(request("sync2"), response) + assert binding.lookup(request("sync2")) == response + assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket").key_prefix == "" + + +async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: + fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + requests: Final = [request("hit"), request("missing"), request("invalid")] + expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} + + assert await binding.async_lookup_batch(requests) == expected + assert binding.lookup_batch(requests) == expected + await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) + assert ("bucket", "cache/first") in fake_gcs.objects + assert ("bucket", "cache/second") in fake_gcs.objects + + +async def test_gcs_facade_binds_only_exact_matching_configuration( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + assert type(facade.cache) is GCSCache + + mismatched_bucket: Final = _native._CacheTestHandle.gcs( + "other", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="buckets must match"): + mismatched_bucket._bind_facade(facade) + mismatched_prefix: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="x", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="key prefixes must match"): + mismatched_prefix._bind_facade(facade) + mismatched_credentials: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + path_service_account="sa.json", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="credentials must match"): + mismatched_credentials._bind_facade(facade) + with pytest.raises(TypeError, match="types must match"): + _native._CacheTestHandle.memory()._bind_facade(facade) + + matching: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + matching._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + await binding.async_store(request("native"), {"value": "native"}) + assert await binding.async_lookup(request("native")) == {"value": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="native") is None + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "key_prefix", "x/"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "path_service_account", "sa.json"): + assert resolver.resolve().kind == "python_callback" + def no_get_cache(*args: object, **kwargs: object) -> None: + return None + + with rebound(facade.cache, "get_cache", no_get_cache): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + class CustomGcs(GCSCache): + pass + + with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + assert resolver.resolve().kind == "python_callback" + custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + with pytest.raises(TypeError, match="types must match"): + matching._bind_facade(custom_facade) + + missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) + with pytest.raises(TypeError, match="requires a configured bucket name"): + matching._bind_facade(missing_bucket) + + +async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + await binding.async_store(request("key"), {"value": "stored"}) + await binding.async_flush() + assert ("bucket", "cache/key") in fake_gcs.objects + assert await binding.async_lookup(request("key")) == {"value": "stored"} + with pytest.raises(NotImplementedError): + await binding.ping() + + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with pytest.raises(AttributeError): + await facade.ping() + assert cast(CacheLookup, facade.cache).flush_cache() is None + + +async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: + wrong_token: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token="wrong-token", + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + wrong_token.lookup(request("missing")) + assert not fake_gcs.objects + + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + binding.lookup(request("server-error")) + assert binding.lookup(request("missing")) is None + + async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively( cluster_nodes: tuple[tuple[str, int], ...], ) -> None: