From eaf210a23c976cde157bb81c8af44275e828ed85 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 4 Sep 2026 15:49:15 -0700 Subject: [PATCH] refactor(rust): share provider callback infrastructure --- .../crates/ai-gateway/src/ocr/handler.rs | 56 +++-- .../crates/ai-gateway/src/ocr/hooks.rs | 2 + litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 14 +- .../crates/ai-gateway/src/ocr/types.rs | 2 + .../crates/ai-gateway/tests/ocr_lifecycle.rs | 130 +++++++++-- litellm-rust/crates/core/src/lib.rs | 1 + litellm-rust/crates/core/src/ocr/handler.rs | 137 ++++++++++-- litellm-rust/crates/core/src/ocr/mod.rs | 1 - litellm-rust/crates/core/src/ocr/observers.rs | 48 ---- .../crates/core/src/provider_callbacks.rs | 206 ++++++++++++++++++ .../python-bridge/src/callback_bindings.rs | 100 +++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 10 +- .../crates/python-bridge/src/ocr_callbacks.rs | 80 ------- .../crates/python-bridge/src/routes/ocr.rs | 4 +- .../messages/handler.py | 6 +- litellm/main.py | 5 +- litellm/ocr/main.py | 10 +- litellm/rust_bridge/callback_adapters.py | 147 +++++++++++++ litellm/rust_bridge/chat_completions.py | 45 +--- litellm/rust_bridge/messages.py | 47 +--- litellm/rust_bridge/ocr.py | 7 +- litellm/rust_bridge/ocr_callbacks.py | 60 +---- .../test_litellm/ocr/sdk_callback_contract.py | 2 +- .../rust_bridge/test_callback_adapters.py | 132 +++++++++++ .../rust_bridge/test_chat_completions.py | 27 ++- 25 files changed, 926 insertions(+), 353 deletions(-) delete mode 100644 litellm-rust/crates/core/src/ocr/observers.rs create mode 100644 litellm-rust/crates/core/src/provider_callbacks.rs create mode 100644 litellm-rust/crates/python-bridge/src/callback_bindings.rs delete mode 100644 litellm-rust/crates/python-bridge/src/ocr_callbacks.rs create mode 100644 litellm/rust_bridge/callback_adapters.py create mode 100644 tests/test_litellm/rust_bridge/test_callback_adapters.py diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index bd2121d335f..952ae8018cb 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::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 litellm_core::provider_callbacks::{ + CallbackDecision, ProviderAttemptObserver, ProviderPreCall, +}; use serde_json::Value; +use std::time::{SystemTime, UNIX_EPOCH}; use super::common_utils::poll_document_intelligence; use super::hooks::OcrLifecycleHooks; @@ -11,13 +13,35 @@ use super::types::PreparedOcrRequest; use crate::client::http_client; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub(crate) async fn execute_ocr_provider_call( +pub(crate) async fn execute_ocr_provider_call( request: PreparedOcrRequest, hooks: &OcrLifecycleHooks, - observer: &mut impl OcrObserver, -) -> Result { + observer: &mut Observer, +) -> Result +where + Observer: ProviderAttemptObserver, + Observer::Error: std::fmt::Display, +{ let request = hooks.prepare_provider_request(request).await?; - let mut request_builder = http_client().post(&request.url).json(&request.body); + let event = ProviderPreCall { + provider: request.custom_llm_provider.clone(), + model: request.model.clone(), + call_id: request.litellm_call_id.clone(), + trace_id: None, + attempt: 1, + started_at: epoch_seconds(), + request: serde_json::from_value(request.body.clone()).map_err(|error| { + Error::InvalidRequest(format!("OCR provider request must be an object: {error}")) + })?, + api_base: request.url.clone(), + headers: request.upstream_headers.iter().cloned().collect(), + }; + let body = match observer.pre_call(&event).await.map_err(callback_error)? { + CallbackDecision::Unchanged => request.body, + CallbackDecision::Replace { payload } => payload, + CallbackDecision::Reject { message, .. } => return Err(Error::InvalidRequest(message)), + }; + let mut request_builder = http_client().post(&request.url).json(&body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); } @@ -25,15 +49,6 @@ pub(crate) async fn execute_ocr_provider_call( request_builder = request_builder.timeout(duration); } - 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; @@ -80,3 +95,14 @@ pub(crate) async fn execute_ocr_provider_call( )? .into_json()) } + +fn callback_error(error: impl std::fmt::Display) -> Error { + Error::InvalidResponse(format!("OCR callback failed: {error}")) +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 446b323db3a..231fffb534c 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -122,6 +122,8 @@ impl OcrLifecycleHooks { }; Ok(ProviderOcrRequest { model, + custom_llm_provider, + litellm_call_id: request.litellm_call_id, config, url, body, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 68131cdc98c..d92491edde0 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,6 +1,6 @@ use litellm_core::Error; use litellm_core::call_lifecycle::CallLifecycle; -use litellm_core::ocr::observers::{NoopOcrObserver, OcrObserver}; +use litellm_core::provider_callbacks::{NoopProviderAttemptObserver, ProviderAttemptObserver}; use serde_json::Value; mod common_utils; @@ -16,7 +16,7 @@ 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 + ocr_with_observer(request, &mut NoopProviderAttemptObserver).await } #[tracing::instrument( @@ -25,10 +25,14 @@ pub async fn ocr(request: OcrRequest<'_>) -> Result { level = "trace", skip_all )] -pub async fn ocr_with_observer( +pub async fn ocr_with_observer( request: OcrRequest<'_>, - observer: &mut impl OcrObserver, -) -> Result { + observer: &mut Observer, +) -> Result +where + Observer: ProviderAttemptObserver, + Observer::Error: std::fmt::Display, +{ let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); CallLifecycle::default() .run_request(request, &hooks, |request| { diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index 75a8e61ddbf..cd8221bcd99 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -50,6 +50,8 @@ impl CallLifecycleRequest for PreparedOcrRequest { pub(crate) struct ProviderOcrRequest { pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) litellm_call_id: String, pub(crate) config: &'static dyn OcrProviderConfig, pub(crate) url: String, pub(crate) body: Value, diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index 4f12d186de6..13fd7539047 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -11,7 +11,9 @@ use litellm_ai_gateway::integrations::custom_logger::{ use litellm_ai_gateway::integrations::types::RequestMetadata; use litellm_ai_gateway::ocr::{OcrRequest, ocr, ocr_with_observer}; use litellm_core::error::Error; -use litellm_core::ocr::observers::{OcrObserver, OcrPostCall, OcrPreCall}; +use litellm_core::provider_callbacks::{ + CallbackDecision, ProviderAttemptObserver, ProviderError, ProviderPostCall, ProviderPreCall, +}; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; @@ -19,16 +21,17 @@ use tokio::net::{TcpListener, TcpStream}; struct ProviderObserver { events: Arc>>, raw_response: Option, - reject: bool, + rejected_callback: Option<&'static str>, + decision: Option<&'static str>, } -impl OcrObserver for ProviderObserver { +impl ProviderAttemptObserver for ProviderObserver { type Error = &'static str; - async fn pre_call(&mut self, input: &OcrPreCall) -> Result<(), Self::Error> { + async fn pre_call(&mut self, input: &ProviderPreCall) -> Result { assert_eq!(input.model, "mistral-ocr-4-1"); assert_eq!( - input.request.data["document"]["document_url"], + input.request["document"]["document_url"], "https://example.com/document.pdf" ); assert!(input.api_base.ends_with("/v1/ocr")); @@ -39,17 +42,50 @@ impl OcrObserver for ProviderObserver { .any(|value| value == "Bearer test-key") ); self.events.lock().unwrap().push("pre"); - if self.reject { - Err("observer failure") - } else { - Ok(()) + match (self.rejected_callback, self.decision) { + (Some("pre"), _) => Err("observer failure"), + (_, Some("replace_pre")) => Ok(CallbackDecision::Replace { + payload: Value::Object( + input + .request + .iter() + .map(|(key, value)| (key.clone(), value.clone())) + .chain(std::iter::once(( + "callback_replaced".to_string(), + json!(true), + ))) + .collect(), + ), + }), + (_, Some("reject_pre")) => Ok(CallbackDecision::Reject { + message: "callback rejected request".to_string(), + status_code: Some(400), + }), + _ => Ok(CallbackDecision::Unchanged), } } - async fn post_call(&mut self, input: &OcrPostCall) -> Result<(), Self::Error> { + async fn post_call( + &mut self, + input: &ProviderPostCall, + ) -> Result { self.events.lock().unwrap().push("post"); - self.raw_response = Some(input.original_response.clone()); - if self.reject { + self.raw_response = input.response.as_str().map(str::to_string); + match (self.rejected_callback, self.decision) { + (Some("post"), _) => Err("observer failure"), + (_, Some("replace_post")) => Ok(CallbackDecision::Replace { + payload: json!({"pages":[{"index":0,"markdown":"masked"}]}), + }), + _ => Ok(CallbackDecision::Unchanged), + } + } + + async fn error(&mut self, input: &ProviderError) -> Result<(), Self::Error> { + assert!(input.committed); + assert!(!input.will_retry); + assert!(!input.message.is_empty()); + self.events.lock().unwrap().push("error"); + if self.rejected_callback == Some("error") { Err("observer failure") } else { Ok(()) @@ -74,7 +110,7 @@ fn observer_request(api_base: &str) -> OcrRequest<'_> { } } -async fn observer_case(status: u16, body: &'static str, reject: bool) { +async fn observer_case(status: u16, body: &'static str, decision: Option<&'static str>) { 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())); @@ -83,6 +119,10 @@ async fn observer_case(status: u16, body: &'static str, reject: bool) { let (mut socket, _) = listener.accept().await.unwrap(); let request = read_http_request(&mut socket).await; assert!(request.starts_with("POST /v1/ocr ")); + assert_eq!( + request.contains(r#""callback_replaced":true"#), + decision == Some("replace_pre") + ); 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}", @@ -93,7 +133,8 @@ async fn observer_case(status: u16, body: &'static str, reject: bool) { let mut observer = ProviderObserver { events: Arc::clone(&events), raw_response: None, - reject, + rejected_callback: None, + decision, }; let result = ocr_with_observer(observer_request(&url), &mut observer).await; tokio::time::timeout(Duration::from_secs(2), server) @@ -102,7 +143,7 @@ async fn observer_case(status: u16, body: &'static str, reject: bool) { .unwrap(); if status != 200 { assert!(matches!(result, Err(Error::Http { status: actual, .. }) if actual == status)); - assert_eq!(*events.lock().unwrap(), ["pre", "http"]); + assert_eq!(*events.lock().unwrap(), ["pre", "http", "error"]); assert_eq!(observer.raw_response, None); } else { assert_eq!(*events.lock().unwrap(), ["pre", "http", "post"]); @@ -110,18 +151,60 @@ async fn observer_case(status: u16, body: &'static str, reject: bool) { if body == "invalid-json" { assert!(matches!(result, Err(Error::InvalidResponse(_)))); } else { - assert_eq!(result.unwrap()["pages"][0]["markdown"], "ok"); + assert_eq!( + result.unwrap()["pages"][0]["markdown"], + if decision == Some("replace_post") { + "masked" + } else { + "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; - } +async fn provider_observers_surround_http_and_can_replace_request_or_response() { + observer_case(200, r#"{"pages":[{"index":0,"markdown":"ok"}]}"#, None).await; + observer_case(200, "invalid-json", None).await; + observer_case(401, r#"{"error":"rejected"}"#, None).await; + observer_case( + 200, + r#"{"pages":[{"index":0,"markdown":"ok"}]}"#, + Some("replace_pre"), + ) + .await; + observer_case( + 200, + r#"{"pages":[{"index":0,"markdown":"ok"}]}"#, + Some("replace_post"), + ) + .await; +} + +#[tokio::test] +async fn provider_callback_rejection_stops_before_http() { + 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, + rejected_callback: None, + decision: Some("reject_pre"), + }; + + let result = ocr_with_observer(observer_request(&url), &mut observer).await; + + assert!( + matches!(result, Err(Error::InvalidRequest(message)) if message == "callback rejected request") + ); + assert_eq!(*events.lock().unwrap(), ["pre"]); + assert!( + tokio::time::timeout(Duration::from_millis(50), listener.accept()) + .await + .is_err() + ); } #[tokio::test] @@ -132,7 +215,8 @@ async fn invalid_ocr_preparation_does_not_call_observers_or_provider() { let mut observer = ProviderObserver { events: Arc::clone(&events), raw_response: None, - reject: false, + rejected_callback: None, + decision: None, }; let request = OcrRequest { document: json!(42), diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 16279a1caf1..12a140f3b0b 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -10,6 +10,7 @@ pub mod messages; #[cfg(any(feature = "observability", test))] pub mod observability; pub mod ocr; +pub mod provider_callbacks; pub mod providers; pub mod realtime; pub mod responses; diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 273c4352b6c..4c1d1b9bfe0 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,8 +1,14 @@ -use reqwest::{RequestBuilder, StatusCode, header::HeaderMap}; +use std::collections::BTreeMap; +use std::time::{SystemTime, UNIX_EPOCH}; + +use reqwest::{RequestBuilder, StatusCode, header::HeaderMap}; +use serde_json::Value; -use super::observers::{OcrObserver, OcrPostCall, OcrPreCall}; use crate::Error; use crate::http_utils::{http_request, truncate_error_body}; +use crate::provider_callbacks::{ + CallbackDecision, ProviderAttemptObserver, ProviderError, ProviderPostCall, ProviderPreCall, +}; pub struct OcrHttpResponse { pub status: StatusCode, @@ -11,37 +17,130 @@ pub struct OcrHttpResponse { } #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub async fn send_ocr_request( +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)?; + event: &ProviderPreCall, + observer: &mut Observer, +) -> Result +where + Observer: ProviderAttemptObserver, + Observer::Error: std::fmt::Display, +{ + let response = match http_request(request).await { + Ok(response) => response, + Err(error) => { + let mapped = transport_error(error); + notify_error(observer, event, &mapped, "provider_request", true).await?; + return Err(mapped); + } + }; let status = response.status(); let headers = response.headers().clone(); - let body = response.text().await.map_err(transport_error)?; + let body = match response.text().await { + Ok(body) => body, + Err(error) => { + let mapped = transport_error(error); + notify_error(observer, event, &mapped, "response_body", true).await?; + return Err(mapped); + } + }; if !status.is_success() { - return Err(Error::Http { + let error = Error::Http { status: status.as_u16(), body: truncate_error_body(&body), - }); + }; + notify_error(observer, event, &error, "provider_response", true).await?; + return Err(error); } - let event = OcrPostCall { - original_response: body, + let post_call = ProviderPostCall { + provider: event.provider.clone(), + model: event.model.clone(), + call_id: event.call_id.clone(), + trace_id: event.trace_id.clone(), + attempt: event.attempt, + started_at: event.started_at, + response: Value::String(body.clone()), + status_code: status.as_u16(), + headers: header_values(&headers), + ended_at: epoch_seconds(), + }; + let body = match observer + .post_call(&post_call) + .await + .map_err(callback_error)? + { + CallbackDecision::Unchanged => body, + CallbackDecision::Replace { + payload: Value::String(replacement), + } => replacement, + CallbackDecision::Replace { payload } => { + serde_json::to_string(&payload).map_err(|error| { + Error::InvalidResponse(format!("callback response is invalid: {error}")) + })? + } + CallbackDecision::Reject { message, .. } => return Err(Error::InvalidResponse(message)), }; - if observer.post_call(&event).await.is_err() { - tracing::warn!("OCR post-call observer failed"); - } Ok(OcrHttpResponse { status, headers, - body: event.original_response, + body, }) } +async fn notify_error( + observer: &mut Observer, + context: &ProviderPreCall, + error: &Error, + stage: &'static str, + committed: bool, +) -> Result<(), Error> +where + Observer: ProviderAttemptObserver, + Observer::Error: std::fmt::Display, +{ + let event = ProviderError { + provider: context.provider.clone(), + model: context.model.clone(), + call_id: context.call_id.clone(), + trace_id: context.trace_id.clone(), + attempt: context.attempt, + started_at: context.started_at, + message: error.to_string(), + stage, + committed, + status_code: match error { + Error::Http { status, .. } => Some(*status), + _ => None, + }, + will_retry: false, + ended_at: epoch_seconds(), + }; + observer.error(&event).await.map_err(callback_error) +} + +fn header_values(headers: &HeaderMap) -> BTreeMap { + headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect() +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + +fn callback_error(error: impl std::fmt::Display) -> Error { + Error::InvalidResponse(format!("OCR callback failed: {error}")) +} + fn transport_error(error: reqwest::Error) -> Error { Error::Network(if error.is_timeout() { "Request timed out".into() diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 15fae1cd960..531a7e14490 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,4 +1,3 @@ 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 deleted file mode 100644 index 4fe050a235a..00000000000 --- a/litellm-rust/crates/core/src/ocr/observers.rs +++ /dev/null @@ -1,48 +0,0 @@ -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/core/src/provider_callbacks.rs b/litellm-rust/crates/core/src/provider_callbacks.rs new file mode 100644 index 00000000000..bd7626a0ead --- /dev/null +++ b/litellm-rust/crates/core/src/provider_callbacks.rs @@ -0,0 +1,206 @@ +use std::collections::BTreeMap; +use std::convert::Infallible; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Deserialize, PartialEq)] +#[serde(tag = "action", rename_all = "snake_case")] +pub enum CallbackDecision { + Unchanged, + Replace { + payload: Value, + }, + Reject { + message: String, + status_code: Option, + }, +} + +#[derive(Clone, Serialize)] +pub struct ProviderPreCall { + pub provider: String, + pub model: String, + pub call_id: String, + pub trace_id: Option, + pub attempt: u32, + pub started_at: f64, + pub request: BTreeMap, + pub api_base: String, + pub headers: BTreeMap, +} + +#[derive(Serialize)] +pub struct ProviderPostCall { + pub provider: String, + pub model: String, + pub call_id: String, + pub trace_id: Option, + pub attempt: u32, + pub started_at: f64, + pub response: Value, + pub status_code: u16, + pub headers: BTreeMap, + pub ended_at: f64, +} + +#[derive(Serialize)] +pub struct ProviderError { + pub provider: String, + pub model: String, + pub call_id: String, + pub trace_id: Option, + pub attempt: u32, + pub started_at: f64, + pub message: String, + pub stage: &'static str, + pub committed: bool, + pub status_code: Option, + pub will_retry: bool, + pub ended_at: f64, +} + +#[derive(Serialize)] +pub struct ProviderStreamEvent { + pub provider: String, + pub model: String, + pub call_id: String, + pub trace_id: Option, + pub attempt: u32, + pub started_at: f64, + pub event: Value, + pub sequence: u64, +} + +#[derive(Serialize)] +pub struct ProviderStreamClose { + pub provider: String, + pub model: String, + pub call_id: String, + pub trace_id: Option, + pub attempt: u32, + pub started_at: f64, + pub outcome: String, + pub ended_at: f64, +} + +#[derive(Serialize)] +pub struct SessionEvent { + pub session_id: String, + pub call_id: String, + pub trace_id: Option, + pub event: Option, + pub response_id: Option, + pub sequence: Option, + pub message: Option, +} + +#[macro_export] +macro_rules! provider_attempt_observer_catalog { + ($consumer:path, $($options:tt)*) => { + $consumer! { + $($options)* + { + pre_call: PreCall($crate::provider_callbacks::ProviderPreCall) -> $crate::provider_callbacks::CallbackDecision = direct; + post_call: PostCall($crate::provider_callbacks::ProviderPostCall) -> $crate::provider_callbacks::CallbackDecision = direct; + error: Error($crate::provider_callbacks::ProviderError) -> () = direct; + } + } + }; +} + +#[macro_export] +macro_rules! streaming_observer_catalog { + ($consumer:path, $($options:tt)*) => { + $consumer! { + $($options)* + { + pre_call: StreamingPreCall($crate::provider_callbacks::ProviderPreCall) -> $crate::provider_callbacks::CallbackDecision = direct; + post_call: StreamingPostCall($crate::provider_callbacks::ProviderPostCall) -> $crate::provider_callbacks::CallbackDecision = direct; + error: StreamingError($crate::provider_callbacks::ProviderError) -> () = direct; + stream_event: StreamEvent($crate::provider_callbacks::ProviderStreamEvent) -> $crate::provider_callbacks::CallbackDecision = direct; + stream_close: StreamClose($crate::provider_callbacks::ProviderStreamClose) -> () = direct; + } + } + }; +} + +#[macro_export] +macro_rules! session_observer_catalog { + ($consumer:path, $($options:tt)*) => { + $consumer! { + $($options)* + { + before_connect: BeforeConnect($crate::provider_callbacks::SessionEvent) -> $crate::provider_callbacks::CallbackDecision = awaitable; + connected: Connected($crate::provider_callbacks::SessionEvent) -> () = awaitable; + before_send: BeforeSend($crate::provider_callbacks::SessionEvent) -> $crate::provider_callbacks::CallbackDecision = awaitable; + after_receive: AfterReceive($crate::provider_callbacks::SessionEvent) -> $crate::provider_callbacks::CallbackDecision = awaitable; + response_complete: ResponseComplete($crate::provider_callbacks::SessionEvent) -> () = awaitable; + response_error: ResponseError($crate::provider_callbacks::SessionEvent) -> () = awaitable; + error: SessionError($crate::provider_callbacks::SessionEvent) -> () = awaitable; + close: Close($crate::provider_callbacks::SessionEvent) -> () = awaitable; + } + } + }; +} + +provider_attempt_observer_catalog!(crate::define_hooks, pub trait ProviderAttemptObserver;); +streaming_observer_catalog!(crate::define_hooks, pub trait StreamingObserver;); +session_observer_catalog!(crate::define_hooks, pub trait SessionObserver;); + +pub struct NoopProviderAttemptObserver; + +impl ProviderAttemptObserver for NoopProviderAttemptObserver { + type Error = Infallible; + + async fn pre_call(&mut self, _input: &ProviderPreCall) -> Result { + Ok(CallbackDecision::Unchanged) + } + + async fn post_call( + &mut self, + _input: &ProviderPostCall, + ) -> Result { + Ok(CallbackDecision::Unchanged) + } + + async fn error(&mut self, _input: &ProviderError) -> Result<(), Infallible> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::CallbackDecision; + + #[test] + fn callback_decisions_have_a_tagged_wire_contract() { + assert_eq!( + serde_json::from_value::(json!({"action": "unchanged"})).unwrap(), + CallbackDecision::Unchanged + ); + assert_eq!( + serde_json::from_value::( + json!({"action": "replace", "payload": {"masked": true}}) + ) + .unwrap(), + CallbackDecision::Replace { + payload: json!({"masked": true}) + } + ); + assert_eq!( + serde_json::from_value::(json!({ + "action": "reject", + "message": "blocked", + "status_code": 400 + })) + .unwrap(), + CallbackDecision::Reject { + message: "blocked".to_string(), + status_code: Some(400) + } + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/callback_bindings.rs b/litellm-rust/crates/python-bridge/src/callback_bindings.rs new file mode 100644 index 00000000000..a8f05e26226 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/callback_bindings.rs @@ -0,0 +1,100 @@ +use std::num::NonZeroUsize; + +use litellm_core::provider_callbacks::{ + CallbackDecision, ProviderAttemptObserver, ProviderError, ProviderPostCall, ProviderPreCall, +}; +use litellm_python_interop::callback_runtime::{AsyncContext, CallbackRuntime, SyncContext}; +use pyo3::prelude::*; + +use crate::constants::OCR_CALLBACK_CAPACITY; +use crate::execution::PythonCallContext; + +litellm_core::provider_attempt_observer_catalog!(crate::bind_python_hooks, + pub(crate) struct PythonProviderSession; + trait ProviderAttemptObserver; +); +litellm_core::streaming_observer_catalog!(crate::bind_python_hooks, + pub struct PythonStreamingSession; + trait litellm_core::provider_callbacks::StreamingObserver; +); +litellm_core::session_observer_catalog!(crate::bind_python_hooks, + pub struct PythonSession; + trait litellm_core::provider_callbacks::SessionObserver; +); + +#[pyclass(frozen)] +struct PythonCallbackRuntime(CallbackRuntime); + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + let _streaming_constructor = PythonStreamingSession::::new; + let _session_constructor = PythonSession::::new; + let capacity = NonZeroUsize::new(OCR_CALLBACK_CAPACITY) + .expect("Python callback capacity is a positive constant"); + module.add( + "__python_callback_runtime__", + PythonCallbackRuntime(CallbackRuntime::new(module, capacity)?), + ) +} + +pub(crate) enum PythonProviderObserver { + Disabled, + Sync(PythonProviderSession), + Async(PythonProviderSession), +} + +impl PythonProviderObserver { + 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("__python_callback_runtime__")? + .extract::>()? + .0 + .clone(); + if context.asynchronous { + Ok(Self::Async(PythonProviderSession::new( + adapter.bind(py), + runtime.async_context(py)?, + )?)) + } else { + Ok(Self::Sync(PythonProviderSession::new( + adapter.bind(py), + runtime.sync_context(py)?, + )?)) + } + } +} + +impl ProviderAttemptObserver for PythonProviderObserver { + type Error = PyErr; + + async fn pre_call(&mut self, input: &ProviderPreCall) -> PyResult { + match self { + Self::Disabled => Ok(CallbackDecision::Unchanged), + Self::Sync(session) => session.pre_call(input).await, + Self::Async(session) => session.pre_call(input).await, + } + } + + async fn post_call(&mut self, input: &ProviderPostCall) -> PyResult { + match self { + Self::Disabled => Ok(CallbackDecision::Unchanged), + Self::Sync(session) => session.post_call(input).await, + Self::Async(session) => session.post_call(input).await, + } + } + + async fn error(&mut self, input: &ProviderError) -> PyResult<()> { + match self { + Self::Disabled => Ok(()), + Self::Sync(session) => session.error(input).await, + Self::Async(session) => session.error(input).await, + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 14123e34392..ad32e0058d8 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,3 +1,4 @@ +mod callback_bindings; mod constants; mod diagnostics; mod errors; @@ -5,7 +6,6 @@ mod execution; #[cfg(feature = "trace-parity")] mod function_trace; mod marshal; -mod ocr_callbacks; mod python_hook_bindings; mod routes; @@ -77,10 +77,14 @@ mod _native { #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + use pyo3::types::{PyDict, PyFrozenSet}; + litellm_python_interop::callback_runtime::register(module)?; - super::ocr_callbacks::register(module)?; + super::callback_bindings::register(module)?; super::errors::register(module)?; - module.add("ready_endpoints", pyo3::types::PyDict::new(module.py()))?; + let ready_endpoints = PyDict::new(module.py()); + ready_endpoints.set_item("ocr", PyFrozenSet::new(module.py(), ["callbacks"])?)?; + module.add("ready_endpoints", ready_endpoints)?; super::routes::register(module)?; module.add_class::()?; super::diagnostics::register(module) diff --git a/litellm-rust/crates/python-bridge/src/ocr_callbacks.rs b/litellm-rust/crates/python-bridge/src/ocr_callbacks.rs deleted file mode 100644 index b93c712f3e2..00000000000 --- a/litellm-rust/crates/python-bridge/src/ocr_callbacks.rs +++ /dev/null @@ -1,80 +0,0 @@ -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; - - 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, - } - } - - 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/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 1ea18a71c4b..1ae739fe9eb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -5,10 +5,10 @@ use litellm_ai_gateway::io::ocr::{OcrRequest, ocr_with_observer}; use pyo3::prelude::*; use serde_json::Value; +use crate::callback_bindings::PythonProviderObserver; 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, @@ -25,7 +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)?; + let mut observer = PythonProviderObserver::new(inputs.callback_adapter, context)?; Ok(async move { let RouteOptions { diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 0ec691cd051..6953527648d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -678,9 +678,11 @@ def anthropic_messages_handler( raw_override: Final = litellm_params.get("rust") request_override: Final = raw_override if isinstance(raw_override, bool) else None - callback_adapter: Final = rust_messages_bridge.MessagesCallbackHandle( + from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter + + callback_adapter: Final = ProviderLoggingAdapter( logging_obj=litellm_logging_obj, - messages=messages, + input=messages, api_key=api_key or "", ) eligible: Final = ( diff --git a/litellm/main.py b/litellm/main.py index 0b1cc41ae65..40530432821 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5129,6 +5129,7 @@ def _dispatch_completion_to_python( # noqa: C901 # one explicit fallback owns def _dispatch_chat_completion(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: """Select one complete native or Python chat provider operation.""" from litellm.rust_bridge import chat_completions as bridge + from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter raw_litellm_metadata: Final = ctx.litellm_params.get("litellm_metadata") litellm_metadata: Final = dict(raw_litellm_metadata) if isinstance(raw_litellm_metadata, dict) else None @@ -5145,9 +5146,9 @@ def _dispatch_chat_completion(ctx: _CompletionDispatchContext) -> _CompletionDis ) raw_request_override: Final = ctx.litellm_params.get("rust") request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - callback_adapter: Final = bridge.ChatCompletionsCallbackHandle( + callback_adapter: Final = ProviderLoggingAdapter( logging_obj=ctx.logging, - messages=ctx.messages, + input=ctx.messages, api_key=ctx.api_key or "", ) optional_params: Final = {**ctx.optional_params, "stream": ctx.stream} # mutable-ok: native request payload diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 018565f0d04..3a13e39f636 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -29,7 +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.rust_bridge.callback_adapters import ProviderLoggingAdapter from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -284,7 +284,13 @@ 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 + ProviderLoggingAdapter( + prepared_request.litellm_logging_obj, + "OCR document processing", + resolved_api_key, + ) + if callbacks_in_rust + else None ), ) diff --git a/litellm/rust_bridge/callback_adapters.py b/litellm/rust_bridge/callback_adapters.py new file mode 100644 index 00000000000..7e2ac2cfa5f --- /dev/null +++ b/litellm/rust_bridge/callback_adapters.py @@ -0,0 +1,147 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import BaseModel, ConfigDict, JsonValue +from typing_extensions import ReadOnly, TypedDict + +from .callbacks import CallbackDecision, CallbackUnchanged, SessionCallbackHandle + + +class PreCallArguments(TypedDict): + complete_input_dict: ReadOnly[Mapping[str, JsonValue]] + api_base: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + + +class ProviderLogging(Protocol): + model_call_details: dict[str, object] + + def pre_call(self, *, input: object, api_key: str | None, additional_args: PreCallArguments) -> None: ... + + def post_call(self, *, original_response: str, input: object, api_key: str | None) -> None: ... + + +class ProviderEvent(BaseModel): + model_config = ConfigDict(extra="forbid") + + provider: str + model: str + call_id: str + trace_id: str | None = None + attempt: int + started_at: float + + +class ProviderPreCall(ProviderEvent): + request: Mapping[str, JsonValue] + api_base: str + headers: Mapping[str, str] + + +class ProviderPostCall(ProviderEvent): + response: JsonValue + status_code: int + headers: Mapping[str, str] + ended_at: float + + +class ProviderError(ProviderEvent): + message: str + stage: str + committed: bool + status_code: int | None + will_retry: bool + ended_at: float + + +class StreamEvent(ProviderEvent): + event: JsonValue + sequence: int + + +class StreamClose(ProviderEvent): + outcome: str + ended_at: float + + +class SessionEvent(BaseModel): + model_config = ConfigDict(extra="forbid") + + session_id: str + call_id: str + trace_id: str | None = None + event: JsonValue | None = None + response_id: str | None = None + sequence: int | None = None + message: str | None = None + + +def _unchanged() -> CallbackUnchanged: + return {"action": "unchanged"} + + +@dataclass(frozen=True, slots=True) +class ProviderLoggingAdapter: + logging_obj: ProviderLogging + input: object + api_key: str | None + + def pre_call(self, payload: object, /) -> CallbackDecision: + event: Final = ProviderPreCall.model_validate(payload) + additional_args: Final[PreCallArguments] = { + "complete_input_dict": event.request, + "api_base": event.api_base, + "headers": event.headers, + } + self.logging_obj.pre_call(input=self.input, api_key=self.api_key, additional_args=additional_args) + return _unchanged() + + def post_call(self, payload: object, /) -> CallbackDecision: + event: Final = ProviderPostCall.model_validate(payload) + response: Final = event.response if isinstance(event.response, str) else json.dumps(event.response) + self.logging_obj.post_call(original_response=response, input=self.input, api_key=self.api_key) + return _unchanged() + + def error(self, payload: object, /) -> None: + event: Final = ProviderError.model_validate(payload) + self.logging_obj.model_call_details["provider_error"] = event.model_dump() + + def stream_event(self, payload: object, /) -> CallbackDecision: + event: Final = StreamEvent.model_validate(payload) + self.logging_obj.model_call_details["provider_stream_event"] = event.model_dump() + return _unchanged() + + def stream_close(self, payload: object, /) -> None: + event: Final = StreamClose.model_validate(payload) + self.logging_obj.model_call_details["provider_stream_close"] = event.model_dump() + + +@dataclass(frozen=True, slots=True) +class SessionCallbackAdapter: + callback: SessionCallbackHandle + + def before_connect(self, payload: object, /) -> CallbackDecision: + return self.callback.before_connect(SessionEvent.model_validate(payload).model_dump()) + + def connected(self, payload: object, /) -> None: + self.callback.connected(SessionEvent.model_validate(payload).model_dump()) + + def before_send(self, payload: object, /) -> CallbackDecision: + return self.callback.before_send(SessionEvent.model_validate(payload).model_dump()) + + def after_receive(self, payload: object, /) -> CallbackDecision: + return self.callback.after_receive(SessionEvent.model_validate(payload).model_dump()) + + def response_complete(self, payload: object, /) -> None: + self.callback.response_complete(SessionEvent.model_validate(payload).model_dump()) + + def response_error(self, payload: object, /) -> None: + self.callback.response_error(SessionEvent.model_validate(payload).model_dump()) + + def error(self, payload: object, /) -> None: + self.callback.error(SessionEvent.model_validate(payload).model_dump()) + + def close(self, payload: object, /) -> None: + self.callback.close(SessionEvent.model_validate(payload).model_dump()) diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 55ae524b967..59936161659 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -2,27 +2,23 @@ from __future__ import annotations -import json from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # runtime typing constructs import httpx -from pydantic import TypeAdapter -from ..litellm_core_utils.litellm_logging import Logging from ..litellm_core_utils.llm_response_utils.convert_dict_to_response import ( convert_to_model_response_object, # pyright: ignore[reportUnknownVariableType] # legacy converter is untyped ) from ..types.utils import ModelResponse from .bindings import UNCHANGED, Unchanged -from .callbacks import CallbackDecision, OneShotCallbackHandle +from .callbacks import OneShotCallbackHandle from .configuration import rust_enabled from .runtime import BridgeErrorContext, EndpointDispatch from .timeouts import timeout_to_seconds RUST_RESPONSE_HEADER: Final = "x-litellm-rust" -_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @dataclass(frozen=True, slots=True) @@ -66,45 +62,6 @@ class RustAchatCompletions(Protocol): ) -> Awaitable[Mapping[str, object]]: ... -@dataclass(frozen=True, slots=True) -class ChatCompletionsCallbackHandle: - logging_obj: Logging - messages: Sequence[object] - api_key: str - - def pre_call(self, payload: object, /) -> CallbackDecision: - event: Final = _EVENT_ADAPTER.validate_python(payload) - request: Final = event.get("request", event) - headers: Final = event.get("headers", {}) - api_base: Final = event.get("api_base", event.get("url", "")) - self.logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # legacy logger is untyped - input=self.messages, - api_key=self.api_key, - additional_args={ # mutable-ok: legacy logger accepts a mutable payload - "complete_input_dict": request if isinstance(request, Mapping) else event, - "api_base": api_base if isinstance(api_base, str) else "", - "headers": headers if isinstance(headers, Mapping) else {}, # mutable-ok: empty logging headers - }, - ) - return {"action": "unchanged"} - - def post_call(self, payload: object, /) -> CallbackDecision: - event: Final = _EVENT_ADAPTER.validate_python(payload) - response: Final = event.get("response", event) - self.logging_obj.post_call( # pyright: ignore[reportUnknownMemberType] # legacy logger is untyped - input=self.messages, - api_key=self.api_key, - original_response=response if isinstance(response, str) else json.dumps(response), - ) - return {"action": "unchanged"} - - def error(self, payload: object, /) -> None: - event: Final = _EVENT_ADAPTER.validate_python(payload) - self.logging_obj.model_call_details["provider_error"] = dict( # mutable-ok: logger stores mutable details - event - ) - - _CHAT: Final = cast( # cast-ok: generic classmethod loses the route Protocol parameters EndpointDispatch[RustChatCompletions, RustAchatCompletions], EndpointDispatch.native( diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 8b6685e874a..8df8f9b8212 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -2,17 +2,13 @@ from __future__ import annotations -import json -from collections.abc import Awaitable, Callable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Awaitable, Callable from typing import Final, Protocol, cast # noqa: TID251 # runtime typing constructs import httpx -from pydantic import TypeAdapter -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.rust_bridge.bindings import UNCHANGED, Unchanged -from litellm.rust_bridge.callbacks import CallbackDecision, OneShotCallbackHandle +from litellm.rust_bridge.callbacks import OneShotCallbackHandle from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.runtime import ( BridgeErrorContext, @@ -22,8 +18,6 @@ from litellm.rust_bridge.runtime import ( ) from litellm.rust_bridge.timeouts import timeout_to_seconds -_CALLBACK_EVENT: Final = TypeAdapter(Mapping[str, object]) - class RustMessages(Protocol): def __call__( @@ -55,43 +49,6 @@ class RustAmessages(Protocol): raise NotImplementedError -@dataclass(frozen=True, slots=True) -class MessagesCallbackHandle: - logging_obj: LiteLLMLoggingObj - messages: Sequence[object] - api_key: str - - def pre_call(self, payload: object, /) -> CallbackDecision: - event: Final = _CALLBACK_EVENT.validate_python(payload) - request: Final = event.get("request", event) - self.logging_obj.pre_call( # pyright: ignore[reportUnknownMemberType] # legacy logger is untyped - input=self.messages, - api_key=self.api_key, - additional_args={ # mutable-ok: legacy logger accepts a mutable payload - "complete_input_dict": request, - "api_base": event.get("api_base", event.get("url", "")), - "headers": event.get("headers", {}), # mutable-ok: empty logging headers - }, - ) - return {"action": "unchanged"} - - def post_call(self, payload: object, /) -> CallbackDecision: - event: Final = _CALLBACK_EVENT.validate_python(payload) - response: Final = event.get("response", event) - self.logging_obj.post_call( # pyright: ignore[reportUnknownMemberType] # legacy logger is untyped - input=self.messages, - api_key=self.api_key, - original_response=response if isinstance(response, str) else json.dumps(response), - ) - return {"action": "unchanged"} - - def error(self, payload: object, /) -> None: - event: Final = _CALLBACK_EVENT.validate_python(payload) - self.logging_obj.model_call_details["provider_error"] = dict( # mutable-ok: logger stores mutable details - event - ) - - _MESSAGES: Final = cast( # cast-ok: generic classmethod loses the route Protocol parameters EndpointDispatch[RustMessages, RustAmessages], EndpointDispatch.native( diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 28fbb2abba3..047d2d9d80b 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -101,9 +101,14 @@ def supports_callback_adapter(*, asynchronous: bool = False) -> bool: if binding.is_overridden(): return False from litellm.rust_bridge import get_native_bridge + from litellm.rust_bridge.loader import native_route_ready native: Final = get_native_bridge() - return native is not None and hasattr(native, "__ocr_callback_runtime__") + return ( + native is not None + and hasattr(native, "__python_callback_runtime__") + and native_route_ready("ocr", frozenset({"callbacks"})) + ) def ocr( diff --git a/litellm/rust_bridge/ocr_callbacks.py b/litellm/rust_bridge/ocr_callbacks.py index b65ca06041f..40b55dc9e06 100644 --- a/litellm/rust_bridge/ocr_callbacks.py +++ b/litellm/rust_bridge/ocr_callbacks.py @@ -1,59 +1,3 @@ -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Final, Protocol +from .callback_adapters import PreCallArguments, ProviderLoggingAdapter -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, - ) +__all__ = ["PreCallArguments", "ProviderLoggingAdapter"] diff --git a/tests/test_litellm/ocr/sdk_callback_contract.py b/tests/test_litellm/ocr/sdk_callback_contract.py index b6dc039c474..817aa1c92c6 100644 --- a/tests/test_litellm/ocr/sdk_callback_contract.py +++ b/tests/test_litellm/ocr/sdk_callback_contract.py @@ -24,7 +24,7 @@ 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 +from litellm.rust_bridge.callback_adapters import PreCallArguments @dataclass(frozen=True, slots=True) diff --git a/tests/test_litellm/rust_bridge/test_callback_adapters.py b/tests/test_litellm/rust_bridge/test_callback_adapters.py new file mode 100644 index 00000000000..d97098eff0d --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_callback_adapters.py @@ -0,0 +1,132 @@ +from dataclasses import dataclass, field +from typing import Final + +from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter, SessionCallbackAdapter +from litellm.rust_bridge.callbacks import CallbackDecision + + +@dataclass +class RecordingLogging: + model_call_details: dict[str, object] = field(default_factory=dict) + calls: list[tuple[str, object]] = field(default_factory=list) + + def pre_call(self, *, input: object, api_key: str | None, additional_args: object) -> None: + self.calls.append(("pre", (input, api_key, additional_args))) + + def post_call(self, *, original_response: str, input: object, api_key: str | None) -> None: + self.calls.append(("post", (original_response, input, api_key))) + + +def provider_event(**updates: object) -> dict[str, object]: + return { + "provider": "mistral", + "model": "mistral-ocr-latest", + "call_id": "call-1", + "trace_id": "trace-1", + "attempt": 1, + "started_at": 10.0, + **updates, + } + + +def test_provider_logging_adapter_preserves_provider_lifecycle() -> None: + logging: Final = RecordingLogging() + adapter: Final = ProviderLoggingAdapter(logging, "OCR document processing", "secret") + + assert adapter.pre_call( + provider_event(request={"model": "mistral-ocr-latest"}, api_base="https://provider.test", headers={}) + ) == {"action": "unchanged"} + assert adapter.post_call(provider_event(response={"pages": []}, status_code=200, headers={}, ended_at=11.0)) == { + "action": "unchanged" + } + adapter.error( + provider_event( + message="retryable", + stage="provider_response", + committed=True, + status_code=429, + will_retry=True, + ended_at=11.0, + ) + ) + assert adapter.stream_event(provider_event(event={"type": "delta"}, sequence=1)) == {"action": "unchanged"} + adapter.stream_close(provider_event(outcome="completed", ended_at=12.0)) + + assert [name for name, _ in logging.calls] == ["pre", "post"] + assert logging.model_call_details["provider_stream_event"] == { + **provider_event(), + "event": {"type": "delta"}, + "sequence": 1, + } + assert logging.model_call_details["provider_stream_close"] == { + **provider_event(), + "outcome": "completed", + "ended_at": 12.0, + } + assert logging.model_call_details["provider_error"] == { + **provider_event(), + "message": "retryable", + "stage": "provider_response", + "committed": True, + "status_code": 429, + "will_retry": True, + "ended_at": 11.0, + } + + +@dataclass +class RecordingSession: + events: list[str] = field(default_factory=list) + + def before_connect(self, payload: object, /) -> CallbackDecision: + self.events.append("before_connect") + return {"action": "unchanged"} + + def connected(self, payload: object, /) -> None: + self.events.append("connected") + + def before_send(self, payload: object, /) -> CallbackDecision: + self.events.append("before_send") + return {"action": "reject", "message": "drop frame", "status_code": None} + + def after_receive(self, payload: object, /) -> CallbackDecision: + self.events.append("after_receive") + return {"action": "replace", "payload": {"type": "masked"}} + + def response_complete(self, payload: object, /) -> None: + self.events.append("response_complete") + + def response_error(self, payload: object, /) -> None: + self.events.append("response_error") + + def error(self, payload: object, /) -> None: + self.events.append("error") + + def close(self, payload: object, /) -> None: + self.events.append("close") + + +def test_session_adapter_preserves_frame_decisions_and_order() -> None: + callback: Final = RecordingSession() + adapter: Final = SessionCallbackAdapter(callback) + event: Final = {"session_id": "session-1", "call_id": "call-1", "event": {"type": "response.create"}} + + assert adapter.before_connect(event) == {"action": "unchanged"} + adapter.connected(event) + assert adapter.before_send(event) == {"action": "reject", "message": "drop frame", "status_code": None} + assert adapter.after_receive(event) == {"action": "replace", "payload": {"type": "masked"}} + adapter.response_complete(event) + adapter.response_error(event) + adapter.error(event) + adapter.close(event) + + assert callback.events == [ + "before_connect", + "connected", + "before_send", + "after_receive", + "response_complete", + "response_error", + "error", + "close", + ] diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 30f5dcbc307..0a1e47a7aef 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -83,9 +83,32 @@ class RecordingCall: } ) callback_adapter.pre_call( - {"request": {"model": model, "messages": messages}, "api_base": api_base, "headers": extra_headers} + { + "provider": custom_llm_provider or "anthropic", + "model": model, + "call_id": "native-test", + "trace_id": None, + "attempt": 1, + "started_at": 1.0, + "request": {"model": model, "messages": messages}, + "api_base": api_base or "", + "headers": extra_headers or {}, + } + ) + callback_adapter.post_call( + { + "provider": custom_llm_provider or "anthropic", + "model": model, + "call_id": "native-test", + "trace_id": None, + "attempt": 1, + "started_at": 1.0, + "response": RUST_RESPONSE, + "status_code": 200, + "headers": {}, + "ended_at": 2.0, + } ) - callback_adapter.post_call({"response": RUST_RESPONSE}) return RUST_RESPONSE