scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(ocr): drive logging callbacks through prepared-request hooks

Browse source
This commit is contained in:
Yujong Lee 2026-09-10 20:57:12 -07:00 • committed by yujonglee
parent 6cc2ae8fcf
commit 23e8cf4825
45 changed files with 1404 additions and 294 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs
View file

@ -32,11 +32,14 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter {
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
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)))

19
litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs
View file

@ -33,13 +33,16 @@ impl OcrAdapter for AzureMistralAdapter {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(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<String> + Sync),
) -> Result<Vec<(String, String)>, 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());
}

2
litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs
View file

@ -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, &params)?;
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(

2
litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs
View file

@ -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, &params)?;
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(

54
litellm-rust/crates/core/src/ocr/hooks.rs
View file

@ -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<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
pub type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + 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<String, Value>,
}
#[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<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse>
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<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse
request: LiteLLMOcrRequest,
) -> 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
})
})

78
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -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<B>(
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::<B>::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::<B>::decode(changed.body)?;
validate(&body.body)?;
build_http_request(client, request, url, headers, &body).await
}
pub(crate) fn build_http_request<B: Serialize>(
pub(crate) async fn build_http_request<B>(
client: &OcrClient,
request: &LiteLLMOcrRequest,
url: &str,
headers: &[(String, String)],
body: &B,
) -> Result<reqwest::Request, OcrError> {
) -> Result<reqwest::Request, OcrError>
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<OcrDocument, OcrError> {
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<B> {
#[serde(flatten)]
body: B,

4
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -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<TokenProviderHandle>,
pub api_key: Option<String>,
pub api_key_source: InputSource,
pub api_base: Option<String>,
@ -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,

3
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -73,6 +73,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
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
})
}

4
litellm-rust/crates/core/tests/azure_ai_ocr.rs
View file

@ -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,

9
litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
View file

@ -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
})
})

245
litellm-rust/crates/core/tests/ocr.rs
View file

@ -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))
);
}

4
litellm-rust/crates/core/tests/reducto_ocr.rs
View file

@ -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,

135
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -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<u16>,
message: String,
},
Internal(String),
Host(PyErr),
}
impl From<BridgeError> 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::<RustUpstreamError>())
}
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::<PyValueError>(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::<RustUpstreamError>(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()));
});
}
}

20
litellm-rust/crates/python-bridge/src/execution.rs
View file

@ -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<T, E, F>(
py: Python<'_>,
future: F,
map_error: fn(E) -> PyErr,
map_error: fn(E) -> BridgeError,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
@ -32,14 +33,14 @@ fn run_sync_on<T, E, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(E) -> PyErr,
map_error: fn(E) -> BridgeError,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
E: Send + 'static,
F: Future<Output = Result<T, E>> + 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<T, E, F>(
py: Python<'_>,
future: F,
map_error: fn(E) -> PyErr,
map_error: fn(E) -> BridgeError,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
@ -67,12 +68,13 @@ where
})
}
fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> PyErr) -> PyResult<T> {
fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> BridgeError) -> PyResult<T> {
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")
}

19
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -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<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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)))
})
}
}

4
litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs
View file

@ -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<f64>,
},
prepare = prepare_transcription,
errors = core_error_to_pyerr,
errors = required_route_error,
}

4
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -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<f64>,
},
prepare = prepare_chat_completions,
errors = chat_completions_error_to_pyerr,
errors = fallback_route_error,
extra = [chat_completions_decline],
}

6
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -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",

4
litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs
View file

@ -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,
)
}

4
litellm-rust/crates/python-bridge/src/routes/messages.rs
View file

@ -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<f64>,
},
prepare = prepare_messages,
errors = core_error_to_pyerr,
errors = fallback_route_error,
}

79
litellm-rust/crates/python-bridge/src/routes/ocr.rs
View file

@ -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<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let document = inputs.document;
) -> PyResult<impl Future<Output = Result<Value, BridgeError>> + 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<PyAny>,
},
optional = {
api_key: Option<String>,
@ -92,9 +148,12 @@ bridge_route! {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
input_sources: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
logging_obj: Option<Py<PyAny>>,
callback_loop: Option<Py<PyAny>>,
token_provider: Option<Py<PyAny>>,
},
prepare = prepare_ocr,
errors = ocr_error_to_pyerr,
errors = std::convert::identity,
}
#[cfg(test)]

207
litellm-rust/crates/python-bridge/src/routes/ocr_callbacks.rs Normal file
View file

@ -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<Py<PyAny>>);
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<Py<PyAny>>,
pub token_provider: Option<Py<PyAny>>,
pub document: Py<PyAny>,
pub document_snapshot: serde_json::Value,
pub api_key: Option<String>,
pub locals: Option<TaskLocals>,
pub error: Mutex<Option<PyErr>>,
}
impl PythonOcrHooks {
pub async fn invoke(&self, callback: Py<PyAny>) -> PyResult<Py<PyAny>> {
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<PyAny>, Py<PyDict>, Py<PyDict>)> {
let body = to_py(py, &request.body)?
.into_bound(py)
.cast_into::<PyDict>()?;
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<OcrPreparedRequest> = 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::<std::collections::BTreeMap<String, String>>()?
.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<String> = 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::<pyo3::exceptions::PyTypeError>(py)
|| !error.is_instance_of::<pyo3::exceptions::PyException>(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::<String>(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(),
))
}
}
})
}
}

15
litellm-rust/crates/python-bridge/src/token_counter.rs
View file

@ -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<Self> {
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<InputTokenCount
counter.count_request(&request)
}
fn token_count_error_to_pyerr(error: Error) -> 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),
}
}

22
litellm-rust/crates/python-interop/src/callback.rs Normal file
View file

@ -0,0 +1,22 @@
use std::cell::Cell;
use pyo3::prelude::*;
thread_local! {
static IN_CALLBACK: Cell<bool> = const { Cell::new(false) };
}
pub fn in_callback() -> bool {
IN_CALLBACK.get()
}
pub fn invoke_callback(callback: &Py<PyAny>) -> PyResult<Py<PyAny>> {
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))
}

2
litellm-rust/crates/python-interop/src/lib.rs
View file

@ -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};

4
litellm/integrations/custom_guardrail.py
View file

@ -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)

6
litellm/litellm_core_utils/guardrail_call_context.py Normal file
View file

@ -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)

21
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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

6
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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]

5
litellm/rust_bridge/_callbacks.py Normal file
View file

@ -0,0 +1,5 @@
from collections.abc import Callable
async def invoke_callback(callback: Callable[[], object]) -> object:
return callback()

9
litellm/rust_bridge/bindings.py
View file

@ -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

27
litellm/rust_bridge/chat_completions.py
View file

@ -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,

68
litellm/rust_bridge/ocr.py
View file

@ -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,
)

10
litellm/rust_bridge/runtime.py
View file

@ -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,

57
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -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()

55
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -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

40
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -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):

14
tests/test_litellm/rust_bridge/test_bindings.py
View file

@ -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

3
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -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):

14
tests/test_litellm_rust/README.md
View file

@ -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

13
tests/test_litellm_rust/conftest.py
View file

@ -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:

339
tests/test_litellm_rust/ocr/test_callbacks.py
View file

@ -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"

22
tests/test_litellm_rust/ocr/test_requests.py
View file

@ -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

2
tests/test_litellm_rust/support/callback_recorder.py
View file

@ -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):

25
tests/test_litellm_rust/support/requests.py
View file

@ -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]: