mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge branch 'main' into litellm_rust_cache_s3
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
0c5bbf8c37
19 changed files with 1292 additions and 14 deletions
16
litellm-rust/Cargo.lock
generated
16
litellm-rust/Cargo.lock
generated
|
|
@ -2704,6 +2704,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"
|
||||
|
|
@ -2924,6 +2939,7 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-azure-blob",
|
||||
"litellm-cache-disk",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-response",
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ 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-s3 = { path = "crates/cache-s3" }
|
||||
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" }
|
||||
|
|
@ -71,6 +72,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"
|
||||
|
|
|
|||
|
|
@ -128,6 +128,14 @@ impl VertexAuth {
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn access_token(
|
||||
&self,
|
||||
config: &VertexConfig,
|
||||
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
|
||||
) -> Result<String, Error> {
|
||||
self.load_provider(config, env_lookup).await?.token().await
|
||||
}
|
||||
|
||||
pub async fn validate_environment(
|
||||
&self,
|
||||
headers: Vec<(String, String)>,
|
||||
|
|
|
|||
20
litellm-rust/crates/cache-gcs/Cargo.toml
Normal file
20
litellm-rust/crates/cache-gcs/Cargo.toml
Normal file
|
|
@ -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"
|
||||
260
litellm-rust/crates/cache-gcs/src/cache.rs
Normal file
260
litellm-rust/crates/cache-gcs/src/cache.rs
Normal file
|
|
@ -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<String>,
|
||||
pub path_service_account: Option<String>,
|
||||
pub endpoint: String,
|
||||
}
|
||||
|
||||
impl GcsConfig {
|
||||
pub fn new(bucket_name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
bucket_name: bucket_name.into(),
|
||||
gcs_path: None,
|
||||
path_service_account: None,
|
||||
endpoint: DEFAULT_ENDPOINT.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GcsCache<S: CacheCodec> {
|
||||
config: GcsConfig,
|
||||
key_prefix: String,
|
||||
client: Client,
|
||||
token: Arc<dyn TokenSource>,
|
||||
codec: S,
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> GcsCache<S> {
|
||||
pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
|
||||
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<dyn TokenSource>,
|
||||
) -> Result<Self, Error> {
|
||||
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<Option<S::Value>, 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<T, F>(future: F) -> Result<T, Error>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>> + 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<S: CacheCodec> BaseCache for GcsCache<S> {
|
||||
type Value = S::Value;
|
||||
type Context = ExactCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
|
||||
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<Option<Self::Value>, 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<Option<Self::Value>, 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<CacheConnectionResult, Error> {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: CacheCodec> BatchCache for GcsCache<S> {
|
||||
async fn async_batch_get_cache(
|
||||
&self,
|
||||
keys: Vec<String>,
|
||||
context: Self::Context,
|
||||
) -> Result<Vec<BatchEntry<Self::Value>>, 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<S: CacheCodec> FlushCache for GcsCache<S> {
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
5
litellm-rust/crates/cache-gcs/src/lib.rs
Normal file
5
litellm-rust/crates/cache-gcs/src/lib.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
mod cache;
|
||||
mod token;
|
||||
|
||||
pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
|
||||
pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};
|
||||
44
litellm-rust/crates/cache-gcs/src/token.rs
Normal file
44
litellm-rust/crates/cache-gcs/src/token.rs
Normal file
|
|
@ -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<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
|
||||
}
|
||||
|
||||
pub struct GcpTokenSource {
|
||||
auth: VertexAuth,
|
||||
config: VertexConfig,
|
||||
}
|
||||
|
||||
impl GcpTokenSource {
|
||||
pub fn new(path_service_account: Option<String>) -> 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<Box<dyn Future<Output = Result<String, Error>> + 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<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
|
||||
Box::pin(async move { Ok(self.0.clone()) })
|
||||
}
|
||||
}
|
||||
324
litellm-rust/crates/cache-gcs/tests/cache.rs
Normal file
324
litellm-rust/crates/cache-gcs/tests/cache.rs
Normal file
|
|
@ -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<JsonCodec<serde_json::Value>> {
|
||||
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<Box<dyn std::future::Future<Output = Result<String, Error>> + 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::<serde_json::Value>::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);
|
||||
}
|
||||
2
litellm-rust/crates/cache/src/error.rs
vendored
2
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -6,6 +6,6 @@ pub enum Error {
|
|||
InvalidEntry,
|
||||
#[error("flushing Redis requires an explicit namespace")]
|
||||
UnscopedFlush,
|
||||
#[error("cache operation is not supported")]
|
||||
#[error("operation is not supported by this cache")]
|
||||
UnsupportedOperation,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ litellm-cache-azure-blob.workspace = true
|
|||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
litellm-cache-s3.workspace = true
|
||||
litellm-cache-gcs.workspace = true
|
||||
litellm-cache-disk.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
serde.workspace = true
|
||||
|
|
|
|||
|
|
@ -81,6 +81,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<String>,
|
||||
}
|
||||
|
||||
pub(super) struct AzureBlobCacheConfig {
|
||||
pub(super) account_url: String,
|
||||
pub(super) container: String,
|
||||
|
|
@ -101,6 +108,7 @@ pub(super) enum CacheBackendConfig {
|
|||
Memory(MemoryCacheConfig),
|
||||
Redis(Box<RedisCacheConfig>),
|
||||
S3(Box<S3CacheConfig>),
|
||||
Gcs(GcsCacheConfig),
|
||||
Disk(DiskCacheConfig),
|
||||
AzureBlob(AzureBlobCacheConfig),
|
||||
}
|
||||
|
|
@ -120,6 +128,7 @@ pub(super) enum UnsupportedCacheConfig {
|
|||
S3Client,
|
||||
S3Credentials,
|
||||
S3Option,
|
||||
GcsBucket,
|
||||
DiskStore,
|
||||
}
|
||||
|
||||
|
|
@ -134,6 +143,7 @@ impl UnsupportedCacheConfig {
|
|||
Self::S3Client => "native S3 client type is not implemented",
|
||||
Self::S3Credentials => "native S3 credentials require Python",
|
||||
Self::S3Option => "native S3 configuration requires Python",
|
||||
Self::GcsBucket => "native GCS cache requires a configured bucket name",
|
||||
Self::DiskStore => "native disk cache requires the built-in diskcache store",
|
||||
}
|
||||
}
|
||||
|
|
@ -184,6 +194,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,
|
||||
|
|
@ -198,10 +215,7 @@ impl NativeCacheConfig {
|
|||
}))
|
||||
}),
|
||||
Some(
|
||||
CacheType::RedisSemantic
|
||||
| CacheType::ValkeySemantic
|
||||
| CacheType::QdrantSemantic
|
||||
| CacheType::Gcs,
|
||||
CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::QdrantSemantic,
|
||||
)
|
||||
| None => Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
|
|
@ -214,8 +228,9 @@ impl NativeCacheConfig {
|
|||
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
|
||||
CacheBackendConfig::S3(_) => None,
|
||||
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");
|
||||
|
|
@ -266,6 +281,31 @@ impl NativeCacheConfig {
|
|||
Some("facade and native backend endpoints must match")
|
||||
}
|
||||
CacheBackendConfig::S3(_) => None,
|
||||
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")
|
||||
}
|
||||
|
|
@ -317,6 +357,23 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
|
|||
})
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_gcs(
|
||||
backend: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Result<GcsCacheConfig, UnsupportedCacheConfig>> {
|
||||
let bucket_name = match backend.getattr("bucket_name")?.extract::<Option<String>>() {
|
||||
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::<String>()?,
|
||||
path_service_account: backend
|
||||
.getattr("path_service_account")?
|
||||
.extract::<Option<String>>()?,
|
||||
}))
|
||||
}
|
||||
|
||||
#[inline(never)]
|
||||
fn project_disk(
|
||||
backend: &Bound<'_, PyAny>,
|
||||
|
|
@ -818,7 +875,7 @@ mod tests {
|
|||
|
||||
use super::{
|
||||
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
|
||||
DiskCacheConfig, NativeCacheConfig, RedisProtocol,
|
||||
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
|
||||
};
|
||||
use crate::cache::native::NativeResponseCache;
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
|
|
@ -894,6 +951,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();
|
||||
|
|
|
|||
|
|
@ -341,6 +341,7 @@ impl FacadeGuard {
|
|||
"redis",
|
||||
),
|
||||
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
|
||||
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
|
||||
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
|
||||
("azure-blob", _) => (
|
||||
"litellm.caching.azure_blob_cache",
|
||||
|
|
@ -392,6 +393,7 @@ impl FacadeGuard {
|
|||
"redis_flush_size",
|
||||
"bucket_name",
|
||||
"key_prefix",
|
||||
"path_service_account",
|
||||
],
|
||||
)?,
|
||||
disk_store: (kind == "disk")
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
|||
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")]
|
||||
|
|
@ -100,6 +102,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<String>,
|
||||
path_service_account: Option<String>,
|
||||
endpoint: Option<String>,
|
||||
token: Option<String>,
|
||||
) -> PyResult<Self> {
|
||||
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<Self> {
|
||||
|
|
|
|||
|
|
@ -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::{
|
||||
|
|
@ -19,6 +20,7 @@ pub(super) enum NativeResponseCache {
|
|||
buffer: Option<Arc<WriteBuffer>>,
|
||||
},
|
||||
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
|
||||
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
|
||||
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
|
||||
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
|
||||
}
|
||||
|
|
@ -64,6 +66,18 @@ impl NativeResponseCache {
|
|||
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
|
||||
}
|
||||
|
||||
pub fn gcs(config: GcsConfig, token: Option<String>) -> Result<Self, Error> {
|
||||
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<Self, Error> {
|
||||
let backend = AzureBlobCache::connect(
|
||||
account_url,
|
||||
|
|
@ -83,7 +97,9 @@ impl NativeResponseCache {
|
|||
cache.backend().account_url(),
|
||||
cache.backend().container_name(),
|
||||
)),
|
||||
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) => None,
|
||||
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::Gcs(_) => {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -94,6 +110,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(_) => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::S3(_) => "s3",
|
||||
Self::Gcs(_) => "gcs",
|
||||
Self::Disk(_) => "disk",
|
||||
Self::AzureBlob(_) => "azure-blob",
|
||||
}
|
||||
|
|
@ -104,6 +121,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.default_ttl(),
|
||||
Self::Redis { cache, .. } => cache.default_ttl(),
|
||||
Self::S3(cache) => cache.default_ttl(),
|
||||
Self::Gcs(cache) => cache.default_ttl(),
|
||||
Self::Disk(cache) => cache.default_ttl(),
|
||||
Self::AzureBlob(cache) => cache.default_ttl(),
|
||||
}
|
||||
|
|
@ -141,13 +159,13 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
|
||||
Self::Redis { cache, .. } => cache.backend().namespace(),
|
||||
Self::S3(_) => None,
|
||||
Self::S3(_) | 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()),
|
||||
Self::S3(_) => None,
|
||||
}
|
||||
|
|
@ -156,14 +174,22 @@ impl NativeResponseCache {
|
|||
pub fn capacity(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
|
||||
Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
|
||||
Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.backend().max_entry_bytes(),
|
||||
Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
|
||||
Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -180,7 +206,11 @@ impl NativeResponseCache {
|
|||
pub fn directory(&self) -> Option<&Path> {
|
||||
match self {
|
||||
Self::Disk(cache) => Some(cache.backend().directory()),
|
||||
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::AzureBlob(_) => None,
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::S3(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -193,6 +223,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.lookup(request, now),
|
||||
Self::Redis { cache, .. } => cache.lookup(request, now),
|
||||
Self::S3(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),
|
||||
}
|
||||
|
|
@ -208,6 +239,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.store(request, response, now),
|
||||
Self::Redis { cache, .. } => cache.store(request, response, now),
|
||||
Self::S3(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),
|
||||
}
|
||||
|
|
@ -222,6 +254,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.lookup_batch(requests, now),
|
||||
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
|
||||
Self::S3(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),
|
||||
}
|
||||
|
|
@ -236,6 +269,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.async_lookup(request, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
|
||||
Self::S3(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,
|
||||
}
|
||||
|
|
@ -258,6 +292,7 @@ impl NativeResponseCache {
|
|||
buffer: Some(buffer),
|
||||
} => buffer.async_store(cache, request, response, now).await,
|
||||
Self::S3(cache) => cache.async_store(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,
|
||||
}
|
||||
|
|
@ -272,6 +307,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
|
||||
Self::S3(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,
|
||||
}
|
||||
|
|
@ -286,6 +322,7 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
|
||||
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
|
||||
Self::S3(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,
|
||||
}
|
||||
|
|
@ -301,6 +338,7 @@ impl NativeResponseCache {
|
|||
cache.async_flush().await
|
||||
}
|
||||
Self::S3(cache) => 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,
|
||||
}
|
||||
|
|
@ -311,8 +349,16 @@ impl NativeResponseCache {
|
|||
Self::Memory(cache) => cache.test_connection().await,
|
||||
Self::Redis { cache, .. } => cache.test_connection().await,
|
||||
Self::S3(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<ResponseCacheCodec>> {
|
||||
match self {
|
||||
Self::Gcs(cache) => Some(cache.backend()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -118,6 +118,7 @@ class XAIChatConfig(OpenAIGPTConfig):
|
|||
base_openai_params: Final = [
|
||||
"logit_bias",
|
||||
"logprobs",
|
||||
"max_completion_tokens",
|
||||
"max_tokens",
|
||||
"n",
|
||||
"parallel_tool_calls",
|
||||
|
|
|
|||
|
|
@ -169,6 +169,15 @@ class _CacheTestHandle:
|
|||
@staticmethod
|
||||
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
|
||||
@staticmethod
|
||||
def gcs(
|
||||
bucket_name: str,
|
||||
*,
|
||||
gcs_path: str | None = None,
|
||||
path_service_account: str | None = None,
|
||||
endpoint: str | None = None,
|
||||
token: str | None = None,
|
||||
) -> _CacheTestHandle: ...
|
||||
@staticmethod
|
||||
def s3(
|
||||
bucket: str,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -90,6 +90,16 @@ class TestXAIReasoningTokenFolding:
|
|||
assert response.usage.total_tokens == 999
|
||||
|
||||
|
||||
def test_max_completion_tokens_is_accepted_and_mapped_to_max_tokens() -> None:
|
||||
optional_params = litellm.get_optional_params(
|
||||
model="grok-4.20",
|
||||
custom_llm_provider="xai",
|
||||
max_completion_tokens=64,
|
||||
)
|
||||
assert optional_params["max_tokens"] == 64, optional_params
|
||||
assert "max_completion_tokens" not in optional_params, optional_params
|
||||
|
||||
|
||||
class TestXAIParallelToolCalls:
|
||||
"""Test suite for XAI parallel tool calls functionality."""
|
||||
|
||||
|
|
|
|||
152
tests/test_litellm_rust/support/fake_gcs.py
Normal file
152
tests/test_litellm_rust/support/fake_gcs.py
Normal file
|
|
@ -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))
|
||||
|
|
@ -27,11 +27,13 @@ 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.disk_cache import DiskCache
|
||||
from litellm.caching.gcs_cache import GCSCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.caching.s3_cache import S3Cache
|
||||
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
|
||||
from tests.test_litellm_rust.support.s3_stub import S3Stub
|
||||
|
||||
|
|
@ -40,6 +42,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]:
|
||||
|
|
@ -59,6 +62,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")
|
||||
|
|
@ -822,6 +834,223 @@ def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -
|
|||
handle._bind_facade(proxied)
|
||||
|
||||
|
||||
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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue