From 0d76359dc9a4e1dba45020626f143e1f1f294bff Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:57:24 -0700 Subject: [PATCH] fix(rust): resolve OCR provider env fallbacks through the secret manager Python reads every provider credential fallback (MISTRAL_API_KEY, AZURE_AI_API_KEY, AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT, Azure AD and Vertex env, ...) through get_secret_str, which consults the configured key_management_system before os.environ. Native OCR read std::env directly, so a key held only in the vault went missing and a stale env copy silently won. OcrClient now carries an injected secret Lookup that the connection exposes to providers and auth crates; the bridge backs it with settings.secret -> get_secret_str, pure Rust keeps the process env. Co-Authored-By: Claude Opus 5 --- litellm-rust/crates/core/src/ocr/handler.rs | 2 +- litellm-rust/crates/core/src/ocr/prepare.rs | 23 +++++-- litellm-rust/crates/core/tests/ocr.rs | 25 +++++++ .../ocr/cohere_parse_transformation.rs | 2 +- .../document_intelligence/transformation.rs | 10 +-- .../llms/src/azure_ai/ocr/transformation.rs | 12 ++-- .../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 | 18 +++-- .../llms/src/cohere/ocr/transformation.rs | 6 +- .../llms/src/mistral/ocr/transformation.rs | 6 +- .../llms/src/reducto/ocr/transformation.rs | 6 +- .../vertex_ai/ocr/deepseek_transformation.rs | 7 +- .../llms/src/vertex_ai/ocr/transformation.rs | 9 +-- .../python-bridge/src/python_settings.rs | 65 ++++++++++++++++++- .../python-bridge/src/routes/ocr/mod.rs | 5 +- litellm/rust_bridge/settings.py | 6 ++ .../test_litellm/rust_bridge/test_settings.py | 46 +++++++++++++ 18 files changed, 226 insertions(+), 41 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 126e79e20e7..19037e49033 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -22,7 +22,7 @@ pub(crate) async fn perform_ocr_request( ) -> Result { request.response_format()?; let config = request.config; - let request = prepare_request(request, caller_document, client.settings()); + let request = prepare_request(request, caller_document, client); 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 72c35469f6d..ed8c7fba503 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,7 +1,7 @@ use litellm_auth::{InputSource, SecretValue, Sourced}; use litellm_llms::base_llm::ocr::{ - settings::OcrSettings, - transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env}, + handler::OcrClient, + transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest}, }; use super::provider_config::OcrProvider; @@ -10,7 +10,7 @@ use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest}; pub(crate) fn prepare_request( request: ResolvedOcrRequest, caller_document: bool, - settings: &OcrSettings, + client: &OcrClient, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); let api_base_env = match request.config.provider() { @@ -23,14 +23,14 @@ pub(crate) fn prepare_request( request .config .get_api_key_env_var() - .and_then(credential_env) + .and_then(|name| client.secrets().get(name)) .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { credentials.api_base.clone().or_else(|| { api_base_env - .and_then(credential_env) + .and_then(|name| client.secrets().get(name)) .map(|value| Sourced::new(value, InputSource::Environment)) }) }); @@ -53,7 +53,12 @@ pub(crate) fn prepare_request( PreparedOcrRequest { model, document, - connection: OcrConnection::new(resolved, transport, settings.clone()), + connection: OcrConnection::new( + resolved, + transport, + client.settings().clone(), + client.secrets().clone(), + ), caller_document, optional_params, input_sources, @@ -63,7 +68,11 @@ pub(crate) fn prepare_request( #[cfg(test)] pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest { - prepare_request(request, true, &OcrSettings::default()) + prepare_request( + request, + true, + &OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()), + ) } #[cfg(test)] diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index f87f16cd033..61d59a38065 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -174,6 +174,30 @@ async fn facade_retains_native_response_when_requested() { ); } +#[tokio::test] +async fn provider_key_fallback_reads_the_injected_secret_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = decode_request(OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: None, + api_base: Some(base.clone()), + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + let client = ocr_client().with_secrets(Arc::new(|name: &str| { + (name == "MISTRAL_API_KEY").then(|| "from-secret-manager".to_string()) + })); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer from-secret-manager")); +} + #[tokio::test] async fn ocr_client_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; @@ -187,6 +211,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { UrlPolicy::default(), VertexAuth::default(), OcrSettings::default(), + Arc::new(litellm_core_utils::settings::ProcessEnvironment), ) .unwrap(); crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) 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 86ee0d96895..045d8744bc9 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 @@ -53,7 +53,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig { ) -> Result { let base = super::transformation::AzureAiOcrConfig::resolve_api_base( request.connection.api_base.as_deref(), - &crate::base_llm::ocr::transformation::credential_env, + &|name: &str| request.connection.secret(name), )?; self.get_complete_url(&base) } 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 5fb20d5900a..8e6182f454f 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 @@ -23,7 +23,7 @@ use crate::base_llm::ocr::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage, OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - ResolvedOcrCredentials, credential_env, decode_and_normalize_response, decode_response, + ResolvedOcrCredentials, decode_and_normalize_response, decode_response, }, }; @@ -181,8 +181,10 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { &request.input_sources, )? }; - self.resolve_headers(&request.connection, &config, &credential_env) - .await + self.resolve_headers(&request.connection, &config, &|name: &str| { + request.connection.secret(name) + }) + .await } fn get_complete_url( @@ -192,7 +194,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { _environment: &Self::Environment, ) -> Result { let endpoint = nonblank(request.connection.api_base.clone()) - .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) + .or_else(|| nonblank(request.connection.secret(AZURE_DI_ENDPOINT_ENV))) .ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; self.build_ocr_url( &endpoint, 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 4a04910aa9a..cd20e75df85 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -10,7 +10,7 @@ use crate::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrRequestContext, - OcrResponseFormat, PreparedOcrRequest, credential_env, + OcrResponseFormat, PreparedOcrRequest, }, }, mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}, @@ -57,8 +57,10 @@ impl BaseOcrConfig for AzureAiOcrConfig { &request.input_sources, )? }; - self.resolve_headers(&request.connection, &config, &credential_env) - .await + self.resolve_headers(&request.connection, &config, &|name: &str| { + request.connection.secret(name) + }) + .await } fn get_complete_url( @@ -67,7 +69,9 @@ impl BaseOcrConfig for AzureAiOcrConfig { _optional_params: &Self::OcrParams, _environment: &Self::Environment, ) -> Result { - self.build_ocr_url(request.connection.api_base.as_deref(), &credential_env) + self.build_ocr_url(request.connection.api_base.as_deref(), &|name: &str| { + request.connection.secret(name) + }) } fn transform_ocr_request( 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 9410f673d29..91fb6461770 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -13,7 +13,7 @@ use serde_json::Value; use crate::base_llm::ocr::{ error::Error, - settings::OcrSettings, + settings::{OcrSettings, Secrets}, transformation::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value, decode_response, @@ -35,6 +35,7 @@ pub struct OcrClient { document_fetcher: MediaFetcher, vertex_auth: VertexAuth, settings: OcrSettings, + secrets: Secrets, } impl OcrClient { @@ -44,6 +45,7 @@ impl OcrClient { url_policy: UrlPolicy, vertex_auth: VertexAuth, settings: OcrSettings, + secrets: Secrets, ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, @@ -51,6 +53,7 @@ impl OcrClient { document_fetcher: MediaFetcher::new(pool, config, url_policy)?, vertex_auth, settings, + secrets, }) } @@ -74,6 +77,10 @@ impl OcrClient { &self.settings } + pub fn secrets(&self) -> &Secrets { + &self.secrets + } + #[cfg(any(test, feature = "test-support"))] pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { Self { @@ -85,6 +92,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), } } @@ -92,6 +100,11 @@ impl OcrClient { pub fn with_settings(self, settings: OcrSettings) -> Self { Self { settings, ..self } } + + #[cfg(any(test, feature = "test-support"))] + pub fn with_secrets(self, secrets: Secrets) -> Self { + Self { secrets, ..self } + } } /// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request, 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 239a5b22000..276b2ca1311 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -1,7 +1,9 @@ -use std::time::Duration; +use std::{sync::Arc, 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 5d1a0c8e0ed..3960282b580 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,9 +1,10 @@ -use std::{collections::BTreeMap, future::Future, time::Duration}; +use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, + settings::ProcessEnvironment, }; use serde::{ Deserialize, Serialize, @@ -15,7 +16,7 @@ use serde_with::serde_as; use crate::base_llm::ocr::{ error::Error, handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, - settings::OcrSettings, + settings::{OcrSettings, Secrets}, }; pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; @@ -160,6 +161,7 @@ pub struct OcrConnection { pub timeout: Duration, pub max_response_bytes: usize, pub settings: OcrSettings, + pub secrets: Secrets, } impl OcrConnection { @@ -167,6 +169,7 @@ impl OcrConnection { credentials: ResolvedOcrCredentials, transport: OcrTransportConfig, settings: OcrSettings, + secrets: Secrets, ) -> Self { let api_key_source = credentials .api_key @@ -191,8 +194,13 @@ impl OcrConnection { .unwrap_or(settings.request_timeout), max_response_bytes: transport.max_response_bytes, settings, + secrets, } } + + pub fn secret(&self, name: &str) -> Option { + self.secrets.get(name) + } } impl Default for OcrConnection { @@ -201,6 +209,7 @@ impl Default for OcrConnection { ResolvedOcrCredentials::default(), OcrTransportConfig::default(), OcrSettings::default(), + Arc::new(ProcessEnvironment), ) } } @@ -563,10 +572,6 @@ pub fn decode_and_normalize_response( }) } -pub fn credential_env(name: &str) -> Option { - std::env::var(name).ok() -} - #[cfg(test)] mod tests { use serde_json::json; @@ -587,6 +592,7 @@ mod tests { ..OcrTransportConfig::default() }, settings.clone(), + Arc::new(ProcessEnvironment), ) .timeout }; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index da6cf90ffcf..d141c68db38 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -13,7 +13,7 @@ use crate::base_llm::ocr::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, - OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, + OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, decode_response_value, }, }; @@ -122,7 +122,9 @@ impl BaseOcrConfig for CohereParseConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - self.resolve_headers(&request.connection, &credential_env) + self.resolve_headers(&request.connection, &|name: &str| { + request.connection.secret(name) + }) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 95658837fc3..2b14372fbec 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -7,7 +7,7 @@ use crate::base_llm::ocr::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, - OcrUsageInfo, PreparedOcrRequest, credential_env, decode_and_normalize_response, + OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, }, }; @@ -84,7 +84,9 @@ impl BaseOcrConfig for MistralOcrConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - self.resolve_headers(&request.connection, &credential_env) + self.resolve_headers(&request.connection, &|name: &str| { + request.connection.secret(name) + }) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 740f0ced090..307ba697316 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -15,7 +15,7 @@ use crate::base_llm::ocr::{ transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, OcrPage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - credential_env, decode_and_normalize_response, + decode_and_normalize_response, }, }; @@ -110,7 +110,9 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - resolve_headers(&request.connection, &credential_env) + resolve_headers(&request.connection, &|name: &str| { + request.connection.secret(name) + }) } fn get_complete_url( 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 8009a65ff77..6fa0b9c5977 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 @@ -9,7 +9,7 @@ use crate::base_llm::ocr::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, - OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, + OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, decode_response_value, }, }; @@ -126,8 +126,9 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { &request.optional_params, &request.input_sources, )?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let location = + vertex::get_vertex_ai_location(&config, &|name: &str| request.connection.secret(name)) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); self.get_complete_url( request.connection.api_base.as_deref(), &environment.project_id, 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 a50e8261aa3..f7941db9364 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -10,7 +10,7 @@ use crate::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrEnvironment, - OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, credential_env, + OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, }, }, mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}, @@ -65,8 +65,9 @@ impl BaseOcrConfig for VertexAiOcrConfig { &request.optional_params, &request.input_sources, )?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let location = + vertex::get_vertex_ai_location(&config, &|name: &str| request.connection.secret(name)) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); self.build_ocr_url( request.connection.api_base.as_deref(), &environment.project_id, @@ -139,7 +140,7 @@ impl VertexAiOcrConfig { .as_ref() .map(litellm_auth::SecretValue::expose), config, - &credential_env, + &|name: &str| connection.secret(name), ) .await .map_err(Error::from) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 79921d67452..272e711ada5 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::settings::Lookup; use pyo3::prelude::*; const MODULE: &str = "litellm.rust_bridge.settings"; @@ -29,6 +30,23 @@ impl PythonSettings { } } +pub(crate) struct PythonSecrets; + +impl Lookup for PythonSecrets { + fn get(&self, name: &str) -> Option { + Python::attach(|py| { + py.import(MODULE) + .and_then(|module| module.getattr("secret")?.call1((name,))) + .and_then(|value| value.extract::>()) + .unwrap_or_else(|error| { + let _ = + PythonSettings::warn(py, &format!("reading secret {name} failed: {error}")); + None + }) + }) + } +} + #[cfg(test)] pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); @@ -36,9 +54,10 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); mod tests { use std::{collections::BTreeSet, ffi::CString}; + use litellm_core_utils::settings::Lookup; use pyo3::{prelude::*, types::PyDict}; - use super::{CONTRACT, PythonSettings}; + use super::{CONTRACT, PythonSecrets, PythonSettings}; #[test] fn every_settings_group_is_in_the_python_contract() { @@ -62,4 +81,48 @@ mod tests { assert_eq!(read, declared); }); } + + #[test] + fn secrets_come_from_the_python_secret_reader_and_a_failed_read_is_unset() { + Python::initialize(); + Python::attach(|py| { + py.run( + c" +import sys +import types +settings = types.ModuleType('litellm.rust_bridge.settings') +settings.warnings = [] +def secret(name): + if name == 'BROKEN': + raise RuntimeError('vault down') + return {'MISTRAL_API_KEY': 'from-vault'}.get(name) +settings.secret = secret +settings.warn = settings.warnings.append +sys.modules.setdefault('litellm', types.ModuleType('litellm')) +sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) +sys.modules['litellm.rust_bridge.settings'] = settings +", + None, + None, + ) + .unwrap(); + }); + assert_eq!( + PythonSecrets.get("MISTRAL_API_KEY").as_deref(), + Some("from-vault") + ); + assert_eq!(PythonSecrets.get("ABSENT"), None); + assert_eq!(PythonSecrets.get("BROKEN"), None); + Python::attach(|py| { + let warnings: Vec = py + .import("litellm.rust_bridge.settings") + .unwrap() + .getattr("warnings") + .unwrap() + .extract() + .unwrap(); + assert_eq!(warnings.len(), 1); + assert!(warnings[0].contains("BROKEN") && warnings[0].contains("vault down")); + }); + } } 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 bb845f48783..966b24a82e7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -3,7 +3,7 @@ mod errors; mod host; mod project; -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; use host::OcrRouteHost; use litellm_auth_gcp::VertexAuth; @@ -16,7 +16,7 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -use crate::{errors::RustBridgeDeclined, http}; +use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSecrets}; const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -45,6 +45,7 @@ fn run_ocr( http::url_policy(py)?, VERTEX_AUTH.clone(), OcrSettings::from_environment(&ProcessEnvironment), + Arc::new(PythonSecrets), ) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index e170f93b198..210ef7ac6b4 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -30,6 +30,12 @@ def warn(message: str) -> None: verbose_logger.warning("%s", message) +def secret(name: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(name) + + def url_policy() -> UrlPolicy: import litellm diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index f75145c2b2c..f3baf463b87 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -3,12 +3,15 @@ import logging from pathlib import Path from typing import Final +import httpx import pytest from pydantic import TypeAdapter import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent from litellm.rust_bridge import settings +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" @@ -73,3 +76,46 @@ def test_warn_reaches_the_litellm_logger(caplog: pytest.LogCaptureFixture) -> No settings.warn("ssl_ecdh_curve 'secp521r1' is not supported") assert [record.getMessage() for record in caplog.records] == ["ssl_ecdh_curve 'secp521r1' is not supported"] + + +class _VaultSecrets(CustomSecretManager): + def __init__(self, secrets: dict[str, str]) -> None: + super().__init__(secret_manager_name="rust_bridge_settings_test") + self.secrets = secrets + + 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 self.secrets.get(secret_name) + + def sync_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return self.secrets.get(secret_name) + + +def test_secret_prefers_the_secret_manager_and_falls_back_to_the_environment_on_a_miss( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "stale-env-key") + monkeypatch.setenv("REDUCTO_API_KEY", "env-only-key") + monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({"MISTRAL_API_KEY": "vault-key"})) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + + assert settings.secret("MISTRAL_API_KEY") == "vault-key" + assert settings.secret("REDUCTO_API_KEY") == "env-only-key" + assert settings.secret("ABSENT_KEY") is None + + +def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "env-key") + monkeypatch.setattr(litellm, "secret_manager_client", None) + + assert settings.secret("MISTRAL_API_KEY") == "env-key"