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

refactor(rust): share provider callback infrastructure

Browse source
This commit is contained in:
Yujong Lee 2026-09-04 15:49:15 -07:00
parent 1a848c473b
commit eaf210a23c
25 changed files with 926 additions and 353 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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<Observer>(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
observer: &mut impl OcrObserver,
) -> Result<Value, Error> {
observer: &mut Observer,
) -> Result<Value, Error>
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)
}

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

@ -122,6 +122,8 @@ impl OcrLifecycleHooks {
};
Ok(ProviderOcrRequest {
model,
custom_llm_provider,
litellm_call_id: request.litellm_call_id,
config,
url,
body,

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

@ -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<Value, Error> {
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<Value, Error> {
level = "trace",
skip_all
)]
pub async fn ocr_with_observer(
pub async fn ocr_with_observer<Observer>(
request: OcrRequest<'_>,
observer: &mut impl OcrObserver,
) -> Result<Value, Error> {
observer: &mut Observer,
) -> Result<Value, Error>
where
Observer: ProviderAttemptObserver,
Observer::Error: std::fmt::Display,
{
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
CallLifecycle::default()
.run_request(request, &hooks, |request| {

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

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

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

@ -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<Mutex<Vec<&'static str>>>,
raw_response: Option<String>,
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<CallbackDecision, Self::Error> {
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<CallbackDecision, Self::Error> {
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),

1
litellm-rust/crates/core/src/lib.rs
View file

@ -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;

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

@ -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<Observer>(
request: RequestBuilder,
event: &OcrPreCall,
observer: &mut impl OcrObserver,
) -> Result<OcrHttpResponse, Error> {
if observer.pre_call(event).await.is_err() {
tracing::warn!("OCR pre-call observer failed");
}
let response = http_request(request).await.map_err(transport_error)?;
event: &ProviderPreCall,
observer: &mut Observer,
) -> Result<OcrHttpResponse, Error>
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>(
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<String, String> {
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()

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

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

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

@ -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<String, String>,
}
#[derive(Serialize)]
pub struct OcrPostCall {
pub original_response: String,
}
#[macro_export]
macro_rules! ocr_observer_catalog {
($consumer:path, $($options:tt)*) => {
$consumer! {
$($options)*
{
pre_call: PreCall($crate::ocr::observers::OcrPreCall) -> () = direct;
post_call: PostCall($crate::ocr::observers::OcrPostCall) -> () = direct;
}
}
};
}
ocr_observer_catalog!(crate::define_hooks, pub trait OcrObserver;);
pub struct NoopOcrObserver;
impl OcrObserver for NoopOcrObserver {
type Error = Infallible;
async fn pre_call(&mut self, _input: &OcrPreCall) -> Result<(), Infallible> {
Ok(())
}
async fn post_call(&mut self, _input: &OcrPostCall) -> Result<(), Infallible> {
Ok(())
}
}

206
litellm-rust/crates/core/src/provider_callbacks.rs Normal file
View file

@ -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<u16>,
},
}
#[derive(Clone, Serialize)]
pub struct ProviderPreCall {
pub provider: String,
pub model: String,
pub call_id: String,
pub trace_id: Option<String>,
pub attempt: u32,
pub started_at: f64,
pub request: BTreeMap<String, Value>,
pub api_base: String,
pub headers: BTreeMap<String, String>,
}
#[derive(Serialize)]
pub struct ProviderPostCall {
pub provider: String,
pub model: String,
pub call_id: String,
pub trace_id: Option<String>,
pub attempt: u32,
pub started_at: f64,
pub response: Value,
pub status_code: u16,
pub headers: BTreeMap<String, String>,
pub ended_at: f64,
}
#[derive(Serialize)]
pub struct ProviderError {
pub provider: String,
pub model: String,
pub call_id: String,
pub trace_id: Option<String>,
pub attempt: u32,
pub started_at: f64,
pub message: String,
pub stage: &'static str,
pub committed: bool,
pub status_code: Option<u16>,
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<String>,
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<String>,
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<String>,
pub event: Option<Value>,
pub response_id: Option<String>,
pub sequence: Option<u64>,
pub message: Option<String>,
}
#[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<CallbackDecision, Infallible> {
Ok(CallbackDecision::Unchanged)
}
async fn post_call(
&mut self,
_input: &ProviderPostCall,
) -> Result<CallbackDecision, Infallible> {
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::<CallbackDecision>(json!({"action": "unchanged"})).unwrap(),
CallbackDecision::Unchanged
);
assert_eq!(
serde_json::from_value::<CallbackDecision>(
json!({"action": "replace", "payload": {"masked": true}})
)
.unwrap(),
CallbackDecision::Replace {
payload: json!({"masked": true})
}
);
assert_eq!(
serde_json::from_value::<CallbackDecision>(json!({
"action": "reject",
"message": "blocked",
"status_code": 400
}))
.unwrap(),
CallbackDecision::Reject {
message: "blocked".to_string(),
status_code: Some(400)
}
);
}
}

100
litellm-rust/crates/python-bridge/src/callback_bindings.rs Normal file
View file

@ -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::<SyncContext>::new;
let _session_constructor = PythonSession::<AsyncContext>::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<SyncContext>),
Async(PythonProviderSession<AsyncContext>),
}
impl PythonProviderObserver {
pub(crate) fn new(
adapter: Option<Py<PyAny>>,
context: PythonCallContext<'_>,
) -> PyResult<Self> {
let Some(adapter) = adapter else {
return Ok(Self::Disabled);
};
let py = context.py;
let module = py.import("litellm.rust_bridge._native")?;
let runtime = module
.getattr("__python_callback_runtime__")?
.extract::<PyRef<'_, PythonCallbackRuntime>>()?
.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<CallbackDecision> {
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<CallbackDecision> {
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,
}
}
}

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

@ -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::ResponsesWebSocketConnection>()?;
super::diagnostics::register(module)

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

@ -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<SyncContext>),
Async(PythonOcrSession<AsyncContext>),
}
impl PythonOcrObserver {
pub(crate) fn new(
adapter: Option<Py<PyAny>>,
context: PythonCallContext<'_>,
) -> PyResult<Self> {
let Some(adapter) = adapter else {
return Ok(Self::Disabled);
};
let py = context.py;
let module = py.import("litellm.rust_bridge._native")?;
let runtime = module
.getattr("__ocr_callback_runtime__")?
.extract::<PyRef<'_, OcrRuntime>>()?
.0
.clone();
if context.asynchronous {
Ok(Self::Async(PythonOcrSession::new(
adapter.bind(py),
runtime.async_context(py)?,
)?))
} else {
Ok(Self::Sync(PythonOcrSession::new(
adapter.bind(py),
runtime.sync_context(py)?,
)?))
}
}
}
impl OcrObserver for PythonOcrObserver {
type Error = PyErr;
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,
}
}
}

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

@ -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 {

6
litellm/llms/anthropic/experimental_pass_through/messages/handler.py
View file

@ -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 = (

5
litellm/main.py
View file

@ -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

10
litellm/ocr/main.py
View file

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

147
litellm/rust_bridge/callback_adapters.py Normal file
View file

@ -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())

45
litellm/rust_bridge/chat_completions.py
View file

@ -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(

47
litellm/rust_bridge/messages.py
View file

@ -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(

7
litellm/rust_bridge/ocr.py
View file

@ -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(

60
litellm/rust_bridge/ocr_callbacks.py
View file

@ -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"]

2
tests/test_litellm/ocr/sdk_callback_contract.py
View file

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

132
tests/test_litellm/rust_bridge/test_callback_adapters.py Normal file
View file

@ -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",
]

27
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -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