scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Wire OCR provider observers through the Python callback runtime and verify SDK parity

Browse source
This commit is contained in:
Yujong Lee 2026-09-02 23:09:18 -07:00
parent 716f71de19
commit 0c61391492
25 changed files with 1112 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
.github/workflows/test-rust.yml vendored
View file

@ -10,6 +10,10 @@ on:
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- "tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py"
- "tests/test_litellm/ocr/**"
- "litellm/ocr/**"
- "litellm/rust_bridge/**"
- ".github/workflows/test-rust.yml"
pull_request:
branches:
@ -25,6 +29,10 @@ on:
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- "tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py"
- "tests/test_litellm/ocr/**"
- "litellm/ocr/**"
- "litellm/rust_bridge/**"
- ".github/workflows/test-rust.yml"
permissions:
@ -133,3 +141,9 @@ jobs:
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- name: Test installed SDK callback parity and fallback
run: |
uv venv /tmp/ocr-callback-sdk --python python
uv pip install --python /tmp/ocr-callback-sdk/bin/python dist/*.whl
/tmp/ocr-callback-sdk/bin/python tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py

2
litellm-rust/crates/ai-gateway/src/io/ocr.rs
View file

@ -1 +1 @@
pub use crate::ocr::{OcrRequest, ocr};
pub use crate::ocr::{OcrRequest, ocr, ocr_with_observer};

38
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
View file

@ -1,9 +1,11 @@
use litellm_core::error::Error;
use litellm_core::http_utils::http_request;
use litellm_core::ocr::handler::send_ocr_request;
use litellm_core::ocr::observers::{OcrObserver, OcrPreCall};
use litellm_core::ocr::transformation::OcrResponseHandling;
use litellm_core::ocr::types::OcrRequestData;
use serde_json::Value;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::common_utils::poll_document_intelligence;
use super::hooks::OcrLifecycleHooks;
use super::types::PreparedOcrRequest;
use crate::client::http_client;
@ -12,6 +14,7 @@ use crate::client::http_client;
pub(crate) async fn execute_ocr_provider_call(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
observer: &mut impl OcrObserver,
) -> Result<Value, Error> {
let request = hooks.prepare_provider_request(request).await?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
@ -22,16 +25,23 @@ pub(crate) async fn execute_ocr_provider_call(
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|err| Error::Network(err.to_string()))?;
let event = OcrPreCall {
model: request.model.clone(),
request: OcrRequestData {
data: request.body,
files: None,
},
api_base: request.url.clone(),
headers: request.upstream_headers.iter().cloned().collect(),
};
let response = send_ocr_request(request_builder, &event, observer).await?;
let status = response.status();
let status = response.status;
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
&& status.as_u16() == 202
{
let operation_url = response
.headers()
.headers
.get("operation-location")
.and_then(|value| value.to_str().ok())
.map(str::to_string)
@ -58,19 +68,7 @@ pub(crate) async fn execute_ocr_provider_call(
.into_json());
}
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text)
let response_json: Value = serde_json::from_str(&response.body)
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
Ok(request

16
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
View file

@ -1,5 +1,6 @@
use litellm_core::Error;
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::ocr::observers::{NoopOcrObserver, OcrObserver};
use serde_json::Value;
mod common_utils;
@ -15,10 +16,23 @@ use prepare::{PreparedOcrCall, prepare_ocr_call};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
ocr_with_observer(request, &mut NoopOcrObserver).await
}
#[tracing::instrument(
name = "ocr",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub async fn ocr_with_observer(
request: OcrRequest<'_>,
observer: &mut impl OcrObserver,
) -> Result<Value, Error> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
CallLifecycle::default()
.run_request(request, &hooks, |request| {
execute_ocr_provider_call(request, &hooks)
execute_ocr_provider_call(request, &hooks, observer)
})
.await
}

140
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -9,12 +9,150 @@ use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_ai_gateway::ocr::{OcrRequest, ocr, ocr_with_observer};
use litellm_core::error::Error;
use litellm_core::http_utils::has_header;
use litellm_core::ocr::observers::{OcrObserver, OcrPostCall, OcrPreCall};
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
struct ProviderObserver {
events: Arc<Mutex<Vec<&'static str>>>,
raw_response: Option<String>,
reject: bool,
}
impl OcrObserver for ProviderObserver {
type Error = &'static str;
async fn pre_call(&mut self, input: &OcrPreCall) -> Result<(), Self::Error> {
assert_eq!(input.model, "mistral-ocr-4-1");
assert_eq!(
input.request.data["document"]["document_url"],
"https://example.com/document.pdf"
);
assert!(input.api_base.ends_with("/v1/ocr"));
assert!(
input
.headers
.values()
.any(|value| value == "Bearer test-key")
);
self.events.lock().unwrap().push("pre");
if self.reject {
Err("observer failure")
} else {
Ok(())
}
}
async fn post_call(&mut self, input: &OcrPostCall) -> Result<(), Self::Error> {
self.events.lock().unwrap().push("post");
self.raw_response = Some(input.original_response.clone());
if self.reject {
Err("observer failure")
} else {
Ok(())
}
}
}
fn observer_request(api_base: &str) -> OcrRequest<'_> {
OcrRequest {
model: "mistral/mistral-ocr-4-1",
document: json!({"type":"document_url","document_url":"https://example.com/document.pdf"}),
api_key: Some("test-key"),
api_base: Some(api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(2)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("observer-test"),
}
}
async fn observer_case(status: u16, body: &'static str, reject: bool) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}/v1", listener.local_addr().unwrap());
let events = Arc::new(Mutex::new(Vec::new()));
let provider_events = Arc::clone(&events);
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut socket).await;
assert!(request.starts_with("POST /v1/ocr "));
provider_events.lock().unwrap().push("http");
let response = format!(
"HTTP/1.1 {status} Test\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
socket.write_all(response.as_bytes()).await.unwrap();
});
let mut observer = ProviderObserver {
events: Arc::clone(&events),
raw_response: None,
reject,
};
let result = ocr_with_observer(observer_request(&url), &mut observer).await;
tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap();
if status != 200 {
assert!(matches!(result, Err(Error::Http { status: actual, .. }) if actual == status));
assert_eq!(*events.lock().unwrap(), ["pre", "http"]);
assert_eq!(observer.raw_response, None);
} else {
assert_eq!(*events.lock().unwrap(), ["pre", "http", "post"]);
assert_eq!(observer.raw_response.as_deref(), Some(body));
if body == "invalid-json" {
assert!(matches!(result, Err(Error::InvalidResponse(_))));
} else {
assert_eq!(result.unwrap()["pages"][0]["markdown"], "ok");
}
}
}
#[tokio::test]
async fn provider_observers_surround_http_and_cannot_replace_its_outcome() {
for reject in [false, true] {
observer_case(200, r#"{"pages":[{"index":0,"markdown":"ok"}]}"#, reject).await;
observer_case(200, "invalid-json", reject).await;
observer_case(401, r#"{"error":"rejected"}"#, reject).await;
}
}
#[tokio::test]
async fn invalid_ocr_preparation_does_not_call_observers_or_provider() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}/v1", listener.local_addr().unwrap());
let events = Arc::new(Mutex::new(Vec::new()));
let mut observer = ProviderObserver {
events: Arc::clone(&events),
raw_response: None,
reject: false,
};
let request = OcrRequest {
document: json!(42),
..observer_request(&url)
};
assert!(
ocr_with_observer(request, &mut observer)
.await
.is_err()
);
assert!(events.lock().unwrap().is_empty());
assert!(
tokio::time::timeout(Duration::from_millis(50), listener.accept())
.await
.is_err()
);
}
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];

51
litellm-rust/crates/core/src/ocr/handler.rs Normal file
View file

@ -0,0 +1,51 @@
use reqwest::{RequestBuilder, StatusCode, header::HeaderMap};
use super::observers::{OcrObserver, OcrPostCall, OcrPreCall};
use crate::Error;
use crate::http_utils::{http_request, truncate_error_body};
pub struct OcrHttpResponse {
pub status: StatusCode,
pub headers: HeaderMap,
pub body: String,
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn send_ocr_request(
request: RequestBuilder,
event: &OcrPreCall,
observer: &mut impl OcrObserver,
) -> Result<OcrHttpResponse, Error> {
if observer.pre_call(event).await.is_err() {
tracing::warn!("OCR pre-call observer failed");
}
let response = http_request(request).await.map_err(transport_error)?;
let status = response.status();
let headers = response.headers().clone();
let body = response.text().await.map_err(transport_error)?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let event = OcrPostCall {
original_response: body,
};
if observer.post_call(&event).await.is_err() {
tracing::warn!("OCR post-call observer failed");
}
Ok(OcrHttpResponse {
status,
headers,
body: event.original_response,
})
}
fn transport_error(error: reqwest::Error) -> Error {
Error::Network(if error.is_timeout() {
"Request timed out".into()
} else {
error.to_string()
})
}

2
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -1,2 +1,4 @@
pub mod handler;
pub mod observers;
pub mod transformation;
pub mod types;

48
litellm-rust/crates/core/src/ocr/observers.rs Normal file
View file

@ -0,0 +1,48 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use serde::Serialize;
use super::types::OcrRequestData;
#[derive(Serialize)]
pub struct OcrPreCall {
pub model: String,
pub request: OcrRequestData,
pub api_base: String,
pub headers: BTreeMap<String, String>,
}
#[derive(Serialize)]
pub struct OcrPostCall {
pub original_response: String,
}
#[macro_export]
macro_rules! ocr_observer_catalog {
($consumer:path, $($options:tt)*) => {
$consumer! {
$($options)*
{
pre_call: PreCall($crate::ocr::observers::OcrPreCall) -> () = direct;
post_call: PostCall($crate::ocr::observers::OcrPostCall) -> () = direct;
}
}
};
}
ocr_observer_catalog!(crate::define_hooks, pub trait OcrObserver;);
pub struct NoopOcrObserver;
impl OcrObserver for NoopOcrObserver {
type Error = Infallible;
async fn pre_call(&mut self, _input: &OcrPreCall) -> Result<(), Infallible> {
Ok(())
}
async fn post_call(&mut self, _input: &OcrPostCall) -> Result<(), Infallible> {
Ok(())
}
}

1
litellm-rust/crates/python-bridge/src/constants.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) const OCR_CALLBACK_CAPACITY: usize = 1024;

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

@ -27,6 +27,24 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr {
}
}
pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr {
let status = match &err {
Error::Http { status, .. } => Some(*status),
_ => None,
};
let exception = core_error_to_pyerr(err);
if let Some(status) = status {
Python::attach(|py| {
if let Err(error) = exception.value(py).setattr("status_code", status) {
return error;
}
exception
})
} else {
exception
}
}
/// 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.

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

@ -10,6 +10,11 @@ use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(crate) struct PythonCallContext<'py> {
pub(crate) py: Python<'py>,
pub(crate) asynchronous: bool,
}
pub(crate) fn run_sync<T, E>(
py: Python<'_>,
future: impl Future<Output = Result<T, E>> + Send + 'static,

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

@ -1,9 +1,11 @@
mod constants;
mod diagnostics;
mod errors;
mod execution;
#[cfg(feature = "trace-parity")]
mod function_trace;
mod marshal;
mod ocr_callbacks;
mod python_hook_bindings;
mod routes;
@ -76,6 +78,7 @@ mod _native {
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
litellm_python_interop::callback_runtime::register(module)?;
super::ocr_callbacks::register(module)?;
super::errors::register(module)?;
module.add("ready_endpoints", pyo3::types::PyDict::new(module.py()))?;
super::routes::register(module)?;

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

@ -0,0 +1,82 @@
use std::num::NonZeroUsize;
use litellm_core::ocr::observers::{OcrObserver, OcrPostCall, OcrPreCall};
use litellm_python_interop::callback_runtime::{AsyncContext, CallbackRuntime, SyncContext};
use pyo3::prelude::*;
use crate::constants::OCR_CALLBACK_CAPACITY;
use crate::execution::PythonCallContext;
litellm_core::ocr_observer_catalog!(crate::bind_python_hooks,
pub(crate) struct PythonOcrSession;
trait OcrObserver;
);
#[pyclass(frozen)]
struct OcrRuntime(CallbackRuntime);
pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
let capacity = NonZeroUsize::new(OCR_CALLBACK_CAPACITY)
.expect("OCR callback capacity is a positive constant");
module.add(
"__ocr_callback_runtime__",
OcrRuntime(CallbackRuntime::new(module, capacity)?),
)
}
pub(crate) enum PythonOcrObserver {
Disabled,
Sync(PythonOcrSession<SyncContext>),
Async(PythonOcrSession<AsyncContext>),
}
impl PythonOcrObserver {
pub(crate) fn new(
adapter: Option<Py<PyAny>>,
context: PythonCallContext<'_>,
) -> PyResult<Self> {
let Some(adapter) = adapter else {
return Ok(Self::Disabled);
};
let py = context.py;
let module = py.import("litellm.rust_bridge._native")?;
let runtime = module
.getattr("__ocr_callback_runtime__")?
.extract::<PyRef<'_, OcrRuntime>>()?
.0
.clone();
if context.asynchronous {
Ok(Self::Async(PythonOcrSession::new(
adapter.bind(py),
runtime.async_context(py)?,
)?))
} else {
Ok(Self::Sync(PythonOcrSession::new(
adapter.bind(py),
runtime.sync_context(py)?,
)?))
}
}
}
impl OcrObserver for PythonOcrObserver {
type Error = PyErr;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn pre_call(&mut self, input: &OcrPreCall) -> PyResult<()> {
match self {
Self::Disabled => Ok(()),
Self::Sync(session) => session.pre_call(input).await,
Self::Async(session) => session.pre_call(input).await,
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn post_call(&mut self, input: &OcrPostCall) -> PyResult<()> {
match self {
Self::Disabled => Ok(()),
Self::Sync(session) => session.post_call(input).await,
Self::Async(session) => session.post_call(input).await,
}
}
}

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

@ -9,43 +9,50 @@ macro_rules! bridge_route {
inputs = $inputs:ident,
required = { $($(#[$required_attr:meta])* $required_name:ident: $required_type:ty),+ $(,)? },
optional = { $($(#[$optional_attr:meta])* $optional_name:ident: $optional_type:ty),* $(,)? },
$(keyword_only = { $($keyword_name:ident: $keyword_type:ty),+ $(,)? },)?
prepare = $prepare:path,
errors = $map_error:path
$(, context = $context:ident)?
$(, extra = [$($extra:ident),* $(,)?])?
$(,)?
) => {
struct $inputs {
$($required_name: $required_type,)*
$($optional_name: $optional_type),*
$($optional_name: $optional_type,)*
$($($keyword_name: $keyword_type,)*)?
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),* $(, *, $($keyword_name=None),*)?))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
$($($keyword_name: $keyword_type,)*)?
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$($optional_name,)*
$($($keyword_name,)*)?
}, $($context { py, asynchronous: false },)?)?;
$crate::execution::run_sync(py, future, $map_error)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),* $(, *, $($keyword_name=None),*)?))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
$($($keyword_name: $keyword_type,)*)?
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$($optional_name,)*
$($($keyword_name,)*)?
}, $($context { py, asynchronous: true },)?)?;
$crate::execution::run_async(py, future, $map_error)
}
@ -61,20 +68,22 @@ macro_rules! bridge_route {
#[cfg(feature = "trace-parity")]
mod trace {
use pyo3::prelude::*;
use super::{$inputs, $map_error, $prepare};
use super::{$inputs, $map_error, $prepare $(, $context)?};
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),* $(, *, $($keyword_name=None),*)?))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
$($($keyword_name: $keyword_type,)*)?
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$($optional_name,)*
$($($keyword_name,)*)?
}, $($context { py, asynchronous: false },)?)?;
$crate::execution::run_sync(
py,
$crate::function_trace::capture(future),
@ -83,17 +92,19 @@ macro_rules! bridge_route {
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),* $(, *, $($keyword_name=None),*)?))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
$($($keyword_name: $keyword_type,)*)?
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$($optional_name,)*
$($($keyword_name,)*)?
}, $($context { py, asynchronous: true },)?)?;
$crate::execution::run_async(
py,
$crate::function_trace::capture(future),
@ -239,7 +250,7 @@ mod tests {
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, callback_adapter=None)",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, *, callback_adapter=None)",
),
(
"transcription",

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

@ -1,15 +1,18 @@
use litellm_core::Error;
use std::future::Future;
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr_with_observer};
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::ocr_error_to_pyerr;
use crate::execution::PythonCallContext;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
use crate::ocr_callbacks::PythonOcrObserver;
fn prepare_ocr(
inputs: OcrInputs,
context: PythonCallContext<'_>,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let _ = inputs.callback_adapter;
let document = inputs.document;
@ -22,6 +25,7 @@ fn prepare_ocr(
timeout_seconds: inputs.timeout_seconds,
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
let mut observer = PythonOcrObserver::new(inputs.callback_adapter, context)?;
Ok(async move {
let RouteOptions {
@ -32,20 +36,23 @@ fn prepare_ocr(
extra_headers,
timeout,
} = options;
run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
ocr_with_observer(
OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
},
&mut observer,
)
.await
})
}
@ -68,8 +75,9 @@ bridge_route! {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
callback_adapter: Option<pyo3::Py<pyo3::PyAny>>,
},
keyword_only = { callback_adapter: Option<Py<PyAny>> },
prepare = prepare_ocr,
errors = ocr_error_to_pyerr,
context = PythonCallContext,
}

34
litellm/ocr/main.py
View file

@ -29,6 +29,7 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.ocr_callbacks import OcrLoggingAdapter
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -233,6 +234,8 @@ def _rust_bridge_api_base(
def _prepare_rust_ocr_call(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
*,
asynchronous: bool,
) -> rust_ocr_bridge.RustOCRRequest:
provider_config: Final = prepared_request.provider_config
api_key_env_var: Final = provider_config.get_api_key_env_var()
@ -254,19 +257,21 @@ def _prepare_rust_ocr_call(
)
rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key)
rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key)
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
callbacks_in_rust: Final = rust_ocr_bridge.supports_callback_adapter(asynchronous=asynchronous)
if not callbacks_in_rust:
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
)
)
return rust_ocr_bridge.RustOCRRequest(
model=prepared_request.model,
document=prepared_request.document,
@ -278,6 +283,9 @@ def _prepare_rust_ocr_call(
),
optional_params=rust_optional_params,
timeout=prepared_request.effective_timeout,
callback_adapter=(
OcrLoggingAdapter(prepared_request.litellm_logging_obj, resolved_api_key) if callbacks_in_rust else None
),
)
@ -289,6 +297,7 @@ def _run_rust_ocr(
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
asynchronous=False,
),
model=prepared_request.model,
provider=prepared_request.custom_llm_provider,
@ -308,6 +317,7 @@ async def _run_rust_aocr(
prepare=lambda: _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
asynchronous=True,
),
model=prepared_request.model,
provider=prepared_request.custom_llm_provider,

3
litellm/rust_bridge/bindings.py
View file

@ -59,6 +59,9 @@ class NativeBinding(Generic[BindingT]):
def reset(self) -> None:
self._override = _UNSET
def is_overridden(self) -> bool:
return not isinstance(self._override, _Unset)
def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None:
native: Final = get_native_bridge()

10
litellm/rust_bridge/ocr.py
View file

@ -96,6 +96,16 @@ def set_rust_ocr(
_OCR.asynchronous.override(aocr)
def supports_callback_adapter(*, asynchronous: bool = False) -> bool:
binding = _OCR.asynchronous if asynchronous else _OCR.sync
if binding.is_overridden():
return False
from litellm.rust_bridge import get_native_bridge
native: Final = get_native_bridge()
return native is not None and hasattr(native, "__ocr_callback_runtime__")
def ocr(
*,
prepare: Callable[[], RustOCRRequest],

59
litellm/rust_bridge/ocr_callbacks.py Normal file
View file

@ -0,0 +1,59 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Protocol
from pydantic import BaseModel, JsonValue
from typing_extensions import ReadOnly, TypedDict
class PreCallArguments(TypedDict):
complete_input_dict: ReadOnly[Mapping[str, JsonValue]]
api_base: ReadOnly[str]
headers: ReadOnly[Mapping[str, str]]
class OcrLogging(Protocol):
def pre_call(self, *, input: str, api_key: str | None, additional_args: PreCallArguments) -> None: ...
def post_call(self, *, original_response: str, input: str, api_key: str | None) -> None: ...
class OcrRequestSnapshot(BaseModel):
data: Mapping[str, JsonValue]
class OcrPreCall(BaseModel):
model: str
request: OcrRequestSnapshot
api_base: str
headers: Mapping[str, str]
class OcrPostCall(BaseModel):
original_response: str
@dataclass(frozen=True, slots=True)
class OcrLoggingAdapter:
logging_obj: OcrLogging
api_key: str | None
def pre_call(self, payload: object) -> None:
event: Final = OcrPreCall.model_validate(payload)
additional_args: Final[PreCallArguments] = {
"complete_input_dict": event.request.data,
"api_base": event.api_base,
"headers": event.headers,
}
self.logging_obj.pre_call(
input="OCR document processing",
api_key=self.api_key,
additional_args=additional_args,
)
def post_call(self, payload: object) -> None:
event: Final = OcrPostCall.model_validate(payload)
self.logging_obj.post_call(
original_response=event.original_response,
input="OCR document processing",
api_key=self.api_key,
)

3
litellm/rust_bridge/runtime.py
View file

@ -94,6 +94,9 @@ class EndpointBinding(Generic[BindingT]):
raise RuntimeError("only native endpoint bindings support resets")
self._native_binding.reset()
def is_overridden(self) -> bool:
return self._native_binding is not None and self._native_binding.is_overridden()
def _attempt(
self,
*,

202
tests/test_litellm/ocr/callback_support.py Normal file
View file

@ -0,0 +1,202 @@
import asyncio
import contextvars
import inspect
import json
import threading
from collections.abc import Generator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, TypedDict
from typing_extensions import ReadOnly
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.ocr.transformation import OCRResponse
MODEL: Final = "mistral/mistral-ocr-4-1"
DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/document.pdf"}
RESPONSE: Final = (
'{"pages":[{"index":0,"markdown":"callback-test"}],"model":"mistral-ocr-4-1","usage_info":{"pages_processed":1}}'
)
REQUEST_CONTEXT: Final = contextvars.ContextVar("ocr_callback_test_context", default="unset")
@dataclass(frozen=True, slots=True)
class CallbackEvent:
name: str
model: object
call_id: object
original_response: object
input: object
additional_args: str
response_type: str
start_time: datetime | None
end_time: datetime | None
context: str
native_provider_hook: bool
class CallbackRecorder(CustomLogger):
def __init__(self, asynchronous: bool, raises: str = "", name: str = "recorder", expected_calls: int = 1) -> None:
super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger exposes untyped keyword arguments
self.asynchronous = asynchronous
self.name = name
self.expected_calls = expected_calls
self.raises = raises
self.events: tuple[CallbackEvent, ...] = ()
self.done = threading.Event()
self.lock = threading.Lock()
def record(
self,
name: str,
kwargs: Mapping[str, object],
response: object = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
) -> None:
event: Final = CallbackEvent(
name,
kwargs.get("model"),
kwargs.get("litellm_call_id"),
kwargs.get("original_response"),
kwargs.get("input"),
json.dumps(kwargs.get("additional_args"), sort_keys=True, default=str),
type(response).__name__,
start_time,
end_time,
REQUEST_CONTEXT.get(),
any(frame.filename == "litellm_callbacks.py" for frame in inspect.stack(context=0)),
)
with self.lock:
self.events += (event,)
terminals: Final = tuple(
item.name for item in self.events if "success" in item.name or "failure" in item.name
)
if len(terminals) >= self.expected_calls * (2 if self.asynchronous and "failure" in name else 1):
self.done.set()
if name == self.raises:
raise ValueError("intentional observer failure")
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
assert model == kwargs.get("model")
assert isinstance(messages, list)
self.record("pre", kwargs)
def log_post_api_call(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime | None
) -> None:
self.record("post", kwargs, response_obj, start_time, end_time)
def log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.record("success", kwargs, response_obj, start_time, end_time)
def log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.record("failure", kwargs, response_obj, start_time, end_time)
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.record("async_success", kwargs, response_obj, start_time, end_time)
async def async_log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.record("async_failure", kwargs, response_obj, start_time, end_time)
async def wait(self) -> tuple[CallbackEvent, ...]:
assert await asyncio.to_thread(self.done.wait, 5), tuple(event.name for event in self.events)
return self.events
class OcrUpstream(ThreadingHTTPServer):
daemon_threads = True
def __init__(self, status: int, body: str, stall: bool = False) -> None:
super().__init__(("127.0.0.1", 0), OcrHandler)
self.status = status
self.body = body.encode()
self.stall = stall
self.release = threading.Event()
self.started = threading.Event()
self.requests: tuple[tuple[str, str], ...] = ()
self.lock = threading.Lock()
@property
def api_base(self) -> str:
return f"http://127.0.0.1:{self.server_port}/v1"
class OcrHandler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
assert isinstance(self.server, OcrUpstream)
body: Final = self.rfile.read(int(self.headers["content-length"])).decode()
with self.server.lock:
self.server.requests += ((self.path, body),)
self.server.started.set()
if self.server.stall:
self.server.release.wait(5)
return
self.send_response(self.server.status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(self.server.body)))
self.end_headers()
self.wfile.write(self.server.body)
def log_message(self, format: str, *args: object) -> None:
pass
@contextmanager
def ocr_upstream(status: int = 200, body: str = RESPONSE, stall: bool = False) -> Generator[OcrUpstream, None, None]:
with OcrUpstream(status, body, stall) as server:
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield server
finally:
server.release.set()
server.shutdown()
thread.join(timeout=5)
assert not thread.is_alive()
def assert_provider_request(server: OcrUpstream) -> None:
assert len(server.requests) == 1
path, body = server.requests[0]
assert path == "/v1/ocr"
assert json.loads(body) == {"model": "mistral-ocr-4-1", "document": DOCUMENT}
class OcrArguments(TypedDict):
model: ReadOnly[str]
document: ReadOnly[dict[str, str]]
api_key: ReadOnly[str]
api_base: ReadOnly[str]
callbacks: ReadOnly[list[CustomLogger]]
rust: ReadOnly[bool]
num_retries: ReadOnly[int]
timeout: ReadOnly[float]
def verify_installed_package() -> None:
root: Final = Path(litellm.__file__).resolve()
assert "site-packages" in root.parts, f"SDK must come from the installed wheel: {root}"
async def call_ocr(arguments: OcrArguments, asynchronous: bool) -> tuple[OCRResponse | None, Exception | None]:
try:
response: Final = (
await litellm.aocr(**arguments) if asynchronous else await asyncio.to_thread(litellm.ocr, **arguments)
)
return OCRResponse.model_validate(response), None
except Exception as error:
return None, error

58
tests/test_litellm/ocr/live_callback_smoke.py Normal file
View file

@ -0,0 +1,58 @@
import asyncio
import os
import sys
from typing import Final
from callback_support import CallbackRecorder, OcrArguments, call_ocr, verify_installed_package
import litellm
async def smoke(*, baseline: bool) -> None:
verify_installed_package()
for asynchronous, rejected in ((False, False), (True, False), (True, True)):
await smoke_case(asynchronous, rejected, baseline=baseline)
async def smoke_case(asynchronous: bool, rejected: bool, *, baseline: bool) -> None:
recorder: Final = CallbackRecorder(asynchronous, name=f"live-{asynchronous}-{rejected}")
arguments: Final[OcrArguments] = {
"model": "mistral/mistral-ocr-latest",
"document": {
"type": "document_url",
"document_url": "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf",
},
"api_key": "invalid-smoke-key" if rejected else os.environ["MISTRAL_API_KEY"],
"callbacks": [recorder],
"rust": True,
"api_base": "https://api.mistral.ai/v1",
"num_retries": 0,
"timeout": 60,
}
response, error = await call_ocr(arguments, asynchronous)
assert (error is not None) is rejected
if error is not None and not baseline:
assert isinstance(error, litellm.AuthenticationError) and error.status_code == 401
if response is not None:
assert len(response.pages) == 1
outcome: Final = (
f"pages={len(response.pages)}"
if response is not None
else f"{type(error).__name__} status={getattr(error, 'status_code', None)}"
)
events: Final = await recorder.wait()
names: Final = tuple(event.name for event in events)
if rejected:
assert names[0] == "pre" and sorted(names[1:]) == ["async_failure", "failure"]
else:
assert names == ("pre", *(() if baseline else ("post",)), "async_success" if asynchronous else "success")
if not baseline:
assert all(event.native_provider_hook for event in events if event.name in ("pre", "post"))
sys.stdout.write(f"async={asynchronous} rejected={rejected} {outcome} events={names}\n")
sys.stdout.flush()
if __name__ == "__main__":
if "--run-live" not in sys.argv:
raise SystemExit("Pass --run-live and set MISTRAL_API_KEY to make real, billable provider calls")
asyncio.run(smoke(baseline="--baseline" in sys.argv))

250
tests/test_litellm/ocr/sdk_callback_contract.py Normal file
View file

@ -0,0 +1,250 @@
import asyncio
import json
import sys
from collections import Counter
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Final
from callback_support import (
DOCUMENT,
MODEL,
REQUEST_CONTEXT,
RESPONSE,
CallbackRecorder,
OcrArguments,
assert_provider_request,
call_ocr,
ocr_upstream,
verify_installed_package,
)
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.rust_bridge.ocr import supports_callback_adapter
from litellm.rust_bridge.ocr_callbacks import PreCallArguments
@dataclass(frozen=True, slots=True)
class Outcome:
response: str | None
error_type: str | None
status: int | None
events: tuple[str, ...]
response_types: tuple[str, ...]
provider_body: str
async def exercise(asynchronous: bool, rust: bool, case: str, *, native_expected: bool | None = None) -> Outcome:
litellm.logging_callback_manager._reset_all_callbacks() # pyright: ignore[reportPrivateUsage] # isolate SDK cases in this dedicated test process
native: Final = rust if native_expected is None else native_expected
if native:
assert supports_callback_adapter(asynchronous=asynchronous), "native callback bridge must be installed"
status: Final = int(case) if case.isdigit() else 500 if case in ("raise_failure", "raise_async_failure") else 200
body: Final = "invalid-json" if case == "malformed" else RESPONSE if status == 200 else '{"error":"rejected"}'
successful: Final = status == 200 and case not in ("malformed", "timeout")
raises: Final = case.removeprefix("raise_") if case.startswith("raise_") else ""
recorder: Final = CallbackRecorder(asynchronous, raises=raises, name=f"first-{asynchronous}-{rust}-{case}")
follower: Final = CallbackRecorder(asynchronous, name=f"second-{asynchronous}-{rust}-{case}")
context: Final = f"request-{asynchronous}-{rust}-{case}"
token: Final = REQUEST_CONTEXT.set(context)
try:
with ocr_upstream(status, body, stall=case == "timeout") as upstream:
params: Final[OcrArguments] = {
"model": MODEL,
"document": DOCUMENT,
"api_key": "test-key",
"api_base": upstream.api_base,
"callbacks": [recorder, follower],
"rust": rust,
"num_retries": 0,
"timeout": 0.2 if case == "timeout" else 5,
}
result, error = await call_ocr(params, asynchronous)
assert (error is None) is successful
if error is not None:
if case.isdigit():
assert getattr(error, "status_code", None) == status
if case == "timeout":
assert isinstance(error, litellm.Timeout)
if result is not None:
assert result.pages[0].markdown == "callback-test"
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), 5)
events: Final = await recorder.wait()
following: Final = await follower.wait()
names: Final = tuple(event.name for event in events)
prefix: Final = ("pre", "post") if status == 200 and case != "timeout" else ("pre",)
assert names[: len(prefix)] == prefix
terminal: Final = "success" if successful else "failure"
expected: Final = (
("async_success",)
if asynchronous and successful
else ((terminal, f"async_{terminal}") if asynchronous else (terminal,))
)
assert sorted(names[len(prefix) :]) == sorted(expected)
assert sorted(event.name for event in following) == sorted(names)
assert len({event.call_id for event in events}) == 1
assert isinstance(events[0].call_id, str) and events[0].call_id
assert all(event.model == "mistral-ocr-4-1" for event in events)
assert events[0].input == "OCR document processing"
pre_arguments: Final = TypeAdapter(PreCallArguments).validate_json(events[0].additional_args)
assert pre_arguments["complete_input_dict"] == {"model": "mistral-ocr-4-1", "document": DOCUMENT}
assert pre_arguments["api_base"] == upstream.api_base + "/ocr"
assert {key.lower(): value for key, value in pre_arguments["headers"].items()}[
"authorization"
] == "Bearer test-key"
for event in events[: len(prefix)]:
assert event.context == context
assert event.native_provider_hook is native
if "post" in prefix:
assert events[1].original_response == body
assert events[1].response_type == "NoneType"
assert events[1].start_time is not None and events[1].end_time is None
for event in events[len(prefix) :]:
assert event.start_time is not None and event.end_time is not None
assert event.end_time >= event.start_time
assert_provider_request(upstream)
return Outcome(
result.model_dump_json() if result is not None else None,
type(error).__name__ if error is not None else None,
getattr(error, "status_code", None),
tuple(sorted(names)),
tuple(sorted(event.response_type for event in events)),
json.dumps(json.loads(upstream.requests[0][1]), sort_keys=True),
)
finally:
REQUEST_CONTEXT.reset(token)
async def verify_parity() -> None:
for asynchronous in (False, True):
for case in (
"success",
"malformed",
"401",
"429",
"500",
"timeout",
"raise_pre",
"raise_post",
"raise_success",
"raise_async_success",
"raise_failure",
"raise_async_failure",
):
await verify_case(asynchronous, case)
await verify_concurrency()
for rust in (False, True):
await verify_delayed_terminal(rust)
for rust in (False, True):
await verify_without_loggers(rust)
async def verify_case(asynchronous: bool, case: str) -> None:
python: Final = await exercise(asynchronous, False, case)
native: Final = await exercise(asynchronous, True, case)
assert python == native, (asynchronous, case, python, native)
sys.stdout.write(f"PASS async={asynchronous} case={case}\n")
sys.stdout.flush()
async def verify_concurrency() -> None:
litellm.logging_callback_manager._reset_all_callbacks() # pyright: ignore[reportPrivateUsage] # isolate the concurrent SDK scenario
recorder: Final = CallbackRecorder(True, name="concurrent", expected_calls=32)
with ocr_upstream() as upstream:
async def one(index: int) -> None:
token: Final = REQUEST_CONTEXT.set(f"concurrent-{index}")
try:
result: Final = await litellm.aocr(
model=MODEL,
document=DOCUMENT,
api_key="test-key",
api_base=upstream.api_base,
callbacks=[recorder],
rust=True,
num_retries=0,
timeout=10,
)
assert result.pages[0].markdown == "callback-test"
finally:
REQUEST_CONTEXT.reset(token)
await asyncio.gather(*(one(index) for index in range(32)))
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), 5)
events: Final = await recorder.wait()
assert len(upstream.requests) == 32
assert Counter(event.name for event in events) == {"pre": 32, "post": 32, "async_success": 32}
assert len({event.call_id for event in events}) == 32
for index in range(32):
assert tuple(event.name for event in events if event.context == f"concurrent-{index}") == (
"pre",
"post",
"async_success",
)
assert all(event.native_provider_hook for event in events if event.name in ("pre", "post"))
class DelayedRecorder(CallbackRecorder):
def __init__(self) -> None:
super().__init__(True, name="delayed")
self.started = asyncio.Event()
self.release = asyncio.Event()
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.started.set()
await self.release.wait()
await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
async def verify_delayed_terminal(rust: bool) -> None:
litellm.logging_callback_manager._reset_all_callbacks() # pyright: ignore[reportPrivateUsage] # isolate delayed delivery
recorder: Final = DelayedRecorder()
with ocr_upstream() as upstream:
result: Final = await asyncio.wait_for(
litellm.aocr(
model=MODEL,
document=DOCUMENT,
api_key="test-key",
api_base=upstream.api_base,
callbacks=[recorder],
rust=rust,
num_retries=0,
),
5,
)
assert result.pages[0].markdown == "callback-test"
await asyncio.wait_for(recorder.started.wait(), 5)
assert tuple(event.name for event in recorder.events) == ("pre", "post")
recorder.release.set()
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), 5)
events: Final = await recorder.wait()
assert tuple(event.name for event in events) == ("pre", "post", "async_success"), events
async def verify_without_loggers(rust: bool) -> None:
litellm.logging_callback_manager._reset_all_callbacks() # pyright: ignore[reportPrivateUsage] # exercise the no-logger configuration
with ocr_upstream() as upstream:
result: Final = await litellm.aocr(
model=MODEL, document=DOCUMENT, api_key="test-key", api_base=upstream.api_base, rust=rust, num_retries=0
)
assert result.pages[0].markdown == "callback-test"
assert_provider_request(upstream)
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), 5)
async def verify_unavailable() -> None:
assert not supports_callback_adapter()
for asynchronous in (False, True):
for case in ("success", "401"):
await exercise(asynchronous, True, case, native_expected=False)
if __name__ == "__main__":
if "--installed" in sys.argv:
verify_installed_package()
asyncio.run(verify_unavailable() if "--without-native" in sys.argv else verify_parity())

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

@ -1,6 +1,10 @@
"""Tests for the optional Rust-backed OCR path."""
import importlib
import os
import subprocess
import types
from pathlib import Path
from typing import Any, Final
from unittest.mock import Mock
@ -25,6 +29,16 @@ DOCUMENT: dict[str, object] = {
"document_url": "https://example.com/doc.pdf",
}
def test_installed_wheel_ocr_callback_parity() -> None:
wheel_python: Final = os.environ.get("LITELLM_OCR_WHEEL_PYTHON")
if wheel_python is None:
pytest.skip("set LITELLM_OCR_WHEEL_PYTHON to the reviewed wheel's interpreter; release-wheel CI requires this lane")
script: Final = Path(__file__).resolve().parents[1] / "rust_bridge" / "sdk_callback_wheel_test.py"
completed: Final = subprocess.run((wheel_python, str(script)), check=False, timeout=240)
assert completed.returncode == 0
FAKE_OCR_RESPONSE: dict[str, object] = {
"pages": [{"index": 0, "markdown": "hello world"}],
"model": "mistral-ocr-2505-completion",

44
tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py Normal file
View file

@ -0,0 +1,44 @@
import importlib.util
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Final
import litellm
def main() -> None:
assert "site-packages" in Path(litellm.__file__).resolve().parts, "install the reviewed wheel first"
spec: Final = importlib.util.find_spec("litellm.rust_bridge._native")
assert spec is not None and spec.origin is not None, "wheel must contain the native extension"
native: Final = Path(spec.origin)
hidden: Final = native.with_suffix(native.suffix + ".disabled")
assert not hidden.exists()
script: Final = Path(__file__).resolve().parents[1] / "ocr" / "sdk_callback_contract.py"
environment: Final = {
**{key: value for key, value in os.environ.items() if key not in ("PYTHONPATH", "PYTHONHOME")},
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
"NO_PROXY": "127.0.0.1,localhost",
}
with tempfile.TemporaryDirectory(prefix="ocr-sdk-callbacks-") as directory:
subprocess.run(
[sys.executable, str(script), "--installed"], cwd=directory, env=environment, check=True, timeout=180
)
native.rename(hidden)
try:
subprocess.run(
[sys.executable, str(script), "--installed", "--without-native"],
cwd=directory,
env=environment,
check=True,
timeout=60,
)
finally:
hidden.rename(native)
sys.stdout.write("Installed-wheel callback parity and unavailable-extension fallback passed\n")
if __name__ == "__main__":
main()