mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge origin/main into litellm_rust_qdrant_semantic_cache
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
3122d21ea6
19 changed files with 1289 additions and 21 deletions
16
litellm-rust/Cargo.lock
generated
16
litellm-rust/Cargo.lock
generated
|
|
@ -2654,6 +2654,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"
|
||||
|
|
@ -2878,6 +2893,7 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-azure-blob",
|
||||
"litellm-cache-disk",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-qdrant-semantic",
|
||||
"litellm-cache-redis",
|
||||
|
|
|
|||
|
|
@ -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-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
|
||||
|
|
@ -73,6 +74,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,7 +6,7 @@ pub enum Error {
|
|||
InvalidEntry,
|
||||
#[error("flushing Redis requires an explicit namespace")]
|
||||
UnscopedFlush,
|
||||
#[error("cache backend does not support this operation")]
|
||||
#[error("operation is not supported by this cache")]
|
||||
UnsupportedOperation,
|
||||
#[error("semantic cache requires request messages")]
|
||||
MissingPrompt,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
litellm-cache-qdrant-semantic.workspace = true
|
||||
|
|
|
|||
|
|
@ -80,6 +80,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,
|
||||
|
|
@ -119,6 +126,7 @@ impl QdrantSemanticCacheConfig {
|
|||
pub(super) enum CacheBackendConfig {
|
||||
Memory(MemoryCacheConfig),
|
||||
Redis(Box<RedisCacheConfig>),
|
||||
Gcs(GcsCacheConfig),
|
||||
Disk(DiskCacheConfig),
|
||||
QdrantSemantic(Box<QdrantSemanticCacheConfig>),
|
||||
AzureBlob(AzureBlobCacheConfig),
|
||||
|
|
@ -136,6 +144,7 @@ pub(super) enum UnsupportedCacheConfig {
|
|||
RedisCredentials,
|
||||
RedisConnection,
|
||||
RedisOption,
|
||||
GcsBucket,
|
||||
DiskStore,
|
||||
QdrantEndpoint,
|
||||
SemanticEmbedding,
|
||||
|
|
@ -149,6 +158,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",
|
||||
Self::QdrantEndpoint => {
|
||||
"native Qdrant requires the default REST port so the gRPC port can be derived"
|
||||
|
|
@ -196,6 +206,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,
|
||||
|
|
@ -216,20 +233,41 @@ impl NativeCacheConfig {
|
|||
backend: CacheBackendConfig::AzureBlob(backend),
|
||||
}))
|
||||
}),
|
||||
Some(
|
||||
CacheType::RedisSemantic
|
||||
| CacheType::ValkeySemantic
|
||||
| CacheType::S3
|
||||
| CacheType::Gcs,
|
||||
)
|
||||
| None => Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
)),
|
||||
Some(CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3) | None => {
|
||||
Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
|
||||
match &self.backend {
|
||||
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::Memory(_) if service.kind() != "memory" => {
|
||||
Some("facade and native backend types must match")
|
||||
}
|
||||
|
|
@ -421,6 +459,23 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConf
|
|||
})
|
||||
}
|
||||
|
||||
#[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_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
|
||||
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
|
||||
|
|
@ -863,7 +918,7 @@ mod tests {
|
|||
|
||||
use super::{
|
||||
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
|
||||
DiskCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
|
||||
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
|
||||
};
|
||||
use crate::cache::native::NativeResponseCache;
|
||||
|
||||
|
|
@ -1067,6 +1122,56 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router
|
|||
});
|
||||
}
|
||||
|
||||
#[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()),
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[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_builtin_disk_configuration_and_rejects_custom_stores() {
|
||||
Python::initialize();
|
||||
|
|
|
|||
|
|
@ -355,6 +355,7 @@ impl FacadeGuard {
|
|||
"QdrantSemanticCache",
|
||||
"qdrant-semantic",
|
||||
),
|
||||
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
|
||||
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
|
||||
("azure-blob", _) => (
|
||||
"litellm.caching.azure_blob_cache",
|
||||
|
|
@ -381,13 +382,16 @@ impl FacadeGuard {
|
|||
return Err(PyTypeError::new_err(message));
|
||||
}
|
||||
let backend_config_names = match kind {
|
||||
"memory" | "redis" | "azure-blob" | "disk" => &[
|
||||
"memory" | "redis" | "azure-blob" | "disk" | "gcs" => &[
|
||||
"namespace",
|
||||
"default_ttl",
|
||||
"max_size_in_memory",
|
||||
"max_size_per_item",
|
||||
"redis_kwargs",
|
||||
"redis_flush_size",
|
||||
"bucket_name",
|
||||
"key_prefix",
|
||||
"path_service_account",
|
||||
][..],
|
||||
"qdrant_semantic" => &[
|
||||
"qdrant_api_base",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use std::env;
|
||||
|
||||
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_http::ClientVariant;
|
||||
|
||||
|
|
@ -74,6 +75,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> {
|
||||
|
|
|
|||
|
|
@ -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()),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ use std::{path::Path, sync::Arc, time::Duration};
|
|||
use litellm_cache::{CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
|
||||
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_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache};
|
||||
use litellm_cache_redis::{RedisCache, RedisTopology};
|
||||
|
|
@ -20,6 +21,7 @@ pub(super) enum NativeResponseCache {
|
|||
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
|
||||
buffer: Option<Arc<WriteBuffer>>,
|
||||
},
|
||||
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
|
||||
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
|
||||
QdrantSemantic(Arc<ResponseCache<QdrantSemanticCache<OpenAiEmbedder, ResponseCacheCodec>>>),
|
||||
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
|
||||
|
|
@ -59,6 +61,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 qdrant_semantic(
|
||||
config: QdrantSemanticCacheConfig,
|
||||
client: reqwest::Client,
|
||||
|
|
@ -103,7 +117,11 @@ impl NativeResponseCache {
|
|||
cache.backend().account_url(),
|
||||
cache.backend().container_name(),
|
||||
)),
|
||||
Self::Memory(_) | Self::Redis { .. } | Self::Disk(_) | Self::QdrantSemantic(_) => None,
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::Gcs(_)
|
||||
| Self::Disk(_)
|
||||
| Self::QdrantSemantic(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -113,6 +131,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(_) => "memory",
|
||||
Self::Redis { .. } => "redis",
|
||||
Self::Gcs(_) => "gcs",
|
||||
Self::Disk(_) => "disk",
|
||||
Self::QdrantSemantic(_) => "qdrant_semantic",
|
||||
Self::AzureBlob(_) => "azure-blob",
|
||||
|
|
@ -123,6 +142,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::QdrantSemantic(_) => None,
|
||||
Self::AzureBlob(cache) => cache.default_ttl(),
|
||||
|
|
@ -131,14 +151,22 @@ impl NativeResponseCache {
|
|||
|
||||
pub fn namespace(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Memory(_) | Self::Disk(_) | Self::QdrantSemantic(_) | Self::AzureBlob(_) => None,
|
||||
Self::Memory(_)
|
||||
| Self::Disk(_)
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
Self::Redis { cache, .. } => cache.backend().namespace(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn topology(&self) -> Option<&RedisTopology> {
|
||||
match self {
|
||||
Self::Memory(_) | Self::Disk(_) | Self::QdrantSemantic(_) | Self::AzureBlob(_) => None,
|
||||
Self::Memory(_)
|
||||
| Self::Disk(_)
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
Self::Redis { cache, .. } => Some(cache.backend().topology()),
|
||||
}
|
||||
}
|
||||
|
|
@ -146,7 +174,7 @@ impl NativeResponseCache {
|
|||
pub fn capacity(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
|
||||
Self::Redis { .. } | Self::Disk(_) => None,
|
||||
Self::Redis { .. } | Self::Gcs(_) | Self::Disk(_) => None,
|
||||
Self::QdrantSemantic(_) => None,
|
||||
Self::AzureBlob(_) => None,
|
||||
}
|
||||
|
|
@ -155,7 +183,7 @@ impl NativeResponseCache {
|
|||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::Memory(cache) => cache.backend().max_entry_bytes(),
|
||||
Self::Redis { .. } | Self::Disk(_) => None,
|
||||
Self::Redis { .. } | Self::Gcs(_) | Self::Disk(_) => None,
|
||||
Self::QdrantSemantic(_) => None,
|
||||
Self::AzureBlob(_) => None,
|
||||
}
|
||||
|
|
@ -174,9 +202,11 @@ impl NativeResponseCache {
|
|||
pub fn directory(&self) -> Option<&Path> {
|
||||
match self {
|
||||
Self::Disk(cache) => Some(cache.backend().directory()),
|
||||
Self::Memory(_) | Self::Redis { .. } | Self::QdrantSemantic(_) | Self::AzureBlob(_) => {
|
||||
None
|
||||
}
|
||||
Self::Memory(_)
|
||||
| Self::Redis { .. }
|
||||
| Self::Gcs(_)
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::AzureBlob(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -216,6 +246,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.lookup(&exact(request), now),
|
||||
Self::Redis { cache, .. } => cache.lookup(&exact(request), now),
|
||||
Self::Gcs(cache) => cache.lookup(&exact(request), now),
|
||||
Self::Disk(cache) => cache.lookup(&exact(request), now),
|
||||
Self::QdrantSemantic(cache) => cache.lookup(request, now),
|
||||
Self::AzureBlob(cache) => cache.lookup(&exact(request), now),
|
||||
|
|
@ -231,6 +262,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.store(&exact(request), response, now),
|
||||
Self::Redis { cache, .. } => cache.store(&exact(request), response, now),
|
||||
Self::Gcs(cache) => cache.store(&exact(request), response, now),
|
||||
Self::Disk(cache) => cache.store(&exact(request), response, now),
|
||||
Self::QdrantSemantic(cache) => cache.store(request, response, now),
|
||||
Self::AzureBlob(cache) => cache.store(&exact(request), response, now),
|
||||
|
|
@ -249,6 +281,9 @@ impl NativeResponseCache {
|
|||
Self::Redis { cache, .. } => {
|
||||
cache.lookup_batch(&requests.iter().map(exact).collect::<Vec<_>>(), now)
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
cache.lookup_batch(&requests.iter().map(exact).collect::<Vec<_>>(), now)
|
||||
}
|
||||
Self::Disk(cache) => {
|
||||
cache.lookup_batch(&requests.iter().map(exact).collect::<Vec<_>>(), now)
|
||||
}
|
||||
|
|
@ -267,6 +302,7 @@ impl NativeResponseCache {
|
|||
match self {
|
||||
Self::Memory(cache) => cache.async_lookup(&exact(request), now).await,
|
||||
Self::Redis { cache, .. } => cache.async_lookup(&exact(request), now).await,
|
||||
Self::Gcs(cache) => cache.async_lookup(&exact(request), now).await,
|
||||
Self::Disk(cache) => cache.async_lookup(&exact(request), now).await,
|
||||
Self::QdrantSemantic(cache) => cache.async_lookup(request, now).await,
|
||||
Self::AzureBlob(cache) => cache.async_lookup(&exact(request), now).await,
|
||||
|
|
@ -292,6 +328,7 @@ impl NativeResponseCache {
|
|||
let request = exact(request);
|
||||
buffer.async_store(cache, &request, response, now).await
|
||||
}
|
||||
Self::Gcs(cache) => cache.async_store(&exact(request), response, now).await,
|
||||
Self::Disk(cache) => cache.async_store(&exact(request), response, now).await,
|
||||
Self::QdrantSemantic(cache) => cache.async_store(request, response, now).await,
|
||||
Self::AzureBlob(cache) => cache.async_store(&exact(request), response, now).await,
|
||||
|
|
@ -312,6 +349,10 @@ impl NativeResponseCache {
|
|||
let requests = requests.iter().map(exact).collect::<Vec<_>>();
|
||||
cache.async_lookup_batch(&requests, now).await
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
let requests = requests.iter().map(exact).collect::<Vec<_>>();
|
||||
cache.async_lookup_batch(&requests, now).await
|
||||
}
|
||||
Self::Disk(cache) => {
|
||||
let requests = requests.iter().map(exact).collect::<Vec<_>>();
|
||||
cache.async_lookup_batch(&requests, now).await
|
||||
|
|
@ -352,6 +393,17 @@ impl NativeResponseCache {
|
|||
)
|
||||
.await
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
cache
|
||||
.async_store_batch(
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(request, value)| (exact(&request), value))
|
||||
.collect(),
|
||||
now,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Self::Disk(cache) => {
|
||||
cache
|
||||
.async_store_batch(
|
||||
|
|
@ -387,6 +439,7 @@ impl NativeResponseCache {
|
|||
}
|
||||
cache.async_flush().await
|
||||
}
|
||||
Self::Gcs(cache) => cache.async_flush().await,
|
||||
Self::Disk(cache) => cache.async_flush().await,
|
||||
Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation),
|
||||
Self::AzureBlob(cache) => cache.async_flush().await,
|
||||
|
|
@ -397,9 +450,17 @@ 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::QdrantSemantic(_) => Err(Error::UnsupportedOperation),
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
@ -26,11 +26,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
|
||||
|
|
@ -38,6 +40,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]:
|
||||
|
|
@ -122,6 +125,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")
|
||||
|
|
@ -699,6 +711,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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue