diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs index 71ca69ddc58..66a02778bef 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs @@ -32,11 +32,14 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { client: &OcrClient, ) -> Result { let params = map_ocr_params(request)?; - let config = AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; + let config = AzureAuthInputs { + azure_ad_token_provider: request.connection.token_provider.clone(), + ..AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)? + }; let headers = validate_environment(&request.connection, &config, &credential_env).await?; let endpoint = nonblank(request.connection.api_base.clone()) .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs index 3107494d39e..6dbc613cdb7 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs @@ -33,13 +33,16 @@ impl OcrAdapter for AzureMistralAdapter { known: params, extra_params: _extra_params, } = _prepare_ocr_request::(request)?; - let config = AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; - let headers = validate_environment(&request.connection, &config, &credential_env).await?; + let config = AzureAuthInputs { + azure_ad_token_provider: request.connection.token_provider.clone(), + ..AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)? + }; let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; + let headers = validate_environment(&request.connection, &config, &credential_env).await?; let document = inline_remote_document( client.document_fetcher(), request.document.clone(), @@ -88,7 +91,9 @@ async fn validate_environment( config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, OcrError> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + if config.azure_ad_token_provider.is_none() + && crate::http_utils::has_header(&connection.extra_headers, "authorization") + { super::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs index 062a0071a34..9e2abd5ccc8 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs @@ -32,7 +32,7 @@ impl OcrAdapter for ReductoLegacyAdapter { super::prepare_document(client, document, &request.connection, &headers).await?; let body = reducto::transform_legacy_ocr_request(&request.model, document, ¶ms)?; let body = merge_extra_params(&body, extra_params)?; - build_http_request(client, request, &url, &headers, &body) + build_http_request(client, request, &url, &headers, &body).await } fn transform_ocr_response( diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs index a49f8105e26..4f398611334 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs @@ -32,7 +32,7 @@ impl OcrAdapter for ReductoV3Adapter { super::prepare_document(client, document, &request.connection, &headers).await?; let body = reducto::transform_v3_ocr_request(&request.model, document, ¶ms)?; let body = merge_extra_params(&body, extra_params)?; - build_http_request(client, request, &url, &headers, &body) + build_http_request(client, request, &url, &headers, &body).await } fn transform_ocr_response( diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 7dd3c6bf8b2..3ad706648d8 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -2,11 +2,12 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use serde::Serialize; +use serde_json::{Map, Value}; + use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument}; use crate::Error; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; -use serde::Serialize; -use serde_json::Value; pub type OcrHookFuture<'a, T> = Pin> + Send + 'a>>; pub type OcrLogFuture<'a> = Pin + Send + 'a>>; @@ -16,7 +17,7 @@ pub struct OcrPreCallRequest { pub model: String, pub custom_llm_provider: String, pub document: OcrDocument, - pub optional_params: Value, + pub optional_params: Map, } #[derive(Clone, Debug, Serialize)] @@ -27,9 +28,19 @@ pub struct OcrDuringCallRequest { pub body: Value, } +pub struct OcrPreparedRequest { + pub model: String, + pub url: String, + pub headers: Vec<(String, String)>, + pub body: Value, +} + pub trait OcrHooks: Send + Sync { - fn has_guardrails(&self) -> bool { - false + fn prepared_request( + &self, + request: OcrPreparedRequest, + ) -> OcrHookFuture<'_, OcrPreparedRequest> { + Box::pin(async move { Ok(request) }) } fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { Box::pin(async move { Ok(request) }) @@ -69,10 +80,22 @@ pub(crate) struct OcrLifecycleHooks { impl CallLifecycleHooks for OcrLifecycleHooks { - type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; - type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; - type SuccessFuture<'a> = OcrLogFuture<'a>; - type FailureFuture<'a> = OcrLogFuture<'a>; + type PreCallFuture<'a> + = OcrHookFuture<'a, LiteLLMOcrRequest> + where + Self: 'a; + type DuringCallFuture<'a> + = OcrHookFuture<'a, LiteLLMOcrRequest> + where + Self: 'a; + type SuccessFuture<'a> + = OcrLogFuture<'a> + where + Self: 'a; + type FailureFuture<'a> + = OcrLogFuture<'a> + where + Self: 'a; fn async_pre_call_hook<'a>( &'a self, @@ -80,27 +103,18 @@ impl CallLifecycleHooks Self::PreCallFuture<'a> { Box::pin(async move { - if !self.hooks.has_guardrails() { - return Ok(request); - } let changed = self .hooks .pre_call(OcrPreCallRequest { model: request.model.clone(), custom_llm_provider: self.provider_name.clone(), document: request.document, - optional_params: Value::Object(request.optional_params), + optional_params: request.optional_params, }) .await?; - let Value::Object(optional_params) = changed.optional_params else { - return Err(super::error::OcrRequestError::RequestField { - path: "guardrail.optional_params".into(), - } - .into()); - }; Ok(LiteLLMOcrRequest { document: changed.document, - optional_params, + optional_params: changed.optional_params, ..request }) }) diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index bf6f924088c..64c4191bb12 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -3,7 +3,7 @@ use serde_json::{Map, Value}; use super::OcrClient; use super::error::{OcrError, OcrRequestError}; -use super::hooks::OcrDuringCallRequest; +use super::hooks::{OcrDuringCallRequest, OcrPreparedRequest}; use super::types::{LiteLLMOcrRequest, OcrDocument}; #[derive(Debug, Deserialize)] @@ -68,55 +68,63 @@ pub(crate) async fn transform_request_body( where B: Serialize + DeserializeOwned, { - let body = if request.hooks.has_guardrails() { - let changed = request - .hooks - .during_call(OcrDuringCallRequest { - model: request.model.clone(), - custom_llm_provider: request.adapter.provider().as_str().into(), - url: url.into(), - body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { - path: "body".into(), - })?, - }) - .await?; - let body = OcrWireBody::::decode(changed.body)?; - validate(&body.body)?; - body - } else { - OcrWireBody { - body, - extra: Map::new(), - } - }; - build_http_request(client, request, url, headers, &body) + let changed = request + .hooks + .during_call(OcrDuringCallRequest { + model: request.model.clone(), + custom_llm_provider: request.adapter.provider().as_str().into(), + url: url.into(), + body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { + path: "body".into(), + })?, + }) + .await?; + let body = OcrWireBody::::decode(changed.body)?; + validate(&body.body)?; + build_http_request(client, request, url, headers, &body).await } -pub(crate) fn build_http_request( +pub(crate) async fn build_http_request( client: &OcrClient, request: &LiteLLMOcrRequest, url: &str, headers: &[(String, String)], body: &B, -) -> Result { +) -> Result +where + B: Serialize + DeserializeOwned, +{ + let prepared = request + .hooks + .prepared_request(OcrPreparedRequest { + model: request.model.clone(), + url: url.into(), + headers: headers.to_vec(), + body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { + path: "body".into(), + })?, + }) + .await?; + let body: B = super::wire::decode_request_value(prepared.body, "guardrail.body")?; let builder = client .provider_http() - .post(url) - .json(body) + .post(&prepared.url) + .json(&body) .timeout(request.connection.timeout); - crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All) - .build() - .map_err(crate::error::TransportError::from) - .map_err(OcrError::from) + crate::http_utils::with_headers( + builder, + &prepared.headers, + crate::http_utils::HeaderPolicy::All, + ) + .build() + .map_err(crate::error::TransportError::from) + .map_err(OcrError::from) } pub(crate) async fn guardrail_document( request: &LiteLLMOcrRequest, url: &str, ) -> Result { - if !request.hooks.has_guardrails() { - return Ok(request.document.clone()); - } let changed = request .hooks .during_call(OcrDuringCallRequest { @@ -133,7 +141,7 @@ pub(crate) async fn guardrail_document( super::wire::decode_request_value(changed.body, "guardrail.document").map_err(OcrError::from) } -#[derive(Serialize)] +#[derive(Serialize, Deserialize)] struct OcrWireBody { #[serde(flatten)] body: B, diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 06519f86c91..f84cb9f02c6 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -8,7 +8,7 @@ use serde_json::{Map, Value}; use super::hooks::{NoopOcrHooks, OcrHooks}; use super::registry::{OcrAdapterKind, resolve_wire_adapter}; use crate::Error; -use crate::auth::InputSource; +use crate::auth::{InputSource, TokenProviderHandle}; use crate::constants::OCR_HTTP_TIMEOUT_SECS; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -60,6 +60,7 @@ pub enum OcrResponseFormat { #[derive(Clone)] pub struct OcrConnection { + pub token_provider: Option, pub api_key: Option, pub api_key_source: InputSource, pub api_base: Option, @@ -74,6 +75,7 @@ pub struct OcrConnection { impl Default for OcrConnection { fn default() -> Self { Self { + token_provider: None, api_key: None, api_key_source: InputSource::Deployment, api_base: None, diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 34d0a7d7b86..215c23735eb 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -73,6 +73,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result wire.optional_params, )?; let connection = OcrConnection { + token_provider: None, api_key: nonblank(wire.api_key), api_key_source, api_base: nonblank(wire.api_base), @@ -150,7 +151,7 @@ pub fn decode_pre_call_result( let changed: Changed = decode_request_value(value, "guardrail")?; Ok(OcrPreCallRequest { document: changed.document, - optional_params: Value::Object(changed.optional_params), + optional_params: changed.optional_params, ..original }) } diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs index d7d532cfef1..24d42fb0bca 100644 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -70,10 +70,6 @@ async fn facade_acquires_supplied_entra_token_for_final_request() { struct ReplaceBodyDocument; impl OcrHooks for ReplaceBodyDocument { - fn has_guardrails(&self) -> bool { - true - } - fn during_call( &self, mut request: OcrDuringCallRequest, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index e4c81dea5a7..083a4fd1e8b 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -361,15 +361,14 @@ async fn pre_call_guardrail_receives_caller_pages_before_mapping() { struct RewritePages; impl OcrHooks for RewritePages { - fn has_guardrails(&self) -> bool { - true - } - fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { Box::pin(async move { assert_eq!(request.optional_params["pages"], json!([0, 2])); Ok(OcrPreCallRequest { - optional_params: json!({"pages": [1]}), + optional_params: serde_json::Map::from_iter([( + "pages".to_string(), + json!([1]), + )]), ..request }) }) diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index cecd8869741..1ceb12df64a 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -3,7 +3,7 @@ use std::sync::{Arc, Mutex}; use serde_json::{Value, json}; use super::OcrClient; -use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest}; +use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest, OcrPreparedRequest}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use super::wire::{OcrWireRequest, decode_request}; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; @@ -123,9 +123,57 @@ struct RecordingHooks { block: bool, } +struct EditPreparedRequest; + +impl OcrHooks for EditPreparedRequest { + fn prepared_request( + &self, + mut request: OcrPreparedRequest, + ) -> OcrHookFuture<'_, OcrPreparedRequest> { + Box::pin(async move { + assert_eq!(request.model, "model"); + assert!(request.url.ends_with("/v1/ocr")); + assert!( + request + .headers + .contains(&("Authorization".into(), "Bearer test-key".into())) + ); + request.body["include_image_base64"] = json!(true); + request + .headers + .push(("x-host-hook".into(), "called".into())); + Ok(request) + }) + } +} + +#[tokio::test] +async fn prepared_request_hook_edits_wire_body_and_headers_without_guardrails() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = wire_request( + "mistral/model", + &base, + json!({"include_image_base64":false}), + ) + .with_host_hooks(Arc::new(EditPreparedRequest), None); + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].contains("x-host-hook: called\r\n")); + let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body["include_image_base64"], true); +} + impl OcrHooks for RecordingHooks { - fn has_guardrails(&self) -> bool { - true + fn prepared_request( + &self, + request: OcrPreparedRequest, + ) -> OcrHookFuture<'_, OcrPreparedRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("prepared"); + Ok(request) + }) } fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { @@ -185,7 +233,10 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { }; perform_ocr(request).await.unwrap(); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]); + assert_eq!( + *events.lock().unwrap(), + ["pre", "during", "prepared", "success"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -224,6 +275,190 @@ async fn upstream_failure_emits_one_terminal_failure() { }; assert!(perform_ocr(request).await.is_err()); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["pre", "during", "prepared", "failure"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } + +#[tokio::test] +async fn every_adapter_runs_the_complete_lifecycle() { + use super::registry::OcrAdapterKind; + + macro_rules! adapter_kinds { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + [$(OcrAdapterKind::$variant,)+] + }; + } + for adapter in super::adapters::for_each_ocr_adapter!(adapter_kinds) { + let (model, response) = match adapter { + OcrAdapterKind::Mistral => ("mistral/model", json!({"pages":[]})), + OcrAdapterKind::AzureMistral => ("azure_ai/model", json!({"pages":[]})), + OcrAdapterKind::AzureDocumentIntelligence => ( + "azure_ai/documentintelligence/prebuilt-read", + json!({"status":"succeeded", "analyzeResult":{"pages":[]}}), + ), + OcrAdapterKind::ReductoLegacy => { + ("reducto/parse-legacy", json!({"result":{"chunks":[]}})) + } + OcrAdapterKind::ReductoV3 => ("reducto/parse-v3", json!({"result":{"chunks":[]}})), + OcrAdapterKind::VertexMistral => ("vertex_ai/mistral-ocr", json!({"pages":[]})), + OcrAdapterKind::VertexDeepSeek => ( + "vertex_ai/deepseek-ocr", + json!({"choices":[{"message":{"content":"text"}}]}), + ), + }; + let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; + let events = Arc::new(Mutex::new(Vec::new())); + let mut request = wire_request( + model, + &base, + json!({"vertex_project":"project", "vertex_location":"us-central1"}), + ) + .with_host_hooks( + Arc::new(RecordingHooks { + events: events.clone(), + block: false, + }), + Some("call-id".into()), + ); + if matches!( + adapter, + OcrAdapterKind::ReductoLegacy | OcrAdapterKind::ReductoV3 + ) { + request.document = request.document.with_source("reducto://ready.pdf".into()); + } + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["pre", "during", "prepared", "success"], + "{model}" + ); + assert_eq!(seen.lock().unwrap().len(), 1, "{model}"); + } +} + +#[derive(Clone, Copy, Debug)] +enum FailureStage { + During, + Prepared, + InvalidBody, + Preparation, + Response, +} + +struct FailingHooks { + recording: RecordingHooks, + stage: FailureStage, +} + +impl OcrHooks for FailingHooks { + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + self.recording.pre_call(request) + } + + fn during_call( + &self, + request: super::hooks::OcrDuringCallRequest, + ) -> OcrHookFuture<'_, super::hooks::OcrDuringCallRequest> { + Box::pin(async move { + let request = self.recording.during_call(request).await?; + if matches!(self.stage, FailureStage::During) { + return Err(crate::Error::InvalidRequest("blocked during call".into())); + } + Ok(request) + }) + } + + fn prepared_request( + &self, + request: OcrPreparedRequest, + ) -> OcrHookFuture<'_, OcrPreparedRequest> { + Box::pin(async move { + let request = self.recording.prepared_request(request).await?; + match self.stage { + FailureStage::Prepared => Err(crate::Error::InvalidRequest( + "blocked prepared request".into(), + )), + FailureStage::InvalidBody => Ok(OcrPreparedRequest { + body: json!({"document":null}), + ..request + }), + _ => Ok(request), + } + }) + } + + fn success<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a super::LiteLLMOcrResponse, + timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + self.recording.success(context, response, timing) + } + + fn failure<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a crate::Error, + timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + assert_eq!(context.litellm_call_id, "call-id"); + assert_eq!(timing.phases.len(), 3); + self.recording.failure(context, error, timing).await; + }) + } +} + +#[rstest::rstest] +#[case(FailureStage::During)] +#[case(FailureStage::Prepared)] +#[case(FailureStage::InvalidBody)] +#[case(FailureStage::Preparation)] +#[case(FailureStage::Response)] +#[tokio::test] +async fn lifecycle_reports_failures_once_at_each_boundary(#[case] stage: FailureStage) { + let (base, seen, server) = mock_server(if matches!(stage, FailureStage::Response) { + vec![MockResponse::json(json!({"pages":"invalid"}))] + } else { + vec![] + }) + .await; + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request( + "mistral/model", + &base, + if matches!(stage, FailureStage::Preparation) { + json!({"pages":"invalid"}) + } else { + json!({}) + }, + ) + .with_host_hooks( + Arc::new(FailingHooks { + recording: RecordingHooks { + events: events.clone(), + block: false, + }, + stage, + }), + Some("call-id".into()), + ); + let result = perform_ocr(request).await; + assert!(result.is_err()); + server.await.unwrap(); + let expected = match stage { + FailureStage::Preparation => vec!["pre", "failure"], + FailureStage::During => vec!["pre", "during", "failure"], + _ => vec!["pre", "during", "prepared", "failure"], + }; + assert_eq!(*events.lock().unwrap(), expected); + assert_eq!( + seen.lock().unwrap().len(), + usize::from(matches!(stage, FailureStage::Response)) + ); +} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 8e86e4713ef..88fba2255f5 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -195,10 +195,6 @@ async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { struct RewriteDocument; impl OcrHooks for RewriteDocument { - fn has_guardrails(&self) -> bool { - true - } - fn during_call( &self, request: OcrDuringCallRequest, diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index e1f458ea0bc..dc04f15ebfc 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -16,24 +16,44 @@ pyo3::create_exception!( "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." ); -pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Auth(message) => PyValueError::new_err(message), - Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) => PyValueError::new_err(err.to_string()), - other => PyRuntimeError::new_err(other.to_string()), +#[derive(Debug)] +pub(crate) enum BridgeError { + InvalidArgument(String), + Declined(String), + Upstream { + status: Option, + message: String, + }, + Internal(String), + Host(PyErr), +} + +impl From for PyErr { + fn from(error: BridgeError) -> Self { + match error { + BridgeError::InvalidArgument(message) => PyValueError::new_err(message), + BridgeError::Declined(reason) => RustBridgeDeclined::new_err(reason), + BridgeError::Upstream { status, message } => { + RustUpstreamError::new_err((status.unwrap_or(0), message)) + } + BridgeError::Internal(message) => PyRuntimeError::new_err(message), + BridgeError::Host(error) => error, + } } } -/// Map a core error for a route whose host keeps a Python implementation. -/// -/// The distinction the host needs is whether the provider was already called. -/// Everything raised before the request goes out is safe for the host to retry -/// on its own path; anything after it is not, because the provider has already -/// done the work and billed for it. -pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { +pub(crate) fn required_route_error(err: Error) -> BridgeError { + match err { + Error::Auth(message) => BridgeError::InvalidArgument(message), + Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => BridgeError::InvalidArgument(err.to_string()), + other => BridgeError::Internal(other.to_string()), + } +} + +pub(crate) fn fallback_route_error(err: Error) -> BridgeError { match err { Error::Unsupported(_) | Error::Auth(_) @@ -46,15 +66,15 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::MissingAzureDocumentIntelligenceCredentials | Error::MissingReductoApiKey | Error::Routing(_) - // Nothing reached the provider, so serving it on Python cannot double - // bill and is the only way the caller gets an answer at all. - | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - Error::Http { status, body } => { - RustUpstreamError::new_err((status, format!("{status}: {body}"))) - } - Error::Network(message) | Error::InvalidResponse(message) => { - RustUpstreamError::new_err((0u16, message)) - } + | Error::Connect(_) => BridgeError::Declined(err.to_string()), + Error::Http { status, body } => BridgeError::Upstream { + status: Some(status), + message: format!("{status}: {body}"), + }, + Error::Network(message) | Error::InvalidResponse(message) => BridgeError::Upstream { + status: None, + message, + }, } } @@ -64,13 +84,16 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustUpstreamError", py.get_type::()) } -pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr { +pub(crate) fn ocr_route_error(err: Error) -> BridgeError { match err { Error::MissingField("document_url" | "image_url") => { - PyValueError::new_err("Document URL is required") + BridgeError::InvalidArgument("Document URL is required".into()) } - Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - other => core_error_to_pyerr(other), + Error::Http { status, .. } => BridgeError::Upstream { + status: Some(status), + message: "OCR provider request failed".into(), + }, + other => required_route_error(other), } } @@ -79,25 +102,71 @@ mod ocr_error_tests { use super::*; #[test] - fn ocr_errors_preserve_python_validation_and_provider_details() { + fn fallback_policy_never_declines_possible_dispatch() { + for error in [ + Error::Network("timeout".into()), + Error::InvalidResponse("malformed".into()), + Error::Http { + status: 429, + body: "limited".into(), + }, + ] { + assert!(matches!( + fallback_route_error(error), + BridgeError::Upstream { .. } + )); + } + for error in [ + Error::Connect("offline".into()), + Error::Unsupported("shape"), + Error::MissingApiKey { provider: "test" }, + ] { + assert!(matches!( + fallback_route_error(error), + BridgeError::Declined(_) + )); + } + } + + #[test] + fn conversion_preserves_absent_status_and_host_exception_identity() { + Python::initialize(); + Python::attach(|py| { + let error: PyErr = BridgeError::Upstream { + status: None, + message: "timeout".into(), + } + .into(); + let args: (u16, String) = error.value(py).getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (0, "timeout".into())); + let original = PyValueError::new_err("host exception"); + let retained = original.value(py).clone(); + let mapped: PyErr = BridgeError::Host(original).into(); + assert!(mapped.value(py).is(&retained)); + }); + } + + #[test] + fn ocr_errors_preserve_status_without_provider_body() { Python::initialize(); Python::attach(|py| { for field in ["document_url", "image_url"] { - let mapped = ocr_error_to_pyerr(Error::MissingField(field)); + let mapped: PyErr = ocr_route_error(Error::MissingField(field)).into(); assert!(mapped.is_instance_of::(py)); assert_eq!(mapped.value(py).to_string(), "Document URL is required"); } - let mapped = ocr_error_to_pyerr(Error::Http { + let mapped: PyErr = ocr_route_error(Error::Http { status: 429, body: r#"{"message":"rate limited"}"#.to_string(), - }); + }) + .into(); assert!(mapped.is_instance_of::(py)); let args: (u16, String) = mapped .value(py) .getattr("args") .and_then(|args| args.extract()) .expect("OCR failures retain status and unprefixed provider message"); - assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); + assert_eq!(args, (429, "OCR provider request failed".to_string())); }); } } diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index b57197b9ddf..72eb85517ee 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -2,6 +2,7 @@ use std::future::Future; use std::panic::AssertUnwindSafe; use std::time::Duration; +use crate::errors::BridgeError; use futures_util::FutureExt; use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; use pyo3::exceptions::PyRuntimeError; @@ -13,7 +14,7 @@ use tokio::time::{self, MissedTickBehavior}; pub(crate) fn run_sync( py: Python<'_>, future: F, - map_error: fn(E) -> PyErr, + map_error: fn(E) -> BridgeError, ) -> PyResult> where T: Serialize + Send + 'static, @@ -32,14 +33,14 @@ fn run_sync_on( py: Python<'_>, runtime: &Runtime, future: F, - map_error: fn(E) -> PyErr, + map_error: fn(E) -> BridgeError, ) -> PyResult> where T: Serialize + Send + 'static, E: Send + 'static, F: Future> + Send + 'static, { - if Handle::try_current().is_ok() { + if Handle::try_current().is_ok() && !litellm_python_interop::in_callback() { return Err(PyRuntimeError::new_err( "synchronous native routes cannot run from a Tokio context; use the async route", )); @@ -53,7 +54,7 @@ where pub(crate) fn run_async( py: Python<'_>, future: F, - map_error: fn(E) -> PyErr, + map_error: fn(E) -> BridgeError, ) -> PyResult> where T: Serialize + Send + 'static, @@ -67,12 +68,13 @@ where }) } -fn map_core_result(result: Result, map_error: fn(E) -> PyErr) -> PyResult { +fn map_core_result(result: Result, map_error: fn(E) -> BridgeError) -> PyResult { match result { Ok(value) => Ok(value), Err(error) => Err( std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error))) - .map_err(panic_to_pyerr)?, + .map_err(panic_to_pyerr)? + .into(), ), } } @@ -124,11 +126,11 @@ mod tests { use super::*; - fn runtime_error(error: Error) -> PyErr { - PyRuntimeError::new_err(error.to_string()) + fn runtime_error(error: Error) -> BridgeError { + BridgeError::Internal(error.to_string()) } - fn panicking_error_mapper(_error: Error) -> PyErr { + fn panicking_error_mapper(_error: Error) -> BridgeError { panic!("error mapper panicked") } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index cf0450a1b30..6f44d5f4cba 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,7 +13,7 @@ use pyo3::prelude::*; use pyo3::types::PyAny; use serde_json::Value; -use crate::errors::core_error_to_pyerr; +use crate::errors::required_route_error; use crate::marshal::{marshal_headers, optional_timeout}; #[pyclass] @@ -37,7 +37,7 @@ impl ResponsesWebSocketConnection { pyo3_async_runtimes::tokio::future_into_py(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await - .map_err(core_error_to_pyerr)?; + .map_err(required_route_error)?; Ok(ResponsesWebSocketConnection { inner }) }) } @@ -45,21 +45,30 @@ impl ResponsesWebSocketConnection { fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult> { let inner = self.inner.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - inner.send_text(text).await.map_err(core_error_to_pyerr) + inner + .send_text(text) + .await + .map_err(|error| PyErr::from(required_route_error(error))) }) } fn recv_text<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - inner.recv_text().await.map_err(core_error_to_pyerr) + inner + .recv_text() + .await + .map_err(|error| PyErr::from(required_route_error(error))) }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - inner.close().await.map_err(core_error_to_pyerr) + inner + .close() + .await + .map_err(|error| PyErr::from(required_route_error(error))) }) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index af60515b0e2..8d722b16240 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -7,7 +7,7 @@ use litellm_core::audio_transcription::{ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::core_error_to_pyerr; +use crate::errors::required_route_error; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; fn prepare_transcription( @@ -67,5 +67,5 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_transcription, - errors = core_error_to_pyerr, + errors = required_route_error, } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 08ab476005c..e41b82c1d8c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -8,7 +8,7 @@ use litellm_core::chat_completions::{ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::chat_completions_error_to_pyerr; +use crate::errors::fallback_route_error; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; fn prepare_chat_completions( @@ -86,6 +86,6 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_chat_completions, - errors = chat_completions_error_to_pyerr, + errors = fallback_route_error, extra = [chat_completions_decline], } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 97313651011..0a2227b0110 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -207,11 +207,11 @@ mod tests { } } - fn map_error(error: Error) -> PyErr { + fn map_error(error: Error) -> crate::errors::BridgeError { if matches!(&error, Error::InvalidRequest(message) if message == "panic in mapper") { panic!("synthetic mapper panic") } - PyLookupError::new_err(error.to_string()) + crate::errors::BridgeError::Host(PyLookupError::new_err(error.to_string())) } } @@ -225,7 +225,7 @@ mod tests { ( "ocr", "aocr", - "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, input_sources=None, timeout_seconds=None)", + "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, input_sources=None, timeout_seconds=None, logging_obj=None, callback_loop=None, token_provider=None)", ), ( "transcription", diff --git a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs index 97ff93f299a..ad3292344b7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs @@ -1,7 +1,7 @@ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::core_error_to_pyerr; +use crate::errors::required_route_error; #[pyfunction] fn gateway_messages<'py>( @@ -20,7 +20,7 @@ fn gateway_messages<'py>( crate::execution::run_async( py, crate::function_trace::capture(future), - core_error_to_pyerr, + required_route_error, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index f69b5e9251d..cdd5cb6c8f2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -5,7 +5,7 @@ use pyo3::prelude::*; use serde_json::Value; use std::future::Future; -use crate::errors::core_error_to_pyerr; +use crate::errors::fallback_route_error; use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; fn prepare_messages( @@ -61,5 +61,5 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_messages, - errors = core_error_to_pyerr, + errors = fallback_route_error, } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index c5def64c2f1..1f7458bcb5a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -1,18 +1,51 @@ -use litellm_core::Error; use std::future::Future; +use std::sync::{Arc, Mutex}; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_core::ocr::wire::{OcrWireRequest, decode_request, is_supported_request}; use pyo3::prelude::*; use serde_json::Value; -use crate::errors::ocr_error_to_pyerr; +use crate::errors::{BridgeError, ocr_route_error}; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; +#[path = "ocr_callbacks.rs"] +mod callbacks; + fn prepare_ocr( inputs: OcrInputs, -) -> PyResult> + Send + 'static> { - let document = inputs.document; +) -> PyResult> + Send + 'static> { + let document: Value = + Python::attach(|py| litellm_python_interop::from_py(inputs.document.bind(py)))?; + let hooks = (inputs.logging_obj.is_some() || inputs.token_provider.is_some()) + .then(|| { + Python::attach(|py| { + Ok::<_, PyErr>(Arc::new(callbacks::PythonOcrHooks { + logger: inputs.logging_obj, + token_provider: match inputs.token_provider { + Some(provider) + if provider.bind(py).is_callable() + && provider.bind(py).is_truthy()? => + { + Some(provider) + } + _ => None, + }, + document: inputs.document, + document_snapshot: document.clone(), + api_key: inputs.api_key.clone(), + locals: inputs + .callback_loop + .map(|event_loop| { + pyo3_async_runtimes::TaskLocals::new(event_loop.into_bound(py)) + .copy_context(py) + }) + .transpose()?, + error: Mutex::new(None), + })) + }) + }) + .transpose()?; let options = RouteOptions::from_python(RouteOptionsInputs { model: inputs.model, api_key: inputs.api_key, @@ -39,7 +72,7 @@ fn prepare_ocr( timeout, } = options; if is_supported_request(&model, custom_llm_provider.as_deref()) { - let request = decode_request(OcrWireRequest { + let mut request = decode_request(OcrWireRequest { model, document, api_key, @@ -49,10 +82,33 @@ fn prepare_ocr( optional_params, input_sources, timeout_seconds: timeout.map(|value| value.as_secs_f64()), - })?; + }) + .map_err(ocr_route_error)?; + if let Some(hooks) = &hooks + && hooks.token_provider.is_some() + { + request.connection.token_provider = + Some(litellm_core::auth::TokenProviderHandle::new(hooks.clone())); + } + let request = match &hooks { + Some(hooks) => request.with_host_hooks(hooks.clone(), None), + None => request, + }; return litellm_core::ocr::ocr(request) .await - .map(|response| response.into_json()); + .map(|response| response.into_json()) + .map_err(|error| { + hooks + .and_then(|hooks| { + hooks + .error + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take() + }) + .map(BridgeError::Host) + .unwrap_or_else(|| ocr_route_error(error)) + }); } run_ocr(OcrRequest { model: &model, @@ -69,6 +125,7 @@ fn prepare_ocr( litellm_call_id: None, }) .await + .map_err(ocr_route_error) }) } @@ -78,8 +135,7 @@ bridge_route! { inputs = OcrInputs, required = { model: String, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - document: serde_json::Value, + document: Py, }, optional = { api_key: Option, @@ -92,9 +148,12 @@ bridge_route! { #[pyo3(from_py_with = litellm_python_interop::from_py)] input_sources: Option, timeout_seconds: Option, + logging_obj: Option>, + callback_loop: Option>, + token_provider: Option>, }, prepare = prepare_ocr, - errors = ocr_error_to_pyerr, + errors = std::convert::identity, } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr_callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr_callbacks.rs new file mode 100644 index 00000000000..0556907d230 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr_callbacks.rs @@ -0,0 +1,207 @@ +use std::sync::Mutex; + +use litellm_core::Error; +use litellm_core::auth::{ResolvedCredential, SecretValue, TokenFuture, TokenProvider}; +use litellm_core::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreparedRequest}; +use litellm_python_interop::{from_py, to_py}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; +use pyo3_async_runtimes::TaskLocals; + +struct PendingCallback(Option>); + +impl Drop for PendingCallback { + fn drop(&mut self) { + if let Some(future) = self.0.take() { + Python::attach(|py| { + if let Err(error) = future.call_method0(py, "cancel") { + error.write_unraisable(py, Some(future.bind(py))); + } + }); + } + } +} + +pub(super) struct PythonOcrHooks { + pub logger: Option>, + pub token_provider: Option>, + pub document: Py, + pub document_snapshot: serde_json::Value, + pub api_key: Option, + pub locals: Option, + pub error: Mutex>, +} + +impl PythonOcrHooks { + pub async fn invoke(&self, callback: Py) -> PyResult> { + if let Some(locals) = &self.locals { + let (mut pending, future) = Python::attach(|py| { + let coroutine = py + .import("litellm.rust_bridge._callbacks")? + .getattr("invoke_callback")? + .call1((callback,))?; + let asyncio = py.import("asyncio")?; + let submitted = locals.context(py).call_method1( + "run", + ( + asyncio.getattr("run_coroutine_threadsafe")?, + coroutine, + locals.event_loop(py), + ), + )?; + let pending = PendingCallback(Some(submitted.clone().unbind())); + let kwargs = PyDict::new(py); + kwargs.set_item("loop", locals.event_loop(py))?; + let wrapped = asyncio + .getattr("wrap_future")? + .call((submitted,), Some(&kwargs))?; + let future = pyo3_async_runtimes::into_future_with_locals(locals, wrapped)?; + Ok::<_, PyErr>((pending, future)) + })?; + let result = future.await; + pending.0.take(); + return result; + } + tokio::task::block_in_place(|| litellm_python_interop::invoke_callback(&callback)) + } + + fn retain_error(&self, error: PyErr) { + *self + .error + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(error); + } + + fn callback( + &self, + py: Python<'_>, + request: &OcrPreparedRequest, + ) -> PyResult<(Py, Py, Py)> { + let body = to_py(py, &request.body)? + .into_bound(py) + .cast_into::()?; + if request.body.get("document") == Some(&self.document_snapshot) { + body.set_item("document", &self.document)?; + } + let headers = PyDict::new(py); + for (name, value) in &request.headers { + headers.set_item(name, value)?; + } + let additional = PyDict::new(py); + additional.set_item("complete_input_dict", &body)?; + additional.set_item("headers", &headers)?; + additional.set_item("api_base", &request.url)?; + let kwargs = PyDict::new(py); + kwargs.set_item("input", "OCR document processing")?; + kwargs.set_item("api_key", &self.api_key)?; + kwargs.set_item("model", &request.model)?; + kwargs.set_item("additional_args", additional)?; + let callback = py.import("functools")?.getattr("partial")?.call( + PyTuple::new( + py, + [self + .logger + .as_ref() + .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("missing OCR logger"))? + .bind(py) + .getattr("pre_call")?], + )?, + Some(&kwargs), + )?; + Ok((callback.unbind(), body.unbind(), headers.unbind())) + } +} + +impl OcrHooks for PythonOcrHooks { + fn prepared_request( + &self, + request: OcrPreparedRequest, + ) -> OcrHookFuture<'_, OcrPreparedRequest> { + Box::pin(async move { + if self.logger.is_none() { + return Ok(request); + } + let result: PyResult = async { + let (callback, body, headers) = Python::attach(|py| self.callback(py, &request))?; + self.invoke(callback).await?; + Python::attach(|py| { + Ok(OcrPreparedRequest { + body: from_py(body.bind(py).as_any())?, + headers: headers + .bind(py) + .extract::>()? + .into_iter() + .collect(), + ..request + }) + }) + } + .await; + result.map_err(|error| { + self.retain_error(error); + Error::InvalidRequest("OCR pre-call hook failed".into()) + }) + }) + } +} + +impl std::fmt::Debug for PythonOcrHooks { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("PythonOcrHooks") + } +} + +impl TokenProvider for PythonOcrHooks { + fn acquire(&self) -> TokenFuture<'_> { + Box::pin(async move { + let result: PyResult = async { + let callback = Python::attach(|py| { + self.token_provider + .as_ref() + .map(|provider| provider.clone_ref(py)) + .ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err("missing token provider") + }) + })?; + let value = self.invoke(callback).await.map_err(|error| { + Python::attach(|py| { + if error.is_instance_of::(py) + || !error.is_instance_of::(py) + { + return error; + } + let wrapped = pyo3::exceptions::PyRuntimeError::new_err(format!( + "Failed to get Azure AD token: {}", + error.value(py) + )); + wrapped.set_cause(py, Some(error)); + wrapped + }) + })?; + Python::attach(|py| { + value.extract::(py).map_err(|_| { + pyo3::exceptions::PyTypeError::new_err("Azure AD token must be a string") + }) + }) + } + .await; + match result { + Ok(token) if token.is_empty() => { + Err(litellm_core::AuthError::AzureTokenAcquisition( + "Missing Azure AI credentials".into(), + )) + } + Ok(token) => Ok(ResolvedCredential::AccessToken { + token: SecretValue::new(token), + expires_on: None, + }), + Err(error) => { + self.retain_error(error); + Err(litellm_core::AuthError::AzureTokenAcquisition( + "host token provider failed".into(), + )) + } + } + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index ee82c170d55..0026173884f 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -6,13 +6,12 @@ use litellm_python_interop::release_gil; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, }; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::PyAny; use tokio::sync::Semaphore; use crate::constants::TOKEN_COUNT_FALLBACK_PARALLELISM; -use crate::errors::RustBridgeDeclined; +use crate::errors::BridgeError; use crate::execution::run_async; /// Counts the input tokens of a raw request body off the Python event loop with @@ -31,7 +30,7 @@ impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json)) - .map_err(token_count_error_to_pyerr)?; + .map_err(token_count_error)?; Ok(Self { inner: Arc::new(inner), encode_slots: Arc::new(Semaphore::new(encode_parallelism())), @@ -53,7 +52,7 @@ impl TokenCounter { .await .map_err(|error| Error::Task(error.to_string()))? }, - token_count_error_to_pyerr, + token_count_error, ) } } @@ -67,18 +66,18 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result PyErr { +fn token_count_error(error: Error) -> BridgeError { let message = error.to_string(); match error { - Error::Load(_) => PyValueError::new_err(message), + Error::Load(_) => BridgeError::InvalidArgument(message), Error::RequestParse(_) | Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems | Error::JsonSerialization(_) - | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), - Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + | Error::JsonUtf8(_) => BridgeError::Declined(message), + Error::Encode(_) | Error::Task(_) => BridgeError::Internal(message), } } diff --git a/litellm-rust/crates/python-interop/src/callback.rs b/litellm-rust/crates/python-interop/src/callback.rs new file mode 100644 index 00000000000..2d2e3629571 --- /dev/null +++ b/litellm-rust/crates/python-interop/src/callback.rs @@ -0,0 +1,22 @@ +use std::cell::Cell; + +use pyo3::prelude::*; + +thread_local! { + static IN_CALLBACK: Cell = const { Cell::new(false) }; +} + +pub fn in_callback() -> bool { + IN_CALLBACK.get() +} + +pub fn invoke_callback(callback: &Py) -> PyResult> { + struct Restore(bool); + impl Drop for Restore { + fn drop(&mut self) { + IN_CALLBACK.set(self.0); + } + } + let _restore = Restore(IN_CALLBACK.replace(true)); + Python::attach(|py| callback.call0(py)) +} diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index 2e562bdae70..bab28befdaa 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -1,5 +1,7 @@ +mod callback; mod gil; mod marshal; +pub use callback::{in_callback, invoke_callback}; pub use gil::{release_count, release_gil}; pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py}; diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 77bf4820a1a..d5f58928ce3 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -851,6 +851,9 @@ class CustomGuardrail(CustomLogger): return None target: Final = self._deployment_hook_target() + from litellm.litellm_core_utils.guardrail_call_context import guardrail_call_type + + context_token: Final = guardrail_call_type.set(call_type) try: if target is not self: request_data["guardrail_to_apply"] = self # rebind-ok: dispatch consumes this key @@ -866,6 +869,7 @@ class CustomGuardrail(CustomLogger): response=response, ) finally: + guardrail_call_type.reset(context_token) if target is not self: request_data.pop("guardrail_to_apply", None) diff --git a/litellm/litellm_core_utils/guardrail_call_context.py b/litellm/litellm_core_utils/guardrail_call_context.py new file mode 100644 index 00000000000..fb0db015d35 --- /dev/null +++ b/litellm/litellm_core_utils/guardrail_call_context.py @@ -0,0 +1,6 @@ +from contextvars import ContextVar +from typing import Final + +from litellm.types.utils import CallTypes + +guardrail_call_type: Final[ContextVar[CallTypes | None]] = ContextVar("guardrail_call_type", default=None) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7109e6942d1..145d9dd2304 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2476,12 +2476,21 @@ class BaseLLMHTTPHandler: extra_headers=headers, timeout=timeout, ) - except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path - verbose_logger.debug( - "Rust Anthropic messages bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return None + except Exception as rust_error: # noqa: BLE001 # only explicit pre-dispatch declines permit fallback + from litellm.rust_bridge.bindings import native_exception_types, upstream_error_details + + exceptions: Final = native_exception_types() + if exceptions is not None and isinstance(rust_error, exceptions[0]): + return None + if exceptions is not None and isinstance(rust_error, exceptions[1]): + status, message = upstream_error_details(rust_error) + raise litellm.APIError( + status_code=status, + message=message, + llm_provider=custom_llm_provider, + model=model, + ) from rust_error + raise if rust_response is None: return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 64a47f4f4ff..05020761207 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -19,6 +19,7 @@ from litellm.cost_calculator import _infer_call_type from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route +from litellm.litellm_core_utils.guardrail_call_context import guardrail_call_type from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks @@ -307,8 +308,9 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response) - call_type: CallTypesLiteral | None = None - if user_api_key_dict.request_route is not None: + context: Final = guardrail_call_type.get() + call_type: CallTypesLiteral | None = context.value if context is not None else None + if call_type is None and user_api_key_dict.request_route is not None: call_types: Final = get_call_types_for_route(user_api_key_dict.request_route) if call_types is not None and len(call_types) > 0: call_type = call_types[0] diff --git a/litellm/rust_bridge/_callbacks.py b/litellm/rust_bridge/_callbacks.py new file mode 100644 index 00000000000..e91eb0c1e44 --- /dev/null +++ b/litellm/rust_bridge/_callbacks.py @@ -0,0 +1,5 @@ +from collections.abc import Callable + + +async def invoke_callback(callback: Callable[[], object]) -> object: + return callback() diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index d16f150a2aa..0f4013cb7b2 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -47,3 +47,12 @@ def native_exception_types() -> tuple[type[BaseException], type[BaseException]] if not isinstance(declined, type) or not isinstance(upstream, type): return None return declined, upstream + + +def upstream_error_details(error: BaseException) -> tuple[int, str]: + args: Final[tuple[object, ...]] = error.args + status_value: Final = args[0] if args else 0 + message_value: Final = args[1] if len(args) > 1 else str(error) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + return status or 500, message diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 674bd8847f7..a75fa0de5f0 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -26,6 +26,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge.bindings import native_exception_types, upstream_error_details from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -274,18 +275,6 @@ def rust_chat_completions_accepts( return True -def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None: - """`(declined, upstream_failed)` from the native module, or None when absent.""" - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - declined: Final = getattr(native_bridge, "RustBridgeDeclined", None) - upstream: Final = getattr(native_bridge, "RustUpstreamError", None) - if declined is None or upstream is None: - return None - return declined, upstream - - def _reraise_or_decline( rust_error: BaseException, *, @@ -299,20 +288,14 @@ def _reraise_or_decline( second attempt bills for it twice. Those surface as an `APIError` carrying the upstream status, which LiteLLM's exception mapping already understands. """ - exceptions: Final = _rust_bridge_exceptions() + exceptions: Final = native_exception_types() if exceptions is None: - verbose_logger.debug( - "Rust chat completions bridge raised %s; falling back to Python path", - type(rust_error).__name__, - ) - return + raise rust_error declined, upstream_failed = exceptions if isinstance(rust_error, upstream_failed): - args: Final = rust_error.args - status: Final = args[0] if args else 0 - message: Final = args[1] if len(args) > 1 else "" + status, message = upstream_error_details(rust_error) raise APIError( - status_code=int(status) or 500, + status_code=status, message=f"litellm rust chat completions: {message}", llm_provider=custom_llm_provider or "", model=model, diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 1fee62c9b27..d5f576e3c7c 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType @@ -14,7 +15,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.call_completion import CallCompletion from litellm.llms.azure_ai.ocr.common_utils import is_azure_cohere_parse_model from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse -from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge.bindings import NativeBinding, native_exception_types, upstream_error_details from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager @@ -69,6 +70,9 @@ class RustOcr(Protocol): optional_params: dict[str, object], input_sources: dict[str, str], timeout_seconds: float | None, + logging_obj: _OCRLogging | None, + callback_loop: asyncio.AbstractEventLoop | None, + token_provider: object, ) -> dict[str, object]: raise NotImplementedError @@ -85,6 +89,9 @@ class RustAocr(Protocol): optional_params: dict[str, object], input_sources: dict[str, str], timeout_seconds: float | None, + logging_obj: _OCRLogging | None, + callback_loop: asyncio.AbstractEventLoop | None, + token_provider: object, ) -> Awaitable[dict[str, object]]: raise NotImplementedError @@ -147,7 +154,6 @@ def supported(request: LiteLLMOcrRequest) -> bool: if request_provider == "azure_ai": return ( not is_azure_cohere_parse_model(request.model) - and not callable(request.kwargs.get("azure_ad_token_provider")) and request.kwargs.get("azure_username") is None and request.kwargs.get("azure_password") is None ) @@ -160,7 +166,8 @@ def _optional_params(request: LiteLLMOcrRequest, resolve_secret: Callable[[str], name: value for name, value in request.kwargs.items() if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) - and name not in ("litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request") + and name + not in ("litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request", "azure_ad_token_provider") } ) request_provider: Final = provider(request) @@ -266,7 +273,7 @@ def _marshal( ) logging_obj.update_from_kwargs( kwargs=dict(logged_kwargs), # mutable-ok: legacy logging mutates its kwargs copy - model=request.model, + model=request.model.removeprefix(f"{request_provider}/"), optional_params=dict(logged_optional_params), # mutable-ok: legacy logging requires concrete dict params litellm_params={ # mutable-ok: legacy logging requires a concrete params dict "litellm_call_id": request.kwargs.get("litellm_call_id"), @@ -274,19 +281,6 @@ def _marshal( }, custom_llm_provider=request_provider, ) - logging_obj.pre_call( - input="OCR document processing", - api_key=api_key, - additional_args={ # mutable-ok: pre_call mutates the additional_args dict - "complete_input_dict": { # mutable-ok: callbacks consume a JSON-serializable request dict - "model": request.model, - "document": document, - **logged_optional_params, - }, - "api_base": request.api_base or "", - "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict - }, - ) return LiteLLMOcrRequest( model=request.model, document=document, @@ -313,17 +307,13 @@ def _map_error(error: Exception, request: LiteLLMOcrRequest) -> Exception: ) if provider_config is None: return error - error_args: Final = cast( # cast-ok: BaseException.args exposes Any while native errors carry scalar args - tuple[object, ...], error.args - ) - status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 - message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) + status, message = upstream_error_details(error) error_factory: Final = cast( # cast-ok: legacy provider error factories have untyped callable parameters Callable[..., Exception], provider_config.get_error_class ) return error_factory( error_message=message, - status_code=status or 500, + status_code=status, headers={}, # mutable-ok: provider error factories require a concrete headers dict ) @@ -349,7 +339,7 @@ def run( try: response: Final = ocr( model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + document=cast(dict[str, object], marshalled.document), # cast-ok: _marshal validates the document dict api_key=marshalled.api_key, api_base=marshalled.api_base, custom_llm_provider=marshalled.custom_llm_provider, @@ -357,9 +347,16 @@ def run( optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict input_sources=marshalled.input_sources, timeout=marshalled.timeout, + logging_obj=cast( + _OCRLogging, request.kwargs["litellm_logging_obj"] + ), # cast-ok: client decorator injects Logging + token_provider=request.kwargs.get("azure_ad_token_provider"), ) except Exception as error: - raise _map_error(error, request) from error + mapped: Final = _map_error(error, request) + if mapped is error: + raise + raise mapped from error return _response(response) if response is not None else None @@ -374,7 +371,7 @@ async def arun( try: response: Final = await aocr( model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + document=cast(dict[str, object], marshalled.document), # cast-ok: _marshal validates the document dict api_key=marshalled.api_key, api_base=marshalled.api_base, custom_llm_provider=marshalled.custom_llm_provider, @@ -382,9 +379,16 @@ async def arun( optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict input_sources=marshalled.input_sources, timeout=marshalled.timeout, + logging_obj=cast( + _OCRLogging, request.kwargs["litellm_logging_obj"] + ), # cast-ok: client decorator injects Logging + token_provider=request.kwargs.get("azure_ad_token_provider"), ) except Exception as error: - raise _map_error(error, request) from error + mapped: Final = _map_error(error, request) + if mapped is error: + raise + raise mapped from error return _response(response) if response is not None else None @@ -399,6 +403,8 @@ def ocr( optional_params: dict[str, object], timeout: float | httpx.Timeout | None, input_sources: Mapping[str, str] | None = None, + logging_obj: _OCRLogging | None = None, + token_provider: object = None, ) -> dict[str, object] | None: rust_ocr: Final = load_rust_ocr() if rust_ocr is None: @@ -413,6 +419,9 @@ def ocr( optional_params=optional_params, input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict timeout_seconds=_timeout_to_seconds(timeout), + logging_obj=logging_obj, + callback_loop=None, + token_provider=token_provider, ) @@ -427,6 +436,8 @@ async def aocr( optional_params: dict[str, object], timeout: float | httpx.Timeout | None, input_sources: Mapping[str, str] | None = None, + logging_obj: _OCRLogging | None = None, + token_provider: object = None, ) -> dict[str, object] | None: rust_aocr: Final = load_rust_aocr() if rust_aocr is None: @@ -441,4 +452,7 @@ async def aocr( optional_params=optional_params, input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict timeout_seconds=_timeout_to_seconds(timeout), + logging_obj=logging_obj, + callback_loop=asyncio.get_running_loop(), + token_provider=token_provider, ) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index d411673439f..9ee6cf10717 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -6,7 +6,7 @@ from enum import Enum from typing import Final, Generic, NoReturn, TypeAlias, TypeVar from litellm.exceptions import APIError -from litellm.rust_bridge.bindings import native_exception_types +from litellm.rust_bridge.bindings import native_exception_types, upstream_error_details NativeT = TypeVar("NativeT") ResultT = TypeVar("ResultT") @@ -137,13 +137,9 @@ def _required_reason(result: RustDeclined | RustUnavailable) -> str: def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: - args: Final[tuple[object, ...]] = error.args - status_value: Final = args[0] if args else 0 - message_value: Final = args[1] if len(args) > 1 else str(error) - status: Final = status_value if isinstance(status_value, int) else 0 - message: Final = message_value if isinstance(message_value, str) else str(message_value) + status, message = upstream_error_details(error) raise APIError( - status_code=status or 500, + status_code=status, message=f"litellm rust {context.route}: {message}", llm_provider=context.provider, model=context.model, diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index a30474245c6..0b68dda027e 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -238,17 +238,66 @@ async def test_gate_invokes_rust_and_marks_response_header(): @pytest.mark.asyncio -async def test_gate_falls_back_to_python_when_bridge_raises(): +async def test_gate_propagates_unclassified_bridge_failure(): bridge = RaisingAsyncMessages() litellm.rust(True) rust_messages.set_rust_messages(amessages=bridge) - response = await _gate() - - assert response is None + with pytest.raises(RuntimeError, match="upstream request failed"): + await _gate() assert bridge.calls == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [0, 400, 429, 500]) +async def test_gate_never_falls_back_after_possible_dispatch(monkeypatch, status): + from types import SimpleNamespace + from litellm.rust_bridge import bindings + + class Declined(Exception): + pass + + class Upstream(Exception): + pass + + error = Upstream(status, "provider failure") + + async def bridge(**kwargs): + raise error + + monkeypatch.setattr( + bindings, "get_native_bridge", lambda: SimpleNamespace(RustBridgeDeclined=Declined, RustUpstreamError=Upstream) + ) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) + with pytest.raises(litellm.APIError) as caught: + await _gate() + assert caught.value.status_code == (status or 500) + assert caught.value.__cause__ is error + + +@pytest.mark.asyncio +async def test_gate_falls_back_only_for_explicit_decline(monkeypatch): + from types import SimpleNamespace + from litellm.rust_bridge import bindings + + class Declined(Exception): + pass + + class Upstream(Exception): + pass + + async def bridge(**kwargs): + raise Declined("unsupported before dispatch") + + monkeypatch.setattr( + bindings, "get_native_bridge", lambda: SimpleNamespace(RustBridgeDeclined=Declined, RustUpstreamError=Upstream) + ) + litellm.rust(True) + rust_messages.set_rust_messages(amessages=bridge) + assert await _gate() is None + + @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages() diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index ddc8439a83a..3e084641632 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2648,6 +2648,61 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: which starved every later callback in litellm.callbacks (notably the lazily-appended VectorStorePreCallHook that attaches provider_specific_fields["search_results"]).""" + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", [CallTypes.aocr, CallTypes.aresponses, CallTypes.anthropic_messages]) + async def test_explicit_call_type_scans_without_proxy_route(self, call_type: CallTypes) -> None: + from fastapi import HTTPException + + from litellm.litellm_core_utils.guardrail_call_context import guardrail_call_type + from litellm.llms.base_llm.ocr.transformation import OCRResponse + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="response-filter", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="secret", action=ContentFilterAction.BLOCK)], + ) + response: Final = ( + OCRResponse(pages=[{"index": 0, "markdown": "secret", "images": [], "dimensions": None}], model="model") + if call_type == CallTypes.aocr + else ResponsesAPIResponse( + id="resp_test", + created_at=0, + model="model", + object="response", + output=[ + { + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "secret", "annotations": []}], + } + ], + ) + if call_type == CallTypes.aresponses + else { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "model", + "content": [{"type": "text", "text": "secret"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + ) + with pytest.raises(HTTPException, match="Content blocked"): + await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": ["response-filter"]}, + response=response, + call_type=call_type, + ) + assert guardrail_call_type.get() is None + @pytest.mark.asyncio async def test_apply_guardrail_retains_request_identity(self) -> None: from litellm.types.guardrails import GuardrailEventHooks diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index dbb4f822d0b..f946d489ecc 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,6 +1,7 @@ """Tests for the optional Rust-backed OCR path.""" import builtins +import asyncio import importlib import types @@ -48,6 +49,7 @@ class RecordingBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.logging_obj: object = None def __call__( self, @@ -60,7 +62,11 @@ class RecordingBridge: optional_params: dict[str, object], input_sources: dict[str, str], timeout_seconds: float | None, + logging_obj: object = None, + callback_loop: asyncio.AbstractEventLoop | None = None, + token_provider: object = None, ) -> dict[str, object]: + self.logging_obj = logging_obj self.calls.append( { "model": model, @@ -94,6 +100,9 @@ class RecordingAsyncBridge: optional_params: dict[str, object], input_sources: dict[str, str], timeout_seconds: float | None, + logging_obj: object = None, + callback_loop: asyncio.AbstractEventLoop | None = None, + token_provider: object = None, ) -> dict[str, object]: self.calls.append( { @@ -123,6 +132,9 @@ class RaisingBridge: optional_params: dict[str, object], input_sources: dict[str, str], timeout_seconds: float | None, + logging_obj: object = None, + callback_loop: asyncio.AbstractEventLoop | None = None, + token_provider: object = None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -139,6 +151,9 @@ class RaisingAsyncBridge: optional_params: dict[str, object], input_sources: dict[str, str], timeout_seconds: float | None, + logging_obj: object = None, + callback_loop: asyncio.AbstractEventLoop | None = None, + token_provider: object = None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -750,18 +765,12 @@ def test_rust_ocr_logging_redacts_azure_credentials(): "azure_ad_token": "****", "client_secret": "****", } - assert logging_obj.pre_call_kwargs is not None - additional_args = logging_obj.pre_call_kwargs["additional_args"] - assert isinstance(additional_args, dict) - complete_input = additional_args["complete_input_dict"] - assert isinstance(complete_input, dict) - assert complete_input["azure_ad_token"] == "****" - assert complete_input["client_secret"] == "****" + assert logging_obj.pre_call_kwargs is None + assert bridge.logging_obj is logging_obj def test_rust_eligibility_rejects_python_only_azure_auth_modes(): for params in ( - {"azure_ad_token_provider": lambda: "token"}, {"azure_username": "user"}, {"azure_password": "password"}, ): @@ -796,7 +805,7 @@ def test_prepare_rust_ocr_call_forwards_global_azure_refresh(monkeypatch: pytest assert bridge.calls[0]["input_sources"] == {"enable_azure_ad_token_refresh": "deployment"} -def test_run_rust_ocr_runs_pre_call_logging(): +def test_run_rust_ocr_passes_retained_logger_to_native_dispatch(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) @@ -813,16 +822,9 @@ def test_run_rust_ocr_runs_pre_call_logging(): resolve_api_key=lambda _name: None, ) - assert logging_obj.pre_call_kwargs is not None - assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" - additional_args = logging_obj.pre_call_kwargs["additional_args"] - complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT - assert complete_input["include_image_base64"] is True - assert additional_args["api_base"] == "https://api.mistral.ai/v1" - assert additional_args["headers"] == { - "x-trace-id": "trace-1", - } + assert logging_obj.pre_call_kwargs is None + assert bridge.logging_obj is logging_obj + assert bridge.calls[0]["document"] is DOCUMENT def test_ocr_routes_to_rust_when_enabled(fake_bridge): diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index 88036a5a556..8f1051e7b21 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -33,3 +33,17 @@ def test_binding_validates_native_attribute( binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) assert binding.load() == expected + + +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ((429, "limited"), (429, "limited")), + ((0, "offline"), (500, "offline")), + ((None, "offline"), (500, "offline")), + (("429", 42), (500, "42")), + ((), (500, "")), + ], +) +def test_upstream_error_details_preserves_protocol(arguments, expected) -> None: + assert bindings.upstream_error_details(Exception(*arguments)) == expected diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index b2fd2e6dcc0..36ec234ba19 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -55,6 +55,9 @@ class _FakeNative: def _fake_native_bridge(monkeypatch): """Expose the bridge's exception classes without the compiled extension.""" monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative()) + from litellm.rust_bridge import bindings + + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) def _hide_native_bridge(monkeypatch): diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md index 4c117fb846b..72bb60d86d1 100644 --- a/tests/test_litellm_rust/README.md +++ b/tests/test_litellm_rust/README.md @@ -4,10 +4,18 @@ This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR be A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions -`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport +`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules use the public SDK with native support required, plus direct `_native` calls for host-hook behavior. `ocr/test_dispatch.py` covers enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports -Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs +Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native request helpers enable Rust and reject unsupported configurations that would fall back to Python. The dispatch test records which OCR entrypoint runs -The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible +## Callback task and cancellation contract + +Synchronous native calls invoke hooks on the caller thread in its current Python context. If called inside an asyncio task, callbacks see that same task and their context changes remain visible to the caller + +Async native calls invoke synchronous hooks on the caller's event loop in separate tasks. Each token-provider or pre-call phase starts with a copy of the context captured at native entry. Context changes remain visible to later callbacks within the same pre-call phase, but do not propagate to another phase or back to the caller. Retained Python objects remain shared regardless of these context boundaries + +Cancelling the native awaitable cancels pending callback delivery and eventually releases its references. A callback-raised `CancelledError` aborts the request before provider transport. Cancellation cannot interrupt a synchronous Python callback that is already executing; it must return or raise before its event loop can process cancellation + +The pending-delivery tests use the event loop's task factory to hold the callback coroutine behind a gate. They wait for cancellation acknowledgement before checking cleanup, rather than relying on a fixed delay diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 17312e4d9d5..370ec4a859a 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -29,11 +29,6 @@ CALLBACK_ATTRIBUTES: Final = ( "_async_success_callback", "_async_failure_callback", ) -EXPECTED_FAILURE_REASONS: Final = { - "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070", - "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070", - "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070", -} def _list_attribute(container: ModuleType, attribute: str) -> list[object]: @@ -95,14 +90,6 @@ def recording_server() -> Generator[RecordingServer]: def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - for item in items: - if "test_litellm_rust" not in item.path.parts: - continue - relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :]) - reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path) - if reason is not None: - item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) - if not _parse_env_bool(os.environ.get("LITELLM_RUST")): skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") for item in items: diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index b08446412c0..a19c24f50ea 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -29,10 +29,319 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer: return recording_server +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_pre_call_retains_context_and_allows_nested_call( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + from contextvars import ContextVar + from litellm.rust_bridge import _native + + context: Final = ContextVar("ocr-callback-context", default="missing") + context.set("caller") + caller_thread: Final = threading.current_thread() + observations: Final = [] + ocr_server.expected_requests = 2 + + class NestedCall(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((context.get(), threading.current_thread())) + nested: Final = _native.ocr( + model="mistral/mistral-ocr-latest", + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + observations.append(nested["pages"][0]["markdown"]) + request_headers(kwargs)["x-outer"] = "retained" + + if asynchronous: + await call_native_aocr(ocr_server, callbacks=[NestedCall()]) + else: + call_native_ocr(ocr_server, callbacks=[NestedCall()]) + + assert observations == [("caller", caller_thread), "native OCR response"] + assert "x-outer" not in ocr_server.requests[0].headers + assert ocr_server.requests[1].headers["x-outer"] == "retained" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_binding_drives_pre_call_and_preserves_host_exception( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + from litellm.rust_bridge import _native + + class Abort(BaseException): + pass + + error: Final = Abort("stop before transport") + original: Final = dict(OCR_DOCUMENT) + observed: Final = [] + ocr_server.expected_requests = 0 + + class HostLogger: + def pre_call(self, *, input, api_key, model, additional_args): + observed.append(additional_args["complete_input_dict"]["document"]) + raise error + + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": original, + "api_key": "test-key", + "api_base": ocr_server.base_url, + "logging_obj": HostLogger(), + } + async def invoke() -> None: + if asynchronous: + await _native.aocr(**arguments, callback_loop=asyncio.get_running_loop()) + else: + _native.ocr(**arguments) + + with pytest.raises(Abort) as caught: + await invoke() + + assert observed == [original] + assert observed[0] is original + assert caught.value is error + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rust_enabled", [False, True], ids=["python", "rust"]) +async def test_callback_reference_and_response_identity_match_python( + ocr_server: RecordingServer, rust_enabled: bool +) -> None: + from tests.test_litellm_rust.support.requests import call_aocr + + original: Final = dict(OCR_DOCUMENT) + token: Final = object() + observed: Final = [] + finished: Final = asyncio.Event() + + class Logger(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(request_body(kwargs)["document"] is original) + kwargs["retained-token"] = token + request_body(kwargs)["include_image_base64"] = True + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + observed.append(kwargs["retained-token"] is token) + observed.append(response_obj) + finished.set() + + litellm.rust(rust_enabled) + response: Final = await call_aocr(ocr_server, document=original, callbacks=[Logger()]) + await asyncio.wait_for(finished.wait(), timeout=10) + assert observed[:2] == [True, True] + assert observed[2] is response + assert ocr_server.requests[0].body["include_image_base64"] is True + + +@pytest.mark.asyncio +async def test_token_provider_document_edit_keeps_live_reference(ocr_server: RecordingServer) -> None: + original: Final = dict(OCR_DOCUMENT) + observed: Final = [] + + def provider() -> str: + original["document_url"] = "data:application/pdf;base64,ZGVm" + return "token" + + class Logger(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(request_body(kwargs)["document"] is original) + + await call_native_aocr( + ocr_server, + model="azure_ai/mistral-ocr-latest", + document=original, + api_key=None, + azure_ad_token_provider=provider, + callbacks=[Logger()], + ) + + assert observed == [True] + assert ocr_server.requests[0].body["document"] == original + + def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): return call_native_ocr(server, callbacks=callbacks, **kwargs) +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_callback_task_and_context_boundaries( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + from contextvars import ContextVar + + context: Final = ContextVar("ocr-task-boundaries", default="missing") + context.set("caller") + caller: Final = asyncio.current_task() + observations: Final = [] + + def provider() -> str: + observations.append(("token", asyncio.current_task(), context.get())) + context.set("token") + return "token" + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append(("pre", asyncio.current_task(), context.get())) + context.set("pre") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append(("next", asyncio.current_task(), context.get())) + + arguments: Final = { + "model": "azure_ai/mistral-ocr-latest", + "api_key": None, + "azure_ad_token_provider": provider, + "callbacks": [Edit(), Observe()], + } + if asynchronous: + await call_native_aocr(ocr_server, **arguments) + else: + call_native_ocr(ocr_server, **arguments) + + assert [phase for phase, _, _ in observations] == ["token", "pre", "next"] + assert [value for _, _, value in observations] == ["caller", "caller" if asynchronous else "token", "pre"] + token_task: Final = observations[0][1] + pre_task: Final = observations[1][1] + assert observations[2][1] is pre_task + if asynchronous: + assert token_task is not None and pre_task is not None + assert token_task is not caller and pre_task is not caller + assert token_task is not pre_task + assert token_task.done() and pre_task.done() + assert context.get() == "caller" + else: + assert token_task is caller and pre_task is caller + assert context.get() == "pre" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["token", "pre_call"]) +async def test_native_cancellation_cancels_pending_callback_and_releases_references( + ocr_server: RecordingServer, phase: str +) -> None: + import gc + import weakref + + from litellm.rust_bridge import _native + + loop: Final = asyncio.get_running_loop() + previous_factory: Final = loop.get_task_factory() + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + tasks: Final = [] + calls: Final = [] + ocr_server.expected_requests = 0 + + async def gate(coroutine): + try: + entered.set() + await release.wait() + return await coroutine + except asyncio.CancelledError: + cancelled.set() + raise + finally: + coroutine.close() + + def factory(event_loop, coroutine, **kwargs): + if getattr(getattr(coroutine, "cr_code", None), "co_name", None) == "invoke_callback": + task: Final = asyncio.Task(gate(coroutine), loop=event_loop, **kwargs) + tasks.append(task) + return task + if previous_factory is not None: + return previous_factory(event_loop, coroutine, **kwargs) + return asyncio.Task(coroutine, loop=event_loop, **kwargs) + + class Host: + def __call__(self): + calls.append("token") + return "token" + + def pre_call(self, **kwargs): + calls.append("pre_call") + + def start(): + host: Final = Host() + reference: Final = weakref.ref(host) + future: Final = _native.aocr( + model="azure_ai/mistral-ocr-latest", + document=OCR_DOCUMENT, + api_key=None if phase == "token" else "test-key", + api_base=ocr_server.base_url, + logging_obj=host, + token_provider=host if phase == "token" else None, + callback_loop=loop, + ) + return future, reference + + loop.set_task_factory(factory) + try: + future, reference = start() + await asyncio.wait_for(entered.wait(), timeout=5) + assert reference() is not None + assert len(tasks) == 1 + future.cancel() + with pytest.raises(asyncio.CancelledError): + await future + await asyncio.wait_for(cancelled.wait(), timeout=5) + await asyncio.gather(*tasks, return_exceptions=True) + assert tasks[0].cancelled() + tasks.clear() + async with asyncio.timeout(5): + while reference() is not None: + gc.collect() + await asyncio.sleep(0.01) + assert calls == [] + assert ocr_server.requests == [] + finally: + loop.set_task_factory(previous_factory) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["token", "pre_call"]) +async def test_native_callback_cancelled_error_aborts_before_provider( + ocr_server: RecordingServer, phase: str +) -> None: + from litellm.rust_bridge import _native + + calls: Final = [] + ocr_server.expected_requests = 0 + + class Host: + def __call__(self): + calls.append("token") + raise asyncio.CancelledError("callback cancelled") + + def pre_call(self, **kwargs): + calls.append("pre_call") + raise asyncio.CancelledError("callback cancelled") + + host: Final = Host() + future: Final = _native.aocr( + model="azure_ai/mistral-ocr-latest", + document=OCR_DOCUMENT, + api_key=None if phase == "token" else "test-key", + api_base=ocr_server.base_url, + logging_obj=host, + token_provider=host if phase == "token" else None, + callback_loop=asyncio.get_running_loop(), + ) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(future, timeout=5) + assert calls == [phase] + assert ocr_server.requests == [] + + async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): return await call_native_aocr(server, callbacks=callbacks, **kwargs) @@ -41,7 +350,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_ observations: Final = [] class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observations.append((model, copy.deepcopy(kwargs["additional_args"]))) call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0]) @@ -64,13 +373,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["include_image_base64"] = True if raise_after_edit: raise RuntimeError("pre-call callback failed") class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(copy.deepcopy(request_body(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False) @@ -83,11 +392,11 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_headers(kwargs)["x-audit-tag"] = "reviewed" class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(dict(request_headers(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()]) @@ -107,12 +416,12 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ aliases: Final = [] class Retain(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): aliases.append(request_body(kwargs)["document"] is original) retained.append(request_body(kwargs)["document"]) class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): original["document_url"] = replacement_url arguments: Final = { @@ -143,7 +452,7 @@ def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_docum retained: Final = [] class RetainAndReplace(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): body = request_body(kwargs) retained.append(body["document"]) body["document"] = replacement @@ -165,11 +474,11 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov observed: Final = [] class Rebind(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(request_body(kwargs)) call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) @@ -182,11 +491,11 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_ queued: Final = [] class QueuePayload(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): queued.append(request_body(kwargs)) class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["queued-edit"] = True call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) @@ -200,7 +509,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o finished: Final = threading.Event() class Stash(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["test-token"] = token def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -280,7 +589,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal observed: Final = [] class TrackInFlightRequest(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["request-token"] = token def log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -364,7 +673,7 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context return "caller-token" class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): assert request_headers(kwargs)["Authorization"] == "Bearer caller-token" observations.append("pre_call") request_headers(kwargs)["Authorization"] = "Bearer edited" diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index d241fe08fc8..5098d2f4ad8 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -164,7 +164,7 @@ def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) - with pytest.raises(RuntimeError, match="OCR transport failed"): + with pytest.raises(litellm.APIConnectionError, match="upstream network error"): call_native_ocr(ocr_server, timeout=0.01) assert len(ocr_server.requests) == 1 @@ -307,12 +307,11 @@ async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callbac ], ids=["oidc-assertion", "document-intelligence-model"], ) -def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks( +def test_native_azure_ocr_caller_token_supports_oidc_override_and_document_intelligence( ocr_server: RecordingServer, isolated_azure_auth: None, configuration: dict[str, object], ) -> None: - ocr_server.expected_requests = 0 calls: Final = [] recorder: Final = RecordingLogger() @@ -327,11 +326,18 @@ def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_call "callbacks": [recorder], **configuration, } - with pytest.raises(NotImplementedError): - call_native_ocr(ocr_server, **arguments) - assert calls == [] - assert recorder.events == () - assert ocr_server.requests == [] + if "model" in configuration: + ocr_server.default_response = ResponseSpec(body={ + "status": "succeeded", + "analyzeResult": {"content": "native OCR response", "pages": [{"pageNumber": 1}]}, + }) + response: Final = call_native_ocr(ocr_server, **arguments) + assert isinstance(response, OCRResponse) + assert calls == ["token"] + assert recorder.names.count("log_pre_api_call") == 1 + assert len(ocr_server.requests) == 1 + assert ocr_server.requests[0].headers["authorization"] == "Bearer unused" + assert "oidc/assertion" not in ocr_server.requests[0].raw_body.decode() @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py index 6de011b1414..d3749ccc095 100644 --- a/tests/test_litellm_rust/support/callback_recorder.py +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -81,7 +81,7 @@ class RecordingLogger(CustomLogger): await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout) return tuple(event for event in self.events if event.name == name) - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): self._record("log_pre_api_call", kwargs) def log_success_event(self, kwargs, response_obj, start_time, end_time): diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index d681d752ffd..5d754a212c1 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -36,11 +36,32 @@ async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return native_ocr.ocr(ocr_arguments(server, **kwargs)) + assert native_ocr.load_rust_ocr() is not None + litellm.rust(True) + assert_native_supported(server, **kwargs) + return call_ocr(server, **kwargs) async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return await native_ocr.aocr(ocr_arguments(server, **kwargs)) + assert native_ocr.load_rust_aocr() is not None + litellm.rust(True) + assert_native_supported(server, **kwargs) + return await call_aocr(server, **kwargs) + + +def assert_native_supported(server: RecordingServer, **kwargs: object) -> None: + arguments: Final = ocr_arguments(server, **kwargs) + request: Final = native_ocr.LiteLLMOcrRequest( + model=arguments["model"], + document=arguments["document"], + api_key=arguments["api_key"], + api_base=arguments["api_base"], + timeout=arguments.get("timeout"), + custom_llm_provider=arguments.get("custom_llm_provider"), + extra_headers=arguments.get("extra_headers"), + kwargs=arguments, + ) + assert native_ocr.supported(request), "test requires native OCR support, not Python fallback" def request_body(kwargs: dict[str, object]) -> dict[str, object]: