scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

test(ocr): enforce Python Rust function trace parity

Browse source
This commit is contained in:
Yujong Lee 2026-09-02 11:02:14 -07:00
parent cd77a2ad86
commit 1ea1d45a5f
19 changed files with 700 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm-rust/Cargo.lock generated
View file

@ -1392,6 +1392,12 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "libc"
version = "0.2.186"
@ -1439,6 +1445,7 @@ dependencies = [
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"tracing",
"url",
]
@ -1457,6 +1464,8 @@ dependencies = [
"serde_json",
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -2284,6 +2293,15 @@ dependencies = [
"digest 0.11.3",
]
[[package]]
name = "sharded-slab"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
dependencies = [
"lazy_static",
]
[[package]]
name = "shlex"
version = "2.0.1"
@ -2422,6 +2440,15 @@ dependencies = [
"syn 3.0.0",
]
[[package]]
name = "thread_local"
version = "1.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070"
dependencies = [
"cfg-if",
]
[[package]]
name = "time"
version = "0.3.53"
@ -2670,6 +2697,17 @@ dependencies = [
"once_cell",
]
[[package]]
name = "tracing-subscriber"
version = "0.3.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
dependencies = [
"sharded-slab",
"thread_local",
"tracing-core",
]
[[package]]
name = "try-lock"
version = "0.2.5"

2
litellm-rust/Cargo.toml
View file

@ -14,6 +14,8 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
tracing = "0.1"
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
litellm-core = { path = "crates/core" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }

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

@ -87,7 +87,14 @@ impl OcrLifecycleHooks {
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_ocr_params(&request.optional_params);
let supported_params = config.get_supported_ocr_params();
let non_default_params = request
.optional_params
.iter()
.filter(|(param, _)| supported_params.contains(&param.as_str()))
.map(|(param, value)| (param.clone(), value.clone()))
.collect();
let filtered_params = config.map_ocr_params(&non_default_params);
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let document = if config.requires_data_uri_document() {

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

@ -18,6 +18,14 @@ use crate::integrations::custom_logger::{
};
use crate::integrations::types::RequestMetadata;
type LifecycleEvents = Arc<Mutex<Vec<&'static str>>>;
fn record_lifecycle_event(events: Option<&LifecycleEvents>, event: &'static str) {
if let Some(events) = events {
events.lock().unwrap().push(event);
}
}
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
@ -80,9 +88,17 @@ struct RecordedLogEvent {
#[derive(Default)]
struct RecordingOcrLogger {
events: Mutex<Vec<RecordedLogEvent>>,
lifecycle_events: Option<LifecycleEvents>,
}
impl RecordingOcrLogger {
fn with_lifecycle_events(lifecycle_events: LifecycleEvents) -> Self {
Self {
events: Mutex::new(Vec::new()),
lifecycle_events: Some(lifecycle_events),
}
}
fn events(&self) -> Vec<RecordedLogEvent> {
self.events.lock().unwrap().clone()
}
@ -96,6 +112,7 @@ impl CustomLogger for RecordingOcrLogger {
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_log_success_event");
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_success_event",
model: model_call_details.model.clone(),
@ -115,6 +132,7 @@ impl CustomLogger for RecordingOcrLogger {
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_log_failure_event");
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_failure_event",
model: model_call_details.model.clone(),
@ -135,22 +153,41 @@ struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
lifecycle_events: Option<LifecycleEvents>,
}
impl RecordingOcrGuardrail {
fn new(hooks: Vec<GuardrailEventHook>) -> Self {
fn with_lifecycle_events(
hooks: Vec<GuardrailEventHook>,
lifecycle_events: LifecycleEvents,
) -> Self {
Self {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
lifecycle_events: Some(lifecycle_events),
}
}
fn blocking_pre_call() -> Self {
fn blocking_pre_call(lifecycle_events: LifecycleEvents) -> Self {
Self {
hooks: vec![GuardrailEventHook::PreCall],
events: Mutex::new(Vec::new()),
block_pre_call: true,
block_during_call: false,
lifecycle_events: Some(lifecycle_events),
}
}
fn blocking_during_call(lifecycle_events: LifecycleEvents) -> Self {
Self {
hooks: vec![GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall],
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: true,
lifecycle_events: Some(lifecycle_events),
}
}
@ -174,6 +211,7 @@ impl CustomGuardrail for RecordingOcrGuardrail {
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_pre_call_hook");
self.events.lock().unwrap().push("async_pre_call_hook");
if self.block_pre_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
@ -191,7 +229,13 @@ impl CustomGuardrail for RecordingOcrGuardrail {
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
record_lifecycle_event(self.lifecycle_events.as_ref(), "async_moderation_hook");
self.events.lock().unwrap().push("async_moderation_hook");
if self.block_during_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked during provider preparation",
)));
}
request.data["body"]["guarded_during"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
@ -242,7 +286,7 @@ fn ocr_dispatch_supports_migrated_providers() {
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.get_supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
@ -281,14 +325,17 @@ fn auth_header_detection_is_case_insensitive() {
#[tokio::test]
async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let provider_events = lifecycle_events.clone();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_request(&mut socket).await;
record_lifecycle_event(Some(&provider_events), "provider_request_received");
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
@ -302,11 +349,13 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
request
});
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
lifecycle_events.clone(),
));
let guardrail = Arc::new(RecordingOcrGuardrail::with_lifecycle_events(
vec![GuardrailEventHook::PreCall, GuardrailEventHook::DuringCall],
lifecycle_events.clone(),
));
let response = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
@ -346,22 +395,81 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
error_kind: None,
}]
);
assert_eq!(
lifecycle_events.lock().unwrap().as_slice(),
[
"async_pre_call_hook",
"async_moderation_hook",
"provider_request_received",
"async_log_success_event",
]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
assert!(request.contains(r#""guarded_during":true"#), "{request}");
}
#[tokio::test]
async fn ocr_lifecycle_during_call_block_skips_provider_and_runs_failure_callback() {
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
lifecycle_events.clone(),
));
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call(
lifecycle_events.clone(),
));
let error = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_millis(100)),
callbacks: vec![logger],
guardrails: vec![guardrail],
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-during-block"),
})
.await
.expect_err("during-call guardrail blocks request");
assert!(matches!(error, Error::InvalidRequest(_)));
assert_eq!(
lifecycle_events.lock().unwrap().as_slice(),
[
"async_pre_call_hook",
"async_moderation_hook",
"async_log_failure_event",
]
);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "provider socket should not be touched");
}
#[tokio::test]
async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let provider_events = lifecycle_events.clone();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let _request = read_http_request(&mut socket).await;
record_lifecycle_event(Some(&provider_events), "provider_request_received");
let response_body = "provider failed";
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
@ -374,7 +482,9 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
.expect("writes response");
});
let logger = Arc::new(RecordingOcrLogger::default());
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
lifecycle_events.clone(),
));
let err = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
@ -408,16 +518,25 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
error_kind: Some("HttpError".to_string()),
}]
);
assert_eq!(
lifecycle_events.lock().unwrap().as_slice(),
["provider_request_received", "async_log_failure_event"]
);
}
#[tokio::test]
async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
let lifecycle_events: LifecycleEvents = Arc::new(Mutex::new(Vec::new()));
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
let logger = Arc::new(RecordingOcrLogger::with_lifecycle_events(
lifecycle_events.clone(),
));
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call(
lifecycle_events.clone(),
));
let err = ocr(OcrRequest {
model: "mistral-ocr-latest",
@ -452,6 +571,10 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
error_kind: Some("InvalidRequest".to_string()),
}]
);
assert_eq!(
lifecycle_events.lock().unwrap().as_slice(),
["async_pre_call_hook", "async_log_failure_event"]
);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "provider socket should not be touched");
}

