From 3106d9c573c6f1898280581e499aade3db8ac13c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:41:04 +0000 Subject: [PATCH] feat(rust-bridge): add cache and secret migration foundations (#42328) * docs(rust): plan Python interop foundation * fix(rust): preserve Python settings coercion at the native boundary * chore(rust): drop interop planning note Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(rust): resolve OCR provider secrets through an async SecretSource before transformation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(rust): project the Python secret manager into the bridge and resolve OCR secrets through it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): drop premium_user from the secret manager snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rust-bridge): read the private key management globals once in the settings snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rust): bound the bridge secret manager state cache to the active snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust): inline coercion unit tests * fix(rust): preserve Python secret manager bindings * refactor(rust-bridge): let settings projectors own their contract specs Each settings group now declares its SettingSpec rows next to the projector that reads them, and the manifest test derives python_settings.json from those tables instead of a hand-copied duplicate. Field carries (group, name) instead of a dotted path, and coercion gains the dict-item reader plus the Redis Boolean, certificate-requirement, non-empty string, and numeric adapters that the cache configuration projection adopts next. Co-Authored-By: Claude Fable 5.1 * refactor(rust-bridge): capture the secret manager binding in one settings read The secret_manager accessor now carries the live client and settings objects, so the bridge classifies the binding from a single snapshot instead of re-reading litellm globals. The unreachable native arm and the service alias go away, the binding-to-state mapping moves next to the snapshot, and the Python callback precomputes its key_manager name. Co-Authored-By: Claude Fable 5.1 * refactor(rust-bridge): execute typed settings field declarations * refactor(rust-bridge): compare cache backends by identity behind one exact trait cache-response gains an object-safe ExactResponseCache so every exact-match backend sits behind one pointer; WriteBuffer flushes through it. The bridge's NativeResponseCache shrinks from nine variants and fifteen per-backend accessors to an exact service plus the three semantic backends, and facade mismatch detection compares BackendIdentity values instead of matching on each backend type. Request projections move next to NativeRequest. Co-Authored-By: Claude Fable 5.1 * refactor(rust-bridge): drive both Python-embedded semantic caches through one execution Redis-semantic and Valkey-semantic operations now share one SemanticExecution body: await the Python embedder, seed the task-local vector, run the native backend, repeat per batch entry. Valkey drops its with_embedder path in favor of the same seeded embedder, and each backend keeps its own embedding-failure policy. PythonEmbedder exposes one call shape. Redis-semantic thresholds are compared at the backend's f32 width, which un-breaks the redis-stack parity tests that a 0.8 facade threshold failed before this branch. Co-Authored-By: Claude Fable 5.1 * wip * feat(rust-bridge): complete response cache runtime surface * fix(rust-bridge): preserve secret manager callback exceptions * refactor(rust-bridge): unify route cache and secret rollout catalog --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 --- litellm-rust/Cargo.lock | 7 + litellm-rust/crates/auth-aws/src/aws.rs | 20 + litellm-rust/crates/auth-aws/src/constants.rs | 13 + litellm-rust/crates/auth-azure/src/lib.rs | 2 +- litellm-rust/crates/auth-azure/src/resolve.rs | 32 +- litellm-rust/crates/auth-gcp/src/lib.rs | 32 + litellm-rust/crates/auth-types/src/secret.rs | 7 + litellm-rust/crates/cache-response/README.md | 2 + .../crates/cache-response/src/buffer.rs | 8 +- .../crates/cache-response/src/exact.rs | 148 +++ litellm-rust/crates/cache-response/src/lib.rs | 2 + .../crates/core-utils/src/serde_compat.rs | 5 + litellm-rust/crates/core/Cargo.toml | 1 + litellm-rust/crates/core/src/ocr/handler.rs | 7 +- litellm-rust/crates/core/src/ocr/prepare.rs | 20 +- .../crates/core/src/ocr/provider_config.rs | 27 + litellm-rust/crates/core/tests/ocr.rs | 85 +- litellm-rust/crates/llms/Cargo.toml | 1 + .../ocr/analyze_transformation.rs | 4 + .../src/aws_textract/ocr/transformation.rs | 4 + .../ocr/cohere_parse_transformation.rs | 4 + .../document_intelligence/transformation.rs | 13 +- .../llms/src/azure_ai/ocr/transformation.rs | 12 + .../crates/llms/src/base_llm/inference/mod.rs | 1 + .../llms/src/base_llm/inference/secrets.rs | 19 + litellm-rust/crates/llms/src/base_llm/mod.rs | 1 + .../crates/llms/src/base_llm/ocr/error.rs | 2 + .../crates/llms/src/base_llm/ocr/handler.rs | 15 +- .../crates/llms/src/base_llm/ocr/settings.rs | 4 +- .../llms/src/base_llm/ocr/transformation.rs | 13 +- .../llms/src/cohere/ocr/transformation.rs | 4 + .../llms/src/mistral/ocr/transformation.rs | 8 + .../llms/src/reducto/ocr/transformation.rs | 8 + .../vertex_ai/ocr/deepseek_transformation.rs | 4 + .../llms/src/vertex_ai/ocr/transformation.rs | 4 + litellm-rust/crates/python-bridge/Cargo.toml | 6 + litellm-rust/crates/python-bridge/README.md | 5 + .../crates/python-bridge/python_settings.json | 154 --- .../crates/python-bridge/src/cache/binding.rs | 31 +- .../crates/python-bridge/src/cache/config.rs | 209 +--- .../python-bridge/src/cache/embedder.rs | 86 +- .../crates/python-bridge/src/cache/facade.rs | 47 +- .../crates/python-bridge/src/cache/handle.rs | 2 +- .../python-bridge/src/cache/identity.rs | 511 +++++++++ .../crates/python-bridge/src/cache/mod.rs | 2 +- .../crates/python-bridge/src/cache/native.rs | 992 +++++------------- .../crates/python-bridge/src/cache/request.rs | 184 +++- .../python-bridge/src/cache/semantic.rs | 212 ++-- .../python-bridge/src/cache/semantic_step.rs | 249 ----- .../crates/python-bridge/src/coercion.rs | 485 +++++++-- .../python-bridge/src/coercion/tests.rs | 372 ------- litellm-rust/crates/python-bridge/src/http.rs | 187 +++- litellm-rust/crates/python-bridge/src/lib.rs | 7 +- .../python-bridge/src/python_settings.rs | 318 ++---- .../python-bridge/src/routes/ocr/host.rs | 5 + .../python-bridge/src/routes/ocr/mod.rs | 94 +- .../python-bridge/src/secrets/callback.rs | 349 ++++++ .../python-bridge/src/secrets/config.rs | 338 ++++++ .../crates/python-bridge/src/secrets/mod.rs | 3 + .../python-bridge/src/secrets/resolved.rs | 252 +++++ .../crates/secrets-types/src/config.rs | 6 +- litellm-rust/crates/secrets/src/error.rs | 2 + litellm-rust/crates/secrets/src/handler.rs | 20 + litellm-rust/crates/secrets/src/lib.rs | 2 +- litellm-rust/crates/secrets/src/resolver.rs | 1 + litellm/chat_completions/dispatch.py | 6 +- .../bedrock/audio_transcription/__init__.py | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/messages/dispatch.py | 6 +- litellm/ocr/dispatch.py | 6 +- litellm/responses/dispatch.py | 6 +- litellm/rust_bridge/_native.pyi | 6 +- litellm/rust_bridge/catalog.py | 80 +- litellm/rust_bridge/configuration.py | 2 +- litellm/rust_bridge/dispatch.py | 6 +- litellm/rust_bridge/response_cache.py | 180 ++++ litellm/rust_bridge/runtime.py | 12 +- litellm/rust_bridge/settings.py | 59 ++ tests/test_litellm/responses/test_dispatch.py | 15 +- .../rust_bridge/ocr/test_secrets.py | 112 ++ .../test_litellm/rust_bridge/test_catalog.py | 147 ++- .../test_litellm/rust_bridge/test_dispatch.py | 48 +- .../test_litellm/rust_bridge/test_runtime.py | 26 +- .../test_litellm/rust_bridge/test_settings.py | 101 +- tests/test_litellm_rust/ocr/test_requests.py | 6 +- tests/test_litellm_rust/test_cache.py | 115 +- tests/unit/chat_completions/test_dispatch.py | 8 +- tests/unit/messages/test_dispatch.py | 4 +- tests/unit/ocr/test_dispatch.py | 10 +- 89 files changed, 4138 insertions(+), 2503 deletions(-) create mode 100644 litellm-rust/crates/cache-response/src/exact.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/inference/mod.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/inference/secrets.rs create mode 100644 litellm-rust/crates/python-bridge/README.md delete mode 100644 litellm-rust/crates/python-bridge/python_settings.json create mode 100644 litellm-rust/crates/python-bridge/src/cache/identity.rs delete mode 100644 litellm-rust/crates/python-bridge/src/cache/semantic_step.rs delete mode 100644 litellm-rust/crates/python-bridge/src/coercion/tests.rs create mode 100644 litellm-rust/crates/python-bridge/src/secrets/callback.rs create mode 100644 litellm-rust/crates/python-bridge/src/secrets/config.rs create mode 100644 litellm-rust/crates/python-bridge/src/secrets/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/secrets/resolved.rs create mode 100644 litellm/rust_bridge/response_cache.py create mode 100644 tests/test_litellm/rust_bridge/ocr/test_secrets.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 464edb3b104..03e0dabbc17 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2894,6 +2894,7 @@ dependencies = [ "litellm-host", "litellm-http", "litellm-llms", + "litellm-secrets", "litellm-types", "mime_guess", "moka", @@ -3006,6 +3007,7 @@ dependencies = [ "litellm-framing", "litellm-host", "litellm-http", + "litellm-secrets", "litellm-types", "reqwest 0.12.28", "rstest", @@ -3024,6 +3026,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "aws-sdk-secretsmanager", "bytes", "criterion", "futures-util", @@ -3047,6 +3050,9 @@ dependencies = [ "litellm-host-python", "litellm-http", "litellm-llms", + "litellm-secrets", + "litellm-secrets-aws", + "litellm-secrets-types", "litellm-token-counter", "litellm-types", "pyo3", @@ -3062,6 +3068,7 @@ dependencies = [ "tokio", "tokio-tungstenite", "url", + "wiremock", ] [[package]] diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index bbcb0f016c8..cb9195ffeb6 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -621,6 +621,26 @@ mod tests { None } + #[test] + fn secret_names_cover_environment_reads() { + let seen = std::sync::Arc::new(std::sync::Mutex::new( + std::collections::BTreeSet::::new(), + )); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + resolve_aws_region(None, &Map::new(), &env); + aws_auth_config(&Map::new(), &env); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| crate::constants::SECRET_NAMES.contains(&name.as_str())) + ); + } + #[test] fn a_region_comes_from_the_call_then_the_model_then_the_environment() { let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]); diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 9e7c6bfab43..26df4f2a350 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -14,6 +14,19 @@ pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; +pub const SECRET_NAMES: &[&str] = &[ + AWS_ACCESS_KEY_ID, + AWS_SECRET_ACCESS_KEY, + AWS_SESSION_TOKEN, + AWS_REGION_NAME, + AWS_REGION, + AWS_SESSION_NAME, + AWS_PROFILE_NAME, + AWS_ROLE_NAME, + AWS_WEB_IDENTITY_TOKEN, + AWS_STS_ENDPOINT, + AWS_EXTERNAL_ID, +]; /// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors /// Python's `_filter_headers_for_aws_signature` allowlist. diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index 5c7c654b69d..9ed505e4160 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -3,5 +3,5 @@ mod native; mod resolve; mod types; -pub use resolve::AzureAuthService; +pub use resolve::{AzureAuthService, SECRET_NAMES}; pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 4d564b6e68a..9a7afe645db 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -19,6 +19,17 @@ const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; +pub const SECRET_NAMES: &[&str] = &[ + AZURE_AD_TOKEN_ENV, + AZURE_TENANT_ID_ENV, + AZURE_CLIENT_ID_ENV, + AZURE_CLIENT_SECRET_ENV, + AZURE_SCOPE_ENV, + AZURE_AUTHORITY_HOST_ENV, + AZURE_CREDENTIAL_ENV, + AZURE_FEDERATED_TOKEN_FILE_ENV, +]; + #[derive(Clone, Debug)] pub(crate) enum AzureCredentialPlan { Supplied(Sourced), @@ -440,13 +451,14 @@ fn non_empty_reference(value: &str, kind: &str) -> Result { #[cfg(test)] mod tests { + use std::collections::BTreeSet; use std::future::Future; use std::sync::{Arc, Mutex}; use serde_json::json; use super::{ - AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference, + AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, SECRET_NAMES, oidc_reference, resolve_reference, select_auth_plan, }; use crate::native::ValidatedAzureRequest; @@ -517,6 +529,24 @@ mod tests { assert!(matches!(plan, AzureCredentialPlan::Native(_))); } + #[test] + fn secret_names_cover_environment_reads() { + let seen = std::sync::Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let inputs = AzureAuthInputs::default(); + select_auth_plan(&inputs, &|name| { + recorded.lock().unwrap().insert(name.to_string()); + None + }) + .unwrap(); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| SECRET_NAMES.contains(&name.as_str())) + ); + } + #[test] fn supplied_token_does_not_require_refresh() { let params = json!({"azure_ad_token": "token"}); diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 534d85acdb0..682f1af5fe1 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -23,6 +23,16 @@ const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; +pub const SECRET_NAMES: &[&str] = &[ + VERTEX_AI_API_KEY_ENV, + VERTEXAI_API_KEY_ENV, + VERTEXAI_CREDENTIALS_ENV, + GOOGLE_APPLICATION_CREDENTIALS_ENV, + VERTEXAI_PROJECT_ENV, + VERTEXAI_LOCATION_ENV, + VERTEX_LOCATION_ENV, +]; + #[derive(Clone, Debug, Default)] pub struct VertexConfig { credentials: Option>, @@ -406,6 +416,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error { #[cfg(test)] mod tests { + use std::collections::BTreeSet; use std::sync::atomic::{AtomicUsize, Ordering}; use serde_json::json; @@ -476,6 +487,27 @@ mod tests { ); } + #[tokio::test] + async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| SECRET_NAMES.contains(&name.as_str())) + ); + } + #[test] fn empty_primary_values_fall_back_to_python_aliases() { let config = config(json!({ diff --git a/litellm-rust/crates/auth-types/src/secret.rs b/litellm-rust/crates/auth-types/src/secret.rs index a07fe3eaad9..7e6789deef7 100644 --- a/litellm-rust/crates/auth-types/src/secret.rs +++ b/litellm-rust/crates/auth-types/src/secret.rs @@ -1,4 +1,5 @@ use serde::Deserialize; +use std::hash::{Hash, Hasher}; use veil::Redact; #[derive(Redact, Clone, Deserialize)] @@ -23,6 +24,12 @@ impl PartialEq for SecretValue { impl Eq for SecretValue {} +impl Hash for SecretValue { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } +} + #[cfg(test)] mod tests { use super::SecretValue; diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 46e561ddad1..d048afb69f8 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -36,6 +36,8 @@ Callers supply Unix time for response freshness. Backend TTL uses its own clock. The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API +The bridge also exposes a production-shaped response cache runtime selected through the Rust catalog. Its shipped rule set is empty, so current SDK, Router, and proxy calls stay on Python and do not construct native cache resources. Tests can inject a rule and build the native memory runtime from an ordinary Python `Cache` configuration without changing the legacy cache classes + Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 606c21410c7..68f2348e28b 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -1,9 +1,9 @@ use std::{sync::Mutex, time::Duration}; -use litellm_cache::{BaseCache, Error, ExactCacheContext}; +use litellm_cache::Error; use serde_json::Value; -use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; +use crate::{ExactResponseCache, ResponseCacheRequest}; pub struct WriteBuffer { flush_size: usize, @@ -18,9 +18,9 @@ impl WriteBuffer { } } - pub async fn async_store>( + pub async fn async_store( &self, - cache: &ResponseCache, + cache: &dyn ExactResponseCache, request: &ResponseCacheRequest, response: Value, now: Duration, diff --git a/litellm-rust/crates/cache-response/src/exact.rs b/litellm-rust/crates/cache-response/src/exact.rs new file mode 100644 index 00000000000..f5e86b2598c --- /dev/null +++ b/litellm-rust/crates/cache-response/src/exact.rs @@ -0,0 +1,148 @@ +use std::{future::Future, pin::Pin, time::Duration}; + +use litellm_cache::{ + BaseCache, BatchCache, CacheConnectionResult, Error, ExactCacheContext, FlushCache, +}; +use serde_json::Value; + +use crate::{CacheEntry, PartialHits, ResponseCache, ResponseCacheRequest}; + +type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// Object-safe view of a `ResponseCache` over an exact-match backend, so hosts can hold every +/// exact backend behind one pointer without erasing which backend it is elsewhere. +pub trait ExactResponseCache: Send + Sync { + fn default_ttl(&self) -> Option; + + fn lookup(&self, request: &ResponseCacheRequest, now: Duration) + -> Result, Error>; + + fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error>; + + fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result; + + fn async_lookup<'a>( + &'a self, + request: &'a ResponseCacheRequest, + now: Duration, + ) -> BoxFuture<'a, Result, Error>>; + + fn async_store<'a>( + &'a self, + request: &'a ResponseCacheRequest, + response: Value, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>>; + + fn async_lookup_batch<'a>( + &'a self, + requests: &'a [ResponseCacheRequest], + now: Duration, + ) -> BoxFuture<'a, Result>; + + fn async_store_batch<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>>; + + fn async_store_entries<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, + ) -> BoxFuture<'a, Result<(), Error>>; + + fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>>; + + fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result>; +} + +impl ExactResponseCache for ResponseCache +where + B: BaseCache + BatchCache + FlushCache, +{ + fn default_ttl(&self) -> Option { + ResponseCache::default_ttl(self) + } + + fn lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + ResponseCache::lookup(self, request, now) + } + + fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + ResponseCache::store(self, request, response, now) + } + + fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + ResponseCache::lookup_batch(self, requests, now) + } + + fn async_lookup<'a>( + &'a self, + request: &'a ResponseCacheRequest, + now: Duration, + ) -> BoxFuture<'a, Result, Error>> { + Box::pin(ResponseCache::async_lookup(self, request, now)) + } + + fn async_store<'a>( + &'a self, + request: &'a ResponseCacheRequest, + response: Value, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_store(self, request, response, now)) + } + + fn async_lookup_batch<'a>( + &'a self, + requests: &'a [ResponseCacheRequest], + now: Duration, + ) -> BoxFuture<'a, Result> { + Box::pin(ResponseCache::async_lookup_batch(self, requests, now)) + } + + fn async_store_batch<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_store_batch(self, entries, now)) + } + + fn async_store_entries<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_store_entries(self, entries)) + } + + fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_flush(self)) + } + + fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result> { + Box::pin(ResponseCache::test_connection(self)) + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 91b36ebe24b..ab9867ac8db 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -2,6 +2,7 @@ mod buffer; mod caching; mod codec; mod embedding; +mod exact; mod response; pub use buffer::WriteBuffer; @@ -11,4 +12,5 @@ pub use caching::{ }; pub use codec::ResponseCacheCodec; pub use embedding::PartialHits; +pub use exact::ExactResponseCache; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/core-utils/src/serde_compat.rs b/litellm-rust/crates/core-utils/src/serde_compat.rs index c767c709f50..fddab1d80e3 100644 --- a/litellm-rust/crates/core-utils/src/serde_compat.rs +++ b/litellm-rust/crates/core-utils/src/serde_compat.rs @@ -17,6 +17,11 @@ pub fn parse_str_bool(value: &str) -> Option { token.eq_ignore_ascii_case("false").then_some(false) } +/// `redis-py` string Booleans: only `1`, `true`, and `yes` (case-insensitive) are true. +pub fn parse_redis_bool(value: &str) -> bool { + value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes") +} + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { deserializer.deserialize_any(Self) diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 69ae8004d46..3bfc5bae925 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -36,6 +36,7 @@ url.workspace = true veil.workspace = true [dev-dependencies] +litellm-secrets.workspace = true litellm-auth-gcp.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 19037e49033..c1265e1e91c 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -22,7 +22,12 @@ pub(crate) async fn perform_ocr_request( ) -> Result { request.response_format()?; let config = request.config; - let request = prepare_request(request, caller_document, client); + let secrets = client + .secret_source() + .resolve(&config.secret_names()) + .await + .map_err(|error| Error::Secret(std::sync::Arc::new(error)))?; + let request = prepare_request(request, caller_document, client, secrets); let hooks = OcrCallHooks::new(host.clone(), &request, config); config.ocr(client, &request, &hooks).await } diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 715aedc69df..54960256faa 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,7 +1,10 @@ use litellm_auth::{InputSource, SecretValue, Sourced}; -use litellm_llms::base_llm::ocr::{ - handler::OcrClient, - transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest}, +use litellm_llms::base_llm::{ + inference::secrets::Secrets, + ocr::{ + handler::OcrClient, + transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest}, + }, }; use super::provider_config::OcrProvider; @@ -11,6 +14,7 @@ pub(crate) fn prepare_request( request: ResolvedOcrRequest, caller_document: bool, client: &OcrClient, + secrets: Secrets, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); let (preferred_api_key_env, api_base_env) = match request.config.provider() { @@ -24,7 +28,7 @@ pub(crate) fn prepare_request( | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None), }; - let secret = |name: &str| client.secrets().truthy(name); + let secret = |name: &str| secrets.truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { credentials.api_key.clone().or_else(|| { preferred_api_key_env @@ -60,12 +64,7 @@ pub(crate) fn prepare_request( PreparedOcrRequest { model, document, - connection: OcrConnection::new( - resolved, - transport, - client.settings().clone(), - client.secrets().clone(), - ), + connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets), caller_document, optional_params, input_sources, @@ -79,6 +78,7 @@ pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedO request, true, &OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()), + std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), ) } diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index d38d87b92cc..0b09e9be354 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -114,6 +114,10 @@ impl OcrConfigKind { with_config!(self, config => config.get_api_key_env_var()) } + pub(crate) fn secret_names(self) -> Vec<&'static str> { + with_config!(self, config => config.secret_names()) + } + pub(crate) fn get_health_check_document(self) -> OcrDocument { with_config!(self, config => config.get_health_check_document()) } @@ -213,6 +217,8 @@ fn is_document_intelligence_model(model: &str) -> bool { #[cfg(test)] mod tests { + use std::collections::HashSet; + use litellm_auth::{InputSource, Sourced}; use litellm_llms::{ base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document, @@ -221,6 +227,27 @@ mod tests { use super::*; + #[rstest] + #[case(OcrConfigKind::AwsTextract)] + #[case(OcrConfigKind::AwsTextractAnalyze)] + #[case(OcrConfigKind::Cohere)] + #[case(OcrConfigKind::Mistral)] + #[case(OcrConfigKind::AzureAi)] + #[case(OcrConfigKind::AzureCohere)] + #[case(OcrConfigKind::AzureDocumentIntelligence)] + #[case(OcrConfigKind::ReductoLegacy)] + #[case(OcrConfigKind::ReductoV3)] + #[case(OcrConfigKind::VertexAi)] + #[case(OcrConfigKind::VertexDeepSeek)] + fn secret_names_include_api_keys_without_duplicates(#[case] config: OcrConfigKind) { + let names = config.secret_names(); + let unique = names.iter().collect::>(); + assert_eq!(names.len(), unique.len()); + if let Some(api_key) = config.get_api_key_env_var() { + assert!(names.contains(&api_key)); + } + } + #[rstest] #[case("cohere")] #[case("mistral")] diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 3aedc7b9023..d376f0df784 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; use litellm_host::{ event::{CallEvent, MachineEvent, WireRequest}, @@ -10,11 +11,14 @@ use litellm_http::{ HttpClientPool, HttpSettings, Resolution, media::{PublicDnsResolver, UrlPolicy}, }; +use litellm_llms::base_llm::inference::secrets::{SecretSource, Secrets}; use litellm_llms::base_llm::ocr::{ error::Error as OcrError, handler::OcrClient, settings::OcrSettings, - transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, + }, }; use rstest::rstest; use serde_json::{Value, json}; @@ -27,6 +31,32 @@ use super::{ }; use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; +struct RecordingSecretSource { + names: Arc>>, + values: &'static [(&'static str, &'static str)], + api_base: String, +} + +impl SecretSource for RecordingSecretSource { + fn resolve<'a>( + &'a self, + names: &'a [&'static str], + ) -> BoxFuture<'a, Result> { + *self.names.lock().unwrap() = names.to_vec(); + let values = self.values; + let api_base = self.api_base.clone(); + Box::pin(async move { + Ok(Arc::new(move |name: &str| match name { + "MISTRAL_AZURE_API_BASE" => Some(api_base.clone()), + _ => values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()), + }) as Secrets) + }) + } +} + #[rstest] #[case::mistral("mistral/model", json!({}))] #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] @@ -184,14 +214,11 @@ async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( #[case] expected_key: &str, ) { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let secret_base = base.clone(); - let client = ocr_client().with_secrets(Arc::new(move |name: &str| match name { - "MISTRAL_AZURE_API_BASE" => Some(secret_base.clone()), - "MISTRAL_API_BASE" => Some("http://127.0.0.1:9/never-read".into()), - _ => secrets - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: secrets, + api_base: base.clone(), })); let request = decode_request(OcrWireRequest { model: "mistral/model".into(), @@ -208,9 +235,47 @@ async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( crate::ocr::client::perform(&client, request).await.unwrap(); server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); } +#[tokio::test] +async fn mistral_ocr_resolves_provider_secrets_before_transformation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: &[("MISTRAL_API_KEY", "source-key")], + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,YWJj" + }), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); +} + #[tokio::test] async fn ocr_client_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; @@ -224,7 +289,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { UrlPolicy::default(), VertexAuth::default(), OcrSettings::default(), - Arc::new(litellm_core_utils::settings::ProcessEnvironment), + Arc::new(litellm_llms::base_llm::inference::secrets::EnvironmentSecrets), ) .unwrap(); crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index f04b78feee1..ed15d9f7cdb 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -18,6 +18,7 @@ litellm-auth-gcp.workspace = true litellm-host.workspace = true litellm-framing.workspace = true litellm-http.workspace = true +litellm-secrets.workspace = true base64.workspace = true bytes.workspace = true data-url = "0.3.2" diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index d476861e6e1..2ce1b0da51b 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -40,6 +40,10 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { type ProviderRequest = AnalyzeDocumentRequest; type Environment = TextractEnvironment; + fn secret_names(&self) -> Vec<&'static str> { + litellm_auth_aws::constants::SECRET_NAMES.to_vec() + } + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["feature_types"] } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index ad630a1ca4c..6eb195defaa 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -29,6 +29,10 @@ impl BaseOcrConfig for TextractDetectTextConfig { type ProviderRequest = DetectDocumentTextRequest; type Environment = TextractEnvironment; + fn secret_names(&self) -> Vec<&'static str> { + litellm_auth_aws::constants::SECRET_NAMES.to_vec() + } + fn get_health_check_document(&self) -> OcrDocument { health_check_document() } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs index 045d8744bc9..09639481cdf 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs @@ -28,6 +28,10 @@ impl BaseOcrConfig for AzureAICohereParseConfig { super::transformation::AzureAiOcrConfig.get_api_key_env_var() } + fn secret_names(&self) -> Vec<&'static str> { + super::transformation::AzureAiOcrConfig.secret_names() + } + fn get_health_check_document(&self) -> OcrDocument { CohereParseConfig.get_health_check_document() } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 9b27fdbb568..bfe0d76aab1 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -2,7 +2,7 @@ use std::{collections::BTreeSet, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_auth::{InputSource, Sourced}; -use litellm_auth_azure::AzureAuthInputs; +use litellm_auth_azure::{AzureAuthInputs, SECRET_NAMES as AZURE_AUTH_SECRET_NAMES}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, @@ -141,6 +141,17 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { Some(AZURE_DI_API_KEY_ENV) } + fn secret_names(&self) -> Vec<&'static str> { + [ + [AZURE_DI_API_KEY_ENV, AZURE_DI_ENDPOINT_ENV].as_slice(), + AZURE_AUTH_SECRET_NAMES, + ] + .into_iter() + .flatten() + .copied() + .collect() + } + fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { ResolvedOcrCredentials { api_key: inputs.api_key.and_then(|key| { diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 6df83e57eab..1fad860f757 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -1,5 +1,6 @@ use litellm_auth::{InputSource, Sourced}; use litellm_auth_azure::AzureAuthInputs; +use litellm_auth_azure::SECRET_NAMES as AZURE_AUTH_SECRET_NAMES; use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl}; use serde_json::Value; @@ -37,6 +38,17 @@ impl BaseOcrConfig for AzureAiOcrConfig { Some(AZURE_AI_API_KEY_ENV) } + fn secret_names(&self) -> Vec<&'static str> { + [ + [AZURE_AI_API_KEY_ENV, AZURE_AI_API_BASE_ENV].as_slice(), + AZURE_AUTH_SECRET_NAMES, + ] + .into_iter() + .flatten() + .copied() + .collect() + } + fn map_ocr_params( &self, non_default_params: &CallArguments, diff --git a/litellm-rust/crates/llms/src/base_llm/inference/mod.rs b/litellm-rust/crates/llms/src/base_llm/inference/mod.rs new file mode 100644 index 00000000000..10c0454f947 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/inference/mod.rs @@ -0,0 +1 @@ +pub mod secrets; diff --git a/litellm-rust/crates/llms/src/base_llm/inference/secrets.rs b/litellm-rust/crates/llms/src/base_llm/inference/secrets.rs new file mode 100644 index 00000000000..eb13fe95116 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/inference/secrets.rs @@ -0,0 +1,19 @@ +use std::sync::Arc; + +use futures_util::future::BoxFuture; +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use litellm_secrets::Error; + +pub type Secrets = Arc; + +pub trait SecretSource: Send + Sync { + fn resolve<'a>(&'a self, names: &'a [&'static str]) -> BoxFuture<'a, Result>; +} + +pub struct EnvironmentSecrets; + +impl SecretSource for EnvironmentSecrets { + fn resolve<'a>(&'a self, _names: &'a [&'static str]) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(Arc::new(ProcessEnvironment) as Secrets) }) + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/mod.rs b/litellm-rust/crates/llms/src/base_llm/mod.rs index 8ed37da4573..9cced64b687 100644 --- a/litellm-rust/crates/llms/src/base_llm/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/mod.rs @@ -2,5 +2,6 @@ pub mod anthropic_messages; pub mod audio_transcription; pub mod base_model_iterator; pub mod chat; +pub mod inference; pub mod ocr; pub mod responses; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index e09842e2856..b3df8fc18c8 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -98,6 +98,8 @@ pub enum Error { "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" )] MissingReductoApiKey, + #[error("secret resolution failed: {0}")] + Secret(#[source] std::sync::Arc), #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 245261d9f92..3ec9de8197f 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; @@ -11,9 +13,10 @@ use litellm_http::{ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; +use crate::base_llm::inference::secrets::SecretSource; use crate::base_llm::ocr::{ error::Error, - settings::{OcrSettings, Secrets}, + settings::OcrSettings, transformation::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value, decode_response, @@ -35,7 +38,7 @@ pub struct OcrClient { document_fetcher: MediaFetcher, vertex_auth: VertexAuth, settings: OcrSettings, - secrets: Secrets, + secrets: Arc, } impl OcrClient { @@ -45,7 +48,7 @@ impl OcrClient { url_policy: UrlPolicy, vertex_auth: VertexAuth, settings: OcrSettings, - secrets: Secrets, + secrets: Arc, ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, @@ -77,7 +80,7 @@ impl OcrClient { &self.settings } - pub fn secrets(&self) -> &Secrets { + pub fn secret_source(&self) -> &Arc { &self.secrets } @@ -92,7 +95,7 @@ impl OcrClient { document_fetcher: MediaFetcher::for_test(document_http), vertex_auth: VertexAuth::default(), settings: OcrSettings::default(), - secrets: std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), + secrets: Arc::new(crate::base_llm::inference::secrets::EnvironmentSecrets), } } @@ -102,7 +105,7 @@ impl OcrClient { } #[cfg(any(test, feature = "test-support"))] - pub fn with_secrets(self, secrets: Secrets) -> Self { + pub fn with_secrets(self, secrets: Arc) -> Self { Self { secrets, ..self } } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs index f5954599b43..87461cd36aa 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -1,9 +1,7 @@ -use std::{sync::Arc, time::Duration}; +use std::time::Duration; use litellm_core_utils::settings::Lookup; -pub type Secrets = Arc; - #[derive(Clone, Debug, PartialEq)] pub struct OcrSettings { pub request_timeout: Duration, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e02a4b7f266..0506ff3d6df 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -14,10 +14,13 @@ use serde::{ use serde_json::{Map, Value}; use serde_with::serde_as; -use crate::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, - settings::{OcrSettings, Secrets}, +use crate::base_llm::{ + inference::secrets::Secrets, + ocr::{ + error::Error, + handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, + settings::OcrSettings, + }, }; pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; @@ -436,6 +439,8 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { None } + fn secret_names(&self) -> Vec<&'static str>; + fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { ResolvedOcrCredentials { api_key: inputs diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index d141c68db38..c0bb4c60563 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -102,6 +102,10 @@ impl BaseOcrConfig for CohereParseConfig { Some(COHERE_API_KEY_ENV) } + fn secret_names(&self) -> Vec<&'static str> { + vec![COHERE_API_KEY_ENV] + } + fn get_health_check_document(&self) -> OcrDocument { OcrDocument::ImageUrl { image_url: COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI.into(), diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 2b14372fbec..149e8056789 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -69,6 +69,14 @@ impl BaseOcrConfig for MistralOcrConfig { Some(MISTRAL_OCR_API_KEY_ENV_VAR) } + fn secret_names(&self) -> Vec<&'static str> { + vec![ + MISTRAL_OCR_API_KEY_ENV_VAR, + "MISTRAL_AZURE_API_KEY", + "MISTRAL_AZURE_API_BASE", + ] + } + fn map_ocr_params( &self, non_default_params: &CallArguments, diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 5272be97c24..f00259984ba 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -92,6 +92,10 @@ impl BaseOcrConfig for ReductoParseV3Config { type ProviderRequest = ReductoV3Request; type Environment = Vec<(String, String)>; + fn secret_names(&self) -> Vec<&'static str> { + vec![REDUCTO_API_KEY_ENV] + } + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["formatting", "retrieval", "settings"] } @@ -180,6 +184,10 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { type ProviderRequest = ReductoLegacyRequest; type Environment = Vec<(String, String)>; + fn secret_names(&self) -> Vec<&'static str> { + vec![REDUCTO_API_KEY_ENV] + } + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["enhance"] } diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 9a23deefb89..86231d50f9c 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -105,6 +105,10 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { VertexAiOcrConfig.get_api_key_env_var() } + fn secret_names(&self) -> Vec<&'static str> { + VertexAiOcrConfig.secret_names() + } + fn map_ocr_params( &self, _arguments: &CallArguments, diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 2d505ba4342..c9342c87e9a 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -34,6 +34,10 @@ impl BaseOcrConfig for VertexAiOcrConfig { Some("VERTEX_AI_API_KEY") } + fn secret_names(&self) -> Vec<&'static str> { + litellm_auth_gcp::SECRET_NAMES.to_vec() + } + fn map_ocr_params( &self, non_default_params: &CallArguments, diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6b4476d897c..4e6c510d104 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true +futures-util.workspace = true litellm-cache.workspace = true litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true @@ -41,6 +42,8 @@ litellm-core-utils.workspace = true litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true +litellm-secrets = { workspace = true, features = ["aws"] } +litellm-secrets-types.workspace = true litellm-types.workspace = true litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } @@ -53,6 +56,7 @@ url.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] +litellm-secrets-aws.workspace = true serde.workspace = true serde_with.workspace = true criterion.workspace = true @@ -60,6 +64,8 @@ futures-util.workspace = true rstest.workspace = true sha2.workspace = true tokio-tungstenite.workspace = true +wiremock = "0.6.5" +aws-sdk-secretsmanager = "1.117.0" [[bench]] name = "serialization" diff --git a/litellm-rust/crates/python-bridge/README.md b/litellm-rust/crates/python-bridge/README.md new file mode 100644 index 00000000000..faaca233f5a --- /dev/null +++ b/litellm-rust/crates/python-bridge/README.md @@ -0,0 +1,5 @@ +Native OCR uses `SecretSource` with `EnvironmentSecrets`, preserving process-environment reads. Readable Python secret managers still make OCR decline to the existing Python implementation. `ResolvedSecrets` and the separate `secret_manager_binding()` snapshot are inactive foundations for a later rollout + +Cache and secret-manager catalog entries remain Python-only, including when `LITELLM_RUST=1`. The new cache runtime is not connected to SDK or gateway caching + +OCR provider requests use the shared `litellm-http` pool. AWS and Google secret-manager SDK clients keep their SDK transports, which do not yet inherit the pool's proxy, TLS, certificate, timeout, or observability configuration. Preserve those SDK transports and configure them equivalently instead of forcing them through reqwest diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json deleted file mode 100644 index ea53d1d2025..00000000000 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ /dev/null @@ -1,154 +0,0 @@ -{ - "http_settings": { - "version": 1, - "fields": { - "ssl_verify": { - "adapter": "SslVerifyInput", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [ - "none", - "bool", - "str" - ], - "unsupported_live": "configuration_error" - }, - "ssl_certificate": { - "adapter": "OptionalStrictString", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "ssl_security_level": { - "adapter": "TuningString", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "ssl_ecdh_curve": { - "adapter": "TuningString", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "force_ipv4": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "http2": { - "adapter": "ExactTrue", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "aiohttp_trust_env": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "disable_aiohttp_trust_env": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "disable_aiohttp_transport": { - "adapter": "ExactTrue", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "user_agent": { - "adapter": "StrictString", - "required": true, - "precedence": "accessor", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - }, - "url_policy": { - "version": 1, - "fields": { - "user_url_validation": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "user_url_allowed_hosts": { - "adapter": "HostCollection", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - }, - "provider_defaults": { - "version": 1, - "fields": { - "vertex_project": { - "adapter": "FalsyOptionalString", - "required": true, - "precedence": "module_global", - "sensitive": true, - "shapes": [], - "unsupported_live": null - }, - "vertex_location": { - "adapter": "FalsyOptionalString", - "required": true, - "precedence": "module_global", - "sensitive": true, - "shapes": [], - "unsupported_live": null - }, - "enable_azure_ad_token_refresh": { - "adapter": "ExactTrue", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - }, - "secret_manager": { - "version": 1, - "fields": { - "readable": { - "adapter": "StrictBool", - "required": true, - "precedence": "accessor", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 2ff73238202..273d3f9ca4e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -11,10 +11,12 @@ use serde_json::Value; use super::{ cache_error, callback::PythonCallback, + config::{CacheBackendConfig, CacheConfigProjection, NativeCacheConfig}, future::{ready_none, ready_value}, native::NativeResponseCache, request::{now, request, requests}, }; +use crate::errors::RustBridgeDeclined; pub(super) enum CacheBinding { Disabled, @@ -22,7 +24,7 @@ pub(super) enum CacheBinding { PythonCallback(PythonCallback), } -#[pyclass(frozen, name = "_CacheTestBinding")] +#[pyclass(frozen, name = "_ResponseCacheRuntime")] pub(crate) struct ResolvedCache { binding: CacheBinding, pid: u32, @@ -66,6 +68,33 @@ impl ResolvedCache { #[pymethods] impl ResolvedCache { + #[staticmethod] + fn from_cache(cache: &Bound<'_, PyAny>) -> PyResult { + let config = match NativeCacheConfig::project(cache)? { + CacheConfigProjection::Native(config) => *config, + CacheConfigProjection::Unsupported(reason) => { + return Err(RustBridgeDeclined::new_err(reason.message())); + } + }; + let service = match config.backend { + CacheBackendConfig::Memory(memory) => NativeResponseCache::memory( + memory.capacity, + memory.default_ttl, + memory.max_entry_bytes, + ), + _ => { + return Err(RustBridgeDeclined::new_err( + "native response cache activation is not implemented for this backend", + )); + } + }; + Ok(Self::new(CacheBinding::Native( + service + .with_scope(config.policy.semantic_cache_scope) + .with_redis_flush_size(config.policy.redis_flush_size), + ))) + } + #[getter] fn kind(&self) -> &'static str { match self.binding { diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 766d526cf5f..b6e08102e18 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -11,7 +11,7 @@ use pyo3::{ types::{PyAny, PyBool, PyDict, PyList, PyString}, }; -use super::{native::NativeResponseCache, request::duration}; +use super::{identity::BackendIdentity, native::NativeResponseCache, request::duration}; #[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct CachePolicy { @@ -293,161 +293,58 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - let default_ttl = match &self.backend { - CacheBackendConfig::Memory(config) => Some(config.default_ttl), - CacheBackendConfig::Redis(config) => Some(config.default_ttl), - CacheBackendConfig::S3(_) => None, - CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO), - CacheBackendConfig::Disk(_) - | CacheBackendConfig::AzureBlob(_) - | CacheBackendConfig::Gcs(_) - | CacheBackendConfig::RedisSemantic(_) - | CacheBackendConfig::QdrantSemantic(_) => None, - }; - if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) - && service.default_ttl() != default_ttl - { - return Some("facade and native backend default TTLs must match"); - } - match &self.backend { - CacheBackendConfig::Memory(config) if service.kind() != "memory" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => { - Some("facade and native backend capacities must match") - } - CacheBackendConfig::Memory(config) - if service.max_entry_bytes() != Some(config.max_entry_bytes) => - { - Some("facade and native backend item limits must match") - } - CacheBackendConfig::Memory(_) => None, - CacheBackendConfig::Redis(_) if service.kind() != "redis" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => { - Some("facade and native backend topologies must match") - } - CacheBackendConfig::Redis(config) => (service.namespace() - != config.namespace.as_deref()) - .then_some("facade and native backend namespaces must match"), - CacheBackendConfig::S3(_) if service.kind() != "s3" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => { - Some("facade and native backend buckets must match") - } - CacheBackendConfig::S3(config) - if service.key_prefix() != Some(config.key_prefix.as_str()) => - { - Some("facade and native backend key prefixes must match") - } - CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => { - Some("facade and native backend regions must match") - } - CacheBackendConfig::S3(config) - if service.endpoint() - != config - .endpoint - .as_ref() - .map(|endpoint| endpoint.url.as_str()) => - { - 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::ValkeySemantic(config) => { - if service.kind() != "valkey-semantic" { - return Some("facade and native backend types must match"); - } - let Some((threshold, index_name)) = service.semantic_config() else { - return Some("facade and native backend types must match"); - }; - (threshold != config.similarity_threshold || index_name != config.index_name) - .then_some("facade and native semantic settings must match") - } - CacheBackendConfig::Disk(_) if service.kind() != "disk" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Disk(config) => { - let Some(directory) = service.directory() else { - return Some("facade and native backend types must match"); - }; - let native = std::fs::canonicalize(directory).ok(); - let facade = std::fs::canonicalize(&config.directory).ok(); - (native != facade).then_some("facade and native backend directories must match") - } - CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::RedisSemantic(config) - if service.index_name() != Some(config.index_name.as_str()) => - { - Some("facade and native backend index names must match") - } - CacheBackendConfig::RedisSemantic(config) - if service.similarity_threshold() - != Some(f64::from(config.similarity_threshold as f32)) => - { - Some("facade and native backend similarity thresholds must match") - } - CacheBackendConfig::RedisSemantic(_) => None, - CacheBackendConfig::QdrantSemantic(config) if service.kind() != "qdrant_semantic" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::QdrantSemantic(config) - if service.collection_name() != Some(config.collection_name.as_str()) => - { - Some("facade and native backend collections must match") - } - CacheBackendConfig::QdrantSemantic(config) - if service.similarity_threshold() != Some(config.similarity_threshold) => - { - Some("facade and native backend similarity thresholds must match") - } - CacheBackendConfig::QdrantSemantic(config) - if service.vector_size() != Some(config.vector_size) => - { - Some("facade and native backend vector sizes must match") - } - CacheBackendConfig::QdrantSemantic(config) - if service.embedding_model() != Some(config.embedding.model.as_str()) => - { - Some("facade and native backend embedding models must match") - } - CacheBackendConfig::QdrantSemantic(_) => None, - CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { - None => Some("facade and native backend types must match"), - Some((account_url, container)) - if account_url != config.account_url || container != config.container => - { - Some("facade and native backend containers must match") - } - Some(_) => None, + self.backend.identity().mismatch(&service.identity()) + } +} + +impl CacheBackendConfig { + /// The identity a native backend must have for this facade configuration to describe it. + pub(super) fn identity(&self) -> BackendIdentity { + match self { + Self::Memory(config) => BackendIdentity::Memory { + capacity: config.capacity, + max_entry_bytes: Some(config.max_entry_bytes), + default_ttl: Some(config.default_ttl), + }, + Self::Redis(config) => BackendIdentity::Redis { + topology: config.topology.clone(), + namespace: config.namespace.clone(), + default_ttl: Some(config.default_ttl), + }, + Self::S3(config) => BackendIdentity::S3 { + bucket: config.bucket.clone(), + key_prefix: config.key_prefix.clone(), + region: config.region.clone(), + endpoint: config + .endpoint + .as_ref() + .map(|endpoint| endpoint.url.clone()), + }, + Self::Gcs(config) => BackendIdentity::Gcs { + bucket_name: config.bucket_name.clone(), + key_prefix: config.key_prefix.clone(), + path_service_account: config.path_service_account.clone(), + }, + Self::ValkeySemantic(config) => BackendIdentity::ValkeySemantic { + index_name: config.index_name.clone(), + similarity_threshold: config.similarity_threshold, + }, + Self::Disk(config) => BackendIdentity::Disk { + directory: config.directory.clone(), + }, + Self::AzureBlob(config) => BackendIdentity::AzureBlob { + account_url: config.account_url.clone(), + container: config.container.clone(), + }, + Self::RedisSemantic(config) => BackendIdentity::RedisSemantic { + index_name: config.index_name.clone(), + similarity_threshold: config.similarity_threshold as f32, + }, + Self::QdrantSemantic(config) => BackendIdentity::QdrantSemantic { + collection_name: config.collection_name.clone(), + similarity_threshold: config.similarity_threshold, + vector_size: config.vector_size, + embedding_model: config.embedding.model.clone(), }, } } diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 9398e5a862b..ffd72e33e1b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -9,6 +9,8 @@ tokio::task_local! { static PREPARED_EMBEDDING: Result, Error>; } +/// Runs `future` with the vector the Python embedder already produced, so the backend's +/// `async_embed` never has to call back into Python from the runtime. pub(super) fn with_prepared_embedding( vector: Result, Error>, future: F, @@ -16,6 +18,7 @@ pub(super) fn with_prepared_embedding( PREPARED_EMBEDDING.scope(vector, future) } +/// The Python object that owns embedding for a semantic backend. pub(super) struct PythonEmbedder(Py); impl Clone for PythonEmbedder { @@ -29,10 +32,6 @@ impl PythonEmbedder { Self(object) } - pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { - Ok(Self(backend.clone().unbind())) - } - pub(super) fn object(&self) -> &Py { &self.0 } @@ -41,18 +40,6 @@ impl PythonEmbedder { visit.call(&self.0) } - pub(super) fn async_embed_awaitable<'py>( - &self, - py: Python<'py>, - prompt: &str, - metadata: &Option, - ) -> PyResult> { - let metadata = to_py(py, metadata)?; - self.0 - .bind(py) - .call_method1("_get_async_embedding", (prompt, metadata)) - } - fn metadata_kwargs<'py>( py: Python<'py>, metadata: Option<&Value>, @@ -62,7 +49,8 @@ impl PythonEmbedder { Ok(kwargs) } - pub(super) fn async_embedding_coroutine( + /// The awaitable of `_get_async_embedding(prompt, metadata=...)`, to run in the caller's loop. + pub(super) fn async_embedding( &self, py: Python<'_>, prompt: &str, @@ -82,35 +70,8 @@ impl PythonEmbedder { .map(|value| value as f32) .collect()) } -} -impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - let result = Python::attach(|py| -> PyResult> { - let metadata = to_py(py, &metadata)?; - self.0 - .bind(py) - .call_method1("_get_embedding", (prompt, metadata))? - .extract() - }) - .map_err(|_| Error::Unavailable)?; - Ok(result.into_iter().map(|value| value as f32).collect()) - } - - fn async_embed( - &self, - _prompt: &str, - _metadata: Option<&Value>, - ) -> impl Future, Error>> + Send { - let seeded = PREPARED_EMBEDDING - .try_with(Clone::clone) - .unwrap_or(Err(Error::Unavailable)); - std::future::ready(seeded) - } -} - -impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + fn embed_sync(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { Python::attach(|py| { let kwargs = Self::metadata_kwargs(py, metadata)?; Self::extract(self.0.bind(py).call_method( @@ -122,15 +83,38 @@ impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { .map_err(|_| Error::Unavailable) } + fn seeded_embedding() -> Result, Error> { + PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)) + } +} + +impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + self.embed_sync(prompt, metadata) + } + fn async_embed( &self, _prompt: &str, _metadata: Option<&Value>, ) -> impl Future, Error>> + Send { - let seeded = PREPARED_EMBEDDING - .try_with(Clone::clone) - .unwrap_or(Err(Error::Unavailable)); - std::future::ready(seeded) + std::future::ready(Self::seeded_embedding()) + } +} + +impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + self.embed_sync(prompt, metadata) + } + + fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + std::future::ready(Self::seeded_embedding()) } } @@ -152,5 +136,9 @@ mod tests { let unscoped = litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await; assert_eq!(unscoped, Err(Error::Unavailable)); + let valkey = with_prepared_embedding(Ok(vec![0.5]), async move { + litellm_cache_valkey_semantic::Embedder::async_embed(&embedder, "prompt", None).await + }); + assert_eq!(valkey.await, Ok(vec![0.5])); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 43b53584943..17fa278ae5e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -11,6 +11,7 @@ use serde_json::Value; use super::{ config::{CacheConfigProjection, NativeCacheConfig}, handle::CacheTestHandle, + identity::BackendIdentity, native::NativeResponseCache, }; @@ -352,47 +353,41 @@ impl FacadeGuard { facade: &Bound<'_, PyAny>, service: &NativeResponseCache, ) -> PyResult { - let kind = service.kind(); + let identity = service.identity(); + let kind = identity.kind(); let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( "only exact built-in Cache facades can be registered", )); } - let cluster = matches!(service.topology(), Some(RedisTopology::Cluster { .. })); - let (module, name, cache_kind) = match (kind, cluster) { - ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), - ("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"), - ("redis_semantic", _) => ( - "litellm.caching.redis_semantic_cache", - "RedisSemanticCache", - "redis-semantic", - ), + let cluster = matches!( + identity, + BackendIdentity::Redis { + topology: RedisTopology::Cluster { .. }, + .. + } + ); + let (module, name) = match (kind, cluster) { + ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache"), + ("redis", false) => ("litellm.caching.redis_cache", "RedisCache"), + ("redis", true) => ("litellm.caching.redis_cluster_cache", "RedisClusterCache"), + ("redis_semantic", _) => ("litellm.caching.redis_semantic_cache", "RedisSemanticCache"), ("qdrant_semantic", _) => ( "litellm.caching.qdrant_semantic_cache", "QdrantSemanticCache", - "qdrant-semantic", ), - ("redis", true) => ( - "litellm.caching.redis_cluster_cache", - "RedisClusterCache", - "redis", - ), - ("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"), - ("valkey-semantic", false) => ( + ("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache"), + ("valkey-semantic", _) => ( "litellm.caching.valkey_semantic_cache", "ValkeySemanticCache", - "valkey-semantic", ), - ("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"), - ("azure-blob", _) => ( - "litellm.caching.azure_blob_cache", - "AzureBlobCache", - "azure-blob", - ), - ("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"), + ("disk", _) => ("litellm.caching.disk_cache", "DiskCache"), + ("azure-blob", _) => ("litellm.caching.azure_blob_cache", "AzureBlobCache"), + ("s3", _) => ("litellm.caching.s3_cache", "S3Cache"), _ => unreachable!(), }; + let cache_kind = identity.cache_type(); let backend = facade.getattr("cache")?; if facade.getattr("type")?.extract::()? != cache_kind || !backend.get_type().is(&py.import(module)?.getattr(name)?) diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 51e1b02c405..61993f42279 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -261,7 +261,7 @@ impl CacheTestHandle { index_name: String, embedder: &Bound<'_, PyAny>, ) -> PyResult { - let python_embedder = PythonEmbedder::from_backend(embedder)?; + let python_embedder = PythonEmbedder::new(embedder.clone().unbind()); let service = NativeResponseCache::valkey_semantic( &url, similarity_threshold, diff --git a/litellm-rust/crates/python-bridge/src/cache/identity.rs b/litellm-rust/crates/python-bridge/src/cache/identity.rs new file mode 100644 index 00000000000..835bafd3ff1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/identity.rs @@ -0,0 +1,511 @@ +use std::{path::PathBuf, time::Duration}; + +use litellm_cache_redis::RedisTopology; + +/// What makes a native backend the one a Python facade describes: the configuration a user can +/// observe on the Python object, captured once so facade projection and native construction +/// compare plain data instead of reaching into each backend type. +#[derive(Clone, Debug, PartialEq)] +pub(super) enum BackendIdentity { + Memory { + capacity: usize, + max_entry_bytes: Option, + default_ttl: Option, + }, + Redis { + topology: RedisTopology, + namespace: Option, + default_ttl: Option, + }, + S3 { + bucket: String, + key_prefix: String, + region: String, + endpoint: Option, + }, + Gcs { + bucket_name: String, + key_prefix: String, + path_service_account: Option, + }, + Disk { + directory: PathBuf, + }, + AzureBlob { + account_url: String, + container: String, + }, + RedisSemantic { + index_name: String, + /// The backend stores the threshold as `f32`; a facade's `f64` is compared at that width. + similarity_threshold: f32, + }, + ValkeySemantic { + index_name: String, + similarity_threshold: f64, + }, + QdrantSemantic { + collection_name: String, + similarity_threshold: f64, + vector_size: u64, + embedding_model: String, + }, +} + +const TYPES: &str = "facade and native backend types must match"; + +impl BackendIdentity { + /// The native backend name reported to Python through `_CacheTestHandle.backend`. + pub(super) fn kind(&self) -> &'static str { + match self { + Self::Memory { .. } => "memory", + Self::Redis { .. } => "redis", + Self::S3 { .. } => "s3", + Self::Gcs { .. } => "gcs", + Self::ValkeySemantic { .. } => "valkey-semantic", + Self::RedisSemantic { .. } => "redis_semantic", + Self::QdrantSemantic { .. } => "qdrant_semantic", + Self::Disk { .. } => "disk", + Self::AzureBlob { .. } => "azure-blob", + } + } + + /// The `LiteLLMCacheType` value a facade of this backend carries in `Cache.type`. + pub(super) fn cache_type(&self) -> &'static str { + match self { + Self::Memory { .. } => "local", + Self::Redis { .. } => "redis", + Self::S3 { .. } => "s3", + Self::Gcs { .. } => "gcs", + Self::ValkeySemantic { .. } => "valkey-semantic", + Self::RedisSemantic { .. } => "redis-semantic", + Self::QdrantSemantic { .. } => "qdrant-semantic", + Self::Disk { .. } => "disk", + Self::AzureBlob { .. } => "azure-blob", + } + } + + /// The first difference between the facade's configuration (`self`) and the native + /// backend (`native`), in the order Python users see the attributes. + pub(super) fn mismatch(&self, native: &Self) -> Option<&'static str> { + let mut differences: Vec<(bool, &'static str)> = Vec::new(); + let mut differs = |condition: bool, message: &'static str| { + differences.push((condition, message)); + }; + match (self, native) { + ( + Self::Memory { + capacity, + max_entry_bytes, + default_ttl, + }, + Self::Memory { + capacity: native_capacity, + max_entry_bytes: native_max_entry_bytes, + default_ttl: native_default_ttl, + }, + ) => { + differs( + default_ttl != native_default_ttl, + "facade and native backend default TTLs must match", + ); + differs( + capacity != native_capacity, + "facade and native backend capacities must match", + ); + differs( + max_entry_bytes != native_max_entry_bytes, + "facade and native backend item limits must match", + ); + } + ( + Self::Redis { + topology, + namespace, + default_ttl, + }, + Self::Redis { + topology: native_topology, + namespace: native_namespace, + default_ttl: native_default_ttl, + }, + ) => { + differs( + default_ttl != native_default_ttl, + "facade and native backend default TTLs must match", + ); + differs( + topology != native_topology, + "facade and native backend topologies must match", + ); + differs( + namespace != native_namespace, + "facade and native backend namespaces must match", + ); + } + ( + Self::S3 { + bucket, + key_prefix, + region, + endpoint, + }, + Self::S3 { + bucket: native_bucket, + key_prefix: native_key_prefix, + region: native_region, + endpoint: native_endpoint, + }, + ) => { + differs( + bucket != native_bucket, + "facade and native backend buckets must match", + ); + differs( + key_prefix != native_key_prefix, + "facade and native backend key prefixes must match", + ); + differs( + region != native_region, + "facade and native backend regions must match", + ); + differs( + endpoint != native_endpoint, + "facade and native backend endpoints must match", + ); + } + ( + Self::Gcs { + bucket_name, + key_prefix, + path_service_account, + }, + Self::Gcs { + bucket_name: native_bucket_name, + key_prefix: native_key_prefix, + path_service_account: native_path_service_account, + }, + ) => { + differs( + bucket_name != native_bucket_name, + "facade and native backend buckets must match", + ); + differs( + key_prefix != native_key_prefix, + "facade and native backend key prefixes must match", + ); + differs( + path_service_account != native_path_service_account, + "facade and native backend credentials must match", + ); + } + ( + Self::Disk { directory }, + Self::Disk { + directory: native_directory, + }, + ) => { + let canonical = |path: &PathBuf| std::fs::canonicalize(path).ok(); + differs( + canonical(directory) != canonical(native_directory), + "facade and native backend directories must match", + ); + } + ( + Self::AzureBlob { + account_url, + container, + }, + Self::AzureBlob { + account_url: native_account_url, + container: native_container, + }, + ) => { + differs( + account_url != native_account_url || container != native_container, + "facade and native backend containers must match", + ); + } + ( + Self::RedisSemantic { + index_name, + similarity_threshold, + }, + Self::RedisSemantic { + index_name: native_index_name, + similarity_threshold: native_similarity_threshold, + }, + ) => { + differs( + index_name != native_index_name, + "facade and native backend index names must match", + ); + differs( + similarity_threshold != native_similarity_threshold, + "facade and native backend similarity thresholds must match", + ); + } + ( + Self::ValkeySemantic { + index_name, + similarity_threshold, + }, + Self::ValkeySemantic { + index_name: native_index_name, + similarity_threshold: native_similarity_threshold, + }, + ) => { + differs( + index_name != native_index_name + || similarity_threshold != native_similarity_threshold, + "facade and native semantic settings must match", + ); + } + ( + Self::QdrantSemantic { + collection_name, + similarity_threshold, + vector_size, + embedding_model, + }, + Self::QdrantSemantic { + collection_name: native_collection_name, + similarity_threshold: native_similarity_threshold, + vector_size: native_vector_size, + embedding_model: native_embedding_model, + }, + ) => { + differs( + collection_name != native_collection_name, + "facade and native backend collections must match", + ); + differs( + similarity_threshold != native_similarity_threshold, + "facade and native backend similarity thresholds must match", + ); + differs( + vector_size != native_vector_size, + "facade and native backend vector sizes must match", + ); + differs( + embedding_model != native_embedding_model, + "facade and native backend embedding models must match", + ); + } + _ => return Some(TYPES), + } + differences + .into_iter() + .find_map(|(condition, message)| condition.then_some(message)) + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_cache_redis::{RedisNode, RedisTopology}; + + use super::BackendIdentity; + + fn memory() -> BackendIdentity { + BackendIdentity::Memory { + capacity: 200, + max_entry_bytes: Some(1024), + default_ttl: Some(Duration::from_secs(60)), + } + } + + fn redis() -> BackendIdentity { + BackendIdentity::Redis { + topology: RedisTopology::Standalone, + namespace: Some("team".into()), + default_ttl: Some(Duration::from_secs(60)), + } + } + + fn s3() -> BackendIdentity { + BackendIdentity::S3 { + bucket: "bucket".into(), + key_prefix: "cache/".into(), + region: "us-east-1".into(), + endpoint: None, + } + } + + fn gcs() -> BackendIdentity { + BackendIdentity::Gcs { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: Some("credentials.json".into()), + } + } + + fn azure() -> BackendIdentity { + BackendIdentity::AzureBlob { + account_url: "https://account.blob.core.windows.net".into(), + container: "cache".into(), + } + } + + fn redis_semantic() -> BackendIdentity { + BackendIdentity::RedisSemantic { + index_name: "idx".into(), + similarity_threshold: 0.8, + } + } + + #[test] + fn redis_semantic_thresholds_compare_at_backend_precision() { + let facade = BackendIdentity::RedisSemantic { + index_name: "idx".into(), + similarity_threshold: 0.8_f64 as f32, + }; + assert_eq!(facade.mismatch(&redis_semantic()), None); + } + + fn valkey_semantic() -> BackendIdentity { + BackendIdentity::ValkeySemantic { + index_name: "idx".into(), + similarity_threshold: 0.8, + } + } + + fn qdrant() -> BackendIdentity { + BackendIdentity::QdrantSemantic { + collection_name: "collection".into(), + similarity_threshold: 0.8, + vector_size: 1536, + embedding_model: "text-embedding-3-small".into(), + } + } + + #[test] + fn identical_identities_have_no_mismatch() { + for identity in [ + memory(), + redis(), + s3(), + gcs(), + azure(), + redis_semantic(), + valkey_semantic(), + qdrant(), + BackendIdentity::Disk { + directory: std::env::temp_dir(), + }, + ] { + assert_eq!(identity.mismatch(&identity), None, "{identity:?}"); + } + } + + #[test] + fn different_kinds_report_a_type_mismatch() { + assert_eq!( + memory().mismatch(&redis()), + Some("facade and native backend types must match") + ); + assert_eq!( + redis_semantic().mismatch(&valkey_semantic()), + Some("facade and native backend types must match") + ); + } + + #[test] + fn the_first_differing_field_names_the_mismatch() { + let BackendIdentity::Memory { capacity, .. } = memory() else { + unreachable!() + }; + assert_eq!( + memory().mismatch(&BackendIdentity::Memory { + capacity: capacity + 1, + max_entry_bytes: Some(1), + default_ttl: Some(Duration::from_secs(60)), + }), + Some("facade and native backend capacities must match") + ); + assert_eq!( + memory().mismatch(&BackendIdentity::Memory { + capacity, + max_entry_bytes: Some(1), + default_ttl: Some(Duration::from_secs(61)), + }), + Some("facade and native backend default TTLs must match") + ); + assert_eq!( + redis().mismatch(&BackendIdentity::Redis { + topology: RedisTopology::Cluster { + startup_nodes: vec![RedisNode { + host: "node".into(), + port: 7000, + }], + }, + namespace: None, + default_ttl: Some(Duration::from_secs(60)), + }), + Some("facade and native backend topologies must match") + ); + assert_eq!( + s3().mismatch(&BackendIdentity::S3 { + bucket: "bucket".into(), + key_prefix: "cache/".into(), + region: "us-east-1".into(), + endpoint: Some("http://localhost:9000".into()), + }), + Some("facade and native backend endpoints must match") + ); + assert_eq!( + gcs().mismatch(&BackendIdentity::Gcs { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: None, + }), + Some("facade and native backend credentials must match") + ); + assert_eq!( + azure().mismatch(&BackendIdentity::AzureBlob { + account_url: "https://account.blob.core.windows.net".into(), + container: "other".into(), + }), + Some("facade and native backend containers must match") + ); + assert_eq!( + valkey_semantic().mismatch(&BackendIdentity::ValkeySemantic { + index_name: "idx".into(), + similarity_threshold: 0.9, + }), + Some("facade and native semantic settings must match") + ); + assert_eq!( + qdrant().mismatch(&BackendIdentity::QdrantSemantic { + collection_name: "collection".into(), + similarity_threshold: 0.8, + vector_size: 1536, + embedding_model: "text-embedding-3-large".into(), + }), + Some("facade and native backend embedding models must match") + ); + } + + #[test] + fn disk_directories_compare_canonically() { + let directory = std::env::temp_dir(); + let mut indirect = directory.clone(); + indirect.push("."); + assert_eq!( + BackendIdentity::Disk { + directory: directory.clone() + } + .mismatch(&BackendIdentity::Disk { + directory: indirect + }), + None + ); + assert_eq!( + BackendIdentity::Disk { directory }.mismatch(&BackendIdentity::Disk { + directory: "/definitely/missing".into() + }), + Some("facade and native backend directories must match") + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index fa028518559..28dd6c3e798 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -5,11 +5,11 @@ mod embedder; mod facade; mod future; mod handle; +mod identity; mod native; mod request; mod resolver; mod semantic; -mod semantic_step; use litellm_cache::Error; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 7fef6f55611..254b9cdea4d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,8 +1,6 @@ -use std::{path::Path, sync::Arc, time::Duration}; +use std::{sync::Arc, time::Duration}; -use litellm_cache::{ - CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext, -}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_disk::DiskCache; use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; @@ -11,8 +9,7 @@ use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCach use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ - CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec, - ResponseCacheRequest, WriteBuffer, + ExactResponseCache, PartialHits, ResponseCache, ResponseCacheCodec, WriteBuffer, }; use litellm_cache_s3::{S3Cache, S3CacheConfig}; use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; @@ -22,69 +19,37 @@ use serde_json::Value; use super::{ config::QdrantSemanticCacheConfig, embedder::PythonEmbedder, - request::NativeRequest, - semantic::{SemanticBody, SemanticOperation, drive}, - semantic_step::{SemanticEmbedExecution, drive_semantic}, + identity::BackendIdentity, + request::{NativeRequest, now}, + semantic::{EmbeddingFailure, SemanticExecution, SemanticOperation, drive}, }; -fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput { - let mut key = request.key.clone(); - if key.preset.is_some() { - return key; +/// What the Python embedder receives for one semantic request. +pub(super) struct EmbeddingInput { + pub(super) prompt: String, + pub(super) metadata: Option, +} + +/// An exact-match backend behind one pointer, with the identity its facade must reproduce. +pub(super) struct ExactService { + cache: Arc, + buffer: Option, + identity: BackendIdentity, +} + +impl ExactService { + fn new(cache: Arc, identity: BackendIdentity) -> Arc { + Arc::new(Self { + cache, + buffer: None, + identity, + }) } - key.fields - .retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input")); - const TENANT: [&str; 3] = [ - "user_api_key", - "user_api_key_team_id", - "user_api_key_org_id", - ]; - let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); - for name in TENANT.into_iter().chain(end_user) { - let sources = [ - request.metadata.as_ref(), - request.litellm_metadata.as_ref(), - request - .litellm_params - .as_ref() - .and_then(|params| params.get("metadata")), - request - .litellm_params - .as_ref() - .and_then(|params| params.get("litellm_metadata")), - ]; - let Some(value) = sources.into_iter().flatten().find_map(|source| { - source - .as_object() - .and_then(|values| values.get(name)) - .filter(|value| !value.is_null()) - }) else { - continue; - }; - let value = match value { - Value::Null => continue, - Value::String(text) => text.clone(), - other => other.to_string(), - }; - key.fields.push(CacheKeyField { - name: name.to_owned(), - value: Some(value), - api_parameter: true, - internal_parameter: false, - }); - } - key } #[derive(Clone)] pub(super) enum NativeResponseCache { - Memory(Arc>>), - Redis { - cache: Arc>>, - buffer: Option>, - }, - S3(Arc>>), - Gcs(Arc>>), + Exact(Arc), ValkeySemantic { cache: Arc>>, embedder: PythonEmbedder, @@ -95,23 +60,25 @@ pub(super) enum NativeResponseCache { embedder: PythonEmbedder, }, QdrantSemantic(Arc>>), - Disk(Arc>>), - AzureBlob(Arc>>), } impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { - Self::Memory(Arc::new(ResponseCache::new(Arc::new( - InMemoryCache::with_clock_and_size_measurement( - Some(capacity), - Some(ttl), - Some(max_entry_bytes), - Some(Arc::new(|entry| { - ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) - })), - super::request::now, - ), - )))) + let backend = InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(ttl), + Some(max_entry_bytes), + Some(Arc::new(|entry| { + ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) + })), + now, + ); + let identity = BackendIdentity::Memory { + capacity: backend.max_size_in_memory(), + max_entry_bytes: backend.max_entry_bytes(), + default_ttl: None, + }; + Self::exact(ResponseCache::new(Arc::new(backend)), identity) } pub fn redis( @@ -122,19 +89,97 @@ impl NativeResponseCache { ) -> Result { let backend = RedisCache::connect(url, topology, ttl, ResponseCacheCodec)?.with_namespace(namespace); - Ok(Self::Redis { - cache: Arc::new(ResponseCache::new(Arc::new(backend))), - buffer: None, - }) + let identity = BackendIdentity::Redis { + topology: backend.topology().clone(), + namespace: backend.namespace().map(str::to_owned), + default_ttl: None, + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) } pub async fn s3(config: S3CacheConfig) -> Self { let runtime = tokio::runtime::Handle::current(); - Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new( - config, + let backend = S3Cache::new(config, ResponseCacheCodec, runtime); + let identity = BackendIdentity::S3 { + bucket: backend.bucket().to_owned(), + key_prefix: backend.key_prefix().to_owned(), + region: backend.region().to_owned(), + endpoint: backend.endpoint().map(str::to_owned), + }; + Self::exact(ResponseCache::new(Arc::new(backend)), identity) + } + + pub fn disk(directory: &str) -> Result { + let backend = DiskCache::open(directory, ResponseCacheCodec)?; + let identity = BackendIdentity::Disk { + directory: backend.directory().to_path_buf(), + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + } + + pub fn gcs(config: GcsConfig, token: Option) -> Result { + let backend = match token { + Some(token) => GcsCache::with_token_source( + config, + ResponseCacheCodec, + Arc::new(StaticTokenSource(token)), + )?, + None => GcsCache::new(config, ResponseCacheCodec)?, + }; + let identity = BackendIdentity::Gcs { + bucket_name: backend.bucket_name().to_owned(), + key_prefix: backend.key_prefix().to_owned(), + path_service_account: backend.path_service_account().map(str::to_owned), + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + } + + pub async fn azure_blob(account_url: &str, container: &str) -> Result { + let backend = AzureBlobCache::connect( + account_url, + container, ResponseCacheCodec, - runtime, - ))))) + tokio::runtime::Handle::current(), + ) + .await?; + let identity = BackendIdentity::AzureBlob { + account_url: backend.account_url().to_owned(), + container: backend.container_name().to_owned(), + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + } + + /// Wraps a built exact backend; the TTL a facade must match comes from the built cache. + fn exact(cache: ResponseCache, identity: BackendIdentity) -> Self + where + ResponseCache: ExactResponseCache + 'static, + B: litellm_cache::BaseCache, + B::Context: Default + PartialEq, + { + let cache: Arc = Arc::new(cache); + let default_ttl = cache.default_ttl(); + let identity = match identity { + BackendIdentity::Memory { + capacity, + max_entry_bytes, + .. + } => BackendIdentity::Memory { + capacity, + max_entry_bytes, + default_ttl, + }, + BackendIdentity::Redis { + topology, + namespace, + .. + } => BackendIdentity::Redis { + topology, + namespace, + default_ttl, + }, + other => other, + }; + Self::Exact(ExactService::new(cache, identity)) } pub fn valkey_semantic( @@ -196,103 +241,39 @@ impl NativeResponseCache { )))) } - pub fn disk(directory: &str) -> Result { - let cache = DiskCache::open(directory, ResponseCacheCodec)?; - Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) - } - - pub fn gcs(config: GcsConfig, token: Option) -> Result { - let backend = match token { - Some(token) => GcsCache::with_token_source( - config, - ResponseCacheCodec, - Arc::new(StaticTokenSource(token)), - )?, - None => GcsCache::new(config, ResponseCacheCodec)?, - }; - Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend))))) - } - - pub async fn azure_blob(account_url: &str, container: &str) -> Result { - let backend = AzureBlobCache::connect( - account_url, - container, - ResponseCacheCodec, - tokio::runtime::Handle::current(), - ) - .await?; - Ok(Self::AzureBlob(Arc::new(ResponseCache::new(Arc::new( - backend, - ))))) - } - - pub fn azure_blob_identity(&self) -> Option<(&str, &str)> { + pub fn identity(&self) -> BackendIdentity { match self { - Self::AzureBlob(cache) => Some(( - cache.backend().account_url(), - cache.backend().container_name(), - )), - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::RedisSemantic { .. } - | Self::QdrantSemantic(_) - | Self::Disk(_) - | Self::Gcs(_) => None, - } - } - - fn exact(request: &NativeRequest) -> ResponseCacheRequest { - ResponseCacheRequest { - key: request.key.clone(), - controls: request.controls, - context: ExactCacheContext { ttl: request.ttl }, - max_age: request.max_age, - } - } - - pub(super) fn semantic_request( - request: &NativeRequest, - ) -> ResponseCacheRequest { - ResponseCacheRequest { - key: request.key.clone(), - controls: request.controls, - context: SemanticCacheContext { - input: request.input.clone(), - messages: request.messages.clone(), - metadata: request.metadata.clone(), - scope: request.scope.clone(), - ttl: request.ttl, + Self::Exact(service) => service.identity.clone(), + Self::ValkeySemantic { cache, .. } => BackendIdentity::ValkeySemantic { + index_name: cache.backend().index_name().to_owned(), + similarity_threshold: cache.backend().similarity_threshold(), + }, + Self::RedisSemantic { cache, .. } => BackendIdentity::RedisSemantic { + index_name: cache.backend().index_name().to_owned(), + similarity_threshold: cache.backend().similarity_threshold(), + }, + Self::QdrantSemantic(cache) => BackendIdentity::QdrantSemantic { + collection_name: cache.backend().collection_name().to_owned(), + similarity_threshold: cache.backend().similarity_threshold(), + vector_size: cache.backend().vector_size(), + embedding_model: cache.backend().embedder().model().to_owned(), }, - max_age: request.max_age, } } - fn semantic( - request: &NativeRequest, - scope: &str, - ) -> ResponseCacheRequest { - ResponseCacheRequest { - key: semantic_key(request, scope), - controls: request.controls, - context: SemanticCacheContext { - input: request.input.clone(), - messages: request.messages.clone(), - metadata: request.metadata.clone(), - scope: Some(scope.to_owned()), - ttl: request.ttl, - }, - max_age: request.max_age, - } + pub fn kind(&self) -> &'static str { + self.identity().kind() } pub fn with_redis_flush_size(self, flush_size: Option) -> Self { match self { - Self::Redis { cache, .. } => Self::Redis { - cache, - buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))), - }, + Self::Exact(service) if matches!(service.identity, BackendIdentity::Redis { .. }) => { + Self::Exact(Arc::new(ExactService { + cache: Arc::clone(&service.cache), + buffer: flush_size.map(WriteBuffer::new), + identity: service.identity.clone(), + })) + } value => value, } } @@ -310,191 +291,6 @@ impl NativeResponseCache { } } - pub fn kind(&self) -> &'static str { - match self { - Self::Memory(_) => "memory", - Self::Redis { .. } => "redis", - Self::S3(_) => "s3", - Self::Gcs(_) => "gcs", - Self::ValkeySemantic { .. } => "valkey-semantic", - Self::RedisSemantic { .. } => "redis_semantic", - Self::QdrantSemantic(_) => "qdrant_semantic", - Self::Disk(_) => "disk", - Self::AzureBlob(_) => "azure-blob", - } - } - - pub fn default_ttl(&self) -> Option { - match self { - 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::ValkeySemantic { cache, .. } => cache.default_ttl(), - Self::RedisSemantic { cache, .. } => cache.default_ttl(), - Self::QdrantSemantic(_) => None, - Self::Disk(cache) => cache.default_ttl(), - Self::AzureBlob(cache) => cache.default_ttl(), - } - } - - pub fn bucket(&self) -> Option<&str> { - match self { - Self::S3(cache) => Some(cache.backend().bucket()), - _ => None, - } - } - - pub fn key_prefix(&self) -> Option<&str> { - match self { - Self::S3(cache) => Some(cache.backend().key_prefix()), - _ => None, - } - } - - pub fn region(&self) -> Option<&str> { - match self { - Self::S3(cache) => Some(cache.backend().region()), - _ => None, - } - } - - pub fn endpoint(&self) -> Option<&str> { - match self { - Self::S3(cache) => cache.backend().endpoint(), - _ => None, - } - } - - pub fn namespace(&self) -> Option<&str> { - match self { - Self::Memory(_) - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::RedisSemantic { .. } - | Self::QdrantSemantic(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - Self::Redis { cache, .. } => cache.backend().namespace(), - } - } - - pub fn topology(&self) -> Option<&RedisTopology> { - match self { - Self::Memory(_) - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::RedisSemantic { .. } - | Self::QdrantSemantic(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - Self::Redis { cache, .. } => Some(cache.backend().topology()), - } - } - - pub fn capacity(&self) -> Option { - match self { - Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::RedisSemantic { .. } - | Self::QdrantSemantic(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - } - } - - pub fn max_entry_bytes(&self) -> Option { - match self { - Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::RedisSemantic { .. } - | Self::QdrantSemantic(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - } - } - - pub fn directory(&self) -> Option<&Path> { - match self { - Self::Disk(cache) => Some(cache.backend().directory()), - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::RedisSemantic { .. } - | Self::QdrantSemantic(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - } - } - - pub fn semantic_config(&self) -> Option<(f64, &str)> { - match self { - Self::ValkeySemantic { cache, .. } => Some(( - cache.backend().similarity_threshold(), - cache.backend().index_name(), - )), - Self::RedisSemantic { cache, .. } => Some(( - f64::from(cache.backend().similarity_threshold()), - cache.backend().index_name(), - )), - _ => None, - } - } - - pub fn index_name(&self) -> Option<&str> { - match self { - Self::RedisSemantic { cache, .. } => Some(cache.backend().index_name()), - _ => None, - } - } - - pub fn similarity_threshold(&self) -> Option { - match self { - Self::RedisSemantic { cache, .. } => { - Some(f64::from(cache.backend().similarity_threshold())) - } - Self::QdrantSemantic(cache) => Some(cache.backend().similarity_threshold()), - _ => None, - } - } - - pub fn collection_name(&self) -> Option<&str> { - match self { - Self::QdrantSemantic(cache) => Some(cache.backend().collection_name()), - _ => None, - } - } - - pub fn vector_size(&self) -> Option { - match self { - Self::QdrantSemantic(cache) => Some(cache.backend().vector_size()), - _ => None, - } - } - - pub fn embedding_model(&self) -> Option<&str> { - match self { - Self::QdrantSemantic(cache) => Some(cache.backend().embedder().model()), - _ => None, - } - } - - pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { - match self { - Self::RedisSemantic { embedder, .. } => Some(embedder), - _ => None, - } - } - pub fn embedder_object(&self) -> Option<&Py> { match self { Self::RedisSemantic { embedder, .. } => Some(embedder.object()), @@ -502,21 +298,49 @@ impl NativeResponseCache { } } + /// The prompt and metadata this backend would embed for `request`, if it has a prompt. + pub(super) fn embedding_input(&self, request: &NativeRequest) -> Option { + let context = match self { + Self::ValkeySemantic { scope, .. } => request.scoped_semantic(scope).context, + Self::RedisSemantic { .. } => request.semantic().context, + Self::Exact(_) | Self::QdrantSemantic(_) => return None, + }; + let prompt = litellm_cache_redis_semantic::prompt_from_context(&context)?; + Some(EmbeddingInput { + prompt, + metadata: context.metadata, + }) + } + + /// Drives a semantic operation whose embedding comes from Python. + fn python_semantic<'py>( + &self, + py: Python<'py>, + operation: SemanticOperation, + ) -> PyResult> { + let (embedder, failure) = match self { + Self::ValkeySemantic { embedder, .. } => (embedder, EmbeddingFailure::Propagate), + Self::RedisSemantic { embedder, .. } => (embedder, EmbeddingFailure::Unavailable), + Self::Exact(_) | Self::QdrantSemantic(_) => { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "semantic execution requires a Python-embedded backend", + )); + } + }; + drive( + py, + SemanticExecution::new(self.clone(), embedder.clone(), failure, operation), + ) + } + pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result, Error> { match self { - Self::Memory(cache) => cache.lookup(&Self::exact(request), now), - Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now), - Self::S3(cache) => cache.lookup(&Self::exact(request), now), + Self::Exact(service) => service.cache.lookup(&request.exact(), now), Self::ValkeySemantic { cache, scope, .. } => { - cache.lookup(&Self::semantic(request, scope), now) + cache.lookup(&request.scoped_semantic(scope), now) } - Self::RedisSemantic { cache, .. } => { - cache.lookup(&Self::semantic_request(request), now) - } - Self::Gcs(cache) => cache.lookup(&Self::exact(request), now), - Self::QdrantSemantic(cache) => cache.lookup(&Self::semantic_request(request), now), - Self::Disk(cache) => cache.lookup(&Self::exact(request), now), - Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now), + Self::RedisSemantic { cache, .. } => cache.lookup(&request.semantic(), now), + Self::QdrantSemantic(cache) => cache.lookup(&request.semantic(), now), } } @@ -527,21 +351,12 @@ impl NativeResponseCache { now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(&Self::exact(request), response, now), - Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), - Self::S3(cache) => cache.store(&Self::exact(request), response, now), + Self::Exact(service) => service.cache.store(&request.exact(), response, now), Self::ValkeySemantic { cache, scope, .. } => { - cache.store(&Self::semantic(request, scope), response, now) + cache.store(&request.scoped_semantic(scope), response, now) } - Self::RedisSemantic { cache, .. } => { - cache.store(&Self::semantic_request(request), response, now) - } - Self::Gcs(cache) => cache.store(&Self::exact(request), response, now), - Self::QdrantSemantic(cache) => { - cache.store(&Self::semantic_request(request), response, now) - } - Self::Disk(cache) => cache.store(&Self::exact(request), response, now), - Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now), + Self::RedisSemantic { cache, .. } => cache.store(&request.semantic(), response, now), + Self::QdrantSemantic(cache) => cache.store(&request.semantic(), response, now), } } @@ -551,29 +366,10 @@ impl NativeResponseCache { now: Duration, ) -> Result { match self { - Self::Memory(cache) => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.lookup_batch(&requests, now) - } - Self::Redis { cache, .. } => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.lookup_batch(&requests, now) - } - Self::S3(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } + Self::Exact(service) => service.cache.lookup_batch(&exact_requests(requests), now), Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } - Self::Gcs(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } - Self::Disk(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } - Self::AzureBlob(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } } } @@ -583,27 +379,14 @@ impl NativeResponseCache { now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await, - Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, - Self::S3(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::Exact(service) => service.cache.async_lookup(&request.exact(), now).await, Self::ValkeySemantic { cache, scope, .. } => { cache - .async_lookup(&Self::semantic(request, scope), now) + .async_lookup(&request.scoped_semantic(scope), now) .await } - Self::RedisSemantic { cache, .. } => { - cache - .async_lookup(&Self::semantic_request(request), now) - .await - } - Self::QdrantSemantic(cache) => { - cache - .async_lookup(&Self::semantic_request(request), now) - .await - } - Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await, - Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await, - Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::RedisSemantic { cache, .. } => cache.async_lookup(&request.semantic(), now).await, + Self::QdrantSemantic(cache) => cache.async_lookup(&request.semantic(), now).await, } } @@ -613,42 +396,16 @@ impl NativeResponseCache { request: NativeRequest, ) -> PyResult> { match self { - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => { + Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); litellm_host_python::run_async( py, - async move { service.async_lookup(&request, super::request::now()).await }, + async move { service.async_lookup(&request, now()).await }, super::cache_error, ) } - Self::ValkeySemantic { - cache, - embedder, - scope, - } => drive_semantic( - py, - SemanticEmbedExecution::lookup( - Arc::clone(cache.backend_arc()), - embedder.clone(), - Self::semantic(&request, scope), - ), - ), - Self::RedisSemantic { .. } => drive( - py, - SemanticBody::new(self.clone(), SemanticOperation::Lookup(request)), - ), - Self::QdrantSemantic(_) => { - let service = self.clone(); - litellm_host_python::run_async( - py, - async move { service.async_lookup(&request, super::request::now()).await }, - super::cache_error, - ) + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::Lookup(request)) } } } @@ -660,61 +417,29 @@ impl NativeResponseCache { now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::Redis { - cache, - buffer: None, - } => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::Redis { - cache, - buffer: Some(buffer), - } => { - buffer - .async_store(cache, &Self::exact(request), response, now) - .await - } - Self::S3(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } + Self::Exact(service) => match &service.buffer { + None => { + service + .cache + .async_store(&request.exact(), response, now) + .await + } + Some(buffer) => { + buffer + .async_store(service.cache.as_ref(), &request.exact(), response, now) + .await + } + }, Self::ValkeySemantic { cache, scope, .. } => { cache - .async_store(&Self::semantic(request, scope), response, now) + .async_store(&request.scoped_semantic(scope), response, now) .await } Self::RedisSemantic { cache, .. } => { - cache - .async_store(&Self::semantic_request(request), response, now) - .await + cache.async_store(&request.semantic(), response, now).await } Self::QdrantSemantic(cache) => { - cache - .async_store(&Self::semantic_request(request), response, now) - .await - } - Self::Gcs(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::Disk(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::AzureBlob(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await + cache.async_store(&request.semantic(), response, now).await } } } @@ -726,51 +451,16 @@ impl NativeResponseCache { response: Value, ) -> PyResult> { match self { - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => { + Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); litellm_host_python::run_async( py, - async move { - service - .async_store(&request, response, super::request::now()) - .await - }, + async move { service.async_store(&request, response, now()).await }, super::cache_error, ) } - Self::ValkeySemantic { - cache, - embedder, - scope, - } => drive_semantic( - py, - SemanticEmbedExecution::store( - Arc::clone(cache.backend_arc()), - embedder.clone(), - Self::semantic(&request, scope), - response, - ), - ), - Self::RedisSemantic { .. } => drive( - py, - SemanticBody::new(self.clone(), SemanticOperation::Store(request, response)), - ), - Self::QdrantSemantic(_) => { - let service = self.clone(); - litellm_host_python::run_async( - py, - async move { - service - .async_store(&request, response, super::request::now()) - .await - }, - super::cache_error, - ) + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::Store(request, response)) } } } @@ -781,37 +471,15 @@ impl NativeResponseCache { now: Duration, ) -> Result { match self { - Self::Memory(cache) => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.async_lookup_batch(&requests, now).await - } - Self::Redis { cache, .. } => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.async_lookup_batch(&requests, now).await - } - Self::S3(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + Self::Exact(service) => { + service + .cache + .async_lookup_batch(&exact_requests(requests), now) .await } Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } - Self::Gcs(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - .await - } - Self::Disk(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - .await - } - Self::AzureBlob(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - .await - } } } @@ -821,31 +489,17 @@ impl NativeResponseCache { now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => { + Self::Exact(service) => { let entries = entries .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) + .map(|(request, value)| (request.exact(), value)) .collect(); - cache.async_store_batch(entries, now).await - } - Self::Redis { cache, .. } => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::S3(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await + service.cache.async_store_batch(entries, now).await } Self::ValkeySemantic { cache, scope, .. } => { let entries = entries .into_iter() - .map(|(request, value)| (Self::semantic(&request, scope), value)) + .map(|(request, value)| (request.scoped_semantic(scope), value)) .collect(); cache.async_store_batch(entries, now).await } @@ -853,28 +507,7 @@ impl NativeResponseCache { Self::QdrantSemantic(cache) => { let entries = entries .into_iter() - .map(|(request, value)| (Self::semantic_request(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::Gcs(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::Disk(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::AzureBlob(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) + .map(|(request, value)| (request.semantic(), value)) .collect(); cache.async_store_batch(entries, now).await } @@ -887,185 +520,54 @@ impl NativeResponseCache { entries: Vec<(NativeRequest, Value)>, ) -> PyResult> { match self { - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => { + Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); litellm_host_python::run_async( py, - async move { - service - .async_store_batch(entries, super::request::now()) - .await - }, + async move { service.async_store_batch(entries, now()).await }, super::cache_error, ) } - Self::ValkeySemantic { - cache, - embedder, - scope, - } => { - let (requests, responses): (Vec<_>, Vec<_>) = entries - .into_iter() - .map(|(request, response)| (Self::semantic(&request, scope), response)) - .unzip(); - drive_semantic( - py, - SemanticEmbedExecution::store_batch( - Arc::clone(cache.backend_arc()), - embedder.clone(), - requests, - responses, - ), - ) - } - Self::RedisSemantic { .. } => drive( - py, - SemanticBody::new(self.clone(), SemanticOperation::StoreBatch(entries.into())), - ), - Self::QdrantSemantic(_) => { - let service = self.clone(); - litellm_host_python::run_async( - py, - async move { - service - .async_store_batch(entries, super::request::now()) - .await - }, - super::cache_error, - ) + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::StoreBatch(entries.into())) } } } pub async fn async_flush(&self) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_flush().await, - Self::Redis { cache, buffer } => { - if let Some(buffer) = buffer { + Self::Exact(service) => { + if let Some(buffer) = &service.buffer { buffer.clear()?; } - cache.async_flush().await + service.cache.async_flush().await } - Self::S3(cache) => cache.async_flush().await, Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } - Self::Gcs(cache) => cache.async_flush().await, - Self::Disk(cache) => cache.async_flush().await, - Self::AzureBlob(cache) => cache.async_flush().await, } } pub async fn test_connection(&self) -> Result { match self { - Self::Memory(cache) => cache.test_connection().await, - Self::Redis { cache, .. } => cache.test_connection().await, - Self::S3(cache) => cache.test_connection().await, + Self::Exact(service) => service.cache.test_connection().await, Self::ValkeySemantic { cache, .. } => cache.test_connection().await, Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } - Self::Gcs(cache) => cache.test_connection().await, - Self::Disk(cache) => cache.test_connection().await, - Self::AzureBlob(cache) => cache.test_connection().await, } } pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { match self { - Self::ValkeySemantic { embedder, .. } => embedder.traverse(visit)?, - Self::RedisSemantic { embedder, .. } => embedder.traverse(visit)?, - _ => {} - } - Ok(()) - } - - pub fn gcs_backend(&self) -> Option<&GcsCache> { - match self { - Self::Gcs(cache) => Some(cache.backend()), - _ => None, + Self::ValkeySemantic { embedder, .. } | Self::RedisSemantic { embedder, .. } => { + embedder.traverse(visit) + } + Self::Exact(_) | Self::QdrantSemantic(_) => Ok(()), } } } -#[cfg(test)] -mod tests { - use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key}; - use serde_json::json; - use sha2::{Digest, Sha256}; - - use super::*; - - fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest { - NativeRequest { - key, - controls: CacheControls::default(), - ttl: None, - max_age: None, - messages: Some(json!([{"role": "user", "content": "prompt"}])), - input: None, - metadata: Some(metadata), - litellm_metadata: None, - litellm_params: None, - scope: None, - } - } - - #[test] - fn semantic_key_matches_python_scope_material() { - let key = CacheKeyInput { - fields: vec![ - CacheKeyField { - name: "model".to_owned(), - value: Some("gpt-4.1".to_owned()), - api_parameter: true, - internal_parameter: false, - }, - CacheKeyField { - name: "messages".to_owned(), - value: Some("prompt".to_owned()), - api_parameter: true, - internal_parameter: false, - }, - ], - ..Default::default() - }; - let request = native_request( - key, - json!({"user_api_key": "k1", "user_api_key_team_id": null}), - ); - let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1")); - assert_eq!(cache_key(&semantic_key(&request, "key")), expected); - - let end_user_request = native_request( - request.key.clone(), - json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}), - ); - let expected = format!( - "{:x}", - Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1") - ); - assert_eq!( - cache_key(&semantic_key(&end_user_request, "end_user")), - expected - ); - - let preset_request = native_request( - CacheKeyInput { - preset: Some("preset-key".to_owned()), - ..Default::default() - }, - json!({"user_api_key": "k1"}), - ); - assert_eq!( - semantic_key(&preset_request, "end_user").preset.as_deref(), - Some("preset-key") - ); - assert!(semantic_key(&preset_request, "end_user").fields.is_empty()); - } +fn exact_requests(requests: &[NativeRequest]) -> Vec { + requests.iter().map(NativeRequest::exact).collect() } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 3b4b910c1f0..627bf9f1840 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,7 +1,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::ExactCacheContext; -use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_cache::{ExactCacheContext, SemanticCacheContext}; +use litellm_cache_response::{CacheControls, CacheKeyField, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; @@ -36,6 +36,99 @@ pub(super) struct NativeRequest { pub(super) scope: Option, } +impl NativeRequest { + pub(super) fn exact(&self) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key.clone(), + controls: self.controls, + context: ExactCacheContext { ttl: self.ttl }, + max_age: self.max_age, + } + } + + /// The request as a semantic backend that keys on the caller's scope sees it. + pub(super) fn semantic(&self) -> ResponseCacheRequest { + self.semantic_with(self.key.clone(), self.scope.clone()) + } + + /// The request keyed the way Python's Valkey semantic cache keys it: prompt fields drop out + /// and the tenant identifiers for `scope` join the key. + pub(super) fn scoped_semantic( + &self, + scope: &str, + ) -> ResponseCacheRequest { + self.semantic_with(semantic_key(self, scope), Some(scope.to_owned())) + } + + fn semantic_with( + &self, + key: CacheKeyInput, + scope: Option, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key, + controls: self.controls, + context: SemanticCacheContext { + input: self.input.clone(), + messages: self.messages.clone(), + metadata: self.metadata.clone(), + scope, + ttl: self.ttl, + }, + max_age: self.max_age, + } + } +} + +fn semantic_key(request: &NativeRequest, scope: &str) -> CacheKeyInput { + let mut key = request.key.clone(); + if key.preset.is_some() { + return key; + } + key.fields + .retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input")); + const TENANT: [&str; 3] = [ + "user_api_key", + "user_api_key_team_id", + "user_api_key_org_id", + ]; + let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); + for name in TENANT.into_iter().chain(end_user) { + let sources = [ + request.metadata.as_ref(), + request.litellm_metadata.as_ref(), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("metadata")), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("litellm_metadata")), + ]; + let Some(value) = sources.into_iter().flatten().find_map(|source| { + source + .as_object() + .and_then(|values| values.get(name)) + .filter(|value| !value.is_null()) + }) else { + continue; + }; + let value = match value { + Value::Null => continue, + Value::String(text) => text.clone(), + other => other.to_string(), + }; + key.fields.push(CacheKeyField { + name: name.to_owned(), + value: Some(value), + api_parameter: true, + internal_parameter: false, + }); + } + key +} + pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; request_input(input) @@ -76,3 +169,90 @@ pub(super) fn now() -> Duration { .duration_since(UNIX_EPOCH) .unwrap_or_default() } + +#[cfg(test)] +mod tests { + use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key}; + use serde_json::json; + use sha2::{Digest, Sha256}; + + use super::*; + + fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest { + NativeRequest { + key, + controls: CacheControls::default(), + ttl: None, + max_age: None, + messages: Some(json!([{"role": "user", "content": "prompt"}])), + input: None, + metadata: Some(metadata), + litellm_metadata: None, + litellm_params: None, + scope: None, + } + } + + #[test] + fn semantic_key_matches_python_scope_material() { + let key = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".to_owned(), + value: Some("gpt-4.1".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "messages".to_owned(), + value: Some("prompt".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + ], + ..Default::default() + }; + let request = native_request( + key, + json!({"user_api_key": "k1", "user_api_key_team_id": null}), + ); + let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1")); + assert_eq!(cache_key(&semantic_key(&request, "key")), expected); + assert_eq!(cache_key(&request.scoped_semantic("key").key), expected); + + let end_user_request = native_request( + request.key.clone(), + json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}), + ); + let expected = format!( + "{:x}", + Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1") + ); + assert_eq!( + cache_key(&semantic_key(&end_user_request, "end_user")), + expected + ); + + let preset_request = native_request( + CacheKeyInput { + preset: Some("preset-key".to_owned()), + ..Default::default() + }, + json!({"user_api_key": "k1"}), + ); + assert_eq!( + semantic_key(&preset_request, "end_user").preset.as_deref(), + Some("preset-key") + ); + assert!(semantic_key(&preset_request, "end_user").fields.is_empty()); + assert_eq!(preset_request.semantic().context.scope, None); + assert_eq!( + preset_request + .scoped_semantic("end_user") + .context + .scope + .as_deref(), + Some("end_user") + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs index 934de01e721..9f4d18d45cd 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -1,7 +1,6 @@ -use std::collections::VecDeque; +use std::{collections::VecDeque, time::Duration}; use litellm_cache::Error; -use litellm_cache_redis_semantic::prompt_from_context; use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; use pyo3::{ PyTraverseError, PyVisit, @@ -23,29 +22,104 @@ pub(super) enum SemanticOperation { StoreBatch(VecDeque<(NativeRequest, Value)>), } +/// What an exception from the Python embedder means for the operation. +#[derive(Clone, Copy)] +pub(super) enum EmbeddingFailure { + /// Raise the Python exception unchanged. + Propagate, + /// Treat the embedding as unavailable and let the backend report that. + Unavailable, +} + enum Phase { Start, AwaitingEmbedding, AwaitingBackend, } -pub(super) struct SemanticBody { +/// Runs a semantic cache operation whose embedding comes from Python: await the Python +/// embedder in the caller's event loop, seed the native backend with the vector, await the +/// backend, and repeat for each entry of a batch. +pub(super) struct SemanticExecution { service: NativeResponseCache, + embedder: PythonEmbedder, + failure: EmbeddingFailure, operation: SemanticOperation, pending: Option<(NativeRequest, Option)>, phase: Phase, + now: Duration, } -impl SemanticBody { - pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self { +impl SemanticExecution { + pub(super) fn new( + service: NativeResponseCache, + embedder: PythonEmbedder, + failure: EmbeddingFailure, + operation: SemanticOperation, + ) -> Self { Self { service, + embedder, + failure, operation, pending: None, phase: Phase::Start, + now: now(), } } + /// Takes the next entry of the operation; `None` once a batch is exhausted. + fn next_pending(&mut self) -> Option<(NativeRequest, Option)> { + match &mut self.operation { + SemanticOperation::Lookup(request) => Some((request.clone(), None)), + SemanticOperation::Store(request, response) => { + Some((request.clone(), Some(std::mem::take(response)))) + } + SemanticOperation::StoreBatch(queue) => queue + .pop_front() + .map(|(request, response)| (request, Some(response))), + } + } + + fn start(&mut self, py: Python<'_>) -> PyResult { + let Some(pending) = self.next_pending() else { + return Ok(ExecutionStep::Return(py.None())); + }; + let (request, response) = &pending; + let enabled = match response { + None => request.controls.reads(), + Some(_) => request.controls.writes(), + }; + let input = enabled + .then(|| self.service.embedding_input(request)) + .flatten(); + self.pending = Some(pending); + let Some(input) = input else { + return self.backend_step(py, Err(Error::Unavailable)); + }; + let awaitable = + self.embedder + .async_embedding(py, &input.prompt, input.metadata.as_ref())?; + self.phase = Phase::AwaitingEmbedding; + Ok(ExecutionStep::Await(awaitable)) + } + + fn embedded(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { + let seed = match result { + Ok(vector) => { + PythonEmbedder::extract(vector.into_bound(py)).map_err(|_| Error::Unavailable) + } + Err(error) => match self.failure { + EmbeddingFailure::Propagate => return Err(error), + EmbeddingFailure::Unavailable if error.is_instance_of::(py) => { + Err(Error::Unavailable) + } + EmbeddingFailure::Unavailable => return Err(error), + }, + }; + self.backend_step(py, seed) + } + fn backend_step( &mut self, py: Python<'_>, @@ -56,11 +130,12 @@ impl SemanticBody { PyRuntimeError::new_err("semantic execution resumed without a pending operation") })?; let service = self.service.clone(); + let now = self.now; let future = async move { match response { - None => service.async_lookup(&request, now()).await, + None => service.async_lookup(&request, now).await, Some(response) => service - .async_store(&request, response, now()) + .async_store(&request, response, now) .await .map(|_| None), } @@ -68,106 +143,45 @@ impl SemanticBody { let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; Ok(ExecutionStep::Await(awaitable.unbind())) } + + fn resume_py( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + match (&self.phase, result) { + (Phase::Start, None) => self.start(py), + (Phase::AwaitingEmbedding, Some(result)) => self.embedded(py, result), + (Phase::AwaitingBackend, Some(Err(error))) => Err(error), + (Phase::AwaitingBackend, Some(Ok(value))) => { + let more = matches!( + &self.operation, + SemanticOperation::StoreBatch(queue) if !queue.is_empty() + ); + if more { + self.phase = Phase::Start; + return self.start(py); + } + Ok(ExecutionStep::Return(value)) + } + _ => Err(PyRuntimeError::new_err( + "invalid semantic cache execution state", + )), + } + } } -impl ExecutionBody for SemanticBody { - fn resume(&mut self, mut result: Option>>) -> PyResult { - Python::attach(|py| { - loop { - match self.phase { - Phase::Start => { - if result.is_some() { - return Err(PyRuntimeError::new_err( - "semantic execution received a result before starting", - )); - } - if self.pending.is_none() { - match &mut self.operation { - SemanticOperation::Lookup(request) => { - self.pending = Some((request.clone(), None)); - } - SemanticOperation::Store(request, response) => { - let response = std::mem::replace(response, Value::Null); - self.pending = Some((request.clone(), Some(response))); - } - SemanticOperation::StoreBatch(queue) => { - let Some((request, response)) = queue.pop_front() else { - return Ok(ExecutionStep::Return(py.None())); - }; - self.pending = Some((request, Some(response))); - } - } - } - let (request, _) = self.pending.as_ref().ok_or_else(|| { - PyRuntimeError::new_err("semantic execution has no pending operation") - })?; - let semantic = NativeResponseCache::semantic_request(request); - let Some(prompt) = prompt_from_context(&semantic.context) else { - return self.backend_step(py, Err(Error::Unavailable)); - }; - let embedder = self.service.semantic_embedder().ok_or_else(|| { - PyRuntimeError::new_err( - "semantic execution requires a redis-semantic backend", - ) - })?; - let coroutine = embedder.async_embedding_coroutine( - py, - &prompt, - semantic.context.metadata.as_ref(), - )?; - self.phase = Phase::AwaitingEmbedding; - return Ok(ExecutionStep::Await(coroutine)); - } - Phase::AwaitingEmbedding => { - let result = result.take().ok_or_else(|| { - PyRuntimeError::new_err( - "semantic execution expected an embedding result", - ) - })?; - let seed = match result { - Ok(value) => PythonEmbedder::extract(value.into_bound(py)) - .map_err(|_| Error::Unavailable), - Err(error) => { - if !error.is_instance_of::(py) { - return Err(error); - } - Err(Error::Unavailable) - } - }; - return self.backend_step(py, seed); - } - Phase::AwaitingBackend => { - let result = result.take().ok_or_else(|| { - PyRuntimeError::new_err("semantic execution expected a backend result") - })?; - let value = match result { - Ok(value) => value, - Err(error) => return Err(error), - }; - let more = matches!( - &self.operation, - SemanticOperation::StoreBatch(queue) if !queue.is_empty() - ); - if more { - self.phase = Phase::Start; - continue; - } - return Ok(ExecutionStep::Return(value)); - } - } - } - }) +impl ExecutionBody for SemanticExecution { + fn resume(&mut self, result: Option>>) -> PyResult { + Python::attach(|py| self.resume_py(py, result)) } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - if let Some(embedder) = self.service.semantic_embedder() { - embedder.traverse(visit)?; - } - Ok(()) + self.embedder.traverse(visit) } } -pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult> { +pub(super) fn drive(py: Python<'_>, body: SemanticExecution) -> PyResult> { let execution = Py::new(py, Execution::new(body))?; py.import("litellm.rust_bridge.lifecycle")? .getattr("drive")? diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs deleted file mode 100644 index 24caf3374d6..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs +++ /dev/null @@ -1,249 +0,0 @@ -use std::{sync::Arc, time::Duration}; - -use litellm_cache::SemanticCacheContext; -use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; -use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context}; -use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; -use serde_json::Value; - -use super::{cache_error, embedder::PythonEmbedder}; - -pub(super) enum Op { - Lookup, - Store(Value), - StoreBatch(Vec), -} - -#[derive(Clone, Copy)] -enum State { - Start, - AwaitingEmbedding, - AwaitingStorage, - Done, -} - -pub(super) struct SemanticEmbedExecution { - backend: Arc>, - embedder: PythonEmbedder, - requests: Vec>, - op: Op, - now: Option, - prepared: Vec>>, - index: usize, - state: State, -} - -impl SemanticEmbedExecution { - pub(super) fn lookup( - backend: Arc>, - embedder: PythonEmbedder, - request: ResponseCacheRequest, - ) -> Self { - Self { - backend, - embedder, - requests: vec![request], - op: Op::Lookup, - now: None, - prepared: vec![None], - index: 0, - state: State::Start, - } - } - - pub(super) fn store( - backend: Arc>, - embedder: PythonEmbedder, - request: ResponseCacheRequest, - response: Value, - ) -> Self { - Self { - backend, - embedder, - requests: vec![request], - op: Op::Store(response), - now: None, - prepared: vec![None], - index: 0, - state: State::Start, - } - } - - pub(super) fn store_batch( - backend: Arc>, - embedder: PythonEmbedder, - requests: Vec>, - responses: Vec, - ) -> Self { - Self { - backend, - embedder, - prepared: vec![None; requests.len()], - requests, - op: Op::StoreBatch(responses), - now: None, - index: 0, - state: State::Start, - } - } - - fn start(&mut self, py: Python<'_>) -> PyResult { - if self.now.is_none() { - self.now = Some(super::request::now()); - } - while self.index < self.requests.len() { - let request = &self.requests[self.index]; - let enabled = match &self.op { - Op::Lookup => request.controls.reads(), - Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(), - }; - if !enabled { - self.index += 1; - continue; - } - let Some(prompt) = prompt_from_context(&request.context) else { - self.index += 1; - continue; - }; - let metadata = request.context.metadata.clone(); - let awaitable = self - .embedder - .async_embed_awaitable(py, &prompt, &metadata)?; - self.state = State::AwaitingEmbedding; - return Ok(ExecutionStep::Await(awaitable.unbind())); - } - self.state = State::AwaitingStorage; - self.storage_step(py) - } - - fn storage_step(&self, py: Python<'_>) -> PyResult { - let requests = self.requests.clone(); - let prepared = self.prepared.clone(); - let backend = Arc::clone(&self.backend); - let now = self - .now - .ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?; - let awaitable = match &self.op { - Op::Lookup => { - let Some(request) = requests.into_iter().next() else { - return Err(PyRuntimeError::new_err( - "semantic lookup requires one request", - )); - }; - match prepared.into_iter().next().flatten() { - Some(values) => { - let backend = backend.with_embedder(PreparedEmbedding(values)); - let cache = Arc::new(ResponseCache::new(Arc::new(backend))); - run_async( - py, - async move { cache.async_lookup(&request, now).await }, - cache_error, - )? - } - None => { - let cache = Arc::new(ResponseCache::new(backend)); - run_async( - py, - async move { cache.async_lookup(&request, now).await }, - cache_error, - )? - } - } - } - Op::Store(response) => { - let Some(request) = requests.into_iter().next() else { - return Err(PyRuntimeError::new_err( - "semantic store requires one request", - )); - }; - let response = response.clone(); - match prepared.into_iter().next().flatten() { - Some(values) => { - let backend = backend.with_embedder(PreparedEmbedding(values)); - let cache = Arc::new(ResponseCache::new(Arc::new(backend))); - run_async( - py, - async move { cache.async_store(&request, response, now).await }, - cache_error, - )? - } - None => { - let cache = Arc::new(ResponseCache::new(backend)); - run_async( - py, - async move { cache.async_store(&request, response, now).await }, - cache_error, - )? - } - } - } - Op::StoreBatch(responses) => { - let responses = responses.clone(); - run_async( - py, - async move { - for ((request, response), prepared) in - requests.into_iter().zip(responses).zip(prepared) - { - let Some(values) = prepared else { - continue; - }; - let backend = backend.with_embedder(PreparedEmbedding(values)); - let cache = ResponseCache::new(Arc::new(backend)); - cache.async_store(&request, response, now).await?; - } - Ok(()) - }, - cache_error, - )? - } - }; - Ok(ExecutionStep::Await(awaitable.unbind())) - } - - fn resume_py( - &mut self, - py: Python<'_>, - result: Option>>, - ) -> PyResult { - match (self.state, result) { - (State::Start, None) => self.start(py), - (State::AwaitingEmbedding, Some(Ok(value))) => { - let values = value.bind(py).extract::>()?; - self.prepared[self.index] = - Some(values.into_iter().map(|value| value as f32).collect()); - self.index += 1; - self.start(py) - } - (State::AwaitingStorage, Some(Ok(value))) => { - self.state = State::Done; - Ok(ExecutionStep::Return(value)) - } - (_, Some(Err(error))) => Err(error), - _ => Err(PyRuntimeError::new_err( - "invalid semantic cache execution state", - )), - } - } -} - -impl ExecutionBody for SemanticEmbedExecution { - fn resume(&mut self, result: Option>>) -> PyResult { - Python::attach(|py| self.resume_py(py, result)) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.embedder.traverse(visit) - } -} - -pub(super) fn drive_semantic<'py>( - py: Python<'py>, - body: SemanticEmbedExecution, -) -> PyResult> { - let execution = Py::new(py, Execution::new(body))?; - py.import("litellm.rust_bridge.lifecycle")? - .getattr("drive")? - .call1((execution,)) -} diff --git a/litellm-rust/crates/python-bridge/src/coercion.rs b/litellm-rust/crates/python-bridge/src/coercion.rs index bb5b8b2d454..1b0b073b3d1 100644 --- a/litellm-rust/crates/python-bridge/src/coercion.rs +++ b/litellm-rust/crates/python-bridge/src/coercion.rs @@ -1,7 +1,4 @@ -use std::collections::BTreeSet; - use litellm_core_utils::serde_compat::parse_str_bool; -use litellm_http::SslVerify; use pyo3::{ exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, prelude::*, @@ -33,36 +30,52 @@ impl From for PyErr { } } -pub(crate) struct Truthy(pub bool); -pub(crate) struct ExactTrue(pub bool); -pub(crate) struct StrBool(pub Option); -pub(crate) struct OptionalStrictString(pub Option); -pub(crate) struct FalsyOptionalString(pub Option); -pub(crate) struct TuningString(pub Option); -pub(crate) struct StringCollection(pub Vec); -pub(crate) struct SslVerifyInput(pub Option); +pub(crate) struct FieldSpec { + name: &'static str, + decode: fn(&Field<'_>) -> Result, +} + +impl FieldSpec { + pub(crate) const fn new( + name: &'static str, + decode: fn(&Field<'_>) -> Result, + ) -> Self { + Self { name, decode } + } + + pub(crate) fn read( + &self, + snapshot: &Bound<'_, PyAny>, + group: &'static str, + ) -> Result { + (self.decode)(&Field::read(snapshot, group, self.name)?) + } +} pub(crate) struct Field<'py> { - path: &'static str, + group: &'static str, + name: &'static str, value: Bound<'py, PyAny>, } impl<'py> Field<'py> { - pub(crate) fn new(path: &'static str, value: Bound<'py, PyAny>) -> Self { - Self { path, value } + pub(crate) fn new(group: &'static str, name: &'static str, value: Bound<'py, PyAny>) -> Self { + Self { group, name, value } } + /// Reads `snapshot.`, distinguishing a field the accessor never declared from a + /// descriptor that raised `AttributeError`. pub(crate) fn read( snapshot: &Bound<'py, PyAny>, - path: &'static str, + group: &'static str, + name: &'static str, ) -> Result { - let name = path.rsplit('.').next().unwrap_or(path); match snapshot.getattr(name) { - Ok(value) => Ok(Self::new(path, value)), + Ok(value) => Ok(Self::new(group, name, value)), Err(error) if error.is_instance_of::(snapshot.py()) => { match Self::missing_field(snapshot, name) { Ok(true) => Err(ProjectionError::InternalSchemaFailure(format!( - "{path}: missing snapshot field" + "{group}.{name}: missing snapshot field" ))), _ => Err(error.into()), } @@ -84,27 +97,49 @@ impl<'py> Field<'py> { && getter.is(object.getattr("__getattribute__")?)) } - fn expected(&self, expected: &'static str) -> Result { + pub(crate) fn path(&self) -> String { + format!("{}.{}", self.group, self.name) + } + + /// A member of this field's collection, reported under the same path. + pub(crate) fn member(&self, value: Bound<'py, PyAny>) -> Self { + Self::new(self.group, self.name, value) + } + + pub(crate) fn expected(&self, expected: &str) -> Result { Ok(format!( "{}: expected {expected}, got {}", - self.path, + self.path(), self.value.get_type().name()? )) } - fn invalid(&self, expected: &'static str) -> ProjectionError { + pub(crate) fn invalid(&self, expected: &str) -> ProjectionError { match self.expected(expected) { Ok(message) => ProjectionError::InvalidConfiguration(message), Err(error) => error, } } - pub(crate) fn truthy(&self) -> Result { - Ok(Truthy(self.value.is_truthy()?)) + pub(crate) fn value(&self) -> &Bound<'py, PyAny> { + &self.value } - pub(crate) fn exact_true(&self) -> ExactTrue { - ExactTrue(self.value.is(PyBool::new(self.value.py(), true))) + pub(crate) fn truthy(&self) -> Result { + Ok(self.value.is_truthy()?) + } + + pub(crate) fn exact_true(&self) -> bool { + self.value.is(PyBool::new(self.value.py(), true)) + } + + pub(crate) fn schema_bool(&self) -> Result { + if !self.value.is_instance_of::() { + return Err(ProjectionError::InternalSchemaFailure( + self.expected("a Boolean")?, + )); + } + Ok(self.exact_true()) } pub(crate) fn strict_string(&self) -> Result { @@ -124,108 +159,366 @@ impl<'py> Field<'py> { self.strict_string() } - pub(crate) fn schema_bool(&self) -> Result { - if !self.value.is_instance_of::() { - return Err(ProjectionError::InternalSchemaFailure( - self.expected("a Boolean")?, - )); - } - Ok(self.exact_true().0) - } - - pub(crate) fn str_bool(&self) -> Result { + pub(crate) fn str_bool(&self) -> Result, ProjectionError> { if self.value.is_none() { - return Ok(StrBool(None)); + return Ok(None); } - Ok(StrBool(parse_str_bool(&self.strict_string()?))) + Ok(parse_str_bool(&self.strict_string()?)) } - pub(crate) fn optional_strict_string(&self) -> Result { + pub(crate) fn optional_strict_string(&self) -> Result, ProjectionError> { if self.value.is_none() { - return Ok(OptionalStrictString(None)); + return Ok(None); } - self.strict_string().map(Some).map(OptionalStrictString) + self.strict_string().map(Some) } - pub(crate) fn falsy_optional_string(&self) -> Result { - if !self.truthy()?.0 { - return Ok(FalsyOptionalString(None)); + pub(crate) fn falsy_optional_string(&self) -> Result, ProjectionError> { + if !self.truthy()? { + return Ok(None); } - self.strict_string().map(Some).map(FalsyOptionalString) + self.strict_string().map(Some) } - pub(crate) fn tuning_string(&self) -> Result { - if !self.truthy()?.0 || !self.value.is_instance_of::() { - return Ok(TuningString(None)); + pub(crate) fn tuning_string(&self) -> Result, ProjectionError> { + if !self.truthy()? || !self.value.is_instance_of::() { + return Ok(None); } - self.strict_string().map(Some).map(TuningString) + self.strict_string().map(Some) } - pub(crate) fn string_collection(&self) -> Result { - if !self.truthy()?.0 { - return Ok(StringCollection(Vec::new())); + pub(crate) fn string_collection(&self) -> Result, ProjectionError> { + if !self.truthy()? { + return Ok(Vec::new()); } if self.value.is_instance_of::() { - return self - .strict_string() - .map(|value| StringCollection(vec![value])); + return self.strict_string().map(|value| vec![value]); } - let values = self - .value + self.value .try_iter()? .filter_map(|item| { let member = match item { - Ok(value) => Self::new(self.path, value), + Ok(value) => self.member(value), Err(error) => return Some(Err(error.into())), }; match member.truthy() { - Ok(Truthy(false)) => None, - Ok(Truthy(true)) => Some(member.strict_string()), + Ok(false) => None, + Ok(true) => Some(member.strict_string()), Err(error) => Some(Err(error)), } }) - .collect::, ProjectionError>>()?; - Ok(StringCollection(values)) + .collect() } - pub(crate) fn host_collection(&self) -> Result { - let values = self - .string_collection()? - .0 - .into_iter() - .map(|host| litellm_http::media::normalize_host(&host)) - .collect::>(); - Ok(StringCollection(values.into_iter().collect())) - } - - pub(crate) fn ssl_verify(&self) -> Result { + pub(crate) fn optional_string_collection( + &self, + ) -> Result>, ProjectionError> { if self.value.is_none() { - return Ok(SslVerifyInput(None)); + return Ok(None); } - if self.value.is_instance_of::() { - return Ok(SslVerifyInput(Some(if self.exact_true().0 { - SslVerify::Enabled - } else { - SslVerify::Disabled - }))); - } - if self.value.is_instance_of::() { - let parsed = match self.str_bool()?.0 { - Some(true) => SslVerify::Enabled, - Some(false) => SslVerify::Disabled, - None => SslVerify::CaBundle(self.strict_string()?.into()), - }; - return Ok(SslVerifyInput(Some(parsed))); - } - let context = self.value.py().import("ssl")?.getattr("SSLContext")?; - if self.value.is_instance(&context)? { - return Err(ProjectionError::UnsupportedLiveObject(self.expected( - "a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported", - )?)); - } - Err(self.invalid("a Boolean, Boolean string, CA path, or None")) + self.string_collection().map(Some) + } + + pub(crate) fn python_binding(&self) -> Option> { + (!self.value.is_none()).then(|| self.value.clone().unbind()) } } #[cfg(test)] -mod tests; +mod tests { + use std::ffi::CString; + + use pyo3::{ + exceptions::{PyLookupError, PyRuntimeError, PyValueError}, + types::PyDict, + }; + use rstest::rstest; + + use super::*; + + fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { + py.eval(&CString::new(source).unwrap(), None, None).unwrap() + } + + #[rstest] + #[case("None", false, false)] + #[case("False", false, false)] + #[case("True", true, true)] + #[case("0", false, false)] + #[case("1", true, false)] + #[case("''", false, false)] + #[case("'false'", true, false)] + #[case("[]", false, false)] + #[case("[0]", true, false)] + #[case("{}", false, false)] + #[case("object()", true, false)] + fn boolean_operations_have_distinct_python_semantics( + #[case] source: &str, + #[case] truth: bool, + #[case] exact: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let value = evaluate(py, source); + let field = Field::new("test", "flag", value.clone()); + assert_eq!(field.truthy().unwrap(), truth); + assert_eq!(field.exact_true(), exact); + assert_eq!( + field.truthy().unwrap(), + py.import("builtins") + .unwrap() + .getattr("bool") + .unwrap() + .call1((value,)) + .unwrap() + .extract::() + .unwrap() + ); + }); + } + + #[rstest] + #[case("None", Ok(None), Ok(None), Ok(None))] + #[case("''", Ok(Some("")), Ok(None), Ok(None))] + #[case( + "' value '", + Ok(Some(" value ")), + Ok(Some(" value ")), + Ok(Some(" value ")) + )] + #[case("[]", Err(()), Ok(None), Ok(None))] + #[case("0", Err(()), Ok(None), Ok(None))] + #[case("1", Err(()), Err(()), Ok(None))] + #[case("object()", Err(()), Err(()), Ok(None))] + fn string_operations_do_not_conflate_absence_and_type_checks( + #[case] source: &str, + #[case] strict: Result, ()>, + #[case] fallback: Result, ()>, + #[case] tuning: Result, ()>, + ) { + Python::initialize(); + Python::attach(|py| { + let field = Field::new("test", "string", evaluate(py, source)); + let owned = + |expected: Result, ()>| expected.map(|value| value.map(str::to_owned)); + assert_eq!( + field.optional_strict_string().map_err(|_| ()), + owned(strict) + ); + assert_eq!( + field.falsy_optional_string().map_err(|_| ()), + owned(fallback) + ); + assert_eq!(field.tuning_string().map_err(|_| ()), owned(tuning)); + }); + } + + #[rstest] + #[case("None", None)] + #[case("' True '", Some(true))] + #[case("' fAlSe '", Some(false))] + #[case("'yes'", None)] + #[case("'1'", None)] + #[case("'unknown'", None)] + fn string_boolean_tokens_remain_separate_from_truthiness( + #[case] source: &str, + #[case] expected: Option, + ) { + Python::initialize(); + Python::attach(|py| { + assert_eq!( + Field::new("test", "flag", evaluate(py, source)) + .str_bool() + .unwrap(), + expected + ); + }); + } + + #[test] + fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +failure = LookupError('protocol failed') +cause = ValueError('cause') +context = RuntimeError('context') +def fail(): + try: + raise context + except RuntimeError: + raise failure from cause +class Bool: + def __bool__(self): return fail() +class Length: + def __len__(self): return fail() +class Iter: + def __iter__(self): return fail() +class Next: + def __iter__(self): return self + def __next__(self): return fail() +class Descriptor: + @property + def flag(self): return fail() +values = (Bool(), Length(), Iter(), Next(), [Bool()]) +descriptor = Descriptor() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let values = locals.get_item("values").unwrap().unwrap(); + for value in values.try_iter().unwrap() { + let error = Field::new("test", "flag", value.unwrap()) + .string_collection() + .err() + .unwrap(); + let error = PyErr::from(error); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + assert!(error.is_instance_of::(py)); + assert!(error.traceback(py).is_some()); + assert!( + error + .value(py) + .getattr("__cause__") + .unwrap() + .is(locals.get_item("cause").unwrap().unwrap()) + ); + assert!( + error + .value(py) + .getattr("__context__") + .unwrap() + .is(locals.get_item("context").unwrap().unwrap()) + ); + } + let error = Field::read( + &locals.get_item("descriptor").unwrap().unwrap(), + "test", + "flag", + ) + .err() + .unwrap(); + assert!( + PyErr::from(error) + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn identity_and_string_contents_do_not_invoke_unrelated_protocols() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +class Hostile: + def __bool__(self): raise AssertionError('bool called') + def __eq__(self, other): raise AssertionError('eq called') + def __str__(self): raise AssertionError('str called') +class Text(str): + def __str__(self): raise AssertionError('str called') + def strip(self): raise AssertionError('strip called') + def lower(self): raise AssertionError('lower called') +hostile = Hostile() +text = Text(' False ') +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let hostile = Field::new("test", "flag", locals.get_item("hostile").unwrap().unwrap()); + assert!(!hostile.exact_true()); + assert!(matches!( + hostile.strict_string(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + let text = Field::new("test", "flag", locals.get_item("text").unwrap().unwrap()); + assert_eq!(text.strict_string().unwrap(), " False "); + assert_eq!(text.str_bool().unwrap(), Some(false)); + }); + } + + #[test] + fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +failure = AttributeError('descriptor failed') +class Snapshot: + @property + def flag(self): raise failure +snapshot = Snapshot() +class Dynamic: + def __getattr__(self, name): raise failure +class Intercepted: + def __getattribute__(self, name): raise failure +dynamic = Dynamic() +intercepted = Intercepted() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let snapshot = locals.get_item("snapshot").unwrap().unwrap(); + let descriptor = PyErr::from(Field::read(&snapshot, "test", "flag").err().unwrap()); + assert!( + descriptor + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + for name in ["dynamic", "intercepted"] { + let value = locals.get_item(name).unwrap().unwrap(); + let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap()); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + } + let missing = PyErr::from(Field::read(&snapshot, "test", "missing").err().unwrap()); + assert!(missing.is_instance_of::(py)); + assert!(missing.to_string().contains("test.missing")); + }); + } + + #[test] + fn configuration_errors_name_fields_without_exposing_values() { + Python::initialize(); + Python::attach(|py| { + for source in [ + "{'secret': 'do-not-print'}", + "['host.test', {'secret': 'do-not-print'}]", + ] { + let field = Field::new("test", "setting", evaluate(py, source)); + let error = PyErr::from(field.falsy_optional_string().err().unwrap()); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("test.setting")); + assert!(!error.to_string().contains("do-not-print")); + } + let hosts = Field::new( + "url_policy", + "user_url_allowed_hosts", + evaluate(py, "['host.test', 1]"), + ); + assert!(matches!( + hosts.string_collection(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + assert!(matches!( + Field::new("test", "flag", evaluate(py, "1")).str_bool(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/coercion/tests.rs b/litellm-rust/crates/python-bridge/src/coercion/tests.rs deleted file mode 100644 index 5ed237c3c64..00000000000 --- a/litellm-rust/crates/python-bridge/src/coercion/tests.rs +++ /dev/null @@ -1,372 +0,0 @@ -use std::ffi::CString; - -use pyo3::{ - exceptions::{PyLookupError, PyRuntimeError, PyValueError}, - types::PyDict, -}; -use rstest::rstest; - -use super::*; - -fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { - py.eval(&CString::new(source).unwrap(), None, None).unwrap() -} - -#[rstest] -#[case("None", false, false)] -#[case("False", false, false)] -#[case("True", true, true)] -#[case("0", false, false)] -#[case("1", true, false)] -#[case("''", false, false)] -#[case("'false'", true, false)] -#[case("[]", false, false)] -#[case("[0]", true, false)] -#[case("{}", false, false)] -#[case("object()", true, false)] -fn boolean_operations_have_distinct_python_semantics( - #[case] source: &str, - #[case] truth: bool, - #[case] exact: bool, -) { - Python::initialize(); - Python::attach(|py| { - let value = evaluate(py, source); - let field = Field::new("test.flag", value.clone()); - assert_eq!(field.truthy().unwrap().0, truth); - assert_eq!(field.exact_true().0, exact); - assert_eq!( - field.truthy().unwrap().0, - py.import("builtins") - .unwrap() - .getattr("bool") - .unwrap() - .call1((value,)) - .unwrap() - .extract::() - .unwrap() - ); - }); -} - -#[rstest] -#[case("None", Ok(None), Ok(None), Ok(None))] -#[case("''", Ok(Some("")), Ok(None), Ok(None))] -#[case( - "' value '", - Ok(Some(" value ")), - Ok(Some(" value ")), - Ok(Some(" value ")) -)] -#[case("[]", Err(()), Ok(None), Ok(None))] -#[case("0", Err(()), Ok(None), Ok(None))] -#[case("1", Err(()), Err(()), Ok(None))] -#[case("object()", Err(()), Err(()), Ok(None))] -fn string_operations_do_not_conflate_absence_and_type_checks( - #[case] source: &str, - #[case] strict: Result, ()>, - #[case] fallback: Result, ()>, - #[case] tuning: Result, ()>, -) { - Python::initialize(); - Python::attach(|py| { - let field = Field::new("test.string", evaluate(py, source)); - let owned = - |expected: Result, ()>| expected.map(|value| value.map(str::to_owned)); - assert_eq!( - field - .optional_strict_string() - .map(|value| value.0) - .map_err(|_| ()), - owned(strict) - ); - assert_eq!( - field - .falsy_optional_string() - .map(|value| value.0) - .map_err(|_| ()), - owned(fallback) - ); - assert_eq!( - field.tuning_string().map(|value| value.0).map_err(|_| ()), - owned(tuning) - ); - }); -} - -#[rstest] -#[case("None", None)] -#[case("' True '", Some(true))] -#[case("' fAlSe '", Some(false))] -#[case("'yes'", None)] -#[case("'1'", None)] -#[case("'unknown'", None)] -fn string_boolean_tokens_remain_separate_from_truthiness( - #[case] source: &str, - #[case] expected: Option, -) { - Python::initialize(); - Python::attach(|py| { - assert_eq!( - Field::new("test.flag", evaluate(py, source)) - .str_bool() - .unwrap() - .0, - expected - ); - }); -} - -#[rstest] -#[case("'EXAMPLE.TEST.'", vec!["example.test"])] -#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])] -#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])] -#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])] -#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])] -#[case("None", vec![])] -#[case("False", vec![])] -fn host_collection_is_owned_normalized_and_deterministic( - #[case] source: &str, - #[case] expected: Vec<&str>, -) { - Python::initialize(); - Python::attach(|py| { - assert_eq!( - Field::new("url_policy.user_url_allowed_hosts", evaluate(py, source)) - .host_collection() - .unwrap() - .0, - expected - ); - }); -} - -#[test] -fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -failure = LookupError('protocol failed') -cause = ValueError('cause') -context = RuntimeError('context') -def fail(): - try: - raise context - except RuntimeError: - raise failure from cause -class Bool: - def __bool__(self): return fail() -class Length: - def __len__(self): return fail() -class Iter: - def __iter__(self): return fail() -class Next: - def __iter__(self): return self - def __next__(self): return fail() -class Descriptor: - @property - def flag(self): return fail() -values = (Bool(), Length(), Iter(), Next(), [Bool()]) -descriptor = Descriptor() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let values = locals.get_item("values").unwrap().unwrap(); - for value in values.try_iter().unwrap() { - let error = Field::new("test.flag", value.unwrap()) - .host_collection() - .err() - .unwrap(); - let error = PyErr::from(error); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - assert!(error.is_instance_of::(py)); - assert!(error.traceback(py).is_some()); - assert!( - error - .value(py) - .getattr("__cause__") - .unwrap() - .is(locals.get_item("cause").unwrap().unwrap()) - ); - assert!( - error - .value(py) - .getattr("__context__") - .unwrap() - .is(locals.get_item("context").unwrap().unwrap()) - ); - } - let error = Field::read( - &locals.get_item("descriptor").unwrap().unwrap(), - "test.flag", - ) - .err() - .unwrap(); - assert!( - PyErr::from(error) - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - }); -} - -#[test] -fn identity_and_string_contents_do_not_invoke_unrelated_protocols() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -class Hostile: - def __bool__(self): raise AssertionError('bool called') - def __eq__(self, other): raise AssertionError('eq called') - def __str__(self): raise AssertionError('str called') -class Text(str): - def __str__(self): raise AssertionError('str called') - def strip(self): raise AssertionError('strip called') - def lower(self): raise AssertionError('lower called') -hostile = Hostile() -text = Text(' False ') -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let hostile = Field::new("test.flag", locals.get_item("hostile").unwrap().unwrap()); - assert!(!hostile.exact_true().0); - assert!(matches!( - hostile.strict_string(), - Err(ProjectionError::InvalidConfiguration(_)) - )); - let text = Field::new("test.flag", locals.get_item("text").unwrap().unwrap()); - assert_eq!(text.strict_string().unwrap(), " False "); - assert_eq!(text.str_bool().unwrap().0, Some(false)); - }); -} - -#[test] -fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -failure = AttributeError('descriptor failed') -class Snapshot: - @property - def flag(self): raise failure -snapshot = Snapshot() -class Dynamic: - def __getattr__(self, name): raise failure -class Intercepted: - def __getattribute__(self, name): raise failure -dynamic = Dynamic() -intercepted = Intercepted() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let snapshot = locals.get_item("snapshot").unwrap().unwrap(); - let descriptor = PyErr::from(Field::read(&snapshot, "test.flag").err().unwrap()); - assert!( - descriptor - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - for name in ["dynamic", "intercepted"] { - let value = locals.get_item(name).unwrap().unwrap(); - let error = PyErr::from(Field::read(&value, "test.flag").err().unwrap()); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - } - let missing = PyErr::from(Field::read(&snapshot, "test.missing").err().unwrap()); - assert!(missing.is_instance_of::(py)); - assert!(missing.to_string().contains("test.missing")); - }); -} - -#[test] -fn configuration_errors_name_fields_without_exposing_values() { - Python::initialize(); - Python::attach(|py| { - for source in [ - "{'secret': 'do-not-print'}", - "['host.test', {'secret': 'do-not-print'}]", - ] { - let field = Field::new("test.setting", evaluate(py, source)); - let error = PyErr::from(field.falsy_optional_string().err().unwrap()); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("test.setting")); - assert!(!error.to_string().contains("do-not-print")); - } - let hosts = Field::new( - "url_policy.user_url_allowed_hosts", - evaluate(py, "['host.test', 1]"), - ); - assert!(matches!( - hosts.host_collection(), - Err(ProjectionError::InvalidConfiguration(_)) - )); - assert!(matches!( - Field::new("test.flag", evaluate(py, "1")).str_bool(), - Err(ProjectionError::InvalidConfiguration(_)) - )); - }); -} - -#[test] -fn projection_releases_the_source_collection() { - Python::initialize(); - Python::attach(|py| { - let source = evaluate(py, "['A.test']"); - let projected = Field::new("test.hosts", source.clone()) - .host_collection() - .unwrap() - .0; - source.call_method1("append", ("b.test",)).unwrap(); - assert_eq!(projected, ["a.test"]); - assert_eq!( - Field::new("test.hosts", source) - .host_collection() - .unwrap() - .0, - ["a.test", "b.test"] - ); - }); -} - -#[rstest] -#[case("True", Some(true))] -#[case("False", Some(false))] -#[case("1", None)] -#[case("None", None)] -#[case("[]", None)] -fn accessor_booleans_are_strict_schema_values( - #[case] source: &str, - #[case] expected: Option, -) { - Python::initialize(); - Python::attach(|py| { - let result = Field::new("secret_manager.readable", evaluate(py, source)).schema_bool(); - match expected { - Some(expected) => assert_eq!(result.unwrap(), expected), - None => { - let error = PyErr::from(result.unwrap_err()); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("secret_manager.readable")); - } - } - }); -} diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 596a89a73d7..2515b409c54 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashSet, + collections::{BTreeSet, HashSet}, path::{Path, PathBuf}, sync::{Arc, LazyLock, Mutex, PoisonError}, }; @@ -10,9 +10,75 @@ use litellm_http::{ TlsSource, Unsupported, media::{PublicDnsResolver, UrlPolicy}, }; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{ + exceptions::PyValueError, + prelude::*, + types::{PyBool, PyDict, PyString}, +}; -use crate::{coercion::Field, python_settings::PythonSettings}; +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const SSL_VERIFY: FieldSpec> = FieldSpec::new("ssl_verify", decode_ssl_verify); +const SSL_CERTIFICATE: FieldSpec> = + FieldSpec::new("ssl_certificate", |field| field.optional_strict_string()); +const SSL_SECURITY_LEVEL: FieldSpec> = + FieldSpec::new("ssl_security_level", |field| field.tuning_string()); +const SSL_ECDH_CURVE: FieldSpec> = + FieldSpec::new("ssl_ecdh_curve", |field| field.tuning_string()); +const FORCE_IPV4: FieldSpec = FieldSpec::new("force_ipv4", |field| field.truthy()); +const HTTP2: FieldSpec = FieldSpec::new("http2", |field| Ok(field.exact_true())); +const AIOHTTP_TRUST_ENV: FieldSpec = + FieldSpec::new("aiohttp_trust_env", |field| field.truthy()); +const DISABLE_AIOHTTP_TRUST_ENV: FieldSpec = + FieldSpec::new("disable_aiohttp_trust_env", |field| field.truthy()); +const DISABLE_AIOHTTP_TRANSPORT: FieldSpec = + FieldSpec::new("disable_aiohttp_transport", |field| Ok(field.exact_true())); +const USER_AGENT: FieldSpec = FieldSpec::new("user_agent", |field| field.schema_string()); +const USER_URL_VALIDATION: FieldSpec = + FieldSpec::new("user_url_validation", |field| field.truthy()); +const USER_URL_ALLOWED_HOSTS: FieldSpec> = + FieldSpec::new("user_url_allowed_hosts", decode_hosts); + +fn decode_hosts(field: &Field<'_>) -> Result, ProjectionError> { + Ok(field + .string_collection()? + .into_iter() + .map(|host| litellm_http::media::normalize_host(&host)) + .collect::>() + .into_iter() + .collect()) +} + +fn decode_ssl_verify(field: &Field<'_>) -> Result, ProjectionError> { + let value = field.value(); + if value.is_none() { + return Ok(None); + } + if value.is_instance_of::() { + return Ok(Some(if field.exact_true() { + SslVerify::Enabled + } else { + SslVerify::Disabled + })); + } + if value.is_instance_of::() { + return Ok(Some(match field.str_bool()? { + Some(true) => SslVerify::Enabled, + Some(false) => SslVerify::Disabled, + None => SslVerify::CaBundle(field.strict_string()?.into()), + })); + } + let context = value.py().import("ssl")?.getattr("SSLContext")?; + if value.is_instance(&context)? { + return Err(ProjectionError::UnsupportedLiveObject(field.expected( + "a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported", + )?)); + } + Err(field.invalid("a Boolean, Boolean string, CA path, or None")) +} static POOL: LazyLock = LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); @@ -80,20 +146,20 @@ pub(crate) fn url_policy(py: Python<'_>) -> PyResult { project_url_policy(&PythonSettings::UrlPolicy.read(py)?) } -fn project_url_policy(value: &Bound<'_, PyAny>) -> PyResult { +fn project_url_policy(snapshot: &Snapshot<'_>) -> PyResult { Ok(UrlPolicy { - validate: Field::read(value, "url_policy.user_url_validation")? - .truthy()? - .0, - allowed_hosts: Field::read(value, "url_policy.user_url_allowed_hosts")? - .host_collection()? - .0, + validate: snapshot.read(&USER_URL_VALIDATION)?, + allowed_hosts: snapshot.read(&USER_URL_ALLOWED_HOSTS)?, }) } fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult> { match kwargs.get_item("ssl_verify")? { - Some(value) => Ok(Field::new("request.ssl_verify", value).ssl_verify()?.0), + Some(value) => Ok(decode_ssl_verify(&Field::new( + "request", + "ssl_verify", + value, + ))?), None => Ok(None), } } @@ -106,39 +172,18 @@ fn for_call(call_ssl_verify: Option, asynchronous: bool) -> HttpSetti } } -fn configured(value: &Bound<'_, PyAny>) -> PyResult { +fn configured(snapshot: &Snapshot<'_>) -> PyResult { Ok(HttpSettingsLayer { - ssl_verify: Field::read(value, "http_settings.ssl_verify")? - .ssl_verify()? - .0, - ssl_certificate: Field::read(value, "http_settings.ssl_certificate")? - .optional_strict_string()? - .0 - .map(PathBuf::from), - ssl_security_level: Field::read(value, "http_settings.ssl_security_level")? - .tuning_string()? - .0, - ssl_ecdh_curve: Field::read(value, "http_settings.ssl_ecdh_curve")? - .tuning_string()? - .0, - force_ipv4: Some(Field::read(value, "http_settings.force_ipv4")?.truthy()?.0), - http2: Some(Field::read(value, "http_settings.http2")?.exact_true().0), - aiohttp_trust_env: Some( - Field::read(value, "http_settings.aiohttp_trust_env")? - .truthy()? - .0, - ), - disable_aiohttp_trust_env: Some( - Field::read(value, "http_settings.disable_aiohttp_trust_env")? - .truthy()? - .0, - ), - disable_aiohttp_transport: Some( - Field::read(value, "http_settings.disable_aiohttp_transport")? - .exact_true() - .0, - ), - user_agent: Some(Field::read(value, "http_settings.user_agent")?.schema_string()?), + ssl_verify: snapshot.read(&SSL_VERIFY)?, + ssl_certificate: snapshot.read(&SSL_CERTIFICATE)?.map(PathBuf::from), + ssl_security_level: snapshot.read(&SSL_SECURITY_LEVEL)?, + ssl_ecdh_curve: snapshot.read(&SSL_ECDH_CURVE)?, + force_ipv4: Some(snapshot.read(&FORCE_IPV4)?), + http2: Some(snapshot.read(&HTTP2)?), + aiohttp_trust_env: Some(snapshot.read(&AIOHTTP_TRUST_ENV)?), + disable_aiohttp_trust_env: Some(snapshot.read(&DISABLE_AIOHTTP_TRUST_ENV)?), + disable_aiohttp_transport: Some(snapshot.read(&DISABLE_AIOHTTP_TRANSPORT)?), + user_agent: Some(snapshot.read(&USER_AGENT)?), ..HttpSettingsLayer::default() }) } @@ -150,12 +195,15 @@ mod tests { use rstest::rstest; use super::*; - use crate::python_settings::CONTRACT; - fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { + fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { + py.eval(&std::ffi::CString::new(source).unwrap(), None, None) + .unwrap() + } + + fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Snapshot<'py> { let source = format!( " -import json import types defaults = dict( ssl_verify=True, @@ -170,14 +218,13 @@ defaults = dict( user_agent='litellm/test', ) defaults.update(dict({overrides})) -settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']['fields']}}) +settings = types.SimpleNamespace(**defaults) " ); let locals = PyDict::new(py); - locals.set_item("contract", CONTRACT).unwrap(); let source = std::ffi::CString::new(source).unwrap(); py.run(&source, Some(&locals), Some(&locals)).unwrap(); - locals.get_item("settings").unwrap().unwrap() + PythonSettings::Http.snapshot(locals.get_item("settings").unwrap().unwrap()) } #[test] @@ -395,7 +442,7 @@ user_agent='litellm/9.9.9', Python::attach(|py| { let value = py.eval(c"__import__('types').SimpleNamespace(user_url_validation=[], user_url_allowed_hosts=['B.test', 'a.test.', 'b.test'])", None, None).unwrap(); assert_eq!( - project_url_policy(&value).unwrap(), + project_url_policy(&PythonSettings::UrlPolicy.snapshot(value)).unwrap(), UrlPolicy { validate: false, allowed_hosts: vec!["a.test".into(), "b.test".into()], @@ -419,4 +466,44 @@ user_agent='litellm/9.9.9', let settings = HttpSettings::from_layers([for_call(None, asynchronous), opted_out]); assert_eq!(settings.trust_proxy_env, expected); } + #[rstest] + #[case("'EXAMPLE.TEST.'", vec!["example.test"])] + #[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])] + #[case("('B.test', 'a.test')", vec!["a.test", "b.test"])] + #[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])] + #[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])] + #[case("None", vec![])] + #[case("False", vec![])] + fn host_collection_is_owned_normalized_and_deterministic( + #[case] source: &str, + #[case] expected: Vec<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + assert_eq!( + decode_hosts(&Field::new( + "url_policy", + "user_url_allowed_hosts", + evaluate(py, source) + )) + .unwrap(), + expected + ); + }); + } + + #[test] + fn projection_releases_the_source_collection() { + Python::initialize(); + Python::attach(|py| { + let source = evaluate(py, "['A.test']"); + let projected = decode_hosts(&Field::new("test", "hosts", source.clone())).unwrap(); + source.call_method1("append", ("b.test",)).unwrap(); + assert_eq!(projected, ["a.test"]); + assert_eq!( + decode_hosts(&Field::new("test", "hosts", source)).unwrap(), + ["a.test", "b.test"] + ); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f13a3ad433f..ed9bc90f650 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -7,6 +7,11 @@ mod http; mod marshal; mod python_settings; mod routes; +#[allow( + dead_code, + reason = "secret-manager foundations await rollout activation" +)] +mod secrets; mod token_counter; #[pymodule(gil_used = true)] @@ -43,7 +48,7 @@ mod _native { let dict = module.dict(); dict.set_item("_CacheTestHandle", py.get_type::())?; dict.set_item("_CacheTestResolver", py.get_type::())?; - dict.set_item("_CacheTestBinding", py.get_type::()) + dict.set_item("_ResponseCacheRuntime", py.get_type::()) } } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index bdc6d14356d..111ac3bc259 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,5 +1,7 @@ use pyo3::prelude::*; +use crate::coercion::{FieldSpec, ProjectionError}; + const MODULE: &str = "litellm.rust_bridge.settings"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -8,28 +10,39 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, + SecretManagerBinding, +} + +pub(crate) struct Snapshot<'py> { + group: PythonSettings, + value: Bound<'py, PyAny>, +} + +impl Snapshot<'_> { + pub(crate) fn read(&self, spec: &FieldSpec) -> Result { + spec.read(&self.value, self.group.name()) + } } impl PythonSettings { - #[cfg(test)] - pub(crate) const ALL: [Self; 4] = [ - Self::Http, - Self::UrlPolicy, - Self::ProviderDefaults, - Self::SecretManager, - ]; - pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", + Self::SecretManagerBinding => "secret_manager_binding", } } - pub(crate) fn read(self, py: Python<'_>) -> PyResult> { - py.import(MODULE)?.getattr(self.name())?.call0() + pub(crate) fn read(self, py: Python<'_>) -> PyResult> { + let value = py.import(MODULE)?.getattr(self.name())?.call0()?; + Ok(Snapshot { group: self, value }) + } + + #[cfg(test)] + pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { + Snapshot { group: self, value } } pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> { @@ -38,209 +51,98 @@ impl PythonSettings { } } -#[cfg(test)] -pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); - #[cfg(test)] mod tests { - use super::{CONTRACT, PythonSettings}; - use pyo3::prelude::*; - use serde_json::{Value, json}; + use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; - struct SettingSpec { - group: &'static str, - name: &'static str, - adapter: &'static str, - precedence: &'static str, - sensitive: bool, - shapes: &'static [&'static str], - unsupported_live: Option<&'static str>, - } - - const SETTINGS: &[SettingSpec] = &[ - SettingSpec { - group: "http_settings", - name: "ssl_verify", - adapter: "SslVerifyInput", - precedence: "module_global", - sensitive: false, - shapes: &["none", "bool", "str"], - unsupported_live: Some("configuration_error"), - }, - SettingSpec { - group: "http_settings", - name: "ssl_certificate", - adapter: "OptionalStrictString", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "ssl_security_level", - adapter: "TuningString", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "ssl_ecdh_curve", - adapter: "TuningString", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "force_ipv4", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "http2", - adapter: "ExactTrue", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "aiohttp_trust_env", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "disable_aiohttp_trust_env", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "disable_aiohttp_transport", - adapter: "ExactTrue", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "user_agent", - adapter: "StrictString", - precedence: "accessor", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "url_policy", - name: "user_url_validation", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "url_policy", - name: "user_url_allowed_hosts", - adapter: "HostCollection", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "provider_defaults", - name: "vertex_project", - adapter: "FalsyOptionalString", - precedence: "module_global", - sensitive: true, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "provider_defaults", - name: "vertex_location", - adapter: "FalsyOptionalString", - precedence: "module_global", - sensitive: true, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "provider_defaults", - name: "enable_azure_ad_token_refresh", - adapter: "ExactTrue", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "secret_manager", - name: "readable", - adapter: "StrictBool", - precedence: "accessor", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - ]; + use super::PythonSettings; + use crate::coercion::FieldSpec; #[test] - fn settings_manifest_matches_the_semantic_contract() { - pyo3::Python::initialize(); - let manifest: Value = pyo3::Python::attach(|py| { - let value = py - .import("json") - .unwrap() - .call_method1("loads", (CONTRACT,)) - .unwrap(); - litellm_host_python::from_py(&value).unwrap() + fn declarations_select_the_decoder_and_read_only_the_requested_field() { + const TRUTHY: FieldSpec = FieldSpec::new("flag", |field| field.truthy()); + const EXACT: FieldSpec = FieldSpec::new("flag", |field| Ok(field.exact_true())); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +reads = [] +class Settings: + value = 1 + @property + def flag(self): + reads.append('flag') + return self.value + @property + def unrelated(self): + raise AssertionError('unrequested field') +settings = Settings() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let value = locals.get_item("settings").unwrap().unwrap(); + let snapshot = PythonSettings::Http.snapshot(value.clone()); + assert!(snapshot.read(&TRUTHY).unwrap()); + assert!(!snapshot.read(&EXACT).unwrap()); + value.setattr("value", true).unwrap(); + assert!(snapshot.read(&EXACT).unwrap()); + assert_eq!( + locals + .get_item("reads") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + ["flag", "flag", "flag"] + ); + }); + } + + #[test] + fn declared_reads_preserve_descriptor_and_decoder_failures_and_name_missing_fields() { + const FLAG: FieldSpec = FieldSpec::new("flag", |field| field.truthy()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +from types import SimpleNamespace +failure = AttributeError('read failed') +class Descriptor: + @property + def flag(self): raise failure +class Truth: + def __bool__(self): raise failure +values = (Descriptor(), SimpleNamespace(flag=Truth())) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let failure = locals.get_item("failure").unwrap().unwrap(); + for value in locals + .get_item("values") + .unwrap() + .unwrap() + .try_iter() + .unwrap() + { + let snapshot = PythonSettings::Http.snapshot(value.unwrap()); + let error = PyErr::from(snapshot.read(&FLAG).unwrap_err()); + assert!(error.value(py).is(&failure)); + assert!(error.traceback(py).is_some()); + } + let missing = PythonSettings::Http.snapshot(py.eval(c"object()", None, None).unwrap()); + let error = PyErr::from(missing.read(&FLAG).unwrap_err()); + assert!(error.is_instance_of::(py)); + assert!( + error + .to_string() + .contains("http_settings.flag: missing snapshot field") + ); }); - let expected: serde_json::Map = PythonSettings::ALL - .into_iter() - .map(|group| { - let fields: serde_json::Map = SETTINGS - .iter() - .filter(|spec| spec.group == group.name()) - .map(|spec| { - ( - spec.name.to_owned(), - json!({ - "adapter": spec.adapter, - "required": true, - "precedence": spec.precedence, - "sensitive": spec.sensitive, - "shapes": spec.shapes, - "unsupported_live": spec.unsupported_live, - }), - ) - }) - .collect(); - ( - group.name().to_owned(), - json!({"version": 1, "fields": fields}), - ) - }) - .collect(); - assert_eq!(manifest, Value::Object(expected)); } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 77c8d5d6641..325377e5285 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -130,6 +130,11 @@ impl RouteHost for OcrRouteHost { } fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + if let Error::Secret(source) = &error + && let Some(original) = crate::secrets::callback::python_error(py, source) + { + return Ok(original); + } Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index d0b13e5056a..2dca6da66cd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -10,16 +10,33 @@ use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_core_utils::settings::ProcessEnvironment; -use litellm_llms::base_llm::ocr::{ - handler::OcrClient, - settings::{OcrSettings, Secrets}, +use litellm_llms::base_llm::{ + inference::secrets::{EnvironmentSecrets, SecretSource}, + ocr::{handler::OcrClient, settings::OcrSettings}, }; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, }; -use crate::{coercion::Field, errors::RustBridgeDeclined, http, python_settings::PythonSettings}; +use crate::{ + coercion::FieldSpec, + errors::RustBridgeDeclined, + http, + python_settings::{PythonSettings, Snapshot}, +}; + +const SECRET_MANAGER_READABLE: FieldSpec = + FieldSpec::new("readable", |field| field.schema_bool()); + +const VERTEX_PROJECT: FieldSpec> = + FieldSpec::new("vertex_project", |field| field.falsy_optional_string()); +const VERTEX_LOCATION: FieldSpec> = + FieldSpec::new("vertex_location", |field| field.falsy_optional_string()); +const ENABLE_AZURE_AD_TOKEN_REFRESH: FieldSpec = + FieldSpec::new("enable_azure_ad_token_refresh", |field| { + Ok(field.exact_true()) + }); const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -62,33 +79,24 @@ fn run_ocr( ) } -fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult { - if Field::read(secret_manager, "secret_manager.readable")?.schema_bool()? { +fn process_environment_secrets(snapshot: &Snapshot<'_>) -> PyResult> { + if snapshot.read(&SECRET_MANAGER_READABLE)? { return Err(RustBridgeDeclined::new_err( "a readable secret manager is configured and the Rust route only reads the process environment", )); } - Ok(Arc::new(ProcessEnvironment)) + Ok(Arc::new(EnvironmentSecrets)) } fn ocr_settings(py: Python<'_>) -> PyResult { project_provider_defaults(&PythonSettings::ProviderDefaults.read(py)?) } -fn project_provider_defaults(value: &Bound<'_, PyAny>) -> PyResult { +fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult { Ok(OcrSettings { - vertex_project: Field::read(value, "provider_defaults.vertex_project")? - .falsy_optional_string()? - .0, - vertex_location: Field::read(value, "provider_defaults.vertex_location")? - .falsy_optional_string()? - .0, - enable_azure_ad_token_refresh: Field::read( - value, - "provider_defaults.enable_azure_ad_token_refresh", - )? - .exact_true() - .0, + vertex_project: snapshot.read(&VERTEX_PROJECT)?, + vertex_location: snapshot.read(&VERTEX_LOCATION)?, + enable_azure_ad_token_refresh: snapshot.read(&ENABLE_AZURE_AD_TOKEN_REFRESH)?, ..OcrSettings::from_environment(&ProcessEnvironment) }) } @@ -120,6 +128,8 @@ mod tests { use super::process_environment_secrets; use crate::errors::RustBridgeDeclined; + use crate::python_settings::PythonSettings; + fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> { let locals = PyDict::new(py); locals.set_item("readable", readable).unwrap(); @@ -132,12 +142,26 @@ mod tests { locals.get_item("manager").unwrap().unwrap() } + #[test] + fn a_readable_secret_manager_sends_the_call_back_to_python() { + Python::initialize(); + Python::attach(|py| { + let declined = process_environment_secrets( + &PythonSettings::SecretManager.snapshot(secret_manager(py, true)), + ) + .err() + .expect("the Rust route declines"); + assert!(declined.is_instance_of::(py)); + }); + } + #[test] fn provider_defaults_distinguish_falsey_values_and_exact_true() { Python::initialize(); Python::attach(|py| { let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap(); - let projected = super::project_provider_defaults(&value).unwrap(); + let snapshot = PythonSettings::ProviderDefaults.snapshot(value.clone()); + let projected = super::project_provider_defaults(&snapshot).unwrap(); assert_eq!(projected.vertex_project, None); assert_eq!(projected.vertex_location, None); assert!(!projected.enable_azure_ad_token_refresh); @@ -146,12 +170,12 @@ mod tests { value .setattr("enable_azure_ad_token_refresh", true) .unwrap(); - let next = super::project_provider_defaults(&value).unwrap(); + let next = super::project_provider_defaults(&snapshot).unwrap(); assert_eq!(next.vertex_project.as_deref(), Some("project")); assert_eq!(next.vertex_location.as_deref(), Some("region")); assert!(next.enable_azure_ad_token_refresh); value.setattr("vertex_project", 1).unwrap(); - let error = super::project_provider_defaults(&value).err().unwrap(); + let error = super::project_provider_defaults(&snapshot).err().unwrap(); assert!(error.is_instance_of::(py)); assert!( error @@ -160,28 +184,4 @@ mod tests { ); }); } - - #[test] - fn a_readable_secret_manager_sends_the_call_back_to_python() { - Python::initialize(); - Python::attach(|py| { - let declined = process_environment_secrets(&secret_manager(py, true)) - .err() - .expect("the Rust route declines"); - assert!(declined.is_instance_of::(py)); - }); - } - - #[test] - fn without_a_readable_secret_manager_secrets_are_the_process_environment() { - Python::initialize(); - Python::attach(|py| { - let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap(); - assert_eq!( - secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"), - None - ); - assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok()); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs new file mode 100644 index 00000000000..2b98081acef --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -0,0 +1,349 @@ +use std::{fmt, future::Future, pin::Pin}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets::{ + Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; + +const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler"; + +struct PythonSecretError(Py); + +impl fmt::Debug for PythonSecretError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("PythonSecretError") + } +} + +impl fmt::Display for PythonSecretError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Python secret manager failed") + } +} + +impl std::error::Error for PythonSecretError {} + +pub(crate) fn python_error(py: Python<'_>, error: &Error) -> Option { + let Error::ExternalManager(source) = error else { + return None; + }; + source + .downcast_ref::() + .map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any())) +} + +/// A secret manager whose reads execute in Python: a custom manager, a legacy compatible +/// client, or a manually assigned SDK client. +pub(crate) struct PythonSecretManager { + client: Py, + system: Option, + /// The `key_manager` name Python's handler dispatches on. + key_manager: &'static str, + settings: Option>, +} + +impl PythonSecretManager { + pub(crate) fn new( + client: Py, + system: Option, + settings: Option>, + ) -> Self { + Self { + client, + system, + key_manager: system.map_or("local", python_name), + settings, + } + } + + fn read(&self, py: Python<'_>, name: &str) -> PyResult> { + let client = self.client.bind(py); + if self.system == Some(KeyManagementSystem::Custom) + || (self.system.is_none() && client.hasattr("sync_read_secret")?) + { + let kwargs = PyDict::new(py); + kwargs.set_item("secret_name", name)?; + if self.system == Some(KeyManagementSystem::Custom) { + let optional_params = self + .settings + .as_ref() + .map(|settings| settings.bind(py).call_method0("model_dump")) + .transpose()?; + kwargs.set_item("optional_params", optional_params)?; + } + return client + .call_method("sync_read_secret", (), Some(&kwargs))? + .extract(); + } + let kwargs = PyDict::new(py); + kwargs.set_item("client", client)?; + kwargs.set_item("key_manager", self.key_manager)?; + kwargs.set_item("secret_name", name)?; + kwargs.set_item( + "key_management_settings", + self.settings + .as_ref() + .map_or_else(|| py.None(), |settings| settings.clone_ref(py)), + )?; + py.import(HANDLER_MODULE)? + .getattr("get_secret_from_manager")? + .call((), Some(&kwargs))? + .extract() + } +} + +/// The `KeyManagementSystem` value as Python spells it. +fn python_name(system: KeyManagementSystem) -> &'static str { + match system { + KeyManagementSystem::GoogleKms => "google_kms", + KeyManagementSystem::AzureKeyVault => "azure_key_vault", + KeyManagementSystem::AwsSecretManager => "aws_secret_manager", + KeyManagementSystem::GoogleSecretManager => "google_secret_manager", + KeyManagementSystem::HashicorpVault => "hashicorp_vault", + KeyManagementSystem::Cyberark => "cyberark", + KeyManagementSystem::Local => "local", + KeyManagementSystem::AwsKms => "aws_kms", + KeyManagementSystem::Custom => "custom", + } +} + +impl ExternalSecretManager for PythonSecretManager { + fn system(&self) -> KeyManagementSystem { + self.system.unwrap_or(KeyManagementSystem::Custom) + } + + fn read_secret<'a>( + &'a self, + name: &'a str, + _settings: &'a KeyManagementSettings, + _environment: &'a (dyn Lookup + Send + Sync), + ) -> Pin, Error>> + Send + 'a>> { + Box::pin(async move { + Python::attach(|py| { + self.read(py, name) + .map(|value| value.map(SecretValue::new).map(Secret::String)) + .map_err(|error| { + Error::ExternalManager(Box::new(PythonSecretError(error.into_value(py)))) + }) + }) + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use litellm_secrets::{ + FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager, + SecretManagerState, SecretResolver, + }; + use pyo3::{prelude::*, types::PyDict}; + + use super::{HANDLER_MODULE, PythonSecretManager, python_error, python_name}; + + #[tokio::test] + async fn callback_failures_preserve_python_exceptions_even_with_environment_fallback() { + Python::initialize(); + for failure_type in ["ValueError", "asyncio.CancelledError"] { + for fallback in [None, Some("environment-key")] { + let (reader, locals) = Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("failure_type", failure_type).unwrap(); + py.run( + c" +import asyncio +failure = eval(failure_type)('secret manager failed') +cause = RuntimeError('original cause') +context = RuntimeError('original context') +failure.__cause__ = cause +failure.__context__ = context +class Manager: + def sync_read_secret(self, secret_name): + raise failure +manager = Manager() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let reader = PythonSecretManager::new( + locals.get_item("manager").unwrap().unwrap().unbind(), + None, + None, + ); + (reader, locals.unbind()) + }); + let resolver = SecretResolver::new( + Arc::new(SecretManagerState::new( + SecretManager::External(Arc::new(reader)), + KeyManagementSettings::default(), + )), + Arc::new(move |_: &str| fallback.map(str::to_owned)), + OidcResolver::default(), + ) + .with_failure_policy(FailurePolicy::EnvironmentFallback); + let error = resolver.get_secret("API_KEY", None).await.unwrap_err(); + Python::attach(|py| { + let original = python_error(py, &error).unwrap(); + let locals = locals.bind(py); + assert!( + original + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + for (attribute, name) in [("__cause__", "cause"), ("__context__", "context")] { + assert!( + original + .value(py) + .getattr(attribute) + .unwrap() + .is(locals.get_item(name).unwrap().unwrap()) + ); + } + assert!(original.traceback(py).is_some()); + }); + } + } + } + + /// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and + /// removes the fake modules again. + fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) { + let locals = PyDict::new(py); + py.run( + c" +import sys, types +calls = [] +def get_secret_from_manager(**kwargs): + calls.append(kwargs) + return 'handled-' + kwargs['secret_name'] +handler = types.ModuleType('litellm.secret_managers.secret_manager_handler') +handler.get_secret_from_manager = get_secret_from_manager +installed = {} +for name in ('litellm', 'litellm.secret_managers'): + if name not in sys.modules: + sys.modules[name] = types.ModuleType(name) + installed[name] = True +sys.modules['litellm.secret_managers.secret_manager_handler'] = handler +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + body(&locals); + py.run( + c" +sys.modules.pop('litellm.secret_managers.secret_manager_handler', None) +for name in installed: + sys.modules.pop(name, None) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + } + + #[test] + fn python_names_round_trip_through_serde() { + for system in [ + KeyManagementSystem::GoogleKms, + KeyManagementSystem::AzureKeyVault, + KeyManagementSystem::AwsSecretManager, + KeyManagementSystem::GoogleSecretManager, + KeyManagementSystem::HashicorpVault, + KeyManagementSystem::Cyberark, + KeyManagementSystem::Local, + KeyManagementSystem::AwsKms, + KeyManagementSystem::Custom, + ] { + assert_eq!( + serde_json::to_value(system).unwrap(), + serde_json::Value::String(python_name(system).to_owned()) + ); + } + } + + #[test] + fn custom_readers_without_a_system_are_called_directly() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +class Manager: + def __init__(self): + self.names = [] + def sync_read_secret(self, secret_name, optional_params=None, timeout=None): + self.names.append(secret_name) + return 'direct-' + secret_name +manager = Manager() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let manager = locals.get_item("manager").unwrap().unwrap(); + let reader = PythonSecretManager::new(manager.clone().unbind(), None, None); + assert_eq!( + reader.read(py, "API_KEY").unwrap().as_deref(), + Some("direct-API_KEY") + ); + assert_eq!( + manager + .getattr("names") + .unwrap() + .extract::>() + .unwrap(), + ["API_KEY"] + ); + }); + } + + #[test] + fn configured_systems_dispatch_through_the_python_handler_with_the_original_settings() { + Python::initialize(); + Python::attach(|py| { + with_fake_handler(py, |locals| { + let client = py.eval(c"object()", None, None).unwrap(); + let settings = py.eval(c"object()", None, None).unwrap(); + let reader = PythonSecretManager::new( + client.clone().unbind(), + Some(KeyManagementSystem::AzureKeyVault), + Some(settings.clone().unbind()), + ); + assert_eq!( + reader.read(py, "API_KEY").unwrap().as_deref(), + Some("handled-API_KEY") + ); + assert!(py.import(HANDLER_MODULE).is_ok()); + let calls = locals.get_item("calls").unwrap().unwrap(); + let call = calls.get_item(0).unwrap().cast_into::().unwrap(); + assert!(call.get_item("client").unwrap().unwrap().is(&client)); + assert!( + call.get_item("key_management_settings") + .unwrap() + .unwrap() + .is(&settings) + ); + assert_eq!( + call.get_item("key_manager") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "azure_key_vault" + ); + assert_eq!( + call.get_item("secret_name") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "API_KEY" + ); + }); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs new file mode 100644 index 00000000000..6fd380fe40c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -0,0 +1,338 @@ +use std::sync::Arc; + +use litellm_secrets::{SecretManager, SecretManagerState}; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; +use pyo3::prelude::*; +use serde_json::Value; + +use super::callback::PythonSecretManager; +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const SYSTEM: FieldSpec> = + FieldSpec::new("system", parse_optional_system); +const ACCESS_MODE: FieldSpec = FieldSpec::new("access_mode", parse_access_mode); +const HOSTED_KEYS: FieldSpec>> = + FieldSpec::new("hosted_keys", |field| field.optional_string_collection()); +const STORE_VIRTUAL_KEYS: FieldSpec = + FieldSpec::new("store_virtual_keys", |field| field.truthy()); +const PREFIX_FOR_STORED_VIRTUAL_KEYS: FieldSpec = + FieldSpec::new("prefix_for_stored_virtual_keys", |field| { + field.strict_string() + }); +const PRIMARY_SECRET_NAME: FieldSpec> = + FieldSpec::new("primary_secret_name", |field| field.falsy_optional_string()); +const KMS_KEY_ID: FieldSpec> = + FieldSpec::new("kms_key_id", |field| field.falsy_optional_string()); +const CUSTOM_SECRET_MANAGER: FieldSpec> = + FieldSpec::new("custom_secret_manager", |field| { + field.falsy_optional_string() + }); +const AWS_REGION_NAME: FieldSpec> = + FieldSpec::new("aws_region_name", |field| field.falsy_optional_string()); +const AWS_ROLE_NAME: FieldSpec> = + FieldSpec::new("aws_role_name", |field| field.falsy_optional_string()); +const AWS_SESSION_NAME: FieldSpec> = + FieldSpec::new("aws_session_name", |field| field.falsy_optional_string()); +const AWS_EXTERNAL_ID: FieldSpec> = + FieldSpec::new("aws_external_id", |field| field.falsy_optional_string()); +const AWS_PROFILE_NAME: FieldSpec> = + FieldSpec::new("aws_profile_name", |field| field.falsy_optional_string()); +const AWS_WEB_IDENTITY_TOKEN: FieldSpec> = + FieldSpec::new("aws_web_identity_token", |field| { + field.falsy_optional_string() + }); +const AWS_STS_ENDPOINT: FieldSpec> = + FieldSpec::new("aws_sts_endpoint", |field| field.falsy_optional_string()); +const REPLICA_REGIONS: FieldSpec>> = + FieldSpec::new("replica_regions", |field| { + field.optional_string_collection() + }); +const CLIENT: FieldSpec>> = + FieldSpec::new("client", |field| Ok(field.python_binding())); +const SETTINGS_OBJECT: FieldSpec>> = + FieldSpec::new("settings_object", |field| Ok(field.python_binding())); + +/// `litellm.secret_manager_client` as the bridge classifies it. +#[derive(Debug)] +pub(crate) enum SecretManagerClient { + /// `None`: reads come from the process environment. + Local, + /// A custom manager, legacy compatible client, or manually assigned SDK client that keeps + /// executing in Python. + PythonCallback(Py), +} + +/// One operation-local capture of the secret manager globals, taken while attached to Python. +#[derive(Debug)] +pub(crate) struct SecretManagerSnapshot { + pub(crate) client: SecretManagerClient, + pub(crate) system: Option, + /// Typed settings that drive native routing: access mode and hosted keys. + pub(crate) settings: KeyManagementSettings, + /// The original `KeyManagementSettings` object, handed back to Python callbacks unchanged. + pub(crate) settings_object: Option>, +} + +impl SecretManagerSnapshot { + pub(crate) fn into_state(self) -> Arc { + match self.client { + SecretManagerClient::Local => Arc::new(SecretManagerState::default()), + SecretManagerClient::PythonCallback(client) => Arc::new(SecretManagerState::new( + SecretManager::External(Arc::new(PythonSecretManager::new( + client, + self.system, + self.settings_object, + ))), + self.settings, + )), + } + } +} + +/// Reads and projects the secret manager settings group in one attached operation. +pub(crate) fn read(py: Python<'_>) -> PyResult { + Ok(project(&PythonSettings::SecretManagerBinding.read(py)?)?) +} + +pub(crate) fn project(snapshot: &Snapshot<'_>) -> Result { + let system = snapshot.read(&SYSTEM)?; + let access_mode = snapshot.read(&ACCESS_MODE)?; + let settings = KeyManagementSettings { + hosted_keys: snapshot.read(&HOSTED_KEYS)?, + store_virtual_keys: Some(snapshot.read(&STORE_VIRTUAL_KEYS)?), + prefix_for_stored_virtual_keys: snapshot.read(&PREFIX_FOR_STORED_VIRTUAL_KEYS)?, + access_mode, + primary_secret_name: snapshot.read(&PRIMARY_SECRET_NAME)?, + kms_key_id: snapshot.read(&KMS_KEY_ID)?, + custom_secret_manager: snapshot.read(&CUSTOM_SECRET_MANAGER)?, + aws_region_name: snapshot.read(&AWS_REGION_NAME)?, + aws_role_name: snapshot.read(&AWS_ROLE_NAME)?, + aws_session_name: snapshot.read(&AWS_SESSION_NAME)?, + aws_external_id: snapshot.read(&AWS_EXTERNAL_ID)?.map(SecretValue::new), + aws_profile_name: snapshot.read(&AWS_PROFILE_NAME)?, + aws_web_identity_token: snapshot + .read(&AWS_WEB_IDENTITY_TOKEN)? + .map(SecretValue::new), + aws_sts_endpoint: snapshot.read(&AWS_STS_ENDPOINT)?, + replica_regions: snapshot.read(&REPLICA_REGIONS)?, + ..KeyManagementSettings::default() + }; + let client = match snapshot.read(&CLIENT)? { + None => SecretManagerClient::Local, + Some(client) => SecretManagerClient::PythonCallback(client), + }; + Ok(SecretManagerSnapshot { + client, + system, + settings, + settings_object: snapshot.read(&SETTINGS_OBJECT)?, + }) +} + +fn parse_optional_system( + field: &Field<'_>, +) -> Result, ProjectionError> { + let Some(value) = field.falsy_optional_string()? else { + return Ok(None); + }; + serde_json::from_value(Value::String(value)) + .map(Some) + .map_err(|error| { + ProjectionError::InvalidConfiguration(format!("secret manager system: {error}")) + }) +} + +fn parse_access_mode(field: &Field<'_>) -> Result { + let value = field.strict_string()?; + serde_json::from_value(Value::String(value)).map_err(|error| { + ProjectionError::InvalidConfiguration(format!("secret manager access mode: {error}")) + }) +} + +#[cfg(test)] +mod tests { + use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, + }; + + use super::{SecretManagerClient, project}; + use crate::python_settings::PythonSettings; + + fn snapshot<'py>( + py: Python<'py>, + system: &str, + access_mode: &str, + store_virtual_keys: Bound<'py, PyAny>, + hosted_keys: Bound<'py, PyAny>, + ) -> crate::python_settings::Snapshot<'py> { + snapshot_with_client( + py, + system, + access_mode, + store_virtual_keys, + hosted_keys, + py.None().into_bound(py), + ) + } + + fn snapshot_with_client<'py>( + py: Python<'py>, + system: &str, + access_mode: &str, + store_virtual_keys: Bound<'py, PyAny>, + hosted_keys: Bound<'py, PyAny>, + client: Bound<'py, PyAny>, + ) -> crate::python_settings::Snapshot<'py> { + let locals = PyDict::new(py); + locals.set_item("client", client).unwrap(); + locals.set_item("system", system).unwrap(); + locals.set_item("access_mode", access_mode).unwrap(); + locals + .set_item("store_virtual_keys", store_virtual_keys) + .unwrap(); + locals.set_item("hosted_keys", hosted_keys).unwrap(); + py.run( + cr#" +from dataclasses import dataclass +from types import SimpleNamespace + +@dataclass(frozen=True, slots=True) +class SecretManager: + system: object + access_mode: object + hosted_keys: object + primary_secret_name: object + store_virtual_keys: object + prefix_for_stored_virtual_keys: object + kms_key_id: object + custom_secret_manager: object + aws_region_name: object + aws_role_name: object + aws_session_name: object + aws_external_id: object + aws_profile_name: object + aws_web_identity_token: object + aws_sts_endpoint: object + replica_regions: object + client: object + settings_object: object + +root = SimpleNamespace(secret_manager=SecretManager( + system=system, + access_mode=access_mode, + hosted_keys=hosted_keys, + primary_secret_name=None, + store_virtual_keys=store_virtual_keys, + prefix_for_stored_virtual_keys="litellm/", + kms_key_id=None, + custom_secret_manager=None, + aws_region_name=None, + aws_role_name=None, + aws_session_name=None, + aws_external_id=None, + aws_profile_name=None, + aws_web_identity_token=None, + aws_sts_endpoint=None, + replica_regions=None, + client=client, + settings_object=None, +)) +"#, + Some(&locals), + Some(&locals), + ) + .unwrap(); + PythonSettings::SecretManagerBinding.snapshot( + locals + .get_item("root") + .unwrap() + .unwrap() + .getattr("secret_manager") + .unwrap(), + ) + } + + #[rstest::rstest] + #[case::string_true(Some("true"), false, true)] + #[case::string_one(Some("1"), false, true)] + #[case::true_value(None, true, true)] + #[case::false_value(None, false, false)] + #[case::string_false(Some("false"), false, true)] + fn python_compatible_boolean_coercion( + #[case] string_value: Option<&str>, + #[case] bool_value: bool, + #[case] expected: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let store_virtual_keys = match string_value { + Some(value) => value.into_pyobject(py).unwrap().into_any(), + None => bool_value.into_pyobject(py).unwrap().to_owned().into_any(), + }; + let hosted_keys = PyTuple::new(py, ["ONE"]).unwrap().into_any(); + let projected = project(&snapshot( + py, + "local", + "read_only", + store_virtual_keys, + hosted_keys, + )) + .unwrap(); + assert_eq!(projected.settings.store_virtual_keys, Some(expected)); + }); + } + + #[test] + fn unknown_system_is_rejected() { + Python::initialize(); + Python::attach(|py| { + let error = project(&snapshot( + py, + "unknown", + "read_only", + false.into_pyobject(py).unwrap().to_owned().into_any(), + PyTuple::empty(py).into_any(), + )) + .unwrap_err(); + let error: PyErr = error.into(); + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + fn client_identity_selects_local_or_python_callback() { + Python::initialize(); + Python::attach(|py| { + let falsy = false.into_pyobject(py).unwrap().to_owned().into_any(); + let local = project(&snapshot( + py, + "local", + "read_only", + falsy.clone(), + PyTuple::empty(py).into_any(), + )) + .unwrap(); + assert!(matches!(local.client, SecretManagerClient::Local)); + assert!(local.settings_object.is_none()); + + let manager = py.eval(c"object()", None, None).unwrap(); + let custom = project(&snapshot_with_client( + py, + "custom", + "read_only", + falsy, + PyTuple::empty(py).into_any(), + manager.clone(), + )) + .unwrap(); + let SecretManagerClient::PythonCallback(client) = custom.client else { + panic!("a live client must stay a Python callback"); + }; + assert!(client.bind(py).is(&manager)); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs new file mode 100644 index 00000000000..f6ca57b08d1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod callback; +pub(crate) mod config; +pub(crate) mod resolved; diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs new file mode 100644 index 00000000000..877c429169a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -0,0 +1,252 @@ +use std::{collections::HashMap, sync::Arc}; + +use futures_util::{future::BoxFuture, future::try_join_all}; +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use litellm_llms::base_llm::inference::secrets::{SecretSource, Secrets}; +use litellm_secrets::{ + Error, FailurePolicy, OidcResolver, Secret, SecretManagerState, SecretResolver, +}; + +use super::config::SecretManagerSnapshot; + +pub(crate) struct ResolvedSecrets { + resolver: SecretResolver, +} + +impl ResolvedSecrets { + pub(crate) fn new(snapshot: SecretManagerSnapshot) -> Self { + Self::from_state(snapshot.into_state()) + } + + fn from_state(state: Arc) -> Self { + Self { + resolver: SecretResolver::new( + state, + Arc::new(ProcessEnvironment), + OidcResolver::default(), + ) + .with_failure_policy(FailurePolicy::EnvironmentFallback), + } + } +} + +impl SecretSource for ResolvedSecrets { + fn resolve<'a>(&'a self, names: &'a [&'static str]) -> BoxFuture<'a, Result> { + Box::pin(async move { + let values = try_join_all(names.iter().map(|name| async move { + self.resolver + .get_secret(name, None) + .await + .map(|secret| secret.map(|secret| ((*name).to_owned(), secret_value(secret)))) + })) + .await? + .into_iter() + .flatten() + .collect::>(); + Ok(Arc::new(ResolvedLookup { values }) as Secrets) + }) + } +} + +struct ResolvedLookup { + values: HashMap, +} + +impl Lookup for ResolvedLookup { + fn get(&self, name: &str) -> Option { + self.values + .get(name) + .cloned() + .or_else(|| ProcessEnvironment.get(name)) + } +} + +fn secret_value(secret: Secret) -> String { + match secret { + Secret::String(value) => value.expose().to_owned(), + Secret::Bool(value) => if value { "True" } else { "False" }.to_owned(), + Secret::Json(value) => value.to_string(), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use aws_sdk_secretsmanager::Client; + use aws_sdk_secretsmanager::config::{ + BehaviorVersion, Credentials, Region, retry::RetryConfig, + }; + use litellm_secrets::{AccessMode, KeyManagementSettings, SecretManager, SecretManagerState}; + use litellm_secrets_aws::AwsSecretsManagerV2; + use serde_json::json; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header}, + }; + + use super::ResolvedSecrets; + use litellm_llms::base_llm::inference::secrets::SecretSource; + + fn state(server: &MockServer, settings: KeyManagementSettings) -> Arc { + let client = Client::from_conf( + aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + Arc::new(SecretManagerState::new( + SecretManager::AwsSecretsManagerV2(AwsSecretsManagerV2::new( + client, + (&settings).into(), + )), + settings, + )) + } + + async fn resolve(state: Arc, name: &'static str) -> Option { + ResolvedSecrets::from_state(state) + .resolve(&[name]) + .await + .unwrap() + .get(name) + } + + #[tokio::test] + async fn hosted_key_miss_falls_back_to_environment() { + let name = "LITELLM_RUST_BRIDGE_HOSTED_KEY_MISS"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId": name}))) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})), + ) + .expect(0) + .mount(&server) + .await; + let result = resolve( + state( + &server, + KeyManagementSettings { + hosted_keys: Some(vec!["OTHER".into()]), + ..Default::default() + }, + ), + name, + ) + .await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 0); + } + + #[tokio::test] + async fn manager_failure_falls_back_to_environment() { + let name = "LITELLM_RUST_BRIDGE_MANAGER_FAILURE"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + let result = resolve(state(&server, KeyManagementSettings::default()), name).await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + + let missing_server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&missing_server) + .await; + let missing = + ResolvedSecrets::from_state(state(&missing_server, KeyManagementSettings::default())) + .resolve(&["LITELLM_RUST_BRIDGE_MANAGER_FAILURE_MISSING"]) + .await; + assert!(matches!(missing, Err(litellm_secrets::Error::Aws(_)))); + } + + #[tokio::test] + async fn write_only_mode_never_consults_the_manager() { + let name = "LITELLM_RUST_BRIDGE_WRITE_ONLY"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})), + ) + .expect(0) + .mount(&server) + .await; + let result = resolve( + state( + &server, + KeyManagementSettings { + access_mode: AccessMode::WriteOnly, + ..Default::default() + }, + ), + name, + ) + .await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 0); + } + + #[tokio::test] + async fn read_only_mode_resolves_from_the_manager() { + let name = "LITELLM_RUST_BRIDGE_READ_ONLY"; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId": name}))) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})), + ) + .expect(1) + .mount(&server) + .await; + assert_eq!( + resolve(state(&server, KeyManagementSettings::default()), name) + .await + .as_deref(), + Some("manager-key") + ); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } + + #[tokio::test] + async fn oidc_failures_are_not_converted_to_missing_secrets() { + let result = ResolvedSecrets::from_state(Arc::new(SecretManagerState::default())) + .resolve(&["oidc/"]) + .await; + assert!(matches!(result, Err(litellm_secrets::Error::InvalidOidc))); + } + + #[tokio::test] + async fn undeclared_names_still_read_the_process_environment() { + let name = "LITELLM_RUST_BRIDGE_UNDECLARED"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + let result = resolve( + state( + &server, + KeyManagementSettings { + hosted_keys: Some(vec!["OTHER".into()]), + ..Default::default() + }, + ), + name, + ) + .await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 0); + } +} diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 36d319311a3..44acf512224 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize}; use crate::SecretValue; -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] pub enum KeyManagementSystem { GoogleKms, @@ -18,7 +18,7 @@ pub enum KeyManagementSystem { Custom, } -#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] pub enum AccessMode { #[default] @@ -33,7 +33,7 @@ impl AccessMode { } } -#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)] #[serde(default)] pub struct KeyManagementSettings { pub hosted_keys: Option>, diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 1be0adc2bf5..de325ff4981 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { OidcFile, #[error("secret cannot be converted to {expected}")] TypeMismatch { expected: &'static str }, + #[error("external secret manager failed")] + ExternalManager(#[source] Box), #[cfg(feature = "aws")] #[error(transparent)] Aws(#[from] litellm_secrets_aws::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 5ab2caa75d4..8762b8e7405 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -1,10 +1,24 @@ +use std::{future::Future, pin::Pin, sync::Arc}; + use litellm_core_utils::settings::Lookup; use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue}; +pub trait ExternalSecretManager: Send + Sync { + fn system(&self) -> KeyManagementSystem; + + fn read_secret<'a>( + &'a self, + name: &'a str, + settings: &'a KeyManagementSettings, + environment: &'a (dyn Lookup + Send + Sync), + ) -> Pin, Error>> + Send + 'a>>; +} + #[derive(Clone)] pub enum SecretManager { Local, + External(Arc), #[cfg(feature = "aws")] AwsKms(crate::aws::AwsKms), #[cfg(feature = "aws")] @@ -25,6 +39,7 @@ impl SecretManager { pub fn system(&self) -> KeyManagementSystem { match self { Self::Local => KeyManagementSystem::Local, + Self::External(manager) => manager.system(), #[cfg(feature = "aws")] Self::AwsKms(_) => KeyManagementSystem::AwsKms, #[cfg(feature = "aws")] @@ -54,6 +69,11 @@ pub async fn get_secret_from_manager( .get(secret_name) .map(SecretValue::new) .map(Secret::String)), + SecretManager::External(manager) => { + manager + .read_secret(secret_name, _settings, environment) + .await + } #[cfg(feature = "aws")] SecretManager::AwsKms(client) => { let ciphertext = environment diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index 1acb5269e66..58aba8494fd 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -7,7 +7,7 @@ mod resolver; mod state; pub use error::Error; -pub use handler::{SecretManager, get_secret_from_manager}; +pub use handler::{ExternalSecretManager, SecretManager, get_secret_from_manager}; pub use litellm_secrets_types::{ AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, }; diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 89439893852..597ca11b171 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -72,6 +72,7 @@ impl SecretResolver { Ok(value) => Ok(value .or_else(|| self.environment_secret(name)) .or(default_value)), + Err(error @ Error::ExternalManager(_)) => Err(error), Err(error) => match self.failure_policy { FailurePolicy::Propagate => Err(error), FailurePolicy::EnvironmentFallback => self diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index d36c0343988..8e4b7f82eb8 100644 --- a/litellm/chat_completions/dispatch.py +++ b/litellm/chat_completions/dispatch.py @@ -4,7 +4,7 @@ from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable from litellm import main -from litellm.rust_bridge.catalog import Context, Delivery, Route +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.chat_completions.entrypoints import ( NATIVE_ACOMPLETION, NATIVE_COMPLETION, @@ -72,8 +72,8 @@ def _public_request( ) -def _context(request: LiteLLMChatCompletionsRequest) -> Context: - return Context( +def _context(request: LiteLLMChatCompletionsRequest) -> RouteContext: + return RouteContext( Route.CHAT_COMPLETIONS, provider=request.custom_llm_provider, model=request.model, diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index 8f35b8eac7a..172316027a8 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -5,7 +5,7 @@ import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.rust_bridge import runtime -from litellm.rust_bridge.catalog import Context, Route +from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.rust_bridge.transcription.native import ( NATIVE_ATRANSCRIPTION, @@ -74,7 +74,7 @@ class BedrockAudioTranscriptionRustDispatch: ) return runtime.run( - Context(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), + RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), binding=NATIVE_TRANSCRIPTION, native=native, python=_no_python_implementation, @@ -107,7 +107,7 @@ class BedrockAudioTranscriptionRustDispatch: ) return await runtime.arun( - Context(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), + RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), binding=NATIVE_ATRANSCRIPTION, native=native, python=_no_async_python_implementation, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2a105301521..18d4c27fec1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -182,10 +182,10 @@ from litellm.utils import ( def _rust_responses_websocket_enabled( custom_llm_provider: str | None, ) -> bool: - from litellm.rust_bridge.catalog import Context, Delivery, Route, decision + from litellm.rust_bridge.catalog import Delivery, Route, RouteContext, decision from litellm.rust_bridge.configuration import Decision - context: Final = Context(Route.RESPONSES, provider=custom_llm_provider, delivery=Delivery.WEBSOCKET) + context: Final = RouteContext(Route.RESPONSES, provider=custom_llm_provider, delivery=Delivery.WEBSOCKET) return decision(context) is not Decision.PYTHON diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index c75f6564d1b..a0c791a136c 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -4,7 +4,7 @@ from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable from litellm.llms.anthropic.experimental_pass_through.messages import handler as main -from litellm.rust_bridge.catalog import Context, Delivery, Route +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook from litellm.rust_bridge.messages.entrypoints import ( NATIVE_AMESSAGES, @@ -71,8 +71,8 @@ def _public_request( ) -def _context(request: LiteLLMMessagesRequest) -> Context: - return Context( +def _context(request: LiteLLMMessagesRequest) -> RouteContext: + return RouteContext( Route.MESSAGES, provider=request.custom_llm_provider, model=request.model, diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 55b19458b7a..b26175c943b 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -6,7 +6,7 @@ import httpx from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.ocr import main from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type -from litellm.rust_bridge.catalog import Context, Route +from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, LiteLLMOcrRequest @@ -52,10 +52,10 @@ _PYTHON_AOCR: Final = cast( # cast-ok: forward the original call shape through ) -def _context(request: LiteLLMOcrRequest) -> Context: +def _context(request: LiteLLMOcrRequest) -> RouteContext: prefix, separator, _ = request.model.partition("/") provider: Final = request.custom_llm_provider or (prefix if separator else None) - return Context(Route.OCR, provider=provider, model=request.model) + return RouteContext(Route.OCR, provider=provider, model=request.model) _DISPATCH: Final = PublicDispatch( diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index b2748fca4b6..d240356805c 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -5,7 +5,7 @@ from typing import Final, TypeAlias, cast # noqa: TID251 # native binding sele from litellm.responses import main from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.rust_bridge.catalog import Context, Delivery, Route +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook from litellm.rust_bridge.public_call import bind, optional_bool, optional_mapping, optional_str, signature from litellm.rust_bridge.responses.entrypoints import ( @@ -64,8 +64,8 @@ def _public_request( ) -def _context(request: LiteLLMResponsesRequest) -> Context: - return Context( +def _context(request: LiteLLMResponsesRequest) -> RouteContext: + return RouteContext( Route.RESPONSES, provider=request.custom_llm_provider, model=request.model, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 31623a53dd0..a033b89a6a8 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -94,7 +94,9 @@ class ResponsesWebSocketConnection: def close(self) -> Future[None]: ... @final -class _CacheTestBinding: +class _ResponseCacheRuntime: + @staticmethod + def from_cache(cache: object) -> _ResponseCacheRuntime: ... @property def kind(self) -> str: ... def lookup( @@ -218,7 +220,7 @@ class _CacheTestHandle: @final class _CacheTestResolver: def __new__(cls, namespace: object) -> _CacheTestResolver: ... - def resolve(self) -> _CacheTestBinding: ... + def resolve(self) -> _ResponseCacheRuntime: ... @final class TokenCounter: diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 8794ff2db95..74ceb1ba123 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -1,9 +1,7 @@ -"""Declarative Rust/Python selection for routes with Rust integration. +"""Ordered rollout policy for routes, cache backends, and secret managers. -Rules are static data matched top to bottom; the first match wins and a -context with no matching rule stays on Python. Whether the Rust core can serve -a specific request body is not decided here: that is Rust admission, which -signals ``RustBridgeDeclined`` before any provider I/O. +The first matching rule wins; unmatched contexts stay on Python. Native +admission separately decides whether the selected implementation can execute. """ from __future__ import annotations @@ -14,6 +12,8 @@ from typing import Final, TypeAlias from litellm.rust_bridge.configuration import Decision, Rollout from litellm.rust_bridge.configuration import decision as _decision +from litellm.types.caching import LiteLLMCacheType +from litellm.types.secret_managers.main import KeyManagementSystem class Route(str, Enum): @@ -31,7 +31,7 @@ class Delivery(Enum): @dataclass(frozen=True, slots=True) -class Context: +class RouteContext: route: Route provider: str | None = None model: str | None = None @@ -39,7 +39,7 @@ class Context: @dataclass(frozen=True, slots=True) -class Rule: +class RouteRule: route: Route rollout: Rollout providers: frozenset[str] | None = None @@ -48,26 +48,76 @@ class Rule: def matches(self, context: Context) -> bool: return ( - context.route is self.route + isinstance(context, RouteContext) + and context.route is self.route and (self.providers is None or context.provider in self.providers) and (self.models is None or context.model in self.models) and (self.deliveries is None or context.delivery in self.deliveries) ) +@dataclass(frozen=True, slots=True) +class CacheContext: + backend: str + + +@dataclass(frozen=True, slots=True) +class CacheRule: + rollout: Rollout + backends: frozenset[str] | None = None + + def matches(self, context: Context) -> bool: + return isinstance(context, CacheContext) and (self.backends is None or context.backend in self.backends) + + +@dataclass(frozen=True, slots=True) +class SecretManagerContext: + system: str + + +@dataclass(frozen=True, slots=True) +class SecretManagerRule: + rollout: Rollout + systems: frozenset[str] | None = None + + def matches(self, context: Context) -> bool: + return isinstance(context, SecretManagerContext) and (self.systems is None or context.system in self.systems) + + +Context: TypeAlias = RouteContext | CacheContext | SecretManagerContext +Rule: TypeAlias = RouteRule | CacheRule | SecretManagerRule Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( - Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), - Rule(Route.OCR, Rollout.RUST_OPT_OUT), - Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), - Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), + RouteRule(Route.OCR, Rollout.RUST_OPT_OUT), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), + RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.LOCAL})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.REDIS})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.REDIS_SEMANTIC})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.VALKEY_SEMANTIC})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.S3})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.DISK})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.QDRANT_SEMANTIC})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.AZURE_BLOB})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.GCS})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_KMS.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AZURE_KEY_VAULT.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AWS_SECRET_MANAGER.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_SECRET_MANAGER.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.HASHICORP_VAULT.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CYBERARK.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.LOCAL.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AWS_KMS.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CUSTOM.value})), ) -def rollout(context: Context, rules: Rules = RULES) -> Rollout: - return next((rule.rollout for rule in rules if rule.matches(context)), Rollout.PYTHON_ONLY) +def rollout(context: Context, rules: Rules | None = None) -> Rollout: + selected_rules: Final = RULES if rules is None else rules + return next((rule.rollout for rule in selected_rules if rule.matches(context)), Rollout.PYTHON_ONLY) -def decision(context: Context, rules: Rules = RULES) -> Decision: +def decision(context: Context, rules: Rules | None = None) -> Decision: return _decision(rollout(context, rules)) diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 791e13a51d0..cd468c80655 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -84,7 +84,7 @@ def reset_rust_configuration() -> None: def rust(enabled: bool | None) -> None: """Set the process override for optional Rust paths. - ``PYTHON_ONLY`` and ``RUST_REQUIRED`` routes in the catalog ignore this switch, + ``PYTHON_ONLY`` and ``RUST_REQUIRED`` entries in the catalog ignore this switch, and an explicit ``LITELLM_RUST`` environment value wins over it. """ _CONFIGURATION.override = enabled diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 7ddc903df58..076b7759c6d 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -6,7 +6,7 @@ from typing import Final, Generic, TypeVar from litellm.rust_bridge import catalog, runtime from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Context, Route, Rules +from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.configuration import decision as rollout_decision @@ -30,12 +30,12 @@ def call_hook( class PublicDispatch(Generic[RequestT]): route: Route request: Callable[[tuple[object, ...], Mapping[str, object]], RequestT | None] - context: Callable[[RequestT], Context] + context: Callable[[RequestT], RouteContext] bypass: Callable[[RequestT], bool] | None = None def _requires_projection(self, rules: Rules) -> bool: for rule in rules: - if rule.route is not self.route: + if not isinstance(rule, RouteRule) or rule.route is not self.route: continue if rule.providers is not None or rule.models is not None or rule.deliveries is not None: if rollout_decision(rule.rollout) is not Decision.PYTHON: diff --git a/litellm/rust_bridge/response_cache.py b/litellm/rust_bridge/response_cache.py new file mode 100644 index 00000000000..82d27fce27b --- /dev/null +++ b/litellm/rust_bridge/response_cache.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import math +from collections.abc import Awaitable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final, Protocol, cast + +from typing_extensions import ReadOnly, Required, TypedDict, assert_never + +from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge.catalog import CacheContext, Rules, decision +from litellm.rust_bridge.configuration import Decision + + +class CacheFacade(Protocol): + @property + def type(self) -> object: ... + + @property + def ttl(self) -> float | None: ... + + @property + def semantic_cache_scope(self) -> str: ... + + def get_cache_key(self, **kwargs: object) -> str: ... # kwargs-ok: mirrors the legacy cache facade contract + + +class NativeCacheKey(TypedDict): + preset: ReadOnly[str] + + +class NativeCacheRequest(TypedDict, total=False): + key: Required[ReadOnly[NativeCacheKey]] + ttl_seconds: ReadOnly[float | None] + max_age_seconds: ReadOnly[float | None] + messages: ReadOnly[object | None] + input: ReadOnly[object | None] + metadata: ReadOnly[object | None] + litellm_metadata: ReadOnly[object | None] + litellm_params: ReadOnly[object | None] + scope: ReadOnly[str] + + +class NativeResponseCacheRuntime(Protocol): + @property + def kind(self) -> str: ... + + def lookup(self, request: NativeCacheRequest) -> object: ... + def store(self, request: NativeCacheRequest, response: object) -> None: ... + def lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: ... + def async_lookup(self, request: NativeCacheRequest) -> Awaitable[object]: ... + def async_store(self, request: NativeCacheRequest, response: object) -> Awaitable[None]: ... + def async_lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> Awaitable[object]: ... + def async_store_batch( + self, + requests: Sequence[NativeCacheRequest], + responses: Sequence[object], + ) -> Awaitable[object]: ... + def async_flush(self) -> Awaitable[None]: ... + def ping(self) -> Awaitable[object]: ... + + +class NativeResponseCacheRuntimeFactory(Protocol): + @staticmethod + def from_cache(cache: CacheFacade) -> NativeResponseCacheRuntime: ... + + +def _runtime_factory(value: object) -> NativeResponseCacheRuntimeFactory | None: + return cast(NativeResponseCacheRuntimeFactory, value) if callable(getattr(value, "from_cache", None)) else None + + +_RUNTIME: Final = NativeBinding("_ResponseCacheRuntime", validate=_runtime_factory) + + +@dataclass(frozen=True, slots=True) +class ResponseCacheRuntime: + native: NativeResponseCacheRuntime + + @property + def kind(self) -> str: + return self.native.kind + + def request(self, cache: CacheFacade, kwargs: Mapping[str, object]) -> NativeCacheRequest | None: + key_value: Final = kwargs.get("cache_key") + key: Final = key_value if isinstance(key_value, str) else cache.get_cache_key(**dict(kwargs)) + if not key: + return None + control_value: Final = kwargs.get("cache") + control: Final = _string_mapping(control_value) + configured_ttl: Final = cache.ttl if cache.ttl is not None else _duration(kwargs.get("ttl")) + control_ttl: Final = _duration(control.get("ttl")) + current_max_age: Final = _duration(control.get("s-max-age")) + legacy_max_age: Final = _duration(control.get("s-maxage")) + ttl: Final = configured_ttl if control_ttl is None else control_ttl + max_age: Final = legacy_max_age if current_max_age is None else current_max_age + return NativeCacheRequest( + key=NativeCacheKey(preset=key), + ttl_seconds=ttl, + max_age_seconds=max_age, + messages=kwargs.get("messages"), + input=kwargs.get("input"), + metadata=kwargs.get("metadata"), + litellm_metadata=kwargs.get("litellm_metadata"), + litellm_params=kwargs.get("litellm_params"), + scope=cache.semantic_cache_scope, + ) + + def lookup(self, request: NativeCacheRequest) -> object: + return self.native.lookup(request) + + def store(self, request: NativeCacheRequest, response: object) -> None: + self.native.store(request, response) + + def lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: + return self.native.lookup_batch(requests) + + async def async_lookup(self, request: NativeCacheRequest) -> object: + return await self.native.async_lookup(request) + + async def async_store(self, request: NativeCacheRequest, response: object) -> None: + await self.native.async_store(request, response) + + async def async_lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: + return await self.native.async_lookup_batch(requests) + + async def async_store_batch( + self, + requests: Sequence[NativeCacheRequest], + responses: Sequence[object], + ) -> object: + return await self.native.async_store_batch(requests, responses) + + async def ping(self) -> object: + return await self.native.ping() + + async def async_flush(self) -> None: + await self.native.async_flush() + + +def resolve_response_cache( + cache: CacheFacade, + rules: Rules | None = None, +) -> ResponseCacheRuntime | None: + backend_value: Final = cache.type + backend: Final = str.__str__(backend_value) if isinstance(backend_value, str) else str(backend_value) + selected: Final = decision(CacheContext(backend=backend), rules) + match selected: + case Decision.PYTHON: + return None + case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: + factory: Final = _RUNTIME.load() + if factory is None: + if selected is Decision.RUST_REQUIRED: + raise RuntimeError("Rust response cache runtime is unavailable") + return None + try: + return ResponseCacheRuntime(factory.from_cache(cache)) + except Exception as error: + exceptions: Final = native_exception_types() + if exceptions is None or not isinstance(error, exceptions[0]): + raise + if selected is Decision.RUST_REQUIRED: + raise RuntimeError(f"Rust response cache runtime declined the cache: {error}") from error + return None + case _: + assert_never(selected) + + +def _duration(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + duration: Final = float(value) + return duration if math.isfinite(duration) and duration >= 0 else None + + +def _string_mapping(value: object) -> Mapping[str, object]: + if not isinstance(value, Mapping): + return {} + source: Final = cast(Mapping[object, object], value) + return {key: item for key, item in source.items() if isinstance(key, str)} diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 1fcde1bf555..cf02a33eab6 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -8,7 +8,7 @@ from typing_extensions import assert_never from litellm.exceptions import APIError from litellm.rust_bridge.bindings import NativeBinding, native_exception_types -from litellm.rust_bridge.catalog import RULES, Context, Rules, decision +from litellm.rust_bridge.catalog import RouteContext, Rules, decision from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.response_metadata import mark_rust_response @@ -42,14 +42,14 @@ class BridgeErrorContext: def run( - context: Context, + context: RouteContext, *, binding: NativeBinding[NativeT], native: Callable[[NativeT], ResultT], python: Callable[[], ResultT], rules: Rules | None = None, ) -> ResultT: - selected: Final = decision(context, RULES if rules is None else rules) + selected: Final = decision(context, rules) match selected: case Decision.PYTHON: return python() @@ -70,14 +70,14 @@ def run( async def arun( - context: Context, + context: RouteContext, *, binding: NativeBinding[NativeT], native: Callable[[NativeT], Awaitable[ResultT]], python: Callable[[], Awaitable[ResultT]], rules: Rules | None = None, ) -> ResultT: - selected: Final = decision(context, RULES if rules is None else rules) + selected: Final = decision(context, rules) match selected: case Decision.PYTHON: return await python() @@ -101,7 +101,7 @@ def _identity(value: ResultT) -> ResultT: return value -def _error_context(context: Context) -> BridgeErrorContext: +def _error_context(context: RouteContext) -> BridgeErrorContext: return BridgeErrorContext(route=context.route.value, provider=context.provider or "", model=context.model or "") diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 9a5cf49f298..866f2fce989 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Final @dataclass(frozen=True, slots=True) @@ -35,6 +36,28 @@ class SecretManager: readable: bool +@dataclass(frozen=True, slots=True) +class SecretManagerBinding: + system: object + access_mode: object + hosted_keys: object + primary_secret_name: object + store_virtual_keys: object + prefix_for_stored_virtual_keys: object + kms_key_id: object + custom_secret_manager: object + aws_region_name: object + aws_role_name: object + aws_session_name: object + aws_external_id: object + aws_profile_name: object + aws_web_identity_token: object + aws_sts_endpoint: object + replica_regions: object + client: object + settings_object: object + + def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -49,6 +72,42 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) +def secret_manager_binding() -> SecretManagerBinding: + import litellm + from litellm.types.secret_managers.main import KeyManagementSettings + + configured_system: Final = ( + litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private + ) + configured_settings: Final = ( + litellm._key_management_settings # pyright: ignore[reportPrivateUsage] # canonical key management globals are private + ) + settings: Final = configured_settings or KeyManagementSettings() + system: Final = ( + configured_system.value if litellm.secret_manager_client is not None and configured_system is not None else None + ) + return SecretManagerBinding( + system=system, + access_mode=settings.access_mode, + hosted_keys=settings.hosted_keys, + primary_secret_name=settings.primary_secret_name, + store_virtual_keys=settings.store_virtual_keys, + prefix_for_stored_virtual_keys=settings.prefix_for_stored_virtual_keys, + kms_key_id=settings.kms_key_id, + custom_secret_manager=settings.custom_secret_manager, + aws_region_name=settings.aws_region_name, + aws_role_name=settings.aws_role_name, + aws_session_name=settings.aws_session_name, + aws_external_id=settings.aws_external_id, + aws_profile_name=settings.aws_profile_name, + aws_web_identity_token=settings.aws_web_identity_token, + aws_sts_endpoint=settings.aws_sts_endpoint, + replica_regions=settings.replica_regions, + client=litellm.secret_manager_client, + settings_object=configured_settings, + ) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/test_litellm/responses/test_dispatch.py b/tests/test_litellm/responses/test_dispatch.py index 2990360d550..45cb5c4f1ad 100644 --- a/tests/test_litellm/responses/test_dispatch.py +++ b/tests/test_litellm/responses/test_dispatch.py @@ -13,7 +13,7 @@ from litellm.responses.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule +from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.responses.entrypoints import ( NATIVE_ARESPONSES, @@ -26,7 +26,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse INPUT: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final = () -RUST_RULES: Final = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),) +RUST_RULES: Final = (RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),) def _response(model: str = "gpt-4o") -> ResponsesAPIResponse: @@ -102,7 +102,8 @@ async def test_async_python_route_forwards_original_call_shape() -> None: response: Final = _response() async def python( - *call_args: object, **call_kwargs: object # kwargs-ok: records call shape + *call_args: object, + **call_kwargs: object, # kwargs-ok: records call shape ) -> ResponsesAPIResponse: captured.append((call_args, call_kwargs)) return response @@ -143,9 +144,7 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: "custom_llm_provider": "anthropic", "litellm_metadata": metadata, } - captured: Final[ - list[tuple[LiteLLMResponsesRequest, tuple[object, ...], Mapping[str, object]]] - ] = [] + captured: Final[list[tuple[LiteLLMResponsesRequest, tuple[object, ...], Mapping[str, object]]]] = [] response: Final = _response("anthropic/claude-sonnet-4-5") def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: rejected fallback @@ -228,9 +227,7 @@ def test_internal_async_marker_bypasses_native() -> None: ((), {}), ), ) -def test_binding_errors_delegate_unchanged_to_python( - args: tuple[object, ...], kwargs: Mapping[str, object] -) -> None: +def test_binding_errors_delegate_unchanged_to_python(args: tuple[object, ...], kwargs: Mapping[str, object]) -> None: captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] response: Final = _response() diff --git a/tests/test_litellm/rust_bridge/ocr/test_secrets.py b/tests/test_litellm/rust_bridge/ocr/test_secrets.py new file mode 100644 index 00000000000..085a42dd373 --- /dev/null +++ b/tests/test_litellm/rust_bridge/ocr/test_secrets.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import configuration +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service + + +class _VaultSecrets(CustomSecretManager): + def __init__(self) -> None: + super().__init__(secret_manager_name="rust_bridge_ocr_test") + + async def async_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return "vault-key" if secret_name == "MISTRAL_API_KEY" else None + + def sync_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return "vault-key" if secret_name == "MISTRAL_API_KEY" else None + + +async def _call(asynchronous: bool, api_base: str) -> OCRResponse: + if asynchronous: + return await litellm.aocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + api_base=api_base, + ) + return litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + api_base=api_base, + ) + + +_RESPONSE: Final = { + "pages": [{"index": 0, "markdown": "parsed document", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("rust_enabled", ("0", "1")) +@pytest.mark.parametrize("access_mode", ("read_only", "read_and_write")) +@pytest.mark.parametrize("system", (None, KeyManagementSystem.CUSTOM)) +async def test_readable_secret_managers_keep_python_ocr_fallback( + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + rust_enabled: str, + access_mode: str, + system: KeyManagementSystem | None, +) -> None: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", rust_enabled) + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets()) + monkeypatch.setattr(litellm, "_key_management_system", system) + monkeypatch.setattr( + litellm, + "_key_management_settings", + KeyManagementSettings(access_mode=access_mode, hosted_keys=["MISTRAL_API_KEY"]), + ) + configuration.reset_rust_configuration() + + with recording_service() as server: + server.default_response = ResponseSpec(body=_RESPONSE) + result: Final = await _call(asynchronous, server.base_url) + + assert result.pages[0].markdown == "parsed document" + assert len(server.requests) == 1 + expected_key: Final = "vault-key" if system is KeyManagementSystem.CUSTOM else "environment-key" + assert server.requests[0].headers["authorization"] == f"Bearer {expected_key}" + assert "x-litellm-rust" not in result._hidden_params.get("additional_headers", {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_no_secret_client_leaves_dormant_binding_settings_unread( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "_key_management_settings", object()) + configuration.reset_rust_configuration() + + with recording_service() as server: + server.default_response = ResponseSpec(body=_RESPONSE) + result: Final = await _call(asynchronous, server.base_url) + + assert result.pages[0].markdown == "parsed document" + assert len(server.requests) == 1 + assert server.requests[0].headers["authorization"] == "Bearer environment-key" + assert result._hidden_params["additional_headers"]["x-litellm-rust"] == "true" diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index 147e863baf5..3fa5f3de7ab 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -6,8 +6,21 @@ from typing import Final import pytest from litellm.rust_bridge import catalog, configuration -from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule +from litellm.rust_bridge.catalog import ( + CacheContext, + CacheRule, + Context, + Delivery, + Route, + RouteContext, + RouteRule, + Rules, + SecretManagerContext, + SecretManagerRule, +) from litellm.rust_bridge.configuration import Decision, Rollout +from litellm.types.caching import LiteLLMCacheType +from litellm.types.secret_managers.main import KeyManagementSystem @pytest.fixture(autouse=True) @@ -34,7 +47,7 @@ def test_shipped_decisions( configuration.rust(process) if environment is not None: monkeypatch.setenv("LITELLM_RUST", environment) - context: Final = Context(route, provider=provider, model="test-model", delivery=delivery) + context: Final = RouteContext(route, provider=provider, model="test-model", delivery=delivery) if route is Route.OCR: enabled: Final = environment == "1" if environment is not None else process is not False @@ -57,31 +70,63 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt configuration.rust(True) monkeypatch.setenv("LITELLM_RUST", "1") - assert catalog.rollout(Context(route), rules=()) is Rollout.PYTHON_ONLY - assert catalog.decision(Context(route), rules=()) is Decision.PYTHON + assert catalog.rollout(RouteContext(route), rules=()) is Rollout.PYTHON_ONLY + assert catalog.decision(RouteContext(route), rules=()) is Decision.PYTHON + + +@pytest.mark.parametrize( + "context", + ( + *(CacheContext(backend.value) for backend in LiteLLMCacheType), + *(SecretManagerContext(system.value) for system in KeyManagementSystem), + CacheContext("custom"), + SecretManagerContext("unknown"), + ), +) +def test_backend_rollouts_stay_on_python_when_global_rust_is_enabled( + monkeypatch: pytest.MonkeyPatch, context: Context +) -> None: + configuration.rust(True) + monkeypatch.setenv("LITELLM_RUST", "1") + + assert catalog.rollout(context) is Rollout.PYTHON_ONLY + assert catalog.decision(context) is Decision.PYTHON + + +def test_response_cache_rules_select_the_whole_backend_runtime() -> None: + rules: Final = ( + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + CacheRule(Rollout.PYTHON_ONLY), + ) + + assert catalog.decision(CacheContext(backend="local"), rules) is Decision.RUST_REQUIRED + assert catalog.decision(CacheContext(backend="redis"), rules) is Decision.PYTHON @pytest.mark.parametrize( ("context", "expected"), ( - (Context(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.RUST_REQUIRED), - (Context(Route.RESPONSES, provider="openai", model="m"), Decision.PYTHON), - (Context(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.STREAMING), Decision.PYTHON), - (Context(Route.RESPONSES, provider="openai", model="other", delivery=Delivery.WEBSOCKET), Decision.PYTHON), - (Context(Route.RESPONSES, provider="anthropic", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), - (Context(Route.MESSAGES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), + ( + RouteContext(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), + Decision.RUST_REQUIRED, + ), + (RouteContext(Route.RESPONSES, provider="openai", model="m"), Decision.PYTHON), + (RouteContext(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.STREAMING), Decision.PYTHON), + (RouteContext(Route.RESPONSES, provider="openai", model="other", delivery=Delivery.WEBSOCKET), Decision.PYTHON), + (RouteContext(Route.RESPONSES, provider="anthropic", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), + (RouteContext(Route.MESSAGES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), ), ) -def test_first_matching_rule_respects_every_constraint(context: Context, expected: Decision) -> None: +def test_first_matching_rule_respects_every_constraint(context: RouteContext, expected: Decision) -> None: rules: Final = ( - Rule( + RouteRule( Route.RESPONSES, Rollout.RUST_REQUIRED, providers=frozenset({"openai"}), models=frozenset({"m"}), deliveries=frozenset({Delivery.WEBSOCKET}), ), - Rule(Route.RESPONSES, Rollout.PYTHON_ONLY), + RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), ) assert catalog.decision(context, rules) is expected @@ -96,4 +141,78 @@ def test_textract_ocr_has_no_python_path_to_opt_out_to( if environment is not None: monkeypatch.setenv("LITELLM_RUST", environment) - assert catalog.decision(Context(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED + assert catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED + + +@pytest.mark.parametrize( + ("context", "expected"), + ( + (RouteContext(Route.OCR, provider="local"), Decision.RUST_REQUIRED), + (RouteContext(Route.OCR, provider="other"), Decision.PYTHON), + (RouteContext(Route.MESSAGES, provider="local"), Decision.PYTHON), + (CacheContext("local"), Decision.RUST_WITH_FALLBACK), + (CacheContext("other"), Decision.PYTHON), + (SecretManagerContext("local"), Decision.PYTHON), + (SecretManagerContext("other"), Decision.RUST_REQUIRED), + ), +) +def test_mixed_rules_select_only_the_matching_domain(context: Context, expected: Decision) -> None: + rules: Final[Rules] = ( + CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({"local"})), + CacheRule(Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + SecretManagerRule(Rollout.RUST_REQUIRED), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"local"})), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + ) + + assert catalog.decision(context, rules) is expected + + +@pytest.mark.parametrize("context", (RouteContext(Route.OCR), CacheContext("local"), SecretManagerContext("local"))) +@pytest.mark.parametrize( + ("rollout", "process", "environment", "expected"), + ( + (Rollout.PYTHON_ONLY, True, "1", Decision.PYTHON), + (Rollout.RUST_REQUIRED, False, "0", Decision.RUST_REQUIRED), + (Rollout.RUST_OPT_IN, None, None, Decision.PYTHON), + (Rollout.RUST_OPT_OUT, None, None, Decision.RUST_WITH_FALLBACK), + (Rollout.RUST_OPT_IN, True, None, Decision.RUST_WITH_FALLBACK), + (Rollout.RUST_OPT_OUT, False, None, Decision.PYTHON), + (Rollout.RUST_OPT_IN, False, "1", Decision.RUST_WITH_FALLBACK), + (Rollout.RUST_OPT_OUT, True, "0", Decision.PYTHON), + ), +) +def test_all_domains_share_rollout_switches_and_first_match( + monkeypatch: pytest.MonkeyPatch, + context: Context, + rollout: Rollout, + process: bool | None, + environment: str | None, + expected: Decision, +) -> None: + configuration.rust(process) + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + rules: Final[Rules] = ( + RouteRule(Route.OCR, rollout), + CacheRule(rollout), + SecretManagerRule(rollout), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED), + CacheRule(Rollout.RUST_REQUIRED), + SecretManagerRule(Rollout.RUST_REQUIRED), + ) + + assert catalog.decision(context, rules) is expected + assert catalog.decision(context, ()) is Decision.PYTHON + + +@pytest.mark.parametrize("context", (RouteContext(Route.OCR), CacheContext("local"), SecretManagerContext("local"))) +def test_empty_constraints_match_nothing(context: Context) -> None: + rules: Final[Rules] = ( + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset()), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset()), + SecretManagerRule(Rollout.RUST_REQUIRED, systems=frozenset()), + ) + + assert catalog.decision(context, rules) is Decision.PYTHON diff --git a/tests/test_litellm/rust_bridge/test_dispatch.py b/tests/test_litellm/rust_bridge/test_dispatch.py index 66f8d114f7a..9a3793a772e 100644 --- a/tests/test_litellm/rust_bridge/test_dispatch.py +++ b/tests/test_litellm/rust_bridge/test_dispatch.py @@ -6,7 +6,7 @@ import pytest from litellm.rust_bridge import configuration from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule, Rules +from litellm.rust_bridge.catalog import CacheRule, Delivery, Route, RouteContext, RouteRule, Rules, SecretManagerRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.dispatch import PublicDispatch @@ -22,14 +22,15 @@ def binding() -> NativeBinding[object]: return bound -def test_route_without_rules_forwards_before_request_projection() -> None: +@pytest.mark.parametrize("rules", ((), (CacheRule(Rollout.RUST_REQUIRED), SecretManagerRule(Rollout.RUST_REQUIRED)))) +def test_route_without_rules_forwards_before_request_projection(rules: Rules) -> None: stream: Final[Iterator[int]] = iter((1, 2)) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: pytest.fail("Python-only routes must not project the request") dispatch: Final = PublicDispatch( - route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: Context(Route.CHAT_COMPLETIONS) + route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: RouteContext(Route.CHAT_COMPLETIONS) ) result: Final = dispatch.run( ("model",), @@ -37,15 +38,15 @@ def test_route_without_rules_forwards_before_request_projection() -> None: python=lambda *args, **kwargs: stream, binding=binding(), native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"), - rules=(), + rules=rules, ) assert result is stream def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None: rules: Final[Rules] = ( - Rule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), - Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED), ) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: @@ -54,7 +55,7 @@ def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None dispatch: Final = PublicDispatch( route=Route.CHAT_COMPLETIONS, request=reject_request, - context=lambda _: Context(Route.CHAT_COMPLETIONS), + context=lambda _: RouteContext(Route.CHAT_COMPLETIONS), ) expected: Final = object() result: Final = dispatch.run( @@ -69,12 +70,12 @@ def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None def test_disabled_optional_rust_rule_forwards_before_projection() -> None: - rules: Final[Rules] = (Rule(Route.OCR, Rollout.RUST_OPT_OUT),) + rules: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: pytest.fail("Disabled optional Rust must not project the request") - dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR)) + dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR)) expected: Final = object() configuration.rust(False) try: @@ -95,12 +96,14 @@ def test_native_stream_result_is_not_consumed_or_wrapped() -> None: request: Final = Request(model="streaming-model") stream: Final[Iterator[int]] = iter((1, 2)) rules: Final[Rules] = ( - Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.STREAMING})), + CacheRule(Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.STREAMING})), ) dispatch: Final = PublicDispatch( route=Route.CHAT_COMPLETIONS, request=lambda args, kwargs: request, - context=lambda value: Context(Route.CHAT_COMPLETIONS, model=value.model, delivery=Delivery.STREAMING), + context=lambda value: RouteContext(Route.CHAT_COMPLETIONS, model=value.model, delivery=Delivery.STREAMING), ) def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Iterator[int]: @@ -122,7 +125,8 @@ def test_native_stream_result_is_not_consumed_or_wrapped() -> None: @pytest.mark.asyncio -async def test_async_route_without_rules_preserves_async_iterator_result() -> None: +@pytest.mark.parametrize("rules", ((), (CacheRule(Rollout.RUST_REQUIRED), SecretManagerRule(Rollout.RUST_REQUIRED)))) +async def test_async_route_without_rules_preserves_async_iterator_result(rules: Rules) -> None: async def chunks() -> AsyncGenerator[int, None]: yield 1 @@ -135,7 +139,7 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No return stream dispatch: Final = PublicDispatch( - route=Route.RESPONSES, request=reject_request, context=lambda _: Context(Route.RESPONSES) + route=Route.RESPONSES, request=reject_request, context=lambda _: RouteContext(Route.RESPONSES) ) result: Final = await dispatch.arun( ("model",), @@ -143,7 +147,7 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No python=python, binding=binding(), native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"), - rules=(), + rules=rules, ) assert result is stream await stream.aclose() @@ -152,11 +156,13 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No @pytest.mark.asyncio async def test_async_dispatch_accepts_websocket_style_none_result() -> None: request: Final = Request(model="realtime-model") - rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})),) + rules: Final[Rules] = ( + RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})), + ) dispatch: Final = PublicDispatch( route=Route.RESPONSES, request=lambda args, kwargs: request, - context=lambda value: Context(Route.RESPONSES, model=value.model, delivery=Delivery.WEBSOCKET), + context=lambda value: RouteContext(Route.RESPONSES, model=value.model, delivery=Delivery.WEBSOCKET), ) async def python(*args: object, **kwargs: object) -> None: # kwargs-ok: public pass-through shape @@ -183,14 +189,14 @@ async def test_async_dispatch_accepts_websocket_style_none_result() -> None: def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() -> None: rules: Final[Rules] = ( - Rule(Route.MESSAGES, Rollout.RUST_REQUIRED), - Rule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})), ) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: pytest.fail("Rules that cannot select Rust must not project the request") - dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR)) + dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR)) expected: Final = object() result: Final = dispatch.run( ("model",), @@ -206,11 +212,11 @@ def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() - @pytest.mark.asyncio async def test_async_bypass_forwards_to_python_without_native() -> None: request: Final = Request(model="bypassed-model") - rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),) + rules: Final[Rules] = (RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),) dispatch: Final = PublicDispatch( route=Route.RESPONSES, request=lambda args, kwargs: request, - context=lambda value: Context(Route.RESPONSES, model=value.model), + context=lambda value: RouteContext(Route.RESPONSES, model=value.model), bypass=lambda value: value.model == "bypassed-model", ) expected: Final = object() diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index fa6c0b30413..bff7ded3114 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -10,7 +10,7 @@ from litellm.exceptions import APIError from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import bindings, configuration, runtime -from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -39,7 +39,7 @@ class NativeFn(Protocol): def __call__(self) -> str: ... -CONTEXT: Final = Context(Route.MESSAGES, provider="anthropic", model="model") +CONTEXT: Final = RouteContext(Route.MESSAGES, provider="anthropic", model="model") RUST: Final = "rust" PYTHON: Final = "python" @@ -50,8 +50,8 @@ def binding(native: NativeFn | None) -> bindings.NativeBinding[NativeFn]: return bound -def rules(rollout: Rollout) -> tuple[Rule, ...]: - return (Rule(Route.MESSAGES, rollout, providers=frozenset({"anthropic"})),) +def rules(rollout: Rollout) -> tuple[RouteRule, ...]: + return (RouteRule(Route.MESSAGES, rollout, providers=frozenset({"anthropic"})),) class Recorder: @@ -74,7 +74,7 @@ def recorder(native_effect: BaseException | None = None) -> Recorder: return Recorder(native_effect) -def run(rollout: Rollout, calls: Recorder, *, native_missing: bool = False, context: Context = CONTEXT) -> str: +def run(rollout: Rollout, calls: Recorder, *, native_missing: bool = False, context: RouteContext = CONTEXT) -> str: return runtime.run( context, binding=binding(None if native_missing else calls.rust), @@ -146,8 +146,8 @@ def test_context_outside_rule_stays_on_python() -> None: calls: Final = recorder() configuration.rust(True) - assert run(Rollout.RUST_REQUIRED, calls, context=Context(Route.MESSAGES, provider="openai")) == "python" - assert run(Rollout.RUST_REQUIRED, calls, context=Context(Route.RESPONSES, provider="anthropic")) == "python" + assert run(Rollout.RUST_REQUIRED, calls, context=RouteContext(Route.MESSAGES, provider="openai")) == "python" + assert run(Rollout.RUST_REQUIRED, calls, context=RouteContext(Route.RESPONSES, provider="anthropic")) == "python" assert calls.calls == (PYTHON, PYTHON) @@ -155,20 +155,20 @@ def test_context_outside_rule_stays_on_python() -> None: @pytest.mark.parametrize( "context", ( - Context(Route.CHAT_COMPLETIONS, provider="anthropic"), - Context(Route.CHAT_COMPLETIONS, provider="bedrock"), - Context(Route.RESPONSES, provider="openai"), - Context(Route.TRANSCRIPTION, provider="openai"), + RouteContext(Route.CHAT_COMPLETIONS, provider="anthropic"), + RouteContext(Route.CHAT_COMPLETIONS, provider="bedrock"), + RouteContext(Route.RESPONSES, provider="openai"), + RouteContext(Route.TRANSCRIPTION, provider="openai"), ), ) @pytest.mark.parametrize("delivery", tuple(Delivery)) async def test_shipped_python_routes_never_load_native( - monkeypatch: pytest.MonkeyPatch, context: Context, delivery: Delivery + monkeypatch: pytest.MonkeyPatch, context: RouteContext, delivery: Delivery ) -> None: monkeypatch.setenv("LITELLM_RUST", "1") configuration.rust(True) calls: Final = recorder() - request: Final = Context(context.route, provider=context.provider, delivery=delivery) + request: Final = RouteContext(context.route, provider=context.provider, delivery=delivery) def reject_load(value: object) -> NativeFn | None: pytest.fail("Python-only dispatch must not load a native binding") diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 023f02cffbb..3a86a69ed8b 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -1,12 +1,8 @@ -import dataclasses import logging -from pathlib import Path from typing import Final import httpx import pytest -from pydantic import TypeAdapter -from typing_extensions import ReadOnly, TypedDict import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager @@ -15,34 +11,6 @@ from litellm.rust_bridge import settings from litellm.secret_managers.main import get_secret_str from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" - - -class SettingSpec(TypedDict): - adapter: ReadOnly[str] - required: ReadOnly[bool] - precedence: ReadOnly[str] - sensitive: ReadOnly[bool] - shapes: ReadOnly[list[str]] - unsupported_live: ReadOnly[str | None] - - -class SettingsGroup(TypedDict): - version: ReadOnly[int] - fields: ReadOnly[dict[str, SettingSpec]] - - -def test_the_rust_contract_matches_the_returned_fields() -> None: - contract: Final = TypeAdapter(dict[str, SettingsGroup]).validate_json(CONTRACT_PATH.read_text()) - - assert {name: tuple(group["fields"]) for name, group in contract.items()} == { - "http_settings": tuple(field.name for field in dataclasses.fields(settings.http_settings())), - "url_policy": tuple(field.name for field in dataclasses.fields(settings.url_policy())), - "provider_defaults": tuple(field.name for field in dataclasses.fields(settings.provider_defaults())), - "secret_manager": tuple(field.name for field in dataclasses.fields(settings.secret_manager())), - } - - def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "user_url_validation", False) monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"]) @@ -140,6 +108,75 @@ def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.Mon assert settings.secret_manager() == settings.SecretManager(readable=False) +def test_secret_manager_projects_custom_settings(monkeypatch: pytest.MonkeyPatch) -> None: + manager_settings: Final = KeyManagementSettings( + access_mode="read_and_write", + hosted_keys=["MISTRAL_API_KEY"], + primary_secret_name="primary", + aws_region_name="us-east-1", + ) + client: Final = _VaultSecrets({"MISTRAL_API_KEY": "vault-key"}) + monkeypatch.setattr(litellm, "secret_manager_client", client) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", manager_settings) + + assert settings.secret_manager_binding() == settings.SecretManagerBinding( + system="custom", + access_mode="read_and_write", + hosted_keys=["MISTRAL_API_KEY"], + primary_secret_name="primary", + store_virtual_keys=manager_settings.store_virtual_keys, + prefix_for_stored_virtual_keys=manager_settings.prefix_for_stored_virtual_keys, + kms_key_id=manager_settings.kms_key_id, + custom_secret_manager=manager_settings.custom_secret_manager, + aws_region_name="us-east-1", + aws_role_name=manager_settings.aws_role_name, + aws_session_name=manager_settings.aws_session_name, + aws_external_id=manager_settings.aws_external_id, + aws_profile_name=manager_settings.aws_profile_name, + aws_web_identity_token=manager_settings.aws_web_identity_token, + aws_sts_endpoint=manager_settings.aws_sts_endpoint, + replica_regions=manager_settings.replica_regions, + client=client, + settings_object=manager_settings, + ) + + +def test_secret_manager_without_a_client_has_no_system(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "secret_manager_client", None) + + assert settings.secret_manager_binding().system is None + + +def test_secret_manager_uses_key_management_defaults(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "_key_management_settings", None) + + defaults: Final = KeyManagementSettings() + result: Final = settings.secret_manager_binding() + + assert result == settings.SecretManagerBinding( + system=None, + access_mode=defaults.access_mode, + hosted_keys=defaults.hosted_keys, + primary_secret_name=defaults.primary_secret_name, + store_virtual_keys=defaults.store_virtual_keys, + prefix_for_stored_virtual_keys=defaults.prefix_for_stored_virtual_keys, + kms_key_id=defaults.kms_key_id, + custom_secret_manager=defaults.custom_secret_manager, + aws_region_name=defaults.aws_region_name, + aws_role_name=defaults.aws_role_name, + aws_session_name=defaults.aws_session_name, + aws_external_id=defaults.aws_external_id, + aws_profile_name=defaults.aws_profile_name, + aws_web_identity_token=defaults.aws_web_identity_token, + aws_sts_endpoint=defaults.aws_sts_endpoint, + replica_regions=defaults.replica_regions, + client=None, + settings_object=None, + ) + + def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "vertex_project", "configured-project") monkeypatch.setattr(litellm, "vertex_location", "europe-west4") diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 51815651eb4..0d3b8ba472d 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -593,7 +593,7 @@ def test_native_projection_errors_never_select_python( import ssl from litellm.rust_bridge import runtime, settings - from litellm.rust_bridge.catalog import Context, Route, Rule + from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR, LiteLLMOcrRequest @@ -621,11 +621,11 @@ def test_native_projection_errors_never_select_python( with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"): runtime.run( - Context(Route.OCR, provider="mistral"), + RouteContext(Route.OCR, provider="mistral"), binding=NATIVE_OCR, native=lambda native: native(request, (), {}), python=python_fallback, - rules=(Rule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), + rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), ) assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 48ca6f5e165..0f389edaa27 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -38,6 +38,9 @@ from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native +from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime, resolve_response_cache from litellm.types.caching import LiteLLMCacheType from litellm.types.llms.custom_llm import CustomLLMItem from litellm.types.utils import EmbeddingResponse @@ -189,6 +192,7 @@ def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) + assert resolve_response_cache(facade) is None with rebound(litellm, "cache", facade): resolver: Final = _CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" @@ -196,6 +200,42 @@ def test_existing_constructor_and_global_are_unchanged() -> None: assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} +async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + + sync_request: Final = runtime.request(facade, {"cache_key": "sync"}) + assert sync_request is not None + runtime.store(sync_request, {"answer": 1}) + assert runtime.lookup(sync_request) == {"answer": 1} + assert facade.cache.get_cache("sync") is None + + async_request: Final = runtime.request(facade, {"cache_key": "async"}) + assert async_request is not None + await runtime.async_store(async_request, {"answer": 2}) + assert await runtime.async_lookup(async_request) == {"answer": 2} + assert await facade.cache.async_get_cache("async") is None + + requests: Final = (sync_request, async_request) + expected: Final = { + "values": [{"answer": 1}, {"answer": 2}], + "missing_indices": [], + } + assert runtime.lookup_batch(requests) == expected + assert await runtime.async_lookup_batch(requests) == expected + + await runtime.async_flush() + assert runtime.lookup(sync_request) is None + assert await runtime.async_lookup(async_request) is None + + def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: resolver: Final = _CacheTestResolver(litellm) @@ -410,9 +450,7 @@ async def test_memory_size_policy_is_applied_by_the_native_host() -> None: await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _CacheTestResolver( - SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0)) - ).resolve() + disabled: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0))).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None @@ -456,9 +494,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: ) -> object: return result, kwargs - binding: Final = _CacheTestResolver( - SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) - ).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve() assert binding.kind == "python_callback" requests: Final = [request("first"), request("second")] kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] @@ -1212,10 +1248,7 @@ def _semantic_embedding(prompt: str) -> list[float]: base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) pivot: Final = min(range(8), key=lambda index: abs(base[index])) direction: Final = _normalized( - [ - (1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] - for index in range(8) - ] + [(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] for index in range(8)] ) # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) @@ -1311,9 +1344,7 @@ def semantic_embedding() -> Generator[DeterministicEmbedding]: [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook ) ) - stack.enter_context( - rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"]) - ) + stack.enter_context(rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])) yield handler @@ -1441,9 +1472,7 @@ async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(url) - await binding.async_store( - semantic_request("async", "name a primary color"), {"answer": "blue"} - ) + await binding.async_store(semantic_request("async", "name a primary color"), {"answer": "blue"}) hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class @@ -1459,9 +1488,7 @@ async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( [{"answer": 1}, {"answer": 2}], ) expected: Final = { - key: json.loads( - cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response")) - ) + key: json.loads(cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))) for key, prompt in ( ("batch-one", "first batch prompt"), ("batch-two", "second batch prompt"), @@ -1471,18 +1498,19 @@ async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( ("batch-one", "first batch prompt"), ("batch-two", "second batch prompt"), ): - assert cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - key, messages=semantic_messages(prompt) - ) == expected[key], key + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) + == expected[key] + ), key cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class "async-python", json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), messages=semantic_messages("python written prompt"), ) - assert await binding.async_lookup( - semantic_request("async-python", "python written prompt") - ) == {"answer": "python"} + assert await binding.async_lookup(semantic_request("async-python", "python written prompt")) == {"answer": "python"} client.close() @@ -1497,13 +1525,9 @@ async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( SEMANTIC_CONTEXT.set("caller-sentinel") response: Final = {"choices": [{"text": "paris"}]} - await binding.async_store( - semantic_request("inline", "what is the capital of france"), response - ) + await binding.async_store(semantic_request("inline", "what is the capital of france"), response) assert ( - await binding.async_lookup( - semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}") - ) + await binding.async_lookup(semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")) == response ) assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None @@ -1540,9 +1564,7 @@ async def test_native_semantic_cancellation_during_embedding_skips_the_backend( semantic_embedding.gate = asyncio.Event() async def lookup() -> object: - return await binding.async_lookup( - semantic_request("cancel", "cancelled prompt") - ) + return await binding.async_lookup(semantic_request("cancel", "cancelled prompt")) task: Final = asyncio.create_task(lookup()) await semantic_embedding.entered.wait() @@ -1587,9 +1609,7 @@ def test_redis_semantic_ttl_is_written_only_when_requested( binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(url) - binding.store( - {**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1} - ) + binding.store({**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}) expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" assert 0 < client.ttl(expiring) <= 12 @@ -1737,16 +1757,13 @@ def test_redis_semantic_handle_rejects_wrong_backends( redis_semantic_cache_index_name=index, ) subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared - redis_url=url, similarity_threshold=0.8, embedding_model=SEMANTIC_EMBEDDING_MODEL, index_name=index, ) with pytest.raises(TypeError): - _CacheTestHandle.redis_semantic( - subclassed_facade.cache - )._bind_facade(subclassed_facade) + _CacheTestHandle.redis_semantic(subclassed_facade.cache)._bind_facade(subclassed_facade) replacement_facade: Final = Cache( type=LiteLLMCacheType.REDIS_SEMANTIC, @@ -1770,9 +1787,7 @@ def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: ) -def test_qdrant_semantic_facade_binds_native_and_shares_entries( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: +def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "shared prompt"}] collection: Final = f"cache_{uuid4().hex}" @@ -1803,9 +1818,7 @@ def test_qdrant_semantic_facade_binds_native_and_shares_entries( assert facade.cache.get_cache("different-key", messages=messages) is None -async def test_qdrant_semantic_async_parity( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: +async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "async prompt"}] collection: Final = f"cache_{uuid4().hex}" @@ -1905,9 +1918,7 @@ async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( await binding.ping() -def test_qdrant_semantic_ignores_request_expiry( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: +def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "persistent prompt"}] collection: Final = f"cache_{uuid4().hex}" @@ -1928,9 +1939,7 @@ def test_qdrant_semantic_ignores_request_expiry( assert python_value["response"] == {"id": "persistent"} -def test_qdrant_semantic_mutation_and_projection_fallback( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: +def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) diff --git a/tests/unit/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py index 63821c74208..2807ed7f8f7 100644 --- a/tests/unit/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -10,7 +10,7 @@ from litellm.chat_completions.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule +from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.chat_completions.entrypoints import ( NATIVE_ACOMPLETION, NATIVE_COMPLETION, @@ -23,7 +23,7 @@ from litellm.types.utils import ModelResponse MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final = () -RUST_RULES: Final = (Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) +RUST_RULES: Final = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) def completion_binding(native: NativeCompletion | None) -> NativeBinding[NativeCompletion]: @@ -117,9 +117,7 @@ def test_native_receives_bound_request_and_original_call_shape() -> None: "custom_llm_provider": "anthropic", "metadata": metadata, } - captured: Final[ - list[tuple[LiteLLMChatCompletionsRequest, tuple[object, ...], Mapping[str, object]]] - ] = [] + captured: Final[list[tuple[LiteLLMChatCompletionsRequest, tuple[object, ...], Mapping[str, object]]]] = [] def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: rejected Rust fallback pytest.fail("Required Rust dispatch must not call Python") diff --git a/tests/unit/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py index 586b77d9a25..48eb1adbf51 100644 --- a/tests/unit/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -12,7 +12,7 @@ from litellm.messages.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule, Rules +from litellm.rust_bridge.catalog import Route, RouteRule, Rules from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.messages.entrypoints import ( NATIVE_AMESSAGES, @@ -25,7 +25,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMe MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final[Rules] = () -RUST_RULES: Final[Rules] = (Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),) +RUST_RULES: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]: diff --git a/tests/unit/ocr/test_dispatch.py b/tests/unit/ocr/test_dispatch.py index e54d4070ba8..2727ff23449 100644 --- a/tests/unit/ocr/test_dispatch.py +++ b/tests/unit/ocr/test_dispatch.py @@ -12,7 +12,7 @@ from litellm.ocr.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule, Rules +from litellm.rust_bridge.catalog import Route, RouteRule, Rules from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.ocr.entrypoints import ( NATIVE_AOCR, @@ -22,8 +22,8 @@ from litellm.rust_bridge.ocr.entrypoints import ( NativeOcr, ) -PYTHON_RULES: Final[Rules] = (Rule(Route.OCR, Rollout.PYTHON_ONLY),) -RUST_RULES: Final[Rules] = (Rule(Route.OCR, Rollout.RUST_REQUIRED),) +PYTHON_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.PYTHON_ONLY),) +RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),) def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]: @@ -403,8 +403,8 @@ def test_provider_scoped_rule_sees_the_provider_named_by_the_model_prefix( model: str, custom_llm_provider: str | None, expected: str ) -> None: rules: Final[Rules] = ( - Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), - Rule(Route.OCR, Rollout.PYTHON_ONLY), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), ) document: Final[Mapping[str, object]] = {"type": "image_url", "image_url": "data:image/png;base64,YQ=="} kwargs: Final[Mapping[str, object]] = (