diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index ed8c7fba503..f1e1dcaaa6d 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -13,24 +13,28 @@ pub(crate) fn prepare_request( client: &OcrClient, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); - let api_base_env = match request.config.provider() { - OcrProvider::Mistral => Some("MISTRAL_API_BASE"), - OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"), - OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => None, + let (preferred_api_key_env, api_base_env) = match request.config.provider() { + OcrProvider::Mistral => ( + Some("MISTRAL_AZURE_API_KEY"), + Some("MISTRAL_AZURE_API_BASE"), + ), + OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")), + OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None), }; + let secret = |name: &str| client.secrets().truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { credentials.api_key.clone().or_else(|| { - request - .config - .get_api_key_env_var() - .and_then(|name| client.secrets().get(name)) + preferred_api_key_env + .into_iter() + .chain(request.config.get_api_key_env_var()) + .find_map(secret) .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(|name| client.secrets().get(name)) + .and_then(secret) .map(|value| Sourced::new(value, InputSource::Environment)) }) }); diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 61d59a38065..3aedc7b9023 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -174,14 +174,30 @@ async fn facade_retains_native_response_when_requested() { ); } +#[rstest] +#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] +#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] +#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] #[tokio::test] -async fn provider_key_fallback_reads_the_injected_secret_source() { +async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( + #[case] secrets: &'static [(&'static str, &'static str)], + #[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 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()), + api_base: None, custom_llm_provider: None, extra_headers: None, optional_params: Default::default(), @@ -189,13 +205,10 @@ async fn provider_key_fallback_reads_the_injected_secret_source() { 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")); + assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); } #[tokio::test]