1
litellm-rust/crates/core/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
tracing.workspace = true
base64.workspace = true
futures-util.workspace = true
bytes.workspace = true

16
litellm-rust/crates/core/src/ocr/transformation.rs
View file

@ -25,16 +25,16 @@ pub enum OcrResponseHandling {
}
pub trait OcrProviderConfig: Sync {
fn supported_ocr_params(&self) -> &'static [&'static str];
fn get_supported_ocr_params(&self) -> &'static [&'static str];
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
let mut mapped_params = Map::new();
for (param, value) in non_default_params {
if self.supported_ocr_params().contains(&param.as_str()) {
mapped_params.insert(param.clone(), value.clone());
}
}
mapped_params
let supported_params = self.get_supported_ocr_params();
non_default_params
.iter()
.filter(|(param, _)| supported_params.contains(&param.as_str()))
.map(|(param, value)| (param.clone(), value.clone()))
.collect()
}
fn transform_ocr_request(

6
litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs
View file

@ -281,8 +281,8 @@ fn page_dimensions(page: &Map<String, Value>) -> Value {
}
impl OcrProviderConfig for AzureAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.get_supported_ocr_params()
}
fn transform_ocr_request(
@ -326,7 +326,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
}
impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS
}

11
litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs
View file

@ -70,10 +70,12 @@ pub struct MistralOcrConfig;
pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
impl OcrProviderConfig for MistralOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
SUPPORTED_OCR_PARAMS
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -100,6 +102,7 @@ impl OcrProviderConfig for MistralOcrConfig {
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_response(
&self,
model: &str,
@ -153,8 +156,8 @@ impl OcrProviderConfig for MistralOcrConfig {
}
}
pub fn supported_ocr_params() -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
pub fn get_supported_ocr_params() -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.get_supported_ocr_params()
}
pub fn map_ocr_params(non_default_params: &Map<String, Value>) -> Map<String, Value> {
@ -181,7 +184,7 @@ mod tests {
#[test]
fn supported_params_match_python_mistral_ocr_config() {
assert_eq!(
supported_ocr_params(),
get_supported_ocr_params(),
&[
"pages",
"include_image_base64",

6
litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs
View file

@ -208,8 +208,8 @@ fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> V
}
impl OcrProviderConfig for VertexAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.get_supported_ocr_params()
}
fn transform_ocr_request(
@ -253,7 +253,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
}
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
fn get_supported_ocr_params(&self) -> &'static [&'static str] {
DEEPSEEK_SUPPORTED_OCR_PARAMS
}

2
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -16,6 +16,8 @@ extension-module = ["pyo3/extension-module"]
panic-test = []
[dependencies]
tracing.workspace = true
tracing-subscriber.workspace = true
futures-util.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }

129
litellm-rust/crates/python-bridge/src/function_trace.rs Normal file
View file

@ -0,0 +1,129 @@
use std::sync::{Arc, Mutex};
use serde::Serialize;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Level, Subscriber};
use tracing_subscriber::filter::{LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
const TARGET: &str = "litellm::function_trace";
#[derive(Clone, Debug, PartialEq, Serialize)]
pub(crate) struct FunctionTraceEvent {
pub(crate) function: &'static str,
pub(crate) depth: usize,
}
#[derive(Clone, Default)]
pub(crate) struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
}
impl FunctionTrace {
pub(crate) fn dispatcher(&self) -> Dispatch {
let filter = filter_fn(|metadata| {
metadata.is_span() && metadata.target() == TARGET && *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE);
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(filter),
),
)
}
pub(crate) fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let depth = context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default();
self.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.push(FunctionTraceEvent {
function: attributes.metadata().name(),
depth,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: TARGET, "wrong_level");
let _second = tracing::trace_span!(target: TARGET, "same_name");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: TARGET, "inner");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0,
},
FunctionTraceEvent {
function: "inner",
depth: 1,
},
]
);
}
}

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

