diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ac9f70e33d7..95dd09282e0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 69728a652bd..a99de014481 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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" } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 6343122b9e9..8025b5f713a 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -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(¶m.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() { diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 8c3f0425149..0dd58b5f9ae 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -18,6 +18,14 @@ use crate::integrations::custom_logger::{ }; use crate::integrations::types::RequestMetadata; +type LifecycleEvents = Arc>>; + +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>, + lifecycle_events: Option, } 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 { 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, events: Mutex>, block_pre_call: bool, + block_during_call: bool, + lifecycle_events: Option, } impl RecordingOcrGuardrail { - fn new(hooks: Vec) -> Self { + fn with_lifecycle_events( + hooks: Vec, + 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"); } diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 3eb9da53277..18917e08ebc 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +tracing.workspace = true base64.workspace = true futures-util.workspace = true bytes.workspace = true diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 3d3c16c8cb6..cb5c58c6fea 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -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) -> Map { - let mut mapped_params = Map::new(); - for (param, value) in non_default_params { - if self.supported_ocr_params().contains(¶m.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(¶m.as_str())) + .map(|(param, value)| (param.clone(), value.clone())) + .collect() } fn transform_ocr_request( diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index b26a7925e8a..f983f698ecf 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -281,8 +281,8 @@ fn page_dimensions(page: &Map) -> 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 } diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 6a8a38204a9..ca4bc8b87af 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -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) -> Map { @@ -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", diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index ee095447028..2fd9bc44630 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -208,8 +208,8 @@ fn ocr_data_from_content(content: Value, usage: Option, 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 } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e89f9cd9aa9..e933237e7d9 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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 } diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs new file mode 100644 index 00000000000..7061e374416 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -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>>, +} + +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 { + self.events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } +} + +struct FunctionTraceLayer { + trace: FunctionTrace, +} + +impl Layer 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, + }, + ] + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 65ce75a7768..6be7ea2e900 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,5 +1,6 @@ mod diagnostics; mod errors; +mod function_trace; mod marshal; mod routes; diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index d2674a80a10..9dc356d2d40 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -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", diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 5588c400972..b2a4238e7b9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -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, timeout_seconds: Option, + trace: Option, }, prepare = prepare_ocr, errors = core_error_to_pyerr, diff --git a/tests/sdk_function_trace/__init__.py b/tests/sdk_function_trace/__init__.py new file mode 100644 index 00000000000..da62b8041f6 --- /dev/null +++ b/tests/sdk_function_trace/__init__.py @@ -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", +] diff --git a/tests/sdk_function_trace/harness.py b/tests/sdk_function_trace/harness.py new file mode 100644 index 00000000000..8f707402449 --- /dev/null +++ b/tests/sdk_function_trace/harness.py @@ -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}") diff --git a/tests/sdk_function_trace/mock_provider.py b/tests/sdk_function_trace/mock_provider.py new file mode 100644 index 00000000000..37eca665586 --- /dev/null +++ b/tests/sdk_function_trace/mock_provider.py @@ -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}") diff --git a/tests/sdk_function_trace/profiler.py b/tests/sdk_function_trace/profiler.py new file mode 100644 index 00000000000..77c78582087 --- /dev/null +++ b/tests/sdk_function_trace/profiler.py @@ -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) diff --git a/tests/test_litellm/ocr/test_function_trace.py b/tests/test_litellm/ocr/test_function_trace.py new file mode 100644 index 00000000000..9d6759a5848 --- /dev/null +++ b/tests/test_litellm/ocr/test_function_trace.py @@ -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