diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 6be74ed2714..8f5fd0d38af 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::future::Future; use std::pin::Pin; @@ -25,6 +26,7 @@ pub(crate) struct OcrLifecycleHooks { logger_runner: CustomLoggerRunner, guardrail_runner: CustomGuardrailRunner, request_metadata: RequestMetadata, + logging_kwargs: HashMap, } type OcrFuture<'a, T> = Pin> + Send + 'a>>; @@ -35,11 +37,13 @@ impl OcrLifecycleHooks { logger_runner: CustomLoggerRunner, guardrail_runner: CustomGuardrailRunner, request_metadata: RequestMetadata, + logging_kwargs: HashMap, ) -> Self { Self { logger_runner, guardrail_runner, request_metadata, + logging_kwargs, } } @@ -167,6 +171,18 @@ impl OcrLifecycleHooks { messages: None, } } + + fn model_call_details( + &self, + context: &CallLifecycleContext, + timing: &CallLifecycleTiming, + ) -> ModelCallDetails { + let mut details = ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ); + details.extra_metadata = self.logging_kwargs.clone(); + details + } } impl CallLifecycleHooks for OcrLifecycleHooks { @@ -202,11 +218,10 @@ impl CallLifecycleHooks for OcrLi return; } let response_obj = CallbackValue::new("ocr", response.clone()); + let model_call_details = self.model_call_details(context, timing); self.logger_runner .async_log_success_event( - &ModelCallDetails::from_standard_logging_payload( - self.standard_logging_payload(context, timing), - ), + &model_call_details, &response_obj, CallbackTiming::new(timing.start_time, timing.end_time), ) @@ -235,12 +250,12 @@ impl CallLifecycleHooks for OcrLi "kind": logging_error.kind, }), ); + let model_call_details = self + .model_call_details(context, timing) + .with_failure_error(logging_error); self.logger_runner .async_log_failure_event( - &ModelCallDetails::from_standard_logging_payload( - self.standard_logging_payload(context, timing), - ) - .with_failure_error(logging_error), + &model_call_details, Some(&response_obj), CallbackTiming::new(timing.start_time, timing.end_time), ) @@ -268,7 +283,7 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { GuardrailContext { call_type: CallType::Ocr, selected_guardrails: Vec::new(), - metadata: std::collections::HashMap::new(), + metadata: HashMap::new(), user_api_key_hash: metadata.user_api_key_hash.clone(), user_api_key_user_id: metadata.user_api_key_user_id.clone(), user_api_key_team_id: metadata.user_api_key_team_id.clone(), diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 5a4b350a4c4..6ff55a4520a 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -1,7 +1,9 @@ +use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; +use serde_json::{json, Value}; use super::hooks::OcrLifecycleHooks; use super::types::{OcrRequest, PreparedOcrRequest}; @@ -25,6 +27,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { }); let model = provider_info.model.to_string(); let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + let logging_kwargs = ocr_logging_kwargs(&request, &model, &custom_llm_provider); PreparedOcrCall { request: PreparedOcrRequest { @@ -42,10 +45,53 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { CustomLoggerRunner::new(request.callbacks), CustomGuardrailRunner::new(request.guardrails), request.request_metadata, + logging_kwargs, ), } } +fn ocr_logging_kwargs( + request: &OcrRequest<'_>, + model: &str, + custom_llm_provider: &str, +) -> HashMap { + let mut kwargs = HashMap::new(); + kwargs.insert("model".to_string(), json!(model)); + kwargs.insert( + "custom_llm_provider".to_string(), + json!(custom_llm_provider), + ); + kwargs.insert("call_type".to_string(), json!("ocr")); + kwargs.insert("document".to_string(), request.document.clone()); + kwargs.insert( + "optional_params".to_string(), + Value::Object(request.optional_params.clone()), + ); + kwargs.insert( + "api_base".to_string(), + request + .api_base + .map(|api_base| json!(api_base)) + .unwrap_or(Value::Null), + ); + kwargs.insert( + "headers".to_string(), + request + .extra_headers + .clone() + .map(Value::Object) + .unwrap_or(Value::Null), + ); + kwargs.insert( + "timeout".to_string(), + request + .timeout + .map(|timeout| json!(timeout.as_secs_f64())) + .unwrap_or(Value::Null), + ); + kwargs +} + fn new_ocr_call_id() -> String { static COUNTER: AtomicU64 = AtomicU64::new(1); let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 35747dc6985..699a0a4066c 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -73,6 +73,8 @@ struct RecordedLogEvent { model: String, call_type: String, user_id: Option, + api_base: Option, + document_type: Option, response_object: Option, error_kind: Option, } @@ -101,6 +103,8 @@ impl CustomLogger for RecordingOcrLogger { model: model_call_details.model.clone(), call_type: model_call_details.call_type.to_string(), user_id: model_call_details.metadata.user_api_key_user_id.clone(), + api_base: log_api_base(model_call_details), + document_type: log_document_type(model_call_details), response_object: Some(response_obj.object.clone()), error_kind: None, }); @@ -120,6 +124,8 @@ impl CustomLogger for RecordingOcrLogger { model: model_call_details.model.clone(), call_type: model_call_details.call_type.to_string(), user_id: model_call_details.metadata.user_api_key_user_id.clone(), + api_base: log_api_base(model_call_details), + document_type: log_document_type(model_call_details), response_object: response_obj.map(|value| value.object.clone()), error_kind: model_call_details .failure_error @@ -131,6 +137,24 @@ impl CustomLogger for RecordingOcrLogger { } } +fn log_api_base(model_call_details: &ModelCallDetails) -> Option { + model_call_details + .extra_metadata + .get("api_base") + .and_then(Value::as_str) + .map(str::to_string) +} + +fn log_document_type(model_call_details: &ModelCallDetails) -> Option { + model_call_details + .extra_metadata + .get("document") + .and_then(Value::as_object) + .and_then(|document| document.get("type")) + .and_then(Value::as_str) + .map(str::to_string) +} + struct RecordingOcrGuardrail { hooks: Vec, events: Mutex>, @@ -338,6 +362,8 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { model: "mistral-ocr-latest".to_string(), call_type: "ocr".to_string(), user_id: Some("user-1".to_string()), + api_base: Some(format!("http://{addr}")), + document_type: Some("document_url".to_string()), response_object: Some("ocr".to_string()), error_kind: None, }] @@ -400,6 +426,8 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { model: "mistral-ocr-latest".to_string(), call_type: "ocr".to_string(), user_id: None, + api_base: Some(format!("http://{addr}")), + document_type: Some("document_url".to_string()), response_object: Some("error".to_string()), error_kind: Some("HttpError".to_string()), }] @@ -444,6 +472,8 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { model: "mistral-ocr-latest".to_string(), call_type: "ocr".to_string(), user_id: None, + api_base: Some(format!("http://{addr}")), + document_type: Some("document_url".to_string()), response_object: Some("error".to_string()), error_kind: Some("InvalidRequest".to_string()), }] diff --git a/litellm-rust/crates/python-bridge/src/integrations_bridge/callback_bridge.rs b/litellm-rust/crates/python-bridge/src/integrations_bridge/callback_bridge.rs new file mode 100644 index 00000000000..2b4aa064ece --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/integrations_bridge/callback_bridge.rs @@ -0,0 +1,60 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyDict}; +use serde_json::Value; + +pub(crate) fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { + let json = py.import("json")?; + let encoded: String = json.call_method1("dumps", (value,))?.extract()?; + serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) +} + +pub(crate) fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { + let json = py.import("json")?; + let encoded = + serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; + Ok(json.call_method1("loads", (encoded,))?.unbind()) +} + +pub(crate) fn py_attr_string(py: Python<'_>, obj: &Py, attr_name: &str) -> Option { + let attr = obj.bind(py).getattr(attr_name).ok()?; + if attr.is_none() { + return None; + } + if let Ok(value_attr) = attr.getattr("value") { + if let Ok(value) = value_attr.extract::() { + return Some(value); + } + } + attr.extract::() + .ok() + .or_else(|| attr.str().ok().map(|value| value.to_string())) +} + +pub(crate) fn py_datetime_from_epoch_seconds( + py: Python<'_>, + timestamp: f64, +) -> PyResult> { + let datetime = py.import("datetime")?.getattr("datetime")?; + Ok(datetime + .call_method1("fromtimestamp", (timestamp,))? + .unbind()) +} + +pub(crate) async fn call_python_awaitable( + obj: Py, + method_name: &'static str, + args: Vec>, + kwargs: Option>, +) -> PyResult> { + let awaitable = Python::with_gil(|py| { + let callable = obj.bind(py).getattr(method_name)?; + let args_tuple = pyo3::types::PyTuple::new(py, args.iter().map(|arg| arg.bind(py)))?; + callable + .call(args_tuple, kwargs.as_ref().map(|dict| dict.bind(py))) + .map(|value| value.unbind()) + })?; + let future = + Python::with_gil(|py| pyo3_async_runtimes::tokio::into_future(awaitable.into_bound(py)))?; + future.await +} diff --git a/litellm-rust/crates/python-bridge/src/integrations_bridge/custom_guardrail_bridge.rs b/litellm-rust/crates/python-bridge/src/integrations_bridge/custom_guardrail_bridge.rs new file mode 100644 index 00000000000..81f251cf933 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/integrations_bridge/custom_guardrail_bridge.rs @@ -0,0 +1,196 @@ +use std::sync::Arc; + +use litellm_ai_gateway::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyDict, PyList}; + +use super::callback_bridge::{call_python_awaitable, json_to_py, py_attr_string, py_to_json}; + +fn py_guardrail_name(py: Python<'_>, obj: &Py) -> String { + py_attr_string(py, obj, "guardrail_name").unwrap_or_else(|| { + obj.bind(py) + .get_type() + .name() + .map(|name| name.to_string()) + .unwrap_or_else(|_| "python_guardrail".to_string()) + }) +} + +fn guardrail_event_hook_from_str(value: &str) -> Option { + match value { + "pre_call" | "GuardrailEventHooks.pre_call" => Some(GuardrailEventHook::PreCall), + "during_call" | "GuardrailEventHooks.during_call" => Some(GuardrailEventHook::DuringCall), + _ => None, + } +} + +fn py_guardrail_hooks(py: Python<'_>, obj: &Py) -> Vec { + let Some(attr) = obj.bind(py).getattr("event_hook").ok() else { + return vec![GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall]; + }; + if attr.is_none() { + return vec![GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall]; + } + if let Ok(list) = attr.downcast::() { + return list + .iter() + .filter_map(|item| { + let value = item + .getattr("value") + .ok() + .and_then(|value_attr| value_attr.extract::().ok()) + .or_else(|| item.extract::().ok()) + .or_else(|| item.str().ok().map(|value| value.to_string()))?; + guardrail_event_hook_from_str(&value) + }) + .collect(); + } + let value = attr + .getattr("value") + .ok() + .and_then(|value_attr| value_attr.extract::().ok()) + .or_else(|| attr.extract::().ok()) + .or_else(|| attr.str().ok().map(|value| value.to_string())); + value + .as_deref() + .and_then(guardrail_event_hook_from_str) + .map(|hook| vec![hook]) + .unwrap_or_default() +} + +pub struct PythonCustomGuardrailAdapter { + obj: Py, + name: String, + hooks: Vec, +} + +impl PythonCustomGuardrailAdapter { + pub fn new(py: Python<'_>, obj: Py) -> Self { + let name = py_guardrail_name(py, &obj); + let hooks = py_guardrail_hooks(py, &obj); + Self { obj, name, hooks } + } + + fn call_guardrail_hook<'a>( + &'a self, + method_name: &'static str, + context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + let obj = Python::with_gil(|py| self.obj.clone_ref(py)); + let call_type = context.call_type.to_string(); + Box::pin(async move { + let (kwargs, data) = Python::with_gil(|py| -> PyResult<(Py, Py)> { + let kwargs = PyDict::new(py); + kwargs.set_item("user_api_key_dict", py.None())?; + kwargs.set_item("cache", py.None())?; + let data = json_to_py(py, request.data)?; + kwargs.set_item("data", data.bind(py))?; + kwargs.set_item("call_type", call_type)?; + Ok((kwargs.unbind(), data)) + }) + .map_err(|err| GuardrailError::blocked(err.to_string()))?; + let result = call_python_awaitable(obj, method_name, Vec::new(), Some(kwargs)) + .await + .map_err(|err| GuardrailError::blocked(err.to_string()))?; + Python::with_gil(|py| -> Result { + let result = result.bind(py); + if result.is_none() { + let value = py_to_json(py, data.bind(py)) + .map_err(|err| GuardrailError::blocked(err.to_string()))?; + return Ok(GuardrailDecision::Allow(GuardrailRequest::new(value))); + } + if let Ok(message) = result.extract::() { + return Err(GuardrailError::blocked(message)); + } + let value = py_to_json(py, result) + .map_err(|err| GuardrailError::blocked(err.to_string()))?; + if value.is_object() { + Ok(GuardrailDecision::Mask(GuardrailRequest::new(value))) + } else { + let fallback = py_to_json(py, data.bind(py)) + .map_err(|err| GuardrailError::blocked(err.to_string()))?; + Ok(GuardrailDecision::Mask(GuardrailRequest::new(fallback))) + } + }) + }) + } +} + +impl CustomGuardrail for PythonCustomGuardrailAdapter { + fn guardrail_name(&self) -> &str { + &self.name + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + self.call_guardrail_hook("async_pre_call_hook", context, request) + } + + fn async_moderation_hook<'a>( + &'a self, + context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + self.call_guardrail_hook("async_moderation_hook", context, request) + } +} + +pub fn py_guardrails_to_rust( + py: Python<'_>, + guardrails: Option>, +) -> PyResult>> { + let Some(guardrails) = guardrails else { + return Ok(Vec::new()); + }; + guardrails + .bind(py) + .try_iter()? + .map(|guardrail| { + let guardrail = guardrail?; + Ok( + Arc::new(PythonCustomGuardrailAdapter::new(py, guardrail.unbind())) + as Arc, + ) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn maps_supported_python_guardrail_hooks() { + assert_eq!( + guardrail_event_hook_from_str("pre_call"), + Some(GuardrailEventHook::PreCall) + ); + assert_eq!( + guardrail_event_hook_from_str("GuardrailEventHooks.during_call"), + Some(GuardrailEventHook::DuringCall) + ); + } + + #[test] + fn leaves_unsupported_python_guardrail_hooks_unmapped() { + assert_eq!(guardrail_event_hook_from_str("post_call"), None); + assert_eq!(guardrail_event_hook_from_str("logging_only"), None); + assert_eq!(guardrail_event_hook_from_str("pre_mcp_call"), None); + assert_eq!(guardrail_event_hook_from_str("during_mcp_call"), None); + assert_eq!( + guardrail_event_hook_from_str("realtime_input_transcription"), + None + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/integrations_bridge/custom_logger_bridge.rs b/litellm-rust/crates/python-bridge/src/integrations_bridge/custom_logger_bridge.rs new file mode 100644 index 00000000000..15956310eda --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/integrations_bridge/custom_logger_bridge.rs @@ -0,0 +1,150 @@ +use std::sync::Arc; + +use litellm_ai_gateway::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, ModelCallDetails, +}; +use pyo3::prelude::*; +use serde_json::{json, Map, Value}; + +use super::callback_bridge::{call_python_awaitable, json_to_py, py_datetime_from_epoch_seconds}; + +fn dropped_log_error(message: impl Into) -> LogError { + LogError { + message: message.into(), + kind: "PythonCallbackError".to_string(), + } +} + +fn model_call_details_json(details: &ModelCallDetails) -> Value { + let mut kwargs: Map = details.extra_metadata.clone().into_iter().collect(); + kwargs.insert("model".to_string(), json!(details.model)); + kwargs.insert( + "custom_llm_provider".to_string(), + json!(details.custom_llm_provider), + ); + kwargs.insert( + "call_type".to_string(), + json!(details.call_type.to_string()), + ); + kwargs.insert( + "litellm_call_id".to_string(), + json!(details.litellm_call_id), + ); + kwargs.insert("request_id".to_string(), json!(details.request_id)); + kwargs.insert("response_cost".to_string(), json!(details.response_cost)); + kwargs.insert( + "metadata".to_string(), + json!({ + "user_api_key_hash": details.metadata.user_api_key_hash, + "user_api_key_user_id": details.metadata.user_api_key_user_id, + "user_api_key_team_id": details.metadata.user_api_key_team_id, + }), + ); + kwargs.insert( + "standard_logging_object".to_string(), + json!(details.standard_logging_payload), + ); + if let Some(error) = &details.failure_error { + kwargs.insert( + "failure_error".to_string(), + json!({ + "message": error.message, + "kind": error.kind, + }), + ); + } + Value::Object(kwargs) +} + +pub struct PythonCustomLoggerAdapter { + obj: Py, +} + +impl PythonCustomLoggerAdapter { + pub fn new(obj: Py) -> Self { + Self { obj } + } +} + +impl CustomLogger for PythonCustomLoggerAdapter { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + let obj = Python::with_gil(|py| self.obj.clone_ref(py)); + let details = model_call_details_json(model_call_details); + let response = json!({ + "object": response_obj.object, + "value": response_obj.value, + }); + Box::pin(async move { + let args = Python::with_gil(|py| -> PyResult>> { + Ok(vec![ + json_to_py(py, details)?, + json_to_py(py, response)?, + py_datetime_from_epoch_seconds(py, timing.start_time)?, + py_datetime_from_epoch_seconds(py, timing.end_time)?, + ]) + }) + .map_err(|err| dropped_log_error(err.to_string()))?; + call_python_awaitable(obj, "async_log_success_event", args, None) + .await + .map_err(|err| dropped_log_error(err.to_string()))?; + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + let obj = Python::with_gil(|py| self.obj.clone_ref(py)); + let details = model_call_details_json(model_call_details); + let response = response_obj.map(|value| { + json!({ + "object": value.object, + "value": value.value, + }) + }); + Box::pin(async move { + let args = Python::with_gil(|py| -> PyResult>> { + Ok(vec![ + json_to_py(py, details)?, + match response { + Some(response) => json_to_py(py, response)?, + None => py.None(), + }, + py_datetime_from_epoch_seconds(py, timing.start_time)?, + py_datetime_from_epoch_seconds(py, timing.end_time)?, + ]) + }) + .map_err(|err| dropped_log_error(err.to_string()))?; + call_python_awaitable(obj, "async_log_failure_event", args, None) + .await + .map_err(|err| dropped_log_error(err.to_string()))?; + Ok(()) + }) + } +} + +pub fn py_callbacks_to_rust( + py: Python<'_>, + callbacks: Option>, +) -> PyResult>> { + let Some(callbacks) = callbacks else { + return Ok(Vec::new()); + }; + callbacks + .bind(py) + .try_iter()? + .map(|callback| { + let callback = callback?; + Ok(Arc::new(PythonCustomLoggerAdapter::new(callback.unbind())) + as Arc) + }) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/integrations_bridge/mod.rs b/litellm-rust/crates/python-bridge/src/integrations_bridge/mod.rs new file mode 100644 index 00000000000..d92462e5d54 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/integrations_bridge/mod.rs @@ -0,0 +1,6 @@ +mod callback_bridge; +mod custom_guardrail_bridge; +mod custom_logger_bridge; + +pub use custom_guardrail_bridge::py_guardrails_to_rust; +pub use custom_logger_bridge::py_callbacks_to_rust; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 946a99f990c..818d9fa8645 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -8,6 +8,7 @@ use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; mod gil; +mod integrations_bridge; type MarshaledOcrInputs = ( Value, @@ -83,7 +84,7 @@ fn marshal_inputs( } #[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, callbacks=None, guardrails=None))] #[allow(clippy::too_many_arguments)] fn ocr( py: Python<'_>, @@ -95,6 +96,8 @@ fn ocr( extra_headers: Option>, optional_params: Option>, timeout_seconds: Option, + callbacks: Option>, + guardrails: Option>, ) -> PyResult> { let (document, extra_headers, optional_params, timeout) = marshal_inputs( py, @@ -103,6 +106,8 @@ fn ocr( optional_params, timeout_seconds, )?; + let callbacks = integrations_bridge::py_callbacks_to_rust(py, callbacks)?; + let guardrails = integrations_bridge::py_guardrails_to_rust(py, guardrails)?; let result = gil::release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { @@ -114,8 +119,8 @@ fn ocr( extra_headers, optional_params, timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), + callbacks, + guardrails, request_metadata: Default::default(), litellm_call_id: None, })) @@ -128,7 +133,7 @@ fn ocr( } #[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, callbacks=None, guardrails=None))] #[allow(clippy::too_many_arguments)] fn aocr( py: Python<'_>, @@ -140,6 +145,8 @@ fn aocr( extra_headers: Option>, optional_params: Option>, timeout_seconds: Option, + callbacks: Option>, + guardrails: Option>, ) -> PyResult> { let (document, extra_headers, optional_params, timeout) = marshal_inputs( py, @@ -148,6 +155,8 @@ fn aocr( optional_params, timeout_seconds, )?; + let callbacks = integrations_bridge::py_callbacks_to_rust(py, callbacks)?; + let guardrails = integrations_bridge::py_guardrails_to_rust(py, guardrails)?; pyo3_async_runtimes::tokio::future_into_py(py, async move { let value = run_ocr(OcrRequest { @@ -159,8 +168,8 @@ fn aocr( extra_headers, optional_params, timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), + callbacks, + guardrails, request_metadata: Default::default(), litellm_call_id: None, }) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 93b3a892659..6cd8b641b34 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -49,6 +49,8 @@ class _PreparedRustOCRCall: api_base: str | None headers: dict[str, object] optional_params: dict[str, object] + callbacks: list[object] + guardrails: list[object] _RUST_OCR_PROVIDERS = { @@ -242,14 +244,33 @@ def _prepare_rust_ocr_call( "headers": resolved_headers, }, ) + callbacks, guardrails = _get_rust_bridge_callbacks() return _PreparedRustOCRCall( api_key=resolved_api_key, api_base=rust_api_base, headers=cast(dict[str, object], resolved_headers), optional_params=rust_optional_params, + callbacks=callbacks, + guardrails=guardrails, ) +def _get_rust_bridge_callbacks() -> tuple[list[object], list[object]]: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_logger import CustomLogger + + active_callbacks = litellm.logging_callback_manager._get_all_callbacks() + guardrails: list[object] = [] + callbacks: list[object] = [] + for callback in active_callbacks: + if isinstance(callback, CustomGuardrail): + if callback not in guardrails: + guardrails.append(callback) + if isinstance(callback, CustomLogger) and callback not in callbacks: + callbacks.append(callback) + return callbacks, guardrails + + def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], @@ -269,6 +290,8 @@ def _run_rust_ocr( extra_headers=prepared.headers, optional_params=prepared.optional_params, timeout=prepared_request.effective_timeout, + callbacks=prepared.callbacks, + guardrails=prepared.guardrails, ) if rust_response is None: return None @@ -294,6 +317,8 @@ async def _run_rust_aocr( extra_headers=prepared.headers, optional_params=prepared.optional_params, timeout=prepared_request.effective_timeout, + callbacks=prepared.callbacks, + guardrails=prepared.guardrails, ) if rust_response is None: return None diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 36a088b6b1a..a592b128f2e 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -19,6 +19,8 @@ class RustOcr(Protocol): extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, + callbacks: list[object] | None, + guardrails: list[object] | None, ) -> dict[str, object]: raise NotImplementedError @@ -34,6 +36,8 @@ class RustAocr(Protocol): extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, + callbacks: list[object] | None, + guardrails: list[object] | None, ) -> Awaitable[dict[str, object]]: raise NotImplementedError @@ -117,6 +121,8 @@ def ocr( extra_headers: dict[str, Any] | None, optional_params: dict[str, object], timeout: Union[float, httpx.Timeout] | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object] | None: rust_ocr = load_rust_ocr() if rust_ocr is None: @@ -130,6 +136,8 @@ def ocr( extra_headers=cast(dict[str, object] | None, extra_headers), optional_params=optional_params, timeout_seconds=_timeout_to_seconds(timeout), + callbacks=callbacks, + guardrails=guardrails, ) @@ -143,6 +151,8 @@ async def aocr( extra_headers: dict[str, Any] | None, optional_params: dict[str, object], timeout: Union[float, httpx.Timeout] | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object] | None: rust_aocr = load_rust_aocr() if rust_aocr is None: @@ -156,4 +166,6 @@ async def aocr( extra_headers=cast(dict[str, object] | None, extra_headers), optional_params=optional_params, timeout_seconds=_timeout_to_seconds(timeout), + callbacks=callbacks, + guardrails=guardrails, ) diff --git a/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py b/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py new file mode 100644 index 00000000000..bd7a7712572 --- /dev/null +++ b/tests/ocr_tests/test_ocr_rust_bridge_callbacks.py @@ -0,0 +1,153 @@ +""" +Tests proving the Rust OCR path executes Python callbacks from the callback +manager. +""" + +import os +from datetime import datetime +from typing import Any, Optional + +import pytest + +import litellm +from base_ocr_unit_tests import TEST_PDF_URL +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import StandardLoggingPayload + +MODEL = "mistral/mistral-ocr-latest" +DOCUMENT: dict[str, object] = { + "type": "document_url", + "document_url": TEST_PDF_URL, +} + + +class OCRCustomLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.standard_logging_payload: Optional[StandardLoggingPayload] = None + self.response_obj: Any = None + self.start_time: Optional[datetime] = None + self.end_time: Optional[datetime] = None + self.kwargs: dict[str, Any] = {} + + async def async_log_success_event( + self, + kwargs: dict[str, Any], + response_obj: Any, + start_time: datetime, + end_time: datetime, + ) -> None: + self.kwargs = kwargs + self.standard_logging_payload = kwargs.get("standard_logging_object") + self.response_obj = response_obj + self.start_time = start_time + self.end_time = end_time + + +class OCRCustomGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="ocr-test-guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + self.calls: list[dict[str, object]] = [] + self.success_log_calls = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data: dict[str, Any], call_type: str + ) -> None: + self.calls.append({"data": data, "call_type": call_type}) + data["document"] = { + **data["document"], + "guardrail_executed": True, + } + data["optional_params"] = { + **data["optional_params"], + "include_image_base64": True, + } + + async def async_log_success_event( + self, + kwargs: dict[str, Any], + response_obj: Any, + start_time: datetime, + end_time: datetime, + ) -> None: + self.success_log_calls += 1 + + +@pytest.fixture(autouse=True) +def reset_litellm_rust_callbacks(): + litellm.logging_callback_manager._reset_all_callbacks() + litellm.use_litellm_rust(False, ocr=None, aocr=None) + yield + litellm.logging_callback_manager._reset_all_callbacks() + litellm.use_litellm_rust(False, ocr=None, aocr=None) + + +def _mistral_api_key() -> str: + api_key = os.getenv("MISTRAL_API_KEY") + if not api_key: + pytest.skip("MISTRAL_API_KEY is required for live Mistral OCR callback tests") + return api_key + + +def _require_native_rust_aocr() -> None: + if rust_ocr_bridge.load_rust_aocr() is None: + pytest.skip("native Rust OCR bridge is not available") + + +@pytest.mark.asyncio +async def test_rust_ocr_executes_custom_logger_from_callback_manager(): + _require_native_rust_aocr() + custom_logger = OCRCustomLogger() + litellm.logging_callback_manager.add_litellm_callback(custom_logger) + litellm.use_litellm_rust(True) + + response = await litellm.aocr( + model=MODEL, + document=DOCUMENT, + api_key=_mistral_api_key(), + ) + + assert len(response.pages) > 0 + assert custom_logger.response_obj["object"] == "ocr" + assert len(custom_logger.response_obj["value"]["pages"]) > 0 + assert isinstance(custom_logger.start_time, datetime) + assert isinstance(custom_logger.end_time, datetime) + assert custom_logger.kwargs["api_base"] == "https://api.mistral.ai/v1" + assert custom_logger.kwargs["document"] == DOCUMENT + assert custom_logger.kwargs["optional_params"] == {} + + logged_payload = custom_logger.standard_logging_payload + assert logged_payload is not None + assert logged_payload["model"] == "mistral-ocr-latest" + assert logged_payload["custom_llm_provider"] == "mistral" + assert logged_payload["call_type"] == "ocr" + assert logged_payload["response_cost"] == 0.0 + + +@pytest.mark.asyncio +async def test_rust_ocr_executes_custom_guardrail_from_callback_manager(): + _require_native_rust_aocr() + custom_guardrail = OCRCustomGuardrail() + litellm.logging_callback_manager.add_litellm_callback(custom_guardrail) + litellm.use_litellm_rust(True) + + response = await litellm.aocr( + model=MODEL, + document=DOCUMENT, + api_key=_mistral_api_key(), + ) + + assert len(response.pages) > 0 + assert custom_guardrail.calls[0]["call_type"] == "ocr" + assert custom_guardrail.calls[0]["data"]["document"]["guardrail_executed"] is True + assert ( + custom_guardrail.calls[0]["data"]["optional_params"]["include_image_base64"] + is True + ) + assert custom_guardrail.success_log_calls == 1 diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index acad249a2bb..7fdd85323ef 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -53,6 +53,8 @@ class RecordingBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object]: self.calls.append( { @@ -64,6 +66,8 @@ class RecordingBridge: "extra_headers": extra_headers, "optional_params": optional_params, "timeout_seconds": timeout_seconds, + "callbacks": callbacks or [], + "guardrails": guardrails or [], } ) return dict(FAKE_OCR_RESPONSE) @@ -85,6 +89,8 @@ class RecordingAsyncBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object]: self.calls.append( { @@ -96,6 +102,8 @@ class RecordingAsyncBridge: "extra_headers": extra_headers, "optional_params": optional_params, "timeout_seconds": timeout_seconds, + "callbacks": callbacks or [], + "guardrails": guardrails or [], } ) return dict(FAKE_OCR_RESPONSE) @@ -112,6 +120,8 @@ class RaisingBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -127,6 +137,8 @@ class RaisingAsyncBridge: extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, + callbacks: list[object] | None = None, + guardrails: list[object] | None = None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -214,9 +226,11 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" + litellm.logging_callback_manager._reset_all_callbacks() rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield + litellm.logging_callback_manager._reset_all_callbacks() rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -400,6 +414,8 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): }, "optional_params": {"include_image_base64": True, "pages": [0]}, "timeout_seconds": 12.5, + "callbacks": [], + "guardrails": [], } @@ -429,6 +445,8 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): "extra_headers": None, "optional_params": {"vertex_project": "project-1"}, "timeout_seconds": 42.0, + "callbacks": [], + "guardrails": [], } @@ -462,6 +480,8 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): }, "optional_params": {"include_image_base64": True}, "timeout_seconds": 12.5, + "callbacks": [], + "guardrails": [], }