test(ocr): prove Rust bridge executes Python callbacks

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<String, Value>,
 }
 
 type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
@@ -35,11 +37,13 @@ impl OcrLifecycleHooks {
         logger_runner: CustomLoggerRunner,
         guardrail_runner: CustomGuardrailRunner,
         request_metadata: RequestMetadata,
+        logging_kwargs: HashMap<String, Value>,
     ) -> 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<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLifecycleHooks {
@@ -202,11 +218,10 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> 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<PreparedOcrRequest, ProviderOcrRequest, Value> 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(),

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<String, Value> {
+    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);

litellm-rust/crates/ai-gateway/src/ocr/tests.rs

@@ -73,6 +73,8 @@ struct RecordedLogEvent {
     model: String,
     call_type: String,
     user_id: Option<String>,
+    api_base: Option<String>,
+    document_type: Option<String>,
     response_object: Option<String>,
     error_kind: Option<String>,
 }
@@ -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<String> {
+    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<String> {
+    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<GuardrailEventHook>,
     events: Mutex<Vec<&'static str>>,
@@ -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()),
         }]

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<Value> {
+    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<Py<PyAny>> {
+    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<PyAny>, attr_name: &str) -> Option<String> {
+    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::<String>() {
+            return Some(value);
+        }
+    }
+    attr.extract::<String>()
+        .ok()
+        .or_else(|| attr.str().ok().map(|value| value.to_string()))
+}
+
+pub(crate) fn py_datetime_from_epoch_seconds(
+    py: Python<'_>,
+    timestamp: f64,
+) -> PyResult<Py<PyAny>> {
+    let datetime = py.import("datetime")?.getattr("datetime")?;
+    Ok(datetime
+        .call_method1("fromtimestamp", (timestamp,))?
+        .unbind())
+}
+
+pub(crate) async fn call_python_awaitable(
+    obj: Py<PyAny>,
+    method_name: &'static str,
+    args: Vec<Py<PyAny>>,
+    kwargs: Option<Py<PyDict>>,
+) -> PyResult<Py<PyAny>> {
+    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
+}

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<PyAny>) -> 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<GuardrailEventHook> {
+    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<PyAny>) -> Vec<GuardrailEventHook> {
+    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::<PyList>() {
+        return list
+            .iter()
+            .filter_map(|item| {
+                let value = item
+                    .getattr("value")
+                    .ok()
+                    .and_then(|value_attr| value_attr.extract::<String>().ok())
+                    .or_else(|| item.extract::<String>().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::<String>().ok())
+        .or_else(|| attr.extract::<String>().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<PyAny>,
+    name: String,
+    hooks: Vec<GuardrailEventHook>,
+}
+
+impl PythonCustomGuardrailAdapter {
+    pub fn new(py: Python<'_>, obj: Py<PyAny>) -> 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<PyDict>, Py<PyAny>)> {
+                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<GuardrailDecision, GuardrailError> {
+                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::<String>() {
+                    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<Py<PyAny>>,
+) -> PyResult<Vec<Arc<dyn CustomGuardrail>>> {
+    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<dyn CustomGuardrail>,
+            )
+        })
+        .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
+        );
+    }
+}

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<String>) -> LogError {
+    LogError {
+        message: message.into(),
+        kind: "PythonCallbackError".to_string(),
+    }
+}
+
+fn model_call_details_json(details: &ModelCallDetails) -> Value {
+    let mut kwargs: Map<String, Value> = 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<PyAny>,
+}
+
+impl PythonCustomLoggerAdapter {
+    pub fn new(obj: Py<PyAny>) -> 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<Vec<Py<PyAny>>> {
+                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<Vec<Py<PyAny>>> {
+                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<Py<PyAny>>,
+) -> PyResult<Vec<Arc<dyn CustomLogger>>> {
+    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<dyn CustomLogger>)
+        })
+        .collect()
+}

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;

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<Py<PyAny>>,
     optional_params: Option<Py<PyAny>>,
     timeout_seconds: Option<f64>,
+    callbacks: Option<Py<PyAny>>,
+    guardrails: Option<Py<PyAny>>,
 ) -> PyResult<Py<PyAny>> {
     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<Py<PyAny>>,
     optional_params: Option<Py<PyAny>>,
     timeout_seconds: Option<f64>,
+    callbacks: Option<Py<PyAny>>,
+    guardrails: Option<Py<PyAny>>,
 ) -> PyResult<Bound<'_, PyAny>> {
     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,
         })

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

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

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

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": [],
     }