@ -1,5 +1,6 @@
mod diagnostics;
mod errors;
mod function_trace;
mod marshal;
mod routes;

2
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -154,7 +154,7 @@ mod tests {
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=None)",
),
(
"transcription",

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

@ -4,8 +4,10 @@ use std::future::Future;
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use pyo3::prelude::*;
use serde_json::Value;
use tracing::instrument::WithSubscriber;
use crate::errors::core_error_to_pyerr;
use crate::function_trace::FunctionTrace;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
fn prepare_ocr(
@ -22,6 +24,7 @@ fn prepare_ocr(
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
let trace = inputs.trace.unwrap_or(false).then(FunctionTrace::default);
Ok(async move {
let RouteOptions {
model,
@ -31,7 +34,7 @@ fn prepare_ocr(
extra_headers,
timeout,
} = options;
run_ocr(OcrRequest {
let future = run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
@ -44,8 +47,14 @@ fn prepare_ocr(
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
});
match trace {
Some(trace) => {
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(serde_json::json!({ "response": response, "trace": trace.events() }))
}
None => future.await,
}
})
}
@ -67,6 +76,7 @@ bridge_route! {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
timeout_seconds: Option<f64>,
trace: Option<bool>,
},
prepare = prepare_ocr,
errors = core_error_to_pyerr,

13
tests/sdk_function_trace/__init__.py Normal file
View file

@ -0,0 +1,13 @@
from tests.sdk_function_trace.harness import (
TraceScenario,
TraceStep,
assert_function_trace_parity,
)
from tests.sdk_function_trace.profiler import FunctionTraceEvent
__all__ = [
"FunctionTraceEvent",
"TraceScenario",
"TraceStep",
"assert_function_trace_parity",
]

39
tests/sdk_function_trace/harness.py Normal file
View file

@ -0,0 +1,39 @@
from __future__ import annotations
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from types import FunctionType
from typing import Final, cast
from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python
@dataclass(frozen=True, slots=True)
class TraceStep:
function: FunctionType
depth: int
@dataclass(frozen=True, slots=True)
class TraceScenario:
steps: tuple[TraceStep, ...]
invoke_python: Callable[[], object]
invoke_rust: Callable[[], Sequence[FunctionTraceEvent]]
def assert_function_trace_parity(scenario: TraceScenario) -> None:
expected: Final = tuple(
FunctionTraceEvent(function=step.function.__name__, depth=step.depth) for step in scenario.steps
)
functions: Final = cast(tuple[FunctionType, ...], tuple(step.function for step in scenario.steps))
with profile_python(functions) as profiler:
scenario.invoke_python()
python_trace: Final = tuple(profiler.events)
rust_trace: Final = tuple(scenario.invoke_rust())
if python_trace != expected:
raise AssertionError(f"Python function trace differs: {python_trace!r} != {expected!r}")
if rust_trace != expected:
raise AssertionError(f"Rust function trace differs: {rust_trace!r} != {expected!r}")
if python_trace != rust_trace:
raise AssertionError(f"Python and Rust function traces differ: {python_trace!r} != {rust_trace!r}")

67
tests/sdk_function_trace/mock_provider.py Normal file
View file

@ -0,0 +1,67 @@
from __future__ import annotations
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from threading import Lock, Thread
from typing import Final, cast
@dataclass(frozen=True, slots=True)
class MockProviderResponse:
status_code: int
headers: tuple[tuple[str, str], ...]
body: bytes
class _MockProviderServer(ThreadingHTTPServer):
def __init__(self, response: MockProviderResponse) -> None:
super().__init__(("127.0.0.1", 0), _MockProviderHandler)
self.response: Final = response
self._request_count = 0
self._request_count_lock: Final = Lock()
def record_request(self) -> None:
with self._request_count_lock:
self._request_count += 1
@property
def request_count(self) -> int:
with self._request_count_lock:
return self._request_count
class _MockProviderHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
content_length: Final = int(self.headers.get("content-length", "0"))
self.rfile.read(content_length)
server: Final = cast(_MockProviderServer, self.server)
server.record_request()
self.send_response(server.response.status_code)
for name, value in server.response.headers:
self.send_header(name, value)
self.send_header("content-length", str(len(server.response.body)))
self.end_headers()
self.wfile.write(server.response.body)
def log_message(self, format: str, *args: object) -> None: # noqa: A002 # matches BaseHTTPRequestHandler
pass
@contextmanager
def mock_provider(response: MockProviderResponse) -> Generator[str]:
server: Final = _MockProviderServer(response)
thread: Final = Thread(target=server.serve_forever, daemon=True)
thread.start()
host, port = cast(tuple[str, int], server.server_address)
try:
yield f"http://{host}:{port}"
finally:
server.shutdown()
server.server_close()
thread.join()
if server.request_count != 1:
raise AssertionError(f"expected one provider request, received {server.request_count}")

49
tests/sdk_function_trace/profiler.py Normal file
View file

@ -0,0 +1,49 @@
from __future__ import annotations
import sys
from collections.abc import Generator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from types import FrameType, FunctionType
from typing import Final
@dataclass(frozen=True, slots=True)
class FunctionTraceEvent:
function: str
depth: int
class PythonProfiler:
def __init__(self, functions: Sequence[FunctionType]) -> None:
self._names_by_code: Final = {function.__code__: function.__name__ for function in functions}
self._seen_frames: Final[set[FrameType]] = set()
self.events: Final[list[FunctionTraceEvent]] = []
def __call__(self, frame: FrameType, event: str, _arg: object) -> None:
if event != "call" or frame in self._seen_frames:
return
function_name: Final = self._names_by_code.get(frame.f_code)
if function_name is None:
return
depth: Final = sum(ancestor.f_code in self._names_by_code for ancestor in _frame_ancestors(frame))
self._seen_frames.add(frame)
self.events.append(FunctionTraceEvent(function=function_name, depth=depth))
def _frame_ancestors(frame: FrameType) -> Generator[FrameType]:
ancestor: Final = frame.f_back
if ancestor is not None:
yield ancestor
yield from _frame_ancestors(ancestor)
@contextmanager
def profile_python(functions: Sequence[FunctionType]) -> Generator[PythonProfiler]:
profiler: Final = PythonProfiler(functions)
previous: Final = sys.getprofile()
sys.setprofile(profiler)
try:
yield profiler
finally:
sys.setprofile(previous)

182
tests/test_litellm/ocr/test_function_trace.py Normal file
View file

@ -0,0 +1,182 @@
from __future__ import annotations
import asyncio
import json
import sys
from collections.abc import Awaitable
from functools import partial
from types import FunctionType
from typing import Final, Protocol, cast
import pytest
from pydantic import StrictInt, StrictStr, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
from litellm.rust_bridge import get_native_bridge
from litellm.rust_bridge import ocr as rust_ocr_bridge
from tests.sdk_function_trace import (
FunctionTraceEvent,
TraceScenario,
TraceStep,
assert_function_trace_parity,
)
from tests.sdk_function_trace.mock_provider import MockProviderResponse, mock_provider
from tests.sdk_function_trace.profiler import profile_python
MODEL: Final = "mistral-ocr-latest"
DOCUMENT: Final[dict[str, str]] = {
"type": "document_url",
"document_url": "https://example.com/document.pdf",
}
OPTIONAL_PARAMS: Final[dict[str, object]] = {"pages": [0], "unsupported": True}
RESPONSE_DATA: Final[dict[str, object]] = {
"pages": [{"index": 0, "markdown": "hello"}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1},
}
PROVIDER_RESPONSE: Final = MockProviderResponse(
status_code=200,
headers=(("content-type", "application/json"),),
body=json.dumps(RESPONSE_DATA).encode(),
)
class _TraceEventPayload(TypedDict):
function: ReadOnly[StrictStr]
depth: ReadOnly[StrictInt]
class _TraceResponsePayload(TypedDict):
response: ReadOnly[object]
trace: ReadOnly[list[_TraceEventPayload]]
_TRACE_RESPONSE: Final = TypeAdapter(_TraceResponsePayload)
class _NativeTraceOcr(Protocol):
def __call__(
self,
*,
model: str,
document: dict[str, str],
api_key: str,
api_base: str,
optional_params: dict[str, object],
trace: bool,
) -> object: ...
def _invoke_python() -> object:
previous_enabled: Final = rust_ocr_bridge.rust_ocr_enabled()
rust_ocr_bridge.use_litellm_rust(False)
try:
with mock_provider(PROVIDER_RESPONSE) as api_base:
return litellm.ocr(
model=f"mistral/{MODEL}",
document=DOCUMENT,
api_key="test-key",
api_base=api_base,
pages=[0],
unsupported=True,
)
finally:
rust_ocr_bridge.use_litellm_rust(previous_enabled)
def _invoke_rust(*, asynchronous: bool = False) -> tuple[FunctionTraceEvent, ...]:
bridge: Final = get_native_bridge()
if bridge is None:
raise AssertionError("The native Rust bridge is required for function-trace parity")
trace_ocr: Final = cast(_NativeTraceOcr, bridge.aocr if asynchronous else bridge.ocr)
with mock_provider(PROVIDER_RESPONSE) as api_base:
invoke: Final = partial(
trace_ocr,
model=f"mistral/{MODEL}",
document=DOCUMENT,
api_key="test-key",
api_base=api_base,
optional_params=OPTIONAL_PARAMS,
trace=True,
)
async def invoke_async() -> object:
return await cast(Awaitable[object], invoke())
raw_result: Final = asyncio.run(invoke_async()) if asynchronous else invoke()
result: Final = _TRACE_RESPONSE.validate_python(raw_result)
return tuple(
FunctionTraceEvent(
function=event["function"],
depth=event["depth"],
)
for event in result["trace"]
)
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
def test_mistral_ocr_transformation_function_trace_parity(asynchronous: bool) -> None:
assert_function_trace_parity(
TraceScenario(
steps=(
TraceStep(cast(FunctionType, MistralOCRConfig.get_supported_ocr_params), depth=0),
TraceStep(cast(FunctionType, MistralOCRConfig.map_ocr_params), depth=0),
TraceStep(cast(FunctionType, MistralOCRConfig.get_supported_ocr_params), depth=1),
TraceStep(cast(FunctionType, MistralOCRConfig.transform_ocr_request), depth=0),
TraceStep(cast(FunctionType, MistralOCRConfig.transform_ocr_response), depth=0),
),
invoke_python=_invoke_python,
invoke_rust=partial(_invoke_rust, asynchronous=asynchronous),
)
)
class First:
@staticmethod
def run() -> None:
return None
class Second:
@staticmethod
def run() -> None:
return None
def test_profiler_matches_code_objects_and_keeps_repeated_calls() -> None:
with profile_python((First.run,)) as profiler:
Second.run()
First.run()
First.run()
assert profiler.events == [
FunctionTraceEvent(function="run", depth=0),
FunctionTraceEvent(function="run", depth=0),
]
def test_profiler_records_selected_function_nesting_depth() -> None:
class Nested:
@staticmethod
def run() -> None:
First.run()
with profile_python((Nested.run, First.run)) as profiler:
Nested.run()
assert profiler.events == [
FunctionTraceEvent(function="run", depth=0),
FunctionTraceEvent(function="run", depth=1),
]
def test_profiler_restores_previous_profiler_after_failure() -> None:
previous: Final = sys.getprofile()
with pytest.raises(RuntimeError, match="stop"):
with profile_python((First.run,)):
raise RuntimeError("stop")
assert sys.getprofile() is previous