diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 9b8b132df62..799812fd03e 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -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 diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 2fc82f0b61f..4f71308f5f2 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -1 +1 @@ -pub use crate::ocr::{OcrRequest, ocr}; +pub use crate::ocr::{OcrRequest, ocr, ocr_with_observer}; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 6c6e12724cd..bd2121d335f 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -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 { 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 diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 2acdd232c80..68131cdc98c 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -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 { + 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 { 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 } diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index c3a89f4394d..be0f9fbaffd 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -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>>, + raw_response: Option, + 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]; diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs new file mode 100644 index 00000000000..273c4352b6c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -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 { + 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() + }) +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index ec2fbb969a6..15fae1cd960 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,2 +1,4 @@ +pub mod handler; +pub mod observers; pub mod transformation; pub mod types; diff --git a/litellm-rust/crates/core/src/ocr/observers.rs b/litellm-rust/crates/core/src/ocr/observers.rs new file mode 100644 index 00000000000..4fe050a235a --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/observers.rs @@ -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, +} + +#[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(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs new file mode 100644 index 00000000000..9308ba185ca --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -0,0 +1 @@ +pub(crate) const OCR_CALLBACK_CAPACITY: usize = 1024; diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 42a07b34d9b..f894bdb35a5 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -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. diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index a08617a2c82..cc037a124ae 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -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( py: Python<'_>, future: impl Future> + Send + 'static, diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 3f377cf0041..14123e34392 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -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)?; diff --git a/litellm-rust/crates/python-bridge/src/ocr_callbacks.rs b/litellm-rust/crates/python-bridge/src/ocr_callbacks.rs new file mode 100644 index 00000000000..74dfc0f1c72 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/ocr_callbacks.rs @@ -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), + Async(PythonOcrSession), +} + +impl PythonOcrObserver { + pub(crate) fn new( + adapter: Option>, + context: PythonCallContext<'_>, + ) -> PyResult { + 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::>()? + .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, + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index d344e948966..64eaf24336e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -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> { 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> { 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> { 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> { 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", diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 1dbc6e56920..1ea18a71c4b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -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> + 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, timeout_seconds: Option, - callback_adapter: Option>, }, + keyword_only = { callback_adapter: Option> }, prepare = prepare_ocr, errors = ocr_error_to_pyerr, + context = PythonCallContext, } diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 125464b7abd..018565f0d04 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -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, diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index e374fcc42b1..fad237a82f4 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -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() diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 45f4d72e863..f0f1a5cb2d2 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -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], diff --git a/litellm/rust_bridge/ocr_callbacks.py b/litellm/rust_bridge/ocr_callbacks.py new file mode 100644 index 00000000000..b65ca06041f --- /dev/null +++ b/litellm/rust_bridge/ocr_callbacks.py @@ -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, + ) diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index df7556b8a09..db91924d70d 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -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, *, diff --git a/tests/test_litellm/ocr/callback_support.py b/tests/test_litellm/ocr/callback_support.py new file mode 100644 index 00000000000..8d535e2e2b8 --- /dev/null +++ b/tests/test_litellm/ocr/callback_support.py @@ -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 diff --git a/tests/test_litellm/ocr/live_callback_smoke.py b/tests/test_litellm/ocr/live_callback_smoke.py new file mode 100644 index 00000000000..599d346e3e1 --- /dev/null +++ b/tests/test_litellm/ocr/live_callback_smoke.py @@ -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)) diff --git a/tests/test_litellm/ocr/sdk_callback_contract.py b/tests/test_litellm/ocr/sdk_callback_contract.py new file mode 100644 index 00000000000..72695290869 --- /dev/null +++ b/tests/test_litellm/ocr/sdk_callback_contract.py @@ -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()) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index f88a9e39121..bb61cfbcdf3 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -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", diff --git a/tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py b/tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py new file mode 100644 index 00000000000..cb5e90bfb46 --- /dev/null +++ b/tests/test_litellm/rust_bridge/sdk_callback_wheel_test.py @@ -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()