diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index c2dff805772..6da5fc07e80 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -96,6 +96,7 @@ jobs: - shard: misc artifact-name: misc test-path: >- + tests/sdk_function_trace tests/test_litellm/batches tests/test_litellm/secret_managers tests/test_litellm/a2a_protocol diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index dd41cf0e84b..71c9bb95ec0 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" @@ -1416,6 +1422,7 @@ dependencies = [ "tokio", "tokio-tungstenite", "tower", + "tracing", ] [[package]] @@ -1435,6 +1442,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "tracing", ] [[package]] @@ -1447,8 +1455,11 @@ dependencies = [ "litellm-python-interop", "pyo3", "pyo3-async-runtimes", + "serde", "serde_json", "tokio", + "tracing", + "tracing-subscriber", ] [[package]] @@ -2276,6 +2287,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" @@ -2414,6 +2434,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" @@ -2662,6 +2691,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 c447d915abe..a096643d0d3 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/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 541beabe170..e3dbdf24ce6 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -14,6 +14,7 @@ path = "src/main.rs" required-features = ["server"] [dependencies] +tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index e0ce165dc93..c1fb328893b 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -32,6 +32,7 @@ pub(super) fn truncate_error_body(body: &str) -> String { format!("{truncated}... (truncated)") } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn ocr_provider_config( provider: &str, model: &str, @@ -73,12 +74,6 @@ pub(super) fn string_headers( .collect() } -pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { - headers - .iter() - .any(|(key, _)| key.eq_ignore_ascii_case(name)) -} - fn document_url_field(document: &Value) -> Result, Error> { let Some(object) = document.as_object() else { return Ok(None); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 815bc84363a..856d9571201 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,12 +1,19 @@ use litellm_core::error::Error; +use litellm_core::http_utils::http_request; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; use super::common_utils::{poll_document_intelligence, truncate_error_body}; -use super::types::ProviderOcrRequest; +use super::hooks::OcrLifecycleHooks; +use super::types::PreparedOcrRequest; use crate::client::http_client; -pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result { +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) async fn execute_ocr_provider_call( + request: PreparedOcrRequest, + hooks: &OcrLifecycleHooks, +) -> Result { + let request = hooks.prepare_provider_request(request).await?; let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -15,8 +22,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Re request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|err| Error::Network(err.to_string()))?; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 401e26d3b29..f8c4f8fe8c5 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,13 +1,10 @@ use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::Error; -use litellm_core::ocr::transformation::OcrAuthStrategy; use serde_json::{Map, Value, json}; use std::future::Future; use std::pin::Pin; -use super::common_utils::{ - convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, -}; +use super::common_utils::{convert_document_url_to_data_uri, string_headers}; use super::types::{PreparedOcrRequest, ProviderOcrRequest}; use crate::integrations::custom_guardrail::{ CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, @@ -62,6 +59,10 @@ impl OcrLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?; + let optional_params = match &request.config { + Ok(config) => config.map_ocr_params(&optional_params), + Err(_) => optional_params, + }; Ok(PreparedOcrRequest { document, optional_params, @@ -69,25 +70,23 @@ impl OcrLifecycleHooks { }) } - async fn prepare_provider_request( + pub(crate) async fn prepare_provider_request( &self, request: PreparedOcrRequest, ) -> Result { - let config = ocr_provider_config(&request.custom_llm_provider, &request.model) - .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?; + let config = request.config?; let env_lookup = |key: &str| std::env::var(key).ok(); - let headers = string_headers(request.extra_headers)?; - let auth_strategy = config.auth_strategy(); - let api_key = (!has_header(&headers, auth_strategy.header_name())) - .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) - .transpose()?; + let upstream_headers = config.validate_environment( + string_headers(request.extra_headers)?, + request.api_key.as_deref(), + &env_lookup, + )?; let url = config.complete_url( request.api_base.as_deref(), &request.model, &request.optional_params, &env_lookup, )?; - let filtered_params = config.map_ocr_params(&request.optional_params); let model = request.model.clone(); let custom_llm_provider = request.custom_llm_provider.clone(); let document = if config.requires_data_uri_document() { @@ -96,9 +95,8 @@ impl OcrLifecycleHooks { request.document }; let body = config - .transform_ocr_request(&request.model, document, filtered_params)? + .transform_ocr_request(&request.model, document, request.optional_params)? .data; - let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); let body = self .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) .await?; @@ -167,9 +165,9 @@ impl OcrLifecycleHooks { } } -impl CallLifecycleHooks for OcrLifecycleHooks { +impl CallLifecycleHooks for OcrLifecycleHooks { type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; - type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>; + type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; type SuccessFuture<'a> = OcrLogFuture<'a>; type FailureFuture<'a> = OcrLogFuture<'a>; @@ -186,7 +184,7 @@ impl CallLifecycleHooks for OcrLi _context: &'a CallLifecycleContext, request: PreparedOcrRequest, ) -> Self::DuringCallFuture<'a> { - Box::pin(async move { self.prepare_provider_request(request).await }) + Box::pin(async move { Ok(request) }) } fn async_log_success_event<'a>( @@ -247,21 +245,6 @@ impl CallLifecycleHooks for OcrLi } } -fn upstream_headers( - headers: &[(String, String)], - auth_strategy: OcrAuthStrategy, - api_key: Option<&str>, -) -> Vec<(String, String)> { - api_key - .map(|api_key| match auth_strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), - }) - .into_iter() - .chain(headers.iter().cloned()) - .collect() -} - fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { GuardrailContext { call_type: CallType::Ocr, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index b59ab626fd3..d9230af1c59 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -13,10 +13,13 @@ pub use types::OcrRequest; use handler::execute_ocr_provider_call; use prepare::{PreparedOcrCall, prepare_ocr_call}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn ocr(request: OcrRequest<'_>) -> Result { let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); CallLifecycle::default() - .run_request(request, &hooks, execute_ocr_provider_call) + .run_request(request, &hooks, |request| { + execute_ocr_provider_call(request, &hooks) + }) .await } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 6231393c889..fedacc62760 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -3,6 +3,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; +use super::common_utils::ocr_provider_config; use super::hooks::OcrLifecycleHooks; use super::types::{OcrRequest, PreparedOcrRequest}; use crate::integrations::custom_guardrail::CustomGuardrailRunner; @@ -13,6 +14,7 @@ pub(crate) struct PreparedOcrCall { pub(crate) hooks: OcrLifecycleHooks, } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { let call_id = request .litellm_call_id @@ -25,9 +27,25 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { }); let model = provider_info.model.to_string(); let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + let config = ocr_provider_config(&custom_llm_provider, &model) + .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone())); + let optional_params = match &config { + Ok(config) => { + let supported = config.supported_ocr_params(); + config.map_ocr_params( + &request + .optional_params + .into_iter() + .filter(|(name, _)| supported.contains(&name.as_str())) + .collect(), + ) + } + Err(_) => request.optional_params, + }; PreparedOcrCall { request: PreparedOcrRequest { + config, model, custom_llm_provider, litellm_call_id: call_id, @@ -35,7 +53,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { api_key: request.api_key.map(str::to_string), api_base: request.api_base.map(str::to_string), extra_headers: request.extra_headers, - optional_params: request.optional_params, + optional_params, timeout: request.timeout, }, hooks: OcrLifecycleHooks::new( diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 8c3f0425149..85e4c408045 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -2,12 +2,13 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use litellm_core::error::Error; +use litellm_core::http_utils::has_header; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; +use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body}; use super::{OcrRequest, ocr}; use crate::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index bde734a4dd1..95e551d79ca 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -25,6 +25,7 @@ pub struct OcrRequest<'a> { } pub(crate) struct PreparedOcrRequest { + pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>, pub(crate) model: String, pub(crate) custom_llm_provider: String, pub(crate) litellm_call_id: String, diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index ab8050734f2..389dbd49505 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -11,6 +11,7 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true +tracing.workspace = true sha2.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 30ba0da5e68..9a96b9d1140 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,11 +1,12 @@ use serde_json::Value; use crate::error::Error; -use crate::http_utils::truncate_error_body; +use crate::http_utils::{http_request, truncate_error_body}; use super::client::http_client; use super::types::ProviderAudioTranscriptionRequest; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { @@ -19,8 +20,7 @@ pub async fn execute_audio_transcription_provider_call( if let Some(duration) = request.timeout { request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|error| Error::Network(error.to_string()))?; let status = response.status(); diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index b71748082bf..31b6de4b3e4 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -11,6 +11,7 @@ pub use handler::execute_audio_transcription_provider_call; pub use prepare::prepare_audio_transcription_provider_call; pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 6288e96b380..bbef97341a9 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -7,6 +7,7 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider} use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { #[cfg(feature = "bedrock-auth")] if provider == "bedrock" { @@ -16,6 +17,7 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv None } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub fn prepare_audio_transcription_provider_call( request: AudioTranscriptionRequest<'_>, ) -> Result { diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index 16a28fbcac0..aa9846427dc 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -15,6 +15,7 @@ pub enum AudioTranscriptionAuth { pub trait AudioTranscriptionProviderConfig: Sync { fn supported_transcription_params(&self) -> &'static [&'static str]; + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn map_transcription_params(&self, params: &Map) -> Map { params .iter() diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index ca51471eb7c..69e5f175ad5 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -7,6 +7,7 @@ use super::transformation::ChatCompletionsProviderConfig; const HEADER_CONTEXT: &str = "chat completions"; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn chat_completions_provider_config( provider: &str, ) -> Option<&'static dyn ChatCompletionsProviderConfig> { diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 7e2731442cc..96d001e2892 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,17 +1,21 @@ use serde_json::Value; use crate::error::Error; -use crate::http_utils::truncate_error_body; +use crate::http_utils::{http_request, truncate_error_body}; use super::client::http_client; +use super::prepare::prepare_provider_request; use super::transformation::ChatCompletionsAuth; use super::types::{ ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData, + ResolvedChatCompletionsRequest, }; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn execute_chat_completions_provider_call( - request: ProviderChatCompletionsRequest, + request: ResolvedChatCompletionsRequest<'_>, ) -> Result { + let request = prepare_provider_request(request)?; let body = serde_json::to_vec(&request.body).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize chat completions request: {err}" @@ -27,7 +31,7 @@ pub(super) async fn execute_chat_completions_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder.send().await.map_err(|err| { + let response = http_request(request_builder).await.map_err(|err| { // Failing to establish the connection means the request never went out, // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 0d009d36d16..32dea17d202 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -19,13 +19,14 @@ pub mod types; use serde_json::{Map, Value}; use handler::execute_chat_completions_provider_call; -use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config}; +use prepare::{parse_messages, resolve_provider_config, resolve_request}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn chat_completions( request: ChatCompletionsRequest<'_>, ) -> Result { - execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await + execute_chat_completions_provider_call(resolve_request(request)?).await } /// Whether the core would accept this request, without resolving credentials or diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 142b2f2aaed..3be2ba21de4 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -6,7 +6,10 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider} use super::common_utils::{chat_completions_provider_config, string_headers}; use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; -use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest}; +use super::types::{ + ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest, + ResolvedChatCompletionsRequest, +}; pub(super) fn resolve_provider_config<'a>( model: &'a str, @@ -34,12 +37,10 @@ pub(super) fn parse_messages(messages: Value) -> Result, Error> .map_err(|err| Error::InvalidRequest(format!("invalid chat completions messages: {err}"))) } -pub(super) fn prepare_chat_completions_call( +pub(super) fn resolve_request( request: ChatCompletionsRequest<'_>, -) -> Result { +) -> Result, Error> { let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; - let env_lookup = |key: &str| std::env::var(key).ok(); - let messages = parse_messages(request.messages)?; if messages.is_empty() { return Err(Error::InvalidRequest( @@ -49,11 +50,29 @@ pub(super) fn prepare_chat_completions_call( if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) { return Err(Error::Unsupported(reason.0)); } + Ok(ResolvedChatCompletionsRequest { + model, + config, + messages, + optional_params: request.optional_params, + api_key: request.api_key, + api_base: request.api_base, + extra_headers: request.extra_headers, + timeout: request.timeout, + }) +} - let mut headers = string_headers(request.extra_headers)?; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +fn validate_environment( + request: &ResolvedChatCompletionsRequest<'_>, + model: &str, + config: &dyn ChatCompletionsProviderConfig, +) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> { + let env_lookup = |key: &str| std::env::var(key).ok(); + let mut headers = string_headers(request.extra_headers.clone())?; let auth = config.auth( request.api_key, - &model, + model, &request.optional_params, &env_lookup, )?; @@ -94,7 +113,16 @@ pub(super) fn prepare_chat_completions_call( headers.push(((*name).to_string(), (*value).to_string())); } } + Ok((headers, auth)) +} +pub(super) fn prepare_provider_request( + request: ResolvedChatCompletionsRequest<'_>, +) -> Result { + let (headers, auth) = validate_environment(&request, &request.model, request.config)?; + let model = request.model; + let config = request.config; + let env_lookup = |key: &str| std::env::var(key).ok(); let url = config.complete_url( request.api_base, &model, @@ -102,7 +130,7 @@ pub(super) fn prepare_chat_completions_call( &env_lookup, )?; let transformed = - config.transform_request(&model, messages, request.optional_params.clone())?; + config.transform_request(&model, request.messages, request.optional_params.clone())?; Ok(ProviderChatCompletionsRequest { model, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index 2858d180e27..f8594dee447 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -2,9 +2,15 @@ use serde_json::{Map, Value, json}; use crate::error::Error; -use super::prepare::prepare_chat_completions_call; +use super::prepare::{prepare_provider_request, resolve_request}; use super::transformation::ChatCompletionsAuth; -use super::types::ChatCompletionsRequest; +use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; + +fn prepare_chat_completions_call( + request: ChatCompletionsRequest<'_>, +) -> Result { + prepare_provider_request(resolve_request(request)?) +} fn request<'a>( model: &'a str, diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index a0868209305..d7b9704c46c 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -62,9 +62,8 @@ pub trait ChatCompletionsProviderConfig: Sync { false } - /// Provider parameter names (post-mapping) the Rust path knows how to place - /// in the upstream body. Anything outside this set declines the request. - fn supported_params(&self) -> &'static [&'static str]; + /// Supported OpenAI parameter names paired with their provider names. + fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)]; /// Parameters consumed as call configuration (credentials, endpoints) /// rather than placed in the body. Accepted, never serialized. @@ -78,7 +77,7 @@ pub trait ChatCompletionsProviderConfig: Sync { optional_params: &Map, ) -> Option { unsupported_param( - self.supported_params(), + self.supported_openai_params(), self.config_params(), optional_params, ) @@ -100,7 +99,7 @@ pub trait ChatCompletionsProviderConfig: Sync { } pub fn unsupported_param( - supported: &'static [&'static str], + supported: &'static [(&'static str, &'static str)], config: &'static [&'static str], optional_params: &Map, ) -> Option { @@ -115,7 +114,9 @@ pub fn unsupported_param( .keys() .any(|key| { key != STREAM_PARAM - && !supported.contains(&key.as_str()) + && !supported + .iter() + .any(|(_, provider_name)| *provider_name == key) && !config.contains(&key.as_str()) }) .then_some(Unsupported("unrecognized request parameter")) diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 35dd543a986..3238d09b6b5 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -22,6 +22,17 @@ pub struct ChatCompletionsRequest<'a> { pub timeout: Option, } +pub(super) struct ResolvedChatCompletionsRequest<'a> { + pub(super) model: String, + pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) messages: Vec, + pub(super) optional_params: Map, + pub(super) api_key: Option<&'a str>, + pub(super) api_base: Option<&'a str>, + pub(super) extra_headers: Option>, + pub(super) timeout: Option, +} + pub(super) struct ProviderChatCompletionsRequest { pub(super) model: String, pub(super) config: &'static dyn ChatCompletionsProviderConfig, diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 10661fadf96..3633130528d 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -5,6 +5,13 @@ use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; use crate::error::{Error, json_type_name}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub async fn http_request( + request: reqwest::RequestBuilder, +) -> Result { + request.send().await +} + /// Bound an upstream error body before it crosses a host boundary, so provider /// bodies stay data-minimized. pub fn truncate_error_body(body: &str) -> String { diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 8dfdb2e361a..8f0f6652fa4 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -10,6 +10,7 @@ pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_b const HEADER_CONTEXT: &str = "messages"; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn messages_provider_config( provider: &str, ) -> Option<&'static dyn AnthropicMessagesProviderConfig> { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 13a65d86131..61ff81bcdc8 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,13 +1,17 @@ use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; use crate::error::Error; +use crate::http_utils::http_request; use super::client::http_client; use super::common_utils::truncate_error_body; -use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest}; +use super::prepare::prepare_provider_request; +use super::types::{AnthropicMessagesResponse, MessagesRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn execute_messages_provider_call( - request: ProviderMessagesRequest, + request: MessagesRequest<'_>, ) -> Result { + let request = prepare_provider_request(request)?; let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -16,8 +20,7 @@ pub(super) async fn execute_messages_provider_call( request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|err| Error::Network(err.to_string()))?; @@ -40,8 +43,9 @@ pub(super) async fn execute_messages_provider_call( } pub(super) async fn execute_messages_provider_stream( - request: ProviderMessagesRequest, + request: MessagesRequest<'_>, ) -> Result { + let request = prepare_provider_request(request)?; if request.provider != ANTHROPIC_MESSAGES_PROVIDER { return Err(Error::InvalidRequest( "streaming messages is not supported for this provider".to_string(), @@ -56,8 +60,7 @@ pub(super) async fn execute_messages_provider_stream( request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index ee2877e61fc..cfa8bda1104 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -16,15 +16,15 @@ pub mod transformation; pub mod types; use handler::{execute_messages_provider_call, execute_messages_provider_stream}; -use prepare::prepare_messages_call; use types::{AnthropicMessagesResponse, MessagesRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn messages(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_call(prepare_messages_call(request)?).await + execute_messages_provider_call(request).await } pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_stream(prepare_messages_call(request)?).await + execute_messages_provider_stream(request).await } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 3b253ac3766..ec83d03f535 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -2,10 +2,11 @@ use crate::error::Error; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; -use super::transformation::MessagesAuthStrategy; +use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use super::types::{MessagesRequest, ProviderMessagesRequest}; +use serde_json::{Map, Value}; -pub(super) fn prepare_messages_call( +pub(super) fn prepare_provider_request( request: MessagesRequest<'_>, ) -> Result { let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) @@ -29,13 +30,46 @@ pub(super) fn prepare_messages_call( .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; let env_lookup = |key: &str| std::env::var(key).ok(); - let mut headers = string_headers(request.extra_headers)?; + let headers = + validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?; + + let typed_request = serde_json::from_value(request.body).map_err(|err| { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = config.transform_request(typed_request)?; + let body = serde_json::to_value(transformed).map_err(|err| { + Error::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) + })?; + + let url = config.complete_url(request.api_base, &model, &env_lookup)?; + + Ok(ProviderMessagesRequest { + provider: provider.to_string(), + model, + config, + url, + body, + upstream_headers: headers, + timeout: request.timeout, + }) +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +fn validate_environment( + config: &dyn AnthropicMessagesProviderConfig, + extra_headers: Option>, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result, Error> { + let mut headers = string_headers(extra_headers)?; let auth_strategy = config.auth_strategy(); let already_authorized = has_header(&headers, auth_strategy.header_name()) || (config.accepts_bearer_auth() && has_bearer_auth(&headers)); if !already_authorized { - let api_key = config.resolve_api_key(request.api_key, &env_lookup)?; + let api_key = config.resolve_api_key(api_key, env_lookup)?; let auth_header = match auth_strategy { MessagesAuthStrategy::Bearer => { ("authorization".to_string(), format!("Bearer {api_key}")) @@ -51,24 +85,5 @@ pub(super) fn prepare_messages_call( } } - let url = config.complete_url(request.api_base, &model, &env_lookup)?; - let typed_request = serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) - })?; - let transformed = config.transform_request(typed_request)?; - let body = serde_json::to_value(transformed).map_err(|err| { - Error::InvalidRequest(format!( - "failed to serialize Anthropic messages request: {err}" - )) - })?; - - Ok(ProviderMessagesRequest { - provider: provider.to_string(), - model, - config, - url, - body, - upstream_headers: headers, - timeout: request.timeout, - }) + Ok(headers) } diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 673a5728aca..a5904c085a0 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -45,6 +45,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { ] } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_request( &self, request: AnthropicMessagesRequest, @@ -52,6 +53,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { Ok(request) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_response( &self, _model: &str, diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 3d3c16c8cb6..ad484c8f968 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -27,6 +27,7 @@ pub enum OcrResponseHandling { pub trait OcrProviderConfig: Sync { fn 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 { @@ -64,6 +65,25 @@ pub trait OcrProviderConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn validate_environment( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result, Error> { + let strategy = self.auth_strategy(); + if crate::http_utils::has_header(&headers, strategy.header_name()) { + return Ok(headers); + } + let api_key = self.resolve_api_key(api_key, env_lookup)?; + let auth_header = match strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(name) => (name.to_string(), api_key), + }; + Ok(std::iter::once(auth_header).chain(headers).collect()) + } + fn auth_strategy(&self) -> OcrAuthStrategy { OcrAuthStrategy::Bearer } diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index 97cc48aa6f2..a7d5a8ad0cf 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -27,7 +27,12 @@ use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage /// per-model gate inside `transform_request`, the function this route replaces. /// Forwarding it would send `top_k` to a model that removed sampling params and /// take a 400 after the call, where Python drops it and succeeds. -const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"]; +const SUPPORTED_PARAMS: &[(&str, &str)] = &[ + ("max_tokens", "max_tokens"), + ("temperature", "temperature"), + ("top_p", "top_p"), + ("stop", "stop_sequences"), +]; pub struct AnthropicChatCompletionsConfig; @@ -112,7 +117,8 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { }) } - fn supported_params(&self) -> &'static [&'static str] { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS } @@ -121,7 +127,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { messages: &[ChatMessage], optional_params: &Map, ) -> Option { - unsupported_param(SUPPORTED_PARAMS, &[], optional_params) + unsupported_param(self.supported_openai_params(), &[], optional_params) .or_else(|| messages.iter().find_map(unsupported_message)) // Anthropic rejects a request whose first turn is not a user turn. // Python only repairs that under `litellm.modify_params`, which the @@ -132,6 +138,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_request( &self, model: &str, @@ -143,6 +150,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_response( &self, _model: &str, diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index 8fcc0f36c7d..f31b961e78a 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -47,6 +47,7 @@ pub fn complete_anthropic_url( } impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 70dad0300f1..b8ca10461fb 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -142,6 +142,7 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess } impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index bb4f6afe5f9..9bf1f73a74d 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -46,10 +46,12 @@ fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a } impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_transcription_params(&self) -> &'static [&'static str] { SUPPORTED_PARAMS } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_transcription_request( &self, _model: &str, @@ -83,6 +85,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_transcription_response( &self, _model: &str, diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index ef5f44b4a14..7be3d108d44 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -23,11 +23,12 @@ use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT /// `additionalModelRequestFields` for Anthropic base models and to /// `inferenceConfig` otherwise, and that branch reads the model catalog the /// core cannot see. -const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"]; - -/// Params that belong in `inferenceConfig`, in the order Python's -/// `AmazonConverseConfig` declares them, so bodies compare cleanly. -const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS; +const SUPPORTED_PARAMS: &[(&str, &str)] = &[ + ("max_tokens", "maxTokens"), + ("temperature", "temperature"), + ("top_p", "topP"), + ("stop", "stopSequences"), +]; const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint"; @@ -66,7 +67,7 @@ fn converse_body(conversation: &Conversation, params: &Map) -> Va }) .collect(); - let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| { + let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| { params .get(*name) .map(|value| ((*name).to_string(), value.clone())) @@ -162,7 +163,8 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { &[("Content-Type", "application/json")] } - fn supported_params(&self) -> &'static [&'static str] { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] { SUPPORTED_PARAMS } @@ -175,32 +177,36 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { messages: &[ChatMessage], optional_params: &Map, ) -> Option { - unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params) - .or_else(|| messages.iter().find_map(unsupported_message)) - // Python's Converse translation drops blank text blocks instead of - // substituting the placeholder the shared conversation builder - // applies, so decline blank text rather than diverge. - .or_else(|| { - messages - .iter() - .any(has_blank_text) - .then_some(Unsupported("blank message text")) - }) - // Converse has no assistant prefill: Python inserts a continue turn - // when a conversation opens or closes on an assistant message, and - // only under `litellm.modify_params`, which the core cannot see. - // Declining both ends also keeps the shared builder's final - // assistant right-strip (an Anthropic rule) unreachable here. - .or_else(|| { - let conversation = build_conversation(messages); - let ends_on_assistant = conversation - .turns - .last() - .is_some_and(|turn| turn.role == TurnRole::Assistant); - (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( - "conversation does not run user turn to user turn", - )) - }) + unsupported_param( + self.supported_openai_params(), + CONFIG_PARAMS, + optional_params, + ) + .or_else(|| messages.iter().find_map(unsupported_message)) + // Python's Converse translation drops blank text blocks instead of + // substituting the placeholder the shared conversation builder + // applies, so decline blank text rather than diverge. + .or_else(|| { + messages + .iter() + .any(has_blank_text) + .then_some(Unsupported("blank message text")) + }) + // Converse has no assistant prefill: Python inserts a continue turn + // when a conversation opens or closes on an assistant message, and + // only under `litellm.modify_params`, which the core cannot see. + // Declining both ends also keeps the shared builder's final + // assistant right-strip (an Anthropic rule) unreachable here. + .or_else(|| { + let conversation = build_conversation(messages); + let ends_on_assistant = conversation + .turns + .last() + .is_some_and(|turn| turn.role == TurnRole::Assistant); + (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( + "conversation does not run user turn to user turn", + )) + }) } fn transform_request( 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..9648321d7ff 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 { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn 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, @@ -134,6 +137,7 @@ impl OcrProviderConfig for MistralOcrConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, @@ -153,6 +157,7 @@ impl OcrProviderConfig for MistralOcrConfig { } } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub fn supported_ocr_params() -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() } @@ -161,6 +166,7 @@ pub fn map_ocr_params(non_default_params: &Map) -> Map Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 498003de149..275e00300d4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -16,6 +16,9 @@ extension-module = ["pyo3/extension-module"] panic-test = [] [dependencies] +serde.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs new file mode 100644 index 00000000000..07b2836b838 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -0,0 +1 @@ +pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs new file mode 100644 index 00000000000..cc153a89b8f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -0,0 +1,23 @@ +use litellm_python_interop::release_count; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +#[pyfunction] +fn gil_stats(py: Python<'_>) -> PyResult> { + let stats = PyDict::new(py); + stats.set_item("releases", release_count())?; + Ok(stats.into_any().unbind()) +} + +#[cfg(feature = "panic-test")] +#[pyfunction] +fn _panic_for_test() { + panic!("intentional PyO3 panic smoke test"); +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + #[cfg(feature = "panic-test")] + module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; + Ok(()) +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs new file mode 100644 index 00000000000..914e2e1e033 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -0,0 +1,61 @@ +use litellm_core::error::Error; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; + +pyo3::create_exception!( + _native, + RustBridgeDeclined, + pyo3::exceptions::PyException, + "The route declined before calling the provider, so the host may retry on its own path." +); + +pyo3::create_exception!( + _native, + RustUpstreamError, + pyo3::exceptions::PyException, + "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." +); + +pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { + match err { + Error::Auth(message) => PyValueError::new_err(message), + Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => PyValueError::new_err(err.to_string()), + other => PyRuntimeError::new_err(other.to_string()), + } +} + +/// Map a core error for a route whose host keeps a Python implementation. +/// +/// The distinction the host needs is whether the provider was already called. +/// Everything raised before the request goes out is safe for the host to retry +/// on its own path; anything after it is not, because the provider has already +/// done the work and billed for it. +pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { + match err { + Error::Unsupported(_) + | Error::Auth(_) + | Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) + | Error::Routing(_) + // Nothing reached the provider, so serving it on Python cannot double + // bill and is the only way the caller gets an answer at all. + | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + Error::Http { status, body } => { + RustUpstreamError::new_err((status, format!("{status}: {body}"))) + } + Error::Network(message) | Error::InvalidResponse(message) => { + RustUpstreamError::new_err((0u16, message)) + } + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + let py = module.py(); + module.add("RustBridgeDeclined", py.get_type::())?; + module.add("RustUpstreamError", py.get_type::()) +} 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..420d237c79d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -0,0 +1,216 @@ +use std::future::Future; +use std::sync::{Arc, Mutex}; + +use serde::Serialize; +use tracing::instrument::WithSubscriber; +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}; + +use crate::constants::FUNCTION_TRACE_TARGET; + +#[derive(Serialize)] +#[serde(untagged)] +pub(crate) enum TraceResponse { + Plain(T), + Traced { + response: T, + trace: Vec, + }, +} + +pub(crate) async fn trace_call( + future: impl Future>, + enabled: bool, +) -> Result, E> { + if !enabled { + return future.await.map(TraceResponse::Plain); + } + let trace = FunctionTrace::default(); + let response = future.with_subscriber(trace.dispatcher()).await?; + Ok(TraceResponse::Traced { + response, + trace: trace.events(), + }) +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct FunctionTraceEvent { + pub function: &'static str, + pub depth: usize, +} + +#[derive(Clone, Default)] +pub struct FunctionTrace { + events: Arc>>, +} + +impl FunctionTrace { + pub fn dispatcher(&self) -> Dispatch { + let filter = filter_fn(|metadata| { + metadata.is_span() + && metadata.target() == FUNCTION_TRACE_TARGET + && *metadata.level() == Level::TRACE + }) + .with_max_level_hint(LevelFilter::TRACE); + Dispatch::new( + Registry::default().with( + FunctionTraceLayer { + trace: self.clone(), + } + .with_filter(filter), + ), + ) + } + + pub 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::*; + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn outer() { + tokio::task::yield_now().await; + inner().await; + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn inner() { + tokio::task::yield_now().await; + } + + #[tokio::test] + async fn concurrent_futures_keep_separate_traces_across_yields() { + use tracing::instrument::WithSubscriber; + + let first = FunctionTrace::default(); + let second = FunctionTrace::default(); + let outside = FunctionTrace::default(); + + async { + tokio::join!( + outer().with_subscriber(first.dispatcher()), + inner().with_subscriber(second.dispatcher()), + ); + inner().await; + } + .with_subscriber(outside.dispatcher()) + .await; + + assert_eq!( + first.events(), + vec![ + FunctionTraceEvent { + function: "outer", + depth: 0 + }, + FunctionTraceEvent { + function: "inner", + depth: 1 + }, + ], + ); + assert_eq!( + second.events(), + vec![FunctionTraceEvent { + function: "inner", + depth: 0 + }], + ); + assert_eq!( + outside.events(), + vec![FunctionTraceEvent { + function: "inner", + depth: 0 + }], + ); + } + + #[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: FUNCTION_TRACE_TARGET, "same_name"); + let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level"); + let _second = tracing::trace_span!(target: FUNCTION_TRACE_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: FUNCTION_TRACE_TARGET, "outer"); + let _outer_guard = outer.enter(); + let _inner = tracing::trace_span!(target: FUNCTION_TRACE_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 2e2624acbe1..1516fcd7d19 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,142 +1,16 @@ -use std::collections::HashMap; -use std::time::Duration; +mod constants; +mod diagnostics; +mod errors; +pub mod function_trace; +mod marshal; +mod routes; -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; -use litellm_core::audio_transcription::{ - AudioTranscriptionRequest, audio_transcription as run_audio_transcription, -}; -use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; -use litellm_core::chat_completions::{ - chat_completions as run_chat_completions, chat_completions_decline_reason, -}; -use litellm_core::error::Error; -use litellm_core::messages::messages as run_messages; -use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; -use litellm_python_interop::{from_py, release_count, release_gil, to_py}; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyAny, PyDict}; -use serde_json::{Map, Value}; +use pyo3::types::PyAny; -pyo3::create_exception!( - _native, - RustBridgeDeclined, - pyo3::exceptions::PyException, - "The route declined before calling the provider, so the host may retry on its own path." -); - -pyo3::create_exception!( - _native, - RustUpstreamError, - pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." -); - -type MarshaledOcrInputs = ( - Value, - Option>, - Map, - Option, -); - -fn messages_response_to_py( - py: Python<'_>, - response: AnthropicMessagesResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn chat_completions_response_to_py( - py: Python<'_>, - response: ChatCompletionsResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn core_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Auth(message) => PyValueError::new_err(message), - Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) => PyValueError::new_err(err.to_string()), - other => PyRuntimeError::new_err(other.to_string()), - } -} - -/// Map a core error for a route whose host keeps a Python implementation. -/// -/// The distinction the host needs is whether the provider was already called. -/// Everything raised before the request goes out is safe for the host to retry -/// on its own path; anything after it is not, because the provider has already -/// done the work and billed for it. -fn chat_completions_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Unsupported(_) - | Error::Auth(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::Routing(_) - // Nothing reached the provider, so serving it on Python cannot double - // bill and is the only way the caller gets an answer at all. - | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - Error::Http { status, body } => { - RustUpstreamError::new_err((status, format!("{status}: {body}"))) - } - Error::Network(message) | Error::InvalidResponse(message) => { - RustUpstreamError::new_err((0u16, message)) - } - } -} - -fn optional_object_to_map( - py: Python<'_>, - name: &'static str, - value: Option>, -) -> PyResult> { - match value { - Some(value) => match from_py(value.bind(py))? { - Value::Object(map) => Ok(map), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), - }, - None => Ok(Map::new()), - } -} - -fn optional_timeout(timeout_seconds: Option) -> Option { - timeout_seconds.and_then(|secs| { - if secs.is_finite() && secs > 0.0 { - Some(Duration::from_secs_f64(secs)) - } else { - None - } - }) -} - -fn marshal_headers( - py: Python<'_>, - headers: Option>, -) -> PyResult> { - let value = match headers { - Some(headers) => from_py(headers.bind(py))?, - None => Value::Object(Map::new()), - }; - let Value::Object(headers) = value else { - return Err(PyValueError::new_err("headers must be a dict")); - }; - headers - .into_iter() - .map(|(name, value)| { - value - .as_str() - .map(|value| (name, value.to_string())) - .ok_or_else(|| PyValueError::new_err("header values must be strings")) - }) - .collect() -} +use crate::errors::core_error_to_pyerr; +use crate::marshal::{marshal_headers, optional_timeout}; #[pyclass] struct ResponsesWebSocketConnection { @@ -186,445 +60,48 @@ impl ResponsesWebSocketConnection { } } -fn marshal_inputs( - py: Python<'_>, - document: Py, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - let document = from_py(document.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - - Ok((document, extra_headers, optional_params, timeout)) -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn ocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (document, extra_headers, optional_params, timeout) = marshal_inputs( - py, - document, - extra_headers, - optional_params, - timeout_seconds, - )?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - })) - }); - - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn aocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (document, extra_headers, optional_params, timeout) = marshal_inputs( - py, - document, - extra_headers, - optional_params, - timeout_seconds, - )?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }) - .await - .map_err(core_error_to_pyerr)?; - - Python::attach(|py| to_py(py, &value)) - }) -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn transcription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( - AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }, - )) - }); - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn atranscription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_audio_transcription(AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }) - .await - .map_err(core_error_to_pyerr)?; - Python::attach(|py| to_py(py, &value)) - }) -} - -type MarshaledMessagesInputs = (Value, Option>, Option); - -fn marshal_messages_inputs( - py: Python<'_>, - body: Py, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let body: Value = from_py(body.bind(py))?; - if !body.is_object() { - return Err(PyValueError::new_err("body must be a dict")); - } - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - Ok((body, extra_headers, optional_timeout(timeout_seconds))) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn messages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (body, extra_headers, timeout) = - marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { - model: &model, - body, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - })) - }); - - match result { - Ok(response) => messages_response_to_py(py, response), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn amessages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (body, extra_headers, timeout) = - marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = run_messages(MessagesRequest { - model: &model, - body, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await - .map_err(core_error_to_pyerr)?; - - Python::attach(|py| messages_response_to_py(py, response)) - }) -} - -type MarshaledChatCompletionsInputs = ( - Value, - Map, - Option>, - Option, -); - -fn marshal_chat_completions_inputs( - py: Python<'_>, - messages: Py, - optional_params: Option>, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let messages: Value = from_py(messages.bind(py))?; - if !messages.is_array() { - return Err(PyValueError::new_err("messages must be a list")); - } - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - Ok(( - messages, - optional_params, - extra_headers, - optional_timeout(timeout_seconds), - )) -} - -/// The decline reason for this request, or `None` when the Rust path accepts -/// it. Resolves no credentials and performs no I/O, so a host can ask before -/// committing to either path. -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] -fn chat_completions_decline( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - custom_llm_provider: Option, -) -> PyResult> { - let messages = from_py(messages.bind(py))?; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - Ok(chat_completions_decline_reason( - &model, - custom_llm_provider.as_deref(), - messages, - &optional_params, - ) - .map(str::to_string)) -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn chat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( - py, - messages, - optional_params, - extra_headers, - timeout_seconds, - )?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( - ChatCompletionsRequest { - model: &model, - messages, - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }, - )) - }); - - match result { - Ok(response) => chat_completions_response_to_py(py, response), - Err(err) => Err(chat_completions_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn achat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( - py, - messages, - optional_params, - extra_headers, - timeout_seconds, - )?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = run_chat_completions(ChatCompletionsRequest { - model: &model, - messages, - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await - .map_err(chat_completions_error_to_pyerr)?; - - Python::attach(|py| chat_completions_response_to_py(py, response)) - }) -} - -#[pyfunction] -fn gil_stats(py: Python<'_>) -> PyResult> { - let stats = PyDict::new(py); - stats.set_item("releases", release_count())?; - Ok(stats.into_any().unbind()) -} - -#[cfg(feature = "panic-test")] -#[pyfunction] -fn _panic_for_test() { - panic!("intentional PyO3 panic smoke test"); -} - #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { - let py = module.py(); - module.add_function(wrap_pyfunction!(ocr, module)?)?; - module.add_function(wrap_pyfunction!(aocr, module)?)?; - module.add_function(wrap_pyfunction!(transcription, module)?)?; - module.add_function(wrap_pyfunction!(atranscription, module)?)?; - module.add_function(wrap_pyfunction!(messages, module)?)?; - module.add_function(wrap_pyfunction!(amessages, module)?)?; - module.add("RustBridgeDeclined", py.get_type::())?; - module.add("RustUpstreamError", py.get_type::())?; - module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; - module.add_function(wrap_pyfunction!(chat_completions, module)?)?; - module.add_function(wrap_pyfunction!(achat_completions, module)?)?; + errors::register(module)?; + routes::register(module)?; module.add_class::()?; - module.add_function(wrap_pyfunction!(gil_stats, module)?)?; - #[cfg(feature = "panic-test")] - module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; - Ok(()) + diagnostics::register(module) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn module_registration_preserves_the_public_surface() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "_native").expect("module should be created"); + _native(&module).expect("module should register"); + + let expected = [ + "RustBridgeDeclined", + "RustUpstreamError", + "ocr", + "aocr", + "transcription", + "atranscription", + "messages", + "amessages", + "chat_completions_decline", + "chat_completions", + "achat_completions", + "ResponsesWebSocketConnection", + "gil_stats", + ]; + + for name in expected { + assert!( + module + .hasattr(name) + .expect("attribute lookup should succeed") + ); + } + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..8d801109dfa --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; +use std::time::Duration; + +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +pub(crate) fn optional_object_to_map( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult> { + match value { + Some(value) => match from_py(value.bind(py))? { + Value::Object(map) => Ok(map), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + }, + None => Ok(Map::new()), + } +} + +pub(crate) fn optional_object( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult>> { + value + .map(|value| optional_object_to_map(py, name, Some(value))) + .transpose() +} + +pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { + timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }) +} + +pub(crate) fn marshal_headers( + py: Python<'_>, + headers: Option>, +) -> PyResult> { + let value = match headers { + Some(headers) => from_py(headers.bind(py))?, + None => Value::Object(Map::new()), + }; + let Value::Object(headers) = value else { + return Err(PyValueError::new_err("headers must be a dict")); + }; + headers + .into_iter() + .map(|(name, value)| { + value + .as_str() + .map(|value| (name, value.to_string())) + .ok_or_else(|| PyValueError::new_err("header values must be strings")) + }) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs new file mode 100644 index 00000000000..c3e290310fa --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -0,0 +1,126 @@ +use std::time::Duration; + +use litellm_core::audio_transcription::{ + AudioTranscriptionRequest, audio_transcription as run_audio_transcription, +}; +use litellm_core::error::Error; +use litellm_python_interop::from_py; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct TranscriptionInputs { + model: String, + audio: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Map, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + Ok(TranscriptionInputs { + model, + audio: from_py(audio.bind(py))?, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + optional_params: optional_object_to_map(py, "optional_params", optional_params)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: TranscriptionInputs) -> Result { + run_audio_transcription(AudioTranscriptionRequest { + model: &inputs.model, + audio: inputs.audio, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + optional_params: inputs.optional_params, + timeout: inputs.timeout, + }) + .await +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn transcription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn atranscription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(transcription, module)?)?; + module.add_function(wrap_pyfunction!(atranscription, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs new file mode 100644 index 00000000000..2ef3810015e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -0,0 +1,156 @@ +use std::time::Duration; + +use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; +use litellm_core::chat_completions::{ + chat_completions as run_chat_completions, chat_completions_decline_reason, +}; +use litellm_core::error::Error; +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::chat_completions_error_to_pyerr; +use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct ChatCompletionsInputs { + model: String, + messages: Value, + optional_params: Map, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let messages: Value = from_py(messages.bind(py))?; + if !messages.is_array() { + return Err(PyValueError::new_err("messages must be a list")); + } + Ok(ChatCompletionsInputs { + model, + messages, + optional_params: optional_object_to_map(py, "optional_params", optional_params)?, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: ChatCompletionsInputs) -> Result { + run_chat_completions(ChatCompletionsRequest { + model: &inputs.model, + messages: inputs.messages, + optional_params: inputs.optional_params, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + timeout: inputs.timeout, + }) + .await +} + +/// The decline reason for this request, or `None` when the Rust path accepts +/// it. Resolves no credentials and performs no I/O, so a host can ask before +/// committing to either path. +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] +fn chat_completions_decline( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + custom_llm_provider: Option, +) -> PyResult> { + let messages = from_py(messages.bind(py))?; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + Ok(chat_completions_decline_reason( + &model, + custom_llm_provider.as_deref(), + messages, + &optional_params, + ) + .map(str::to_string)) +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn chat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + messages, + optional_params, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, chat_completions_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn achat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + messages, + optional_params, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, chat_completions_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; + module.add_function(wrap_pyfunction!(chat_completions, module)?)?; + module.add_function(wrap_pyfunction!(achat_completions, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs new file mode 100644 index 00000000000..a29fbdd9fdc --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -0,0 +1,122 @@ +use std::time::Duration; + +use litellm_core::error::Error; +use litellm_core::messages::messages as run_messages; +use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct MessagesInputs { + model: String, + body: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let body: Value = from_py(body.bind(py))?; + if !body.is_object() { + return Err(PyValueError::new_err("body must be a dict")); + } + Ok(MessagesInputs { + model, + body, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: MessagesInputs) -> Result { + run_messages(MessagesRequest { + model: &inputs.model, + body: inputs.body, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + timeout: inputs.timeout, + }) + .await +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn messages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + body, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn amessages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + body, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(messages, module)?)?; + module.add_function(wrap_pyfunction!(amessages, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs new file mode 100644 index 00000000000..ab31702d70f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -0,0 +1,53 @@ +use std::future::Future; + +use litellm_core::error::Error; +use litellm_python_interop::{release_gil, to_py}; +use pyo3::prelude::*; +use serde::Serialize; + +use crate::function_trace::trace_call; + +mod audio_transcription; +mod chat_completions; +mod messages; +mod ocr; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + ocr::register(module)?; + audio_transcription::register(module)?; + messages::register(module)?; + chat_completions::register(module) +} + +fn block_on( + py: Python<'_>, + call: impl Future> + Send, + trace: bool, + map_err: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send, +{ + let result = release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(trace_call(call, trace)) + }); + match result { + Ok(response) => to_py(py, &response), + Err(err) => Err(map_err(err)), + } +} + +fn into_py_future<'py, T>( + py: Python<'py>, + call: impl Future> + Send + 'static, + trace: bool, + map_err: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let response = trace_call(call, trace).await.map_err(map_err)?; + Python::attach(|py| to_py(py, &response)) + }) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs new file mode 100644 index 00000000000..ad2a839de07 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -0,0 +1,128 @@ +use std::time::Duration; + +use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; +use litellm_core::error::Error; +use litellm_python_interop::from_py; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; + +use super::{block_on, into_py_future}; + +struct OcrInputs { + model: String, + document: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Map, + timeout: Option, +} + +#[allow(clippy::too_many_arguments)] +fn marshal_inputs( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + Ok(OcrInputs { + model, + document: from_py(document.bind(py))?, + api_key, + api_base, + custom_llm_provider, + extra_headers: optional_object(py, "extra_headers", extra_headers)?, + optional_params: optional_object_to_map(py, "optional_params", optional_params)?, + timeout: optional_timeout(timeout_seconds), + }) +} + +async fn call(inputs: OcrInputs) -> Result { + run_ocr(OcrRequest { + model: &inputs.model, + document: inputs.document, + api_key: inputs.api_key.as_deref(), + api_base: inputs.api_base.as_deref(), + custom_llm_provider: inputs.custom_llm_provider.as_deref(), + extra_headers: inputs.extra_headers, + optional_params: inputs.optional_params, + timeout: inputs.timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await +} + +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn ocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + block_on(py, call(inputs), trace, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn aocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + let inputs = marshal_inputs( + py, + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + )?; + into_py_future(py, call(inputs), trace, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(aocr, module)?) +} diff --git a/tests/sdk_function_trace/README.md b/tests/sdk_function_trace/README.md new file mode 100644 index 00000000000..d3a3b654aea --- /dev/null +++ b/tests/sdk_function_trace/README.md @@ -0,0 +1,30 @@ +# SDK function tracing + +The compare runner executes the same SDK calls through the Python engine and the Rust native bridge against a local HTTP provider fixture, then prints their pipeline trees side by side. Matching calls align on the same row in green; Python-only calls are blue, Rust-only calls yellow, and reordered calls red. Gaps preserve execution order and each column retains its own nesting. A comparison column labels every row even without color. Colors are enabled in terminals unless `NO_COLOR` is set. A difference summary follows (shared step order, python-only steps, rust-only steps). Each invocation must issue exactly one HTTP request. It requires the LiteLLM Python dependencies and the native extension built with tracing support + +From the repository root, using the project's Python environment: + +```bash +uv run python -m tests.sdk_function_trace.compare +uv run python -m tests.sdk_function_trace.compare --route ocr +uv run python -m tests.sdk_function_trace.compare --route ocr --sync +uv run python -m tests.sdk_function_trace.compare --route all --both --check +``` + +Calls default to async; use `--sync` for synchronous calls or `--both` for the complete matrix. Python sync Messages raises `not implemented for sync calls`; only that exact failure is marked `SKIP`, and the runner still executes Rust sync Messages and subsequent routes. Bedrock transcription has no independent Python provider implementation: its Python trace covers SDK dispatch into Rust + +Both engines are projected onto a shared per-route step table (`steps.py`): canonical names such as `transform_ocr_request` map Python functions (`MistralOCRConfig.transform_ocr_request`) and Rust spans (`transform_ocr_request`) to the same label. Only the first occurrence of each step is kept. Python indentation uses each event's actual frame ancestors and the nearest already displayed ancestor, so returned helpers and coroutine resumptions do not create false parents. Rust indentation uses instrumented span ancestry. Unmatched Rust span names pass through unchanged. `--full` prints every captured runtime event; validation still uses projected steps + +Every report checks required stage presence and dependency order. Provider lookup must precede request transformation, which must precede HTTP, followed by response transformation. The handler must precede HTTP; parameter mapping and supported-parameter checks must precede request transformation. Environment validation and URL construction, where mapped, must precede HTTP. Python transcription is checked only through native dispatch. `--check` also requires identical canonical step sequences for comparable routes and exits nonzero for missing, extra, or reordered steps, or an unexpected call failure, after finishing all selected cases + +Individual stage checks are separate from cross-language `step parity`. Passing stage checks cannot override a failing step comparison. Bedrock transcription and Python sync Messages report `UNAVAILABLE` for cross-language parity because they lack an independent Python execution to compare. Absolute nesting depth is not a cross-language gate: async Python Messages dispatches its handler onto another thread. See `route-comparison.md` for the audited matrix and remaining contract limitations + +The Python runner uses the existing `profile_python` / `sys.setprofile` collector, selecting executed code under the installed `litellm` source directory instead of maintaining a function-name allowlist. It prints source locations and qualified function names, including repeated calls. Coroutine resumptions are counted once per invocation. It profiles the current thread and threads created during the call, including the fresh async executor. Existing worker threads are not retroactively profiled; background Python calls may appear, and indentation follows selected Python stack ancestors within each thread + +The Rust runner calls the compiled PyO3 SDK entrypoints with `trace=True`. The existing `FunctionTrace` subscriber collects `#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]` spans for the route entrypoint, preparation, provider lookup, HTTP handler, and selected provider transformations. The shared `http_request` helper instruments the existing Rust send operation without changing clients, timeouts, signing, or error mapping. Function names come from the actual functions. `WithSubscriber` attaches the collector to each future across async polls. Arguments and provider payloads are not recorded in trace events. Uninstrumented functions do not appear; this is scoped instrumentation, not an exhaustive native call graph + +Tracing is opt-in: native calls without `trace=True` keep their original response shape. Traced calls return `{"response": ..., "trace": [{"function": ..., "depth": ...}]}`. The runners print only trace events. Missing native support or empty traces fail instead of falling back to source searching. The old `--repo`, `--signatures`, and `--calls` options are removed + +`profile_python(functions)` still supports direct function references for focused parity checks. `assert_function_trace_parity` compares selected Python events with Rust events supplied by an executable scenario. Successful stage checks prove the declared pipeline ran in a valid dependency order for this fixture; they do not assert identical function contracts, request bodies, responses, streaming behavior, or live-provider correctness + +Build the extension with `maturin develop` in the project's virtual environment. Then run either command above to get the executed function order 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/compare.py b/tests/sdk_function_trace/compare.py new file mode 100644 index 00000000000..941c1b6e067 --- /dev/null +++ b/tests/sdk_function_trace/compare.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import argparse +import os +import sys +from typing import Final + +from tests.sdk_function_trace.fixtures import ROUTES +from tests.sdk_function_trace.report import compare, render + + +def _run(route: str, asynchronous: bool, *, full: bool, colorize: bool) -> bool: + comparison: Final = compare(route, asynchronous=asynchronous) + sys.stdout.write(render(comparison, full=full, colorize=colorize)) + return comparison.passed + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description="Compare Python and Rust SDK pipeline steps per route") + parser.add_argument("--route", choices=("all", *ROUTES), default="all") + mode: Final = parser.add_mutually_exclusive_group() + mode.add_argument("--async", dest="asynchronous", action="store_true", default=True) + mode.add_argument("--sync", dest="asynchronous", action="store_false") + mode.add_argument("--both", action="store_true", help="run async and sync for every selected route") + parser.add_argument( + "--check", action="store_true", help="exit nonzero for missing, extra, or reordered comparable steps" + ) + parser.add_argument( + "--full", action="store_true", help="print every captured runtime event instead of pipeline steps" + ) + args: Final = parser.parse_args() + os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") + colorize: Final = sys.stdout.isatty() and "NO_COLOR" not in os.environ + results: Final = tuple( + _run(selected, selected_mode, full=args.full, colorize=colorize) + for selected in ROUTES + if args.route in ("all", selected) + for selected_mode in ((True, False) if args.both else (args.asynchronous,)) + ) + if args.check and not all(results): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/sdk_function_trace/fixtures.py b/tests/sdk_function_trace/fixtures.py new file mode 100644 index 00000000000..47bbe839627 --- /dev/null +++ b/tests/sdk_function_trace/fixtures.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import base64 +import io +import json +import wave +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final, Protocol, cast + +from tests.sdk_function_trace.mock_provider import MockProviderResponse +from tests.sdk_function_trace.steps import Engine + +ANTHROPIC_MODEL: Final = "claude-sonnet-5" +OCR_MODEL: Final = "mistral-ocr-latest" +AUDIO_MODEL: Final = "mistral.voxtral-mini-3b-2507" + + +class SdkCall(Protocol): + def __call__(self, **kwargs: object) -> object: ... + + +@dataclass(frozen=True, slots=True) +class Fixture: + kwargs: dict[str, object] + provider_response: MockProviderResponse + + +@dataclass(frozen=True, slots=True) +class RouteSpec: + label: str + python_entrypoints: tuple[str, str] + rust_entrypoints: tuple[str, str] + fixture: Callable[[Engine], Fixture] + + +@dataclass(frozen=True, slots=True) +class Invocation: + function: SdkCall + kwargs: dict[str, object] + provider_response: MockProviderResponse + label: str + + +def audio_bytes() -> bytes: + with io.BytesIO() as buffer: + with wave.open(buffer, "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(16000) + audio.writeframes(b"\x00\x00" * 1600) + return buffer.getvalue() + + +def _anthropic_message_response() -> MockProviderResponse: + body: Final = { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + return MockProviderResponse(200, (("content-type", "application/json"),), json.dumps(body).encode()) + + +def _conversation() -> dict[str, object]: + return {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + + +def _ocr_fixture(engine: Engine) -> Fixture: + return Fixture( + kwargs={ + "model": f"mistral/{OCR_MODEL}", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + **({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}), + }, + provider_response=MockProviderResponse( + 200, + (("content-type", "application/json"),), + json.dumps( + { + "pages": [{"index": 0, "markdown": "hello"}], + "model": OCR_MODEL, + "usage_info": {"pages_processed": 1}, + } + ).encode(), + ), + ) + + +def _chat_completions_fixture(engine: Engine) -> Fixture: + conversation: Final = _conversation() + payload: Final = ( + {"messages": conversation["messages"], "optional_params": {"max_tokens": 16}} + if engine == "rust" + else conversation + ) + return Fixture( + kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, + provider_response=_anthropic_message_response(), + ) + + +def _messages_fixture(engine: Engine) -> Fixture: + conversation: Final = _conversation() + payload: Final = {"body": {**conversation, "model": ANTHROPIC_MODEL}} if engine == "rust" else conversation + return Fixture( + kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, + provider_response=_anthropic_message_response(), + ) + + +def _transcription_fixture(engine: Engine) -> Fixture: + credentials: Final = { + "aws_access_key_id": "test-access", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + payload: Final = ( + { + "audio": {"data": base64.b64encode(audio_bytes()).decode(), "format": "wav"}, + "optional_params": credentials, + } + if engine == "rust" + else {"file": ("sample.wav", audio_bytes(), "audio/wav"), **credentials} + ) + return Fixture( + kwargs={"model": f"bedrock/{AUDIO_MODEL}", **payload}, + provider_response=MockProviderResponse( + 200, + (("content-type", "application/json"),), + json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + } + ).encode(), + ), + ) + + +ROUTE_SPECS: Final[dict[str, RouteSpec]] = { + "chat_completions": RouteSpec( + label="anthropic", + python_entrypoints=("completion", "acompletion"), + rust_entrypoints=("chat_completions", "achat_completions"), + fixture=_chat_completions_fixture, + ), + "audio_transcription": RouteSpec( + label="bedrock (Rust-only provider; Python trace covers SDK dispatch)", + python_entrypoints=("transcription", "atranscription"), + rust_entrypoints=("transcription", "atranscription"), + fixture=_transcription_fixture, + ), + "messages": RouteSpec( + label="anthropic", + python_entrypoints=("create", "acreate"), + rust_entrypoints=("messages", "amessages"), + fixture=_messages_fixture, + ), + "ocr": RouteSpec( + label="mistral", + python_entrypoints=("ocr", "aocr"), + rust_entrypoints=("ocr", "aocr"), + fixture=_ocr_fixture, + ), +} + +ROUTES: Final = tuple(ROUTE_SPECS) + + +def sdk_invocation(route: str, *, engine: Engine, asynchronous: bool) -> Invocation: + import litellm + from litellm.anthropic_interface import messages as sdk_messages + from litellm.rust_bridge import get_native_bridge + + rust: Final = engine == "rust" + bridge: Final = get_native_bridge() if rust else None + if rust and bridge is None: + raise RuntimeError("Build the native extension first: maturin develop") + spec: Final = ROUTE_SPECS.get(route) + if spec is None: + raise ValueError(f"Unknown route: {route}") + fixture: Final = spec.fixture(engine) + owner: Final = bridge if rust else (sdk_messages if route == "messages" else litellm) + entrypoint: Final = (spec.rust_entrypoints if rust else spec.python_entrypoints)[int(asynchronous)] + return Invocation( + function=cast(SdkCall, getattr(owner, entrypoint)), + kwargs={ + **fixture.kwargs, + "api_key": "test-key", + **({"trace": True, "timeout_seconds": 5} if rust else {"timeout": 5}), + }, + provider_response=fixture.provider_response, + label=spec.label, + ) 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/ocr-comparison.md b/tests/sdk_function_trace/ocr-comparison.md new file mode 100644 index 00000000000..d252480e218 --- /dev/null +++ b/tests/sdk_function_trace/ocr-comparison.md @@ -0,0 +1,59 @@ +# OCR Python and Rust comparison + +Audited implementation revision: `edcba483b2`. The implementations do not match in function contracts, call structure, or all tested response behavior. This audit changes the source listing coverage, not OCR runtime behavior + +Run both source listings from the repository root: + +```bash +python3 tests/sdk_function_trace/list_python_steps.py --route ocr --signatures --calls +uv run tests/sdk_function_trace/list_rust_steps.py --route ocr --signatures --calls +``` + +Both cover Mistral, Azure AI Mistral, Azure Document Intelligence, Vertex Mistral, and Vertex DeepSeek. Listings show declarations and source call sites, not executed traces + +## Function contracts + +Comparing Python `BaseOCRConfig` with Rust `OcrProviderConfig`, omitting `self` and language-specific ownership details: + +| Python | Rust | Difference | +| --- | --- | --- | +| `get_supported_ocr_params(model)` | `supported_ocr_params()` | Name and model argument | +| `get_api_key_env_var()` | No corresponding method | Missing contract | +| `map_ocr_params(non_default_params, optional_params, model)` | `map_ocr_params(non_default_params)` | Missing accumulator and model | +| `validate_environment(headers, model, api_key, api_base, litellm_params, **kwargs)` | Separate auth/key/header helpers | Different contract | +| `get_complete_url(api_base, model, optional_params, litellm_params, **kwargs)` | `complete_url(api_base, model, optional_params, env_lookup)` | Name and context | +| `transform_ocr_request(model, document, optional_params, headers, **kwargs)` | `transform_ocr_request(model, document, optional_params)` | Missing headers and extra context | +| `async_transform_ocr_request(...)` | No corresponding method | Missing async override | +| `transform_ocr_response(model, raw_response, logging_obj, **kwargs)` | `transform_ocr_response(model, response_json)` | Missing HTTP metadata, logging and extra context | +| `async_transform_ocr_response(...)` | No corresponding method | Missing async override | +| `get_error_class(error_message, status_code, headers)` | Central Rust error mapping | Different contract | + +Python's default mapper returns the supplied `optional_params`; Rust's filters `non_default_params`. Provider overrides must also be compared + +Python maps parameters during SDK preparation, before HTTP-handler environment validation and URL construction. Rust resolves auth and URL before mapping parameters in `prepare_provider_request`. Python has async provider transforms; both native entrypoints execute the same Rust async route using synchronous transform hooks, with polling and document downloading in gateway helpers + +The native bindings also accept `optional_params` and `timeout_seconds`, while the Python SDK accepts `**kwargs` and `timeout`. Public SDK calls with Rust enabled still execute Python preparation before entering Rust, so matching SDK responses would not prove matching standalone Rust steps + +## Runtime results + +Built the native extension from the audited source using `cargo build -p litellm-python-bridge --features extension-module --offline`. Supplied that build's functions through `use_litellm_rust` dependency injection. Ran public `litellm.ocr` and `litellm.aocr` with Rust disabled and enabled against identical local HTTP response fixtures, requiring one request per invocation + +Successful `model_dump()` results and failure exception classes were compared. These checks cover Mistral response outcomes only, not request equality, error messages, live providers, or every execution branch + +| Mistral response fixture | Sync | Async | Observation | +| --- | --- | --- | --- | +| Valid page/model/usage | Match | Match | Same normalized response | +| Model omitted | Match | Match | Both use the requested model | +| `model: null` | Different | Different | Python rejects; Rust uses the requested model | +| `pages: null` | Different | Different | Python rejects; Rust returns an empty array | +| Invalid page element | Match | Match | Both reject during response validation | + +Six of ten fixture/mode comparisons match, four differ. Rust's Mistral response transform conflates missing values with explicit nulls through `as_array`/`as_str` fallbacks. Python preserves explicit nulls into response validation, which rejects them + +## Other provider gaps found in source + +Azure Document Intelligence's Python configuration supports `pages`, `features`, and `req_format`; Rust lists only `pages`. Python normalizes parameters before URL construction; Rust normalizes pages during URL construction + +Python preserves Azure `content`, `tables`, and `keyValuePairs`, and supports retaining the native operation payload. Rust's `OcrResponseData` has no corresponding fields, and its Azure transform does not preserve those values + +Azure and Vertex async document transforms and Azure polling also use different helper contracts. Their runtime equivalence was not tested in this audit diff --git a/tests/sdk_function_trace/profiler.py b/tests/sdk_function_trace/profiler.py new file mode 100644 index 00000000000..c71c74ab0d3 --- /dev/null +++ b/tests/sdk_function_trace/profiler.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import sys +import threading +from collections.abc import Generator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import CodeType, FrameType, FunctionType +from typing import Final + + +@dataclass(frozen=True, slots=True) +class FunctionTraceEvent: + function: str + depth: int + ancestors: tuple[str, ...] | None = None + + +class PythonProfiler: + def __init__(self, functions: Sequence[FunctionType], source_root: Path | None = None) -> None: + self._source_root: Final = str(source_root.resolve()) + "/" if source_root is not None else 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.function_name(frame.f_code) + if function_name is None: + return + ancestors: Final = tuple( + name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor.f_code)) is not None + ) + self._seen_frames.add(frame) + self.events.append( + FunctionTraceEvent( + function=function_name, + depth=len(ancestors), + ancestors=ancestors if self._source_root is not None else None, + ) + ) + + def function_name(self, code: CodeType) -> str | None: + if self._source_root is None: + return self._names_by_code.get(code) + if not code.co_filename.startswith(self._source_root): + return None + relative: Final = code.co_filename.removeprefix(self._source_root) + return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" + + +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] = (), *, source_root: Path | None = None, threads: bool = False +) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(functions, source_root) + previous_thread: Final = threading.getprofile() + if threads: + threading.setprofile(profiler) + previous: Final = sys.getprofile() + sys.setprofile(profiler) + try: + yield profiler + finally: + sys.setprofile(previous) + if threads: + threading.setprofile(previous_thread) diff --git a/tests/sdk_function_trace/report.py b/tests/sdk_function_trace/report.py new file mode 100644 index 00000000000..9b654e571f8 --- /dev/null +++ b/tests/sdk_function_trace/report.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from tests.sdk_function_trace.fixtures import ROUTE_SPECS +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.runtime import ( + TraceDiff, + TraceFailed, + TraceOk, + TraceRun, + TraceSkipped, + attempt_trace, + trace_diff, +) +from tests.sdk_function_trace.steps import Engine, pipeline_issues, pipeline_steps +from tests.sdk_function_trace.table import format_trace_table + +_PYTHON_ONLY_COLOR: Final = "\033[34m" +_RUST_ONLY_COLOR: Final = "\033[33m" +_RESET: Final = "\033[0m" + +_ENGINE_COLOR: Final[dict[Engine, str]] = {"python": _PYTHON_ONLY_COLOR, "rust": _RUST_ONLY_COLOR} + + +@dataclass(frozen=True, slots=True) +class EngineReport: + engine: Engine + run: TraceRun + events: tuple[FunctionTraceEvent, ...] + steps: tuple[FunctionTraceEvent, ...] + issues: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Comparison: + route: str + label: str + asynchronous: bool + engines: tuple[EngineReport, ...] + diff: TraceDiff + + @property + def comparable(self) -> bool: + return self.route != "audio_transcription" and all(isinstance(report.run, TraceOk) for report in self.engines) + + @property + def passed(self) -> bool: + return ( + (not self.comparable or self.diff.matches) + and not any(report.issues for report in self.engines) + and all(not isinstance(report.run, TraceFailed) for report in self.engines) + ) + + +def _events(run: TraceRun) -> tuple[FunctionTraceEvent, ...]: + match run: + case TraceOk(events=events): + return events + case TraceSkipped() | TraceFailed(): + return () + + +def _engine_report(route: str, engine: Engine, run: TraceRun) -> EngineReport: + events: Final = _events(run) + steps: Final = pipeline_steps(route, engine, events) + issues: Final = pipeline_issues(route, engine, steps) if isinstance(run, TraceOk) else () + return EngineReport(engine=engine, run=run, events=events, steps=steps, issues=issues) + + +def compare(route: str, *, asynchronous: bool) -> Comparison: + runs: Final = { + engine: attempt_trace(route, engine=engine, asynchronous=asynchronous) for engine in ("python", "rust") + } + engines: Final = tuple(_engine_report(route, engine, run) for engine, run in runs.items()) + return Comparison( + route=route, + label=ROUTE_SPECS[route].label, + asynchronous=asynchronous, + engines=engines, + diff=trace_diff(engines[0].steps, engines[1].steps), + ) + + +def _tree_line(event: FunctionTraceEvent, only: frozenset[str], marker: str, color: str, *, colorize: bool) -> str: + line: Final = f"{' ' * event.depth}{event.function}" + (f" {marker}" if event.function in only else "") + return f"{color}{line}{_RESET}\n" if colorize and event.function in only else f"{line}\n" + + +def _tree_lines( + events: tuple[FunctionTraceEvent, ...], + only: frozenset[str], + marker: str, + color: str, + *, + colorize: bool, +) -> tuple[str, ...]: + return tuple(_tree_line(event, only, marker, color, colorize=colorize) for event in events) + + +def _engine_lines( + report: EngineReport, diff: TraceDiff, *, comparable: bool, full: bool, colorize: bool +) -> tuple[str, ...]: + match report.run: + case TraceSkipped(reason=reason): + return (f"{report.engine}: SKIP ({reason})\n\n",) + case TraceFailed(reason=reason): + return (f"{report.engine}: FAIL ({reason})\n\n",) + case TraceOk(): + shown: Final = report.events if full else report.steps + only: Final = ( + () if full or not comparable else (diff.python_only if report.engine == "python" else diff.rust_only) + ) + return ( + f"{report.engine} ({len(shown)} steps)\n\n", + *_tree_lines( + shown, + frozenset(only), + f"<- {report.engine} only", + _ENGINE_COLOR[report.engine], + colorize=colorize, + ), + "\n", + ) + + +def _parity_lines(comparison: Comparison) -> tuple[str, ...]: + if not comparison.comparable: + if comparison.route == "audio_transcription": + return ("step parity: UNAVAILABLE (Bedrock transcription has no independent Python implementation)\n",) + return ("step parity: UNAVAILABLE (both engines must complete)\n",) + diff: Final = comparison.diff + order: Final = "the same" if diff.shared_order_matches else "a different" + return ( + "diff\n\n", + f"shared steps appear in {order} order\n", + f"python-only: {', '.join(diff.python_only) or 'none'}\n", + f"rust-only: {', '.join(diff.rust_only) or 'none'}\n\n", + f"step parity: {'PASS' if diff.matches else 'FAIL'}\n", + ) + + +def _stage_lines(comparison: Comparison) -> tuple[str, ...]: + return tuple( + f"{report.engine} " + f"{'SDK dispatch only' if comparison.route == 'audio_transcription' and report.engine == 'python' else 'pipeline'}: " + f"{'FAIL: ' + '; '.join(report.issues) if report.issues else 'PASS'}\n" + for report in comparison.engines + if isinstance(report.run, TraceOk) + ) + + +def render(comparison: Comparison, *, full: bool, colorize: bool) -> str: + mode: Final = "async" if comparison.asynchronous else "sync" + traces: Final = ( + (format_trace_table(comparison.engines[0].steps, comparison.engines[1].steps, colorize=colorize) + "\n\n",) + if not full and all(isinstance(report.run, TraceOk) for report in comparison.engines) + else tuple( + line + for report in comparison.engines + for line in _engine_lines( + report, comparison.diff, comparable=comparison.comparable, full=full, colorize=colorize + ) + ) + ) + return "".join( + ( + f"route: {comparison.route} provider: {comparison.label} mode: {mode}\n\n", + *traces, + *_parity_lines(comparison), + *_stage_lines(comparison), + "Each successful invocation issued exactly one local provider request\n\n", + ) + ) diff --git a/tests/sdk_function_trace/route-comparison.md b/tests/sdk_function_trace/route-comparison.md new file mode 100644 index 00000000000..009d3544d05 --- /dev/null +++ b/tests/sdk_function_trace/route-comparison.md @@ -0,0 +1,26 @@ +# SDK route trace audit + +Run the four native HTTP route families in both modes from the repository root: + +```bash +uv run python -m tests.sdk_function_trace.compare --route all --both --check +``` + +The local fixture matrix on 2026-09-02 completed 15 successful engine invocations and one expected skip. Every successful invocation issued exactly one local HTTP request. All five comparable route/mode pairs have identical canonical steps in the same order, with no Python-only or Rust-only steps + +| Route | Python async | Python sync | Rust async | Rust sync | +| --- | --- | --- | --- | --- | +| Chat completions, Anthropic | Pass | Pass | Pass | Pass | +| Messages, Anthropic | Pass | Unsupported, skipped | Pass | Pass | +| OCR, Mistral | Pass | Pass | Pass | Pass | +| Audio transcription, Bedrock | Dispatch only | Dispatch only | Pass | Pass | + +The same canonical step sequence ran in sync and async for each engine with both modes available. Bedrock transcription's Python SDK delegates to Rust, so its two successful calls do not establish independent provider parity. Realtime and Responses WebSockets are outside this HTTP fixture runner + +Chat and OCR also have identical projected nesting in both modes. Async Messages has the same helper nesting beneath its handler, but Python starts that handler on a worker thread, so it appears as a second root. The comparison preserves this physical thread boundary and checks step order independently of absolute depth + +Rust now resolves chat providers and supported parameters before entering its handler. Chat and Messages validate the environment and transform requests inside their handlers. Messages builds the final URL after transformation. OCR resolves its config and maps supported parameters during preparation, then validates credentials, builds the URL, and transforms the request inside its handler. Its during-call guardrails still run before HTTP, within the provider-call lifecycle phase + +The environment hooks execute credential and header validation. Chat's supported-parameter hooks return OpenAI names paired with provider names and feed the existing request acceptance checks. The direct Rust API still accepts provider-mapped parameters, and its supported subset is smaller than Python's. Matching the pipeline does not establish identical parameter contracts + +`--check` now fails if either comparable engine has missing, extra, or reordered canonical steps, even if its individual stage checks pass. Bedrock transcription and sync Messages report `UNAVAILABLE` for cross-language parity; native execution is still checked. Passing establishes step coverage and order for one non-streaming fixture per route, not complete request, response, error, or provider parity. The previously recorded OCR response gaps remain in `ocr-comparison.md` diff --git a/tests/sdk_function_trace/runtime.py b/tests/sdk_function_trace/runtime.py new file mode 100644 index 00000000000..d5bf15694bc --- /dev/null +++ b/tests/sdk_function_trace/runtime.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +import os +from collections.abc import Awaitable, Generator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final, cast +from unittest.mock import patch + +from pydantic import BaseModel, ConfigDict + +from tests.sdk_function_trace.fixtures import Invocation, sdk_invocation +from tests.sdk_function_trace.mock_provider import mock_provider +from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python +from tests.sdk_function_trace.steps import Engine + + +class TraceEventPayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + function: str + depth: int + + +class TraceResponsePayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + response: object + trace: tuple[TraceEventPayload, ...] | list[TraceEventPayload] + + +@contextmanager +def _python_engine() -> Generator[None]: + from litellm.rust_bridge import ocr as ocr_bridge + + previous_ocr: Final = ocr_bridge.rust_ocr_enabled() + with patch.dict(os.environ, {"LITELLM_RUST": "false"}): + ocr_bridge.use_litellm_rust(False) + try: + yield + finally: + ocr_bridge.use_litellm_rust(previous_ocr) + + +def _invoke(case: Invocation, api_base: str, *, asynchronous: bool) -> object: + async def invoke_async() -> object: + return await cast("Awaitable[object]", case.function(**case.kwargs, api_base=api_base)) + + if asynchronous: + return asyncio.run(invoke_async()) + return case.function(**case.kwargs, api_base=api_base) + + +def collect(case: Invocation, api_base: str, *, engine: Engine, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: + import litellm + + if engine == "rust": + payload: Final = TraceResponsePayload.model_validate(_invoke(case, api_base, asynchronous=asynchronous)) + return tuple(FunctionTraceEvent(event.function, event.depth) for event in payload.trace) + with profile_python(source_root=Path(litellm.__file__).parent, threads=True) as profiler: + _invoke(case, api_base, asynchronous=asynchronous) + return tuple(profiler.events) + + +def run_trace(route: str, *, engine: Engine, asynchronous: bool = False) -> tuple[FunctionTraceEvent, ...]: + case: Final = sdk_invocation(route, engine=engine, asynchronous=asynchronous) + with _python_engine(), mock_provider(case.provider_response) as api_base: + events: Final = collect(case, api_base, engine=engine, asynchronous=asynchronous) + if not events: + raise RuntimeError(f"No runtime events for {route}; rebuild the native extension with tracing support") + return events + + +@dataclass(frozen=True, slots=True) +class TraceOk: + events: tuple[FunctionTraceEvent, ...] + + +@dataclass(frozen=True, slots=True) +class TraceSkipped: + reason: str + + +@dataclass(frozen=True, slots=True) +class TraceFailed: + reason: str + + +TraceRun = TraceOk | TraceSkipped | TraceFailed + + +def attempt_trace(route: str, *, engine: Engine, asynchronous: bool) -> TraceRun: + try: + return TraceOk(run_trace(route, engine=engine, asynchronous=asynchronous)) + except Exception as error: + reason: Final = f"{type(error).__name__}: {error}" + if ( + route == "messages" + and engine == "python" + and not asynchronous + and isinstance(error, ValueError) + and str(error) == "anthropic_messages_handler is not implemented for sync calls" + ): + return TraceSkipped(reason) + return TraceFailed(reason) + + +@dataclass(frozen=True, slots=True) +class TraceDiff: + python_only: tuple[str, ...] + rust_only: tuple[str, ...] + shared_order_matches: bool + + @property + def matches(self) -> bool: + return not self.python_only and not self.rust_only and self.shared_order_matches + + +def trace_diff(python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...]) -> TraceDiff: + python_names: Final = {event.function for event in python} + rust_names: Final = {event.function for event in rust} + shared_python: Final = tuple(event.function for event in python if event.function in rust_names) + shared_rust: Final = tuple(event.function for event in rust if event.function in python_names) + return TraceDiff( + python_only=tuple(event.function for event in python if event.function not in rust_names), + rust_only=tuple(event.function for event in rust if event.function not in python_names), + shared_order_matches=bool(shared_python) and shared_python == shared_rust, + ) diff --git a/tests/sdk_function_trace/steps.py b/tests/sdk_function_trace/steps.py new file mode 100644 index 00000000000..bb50d4ebe57 --- /dev/null +++ b/tests/sdk_function_trace/steps.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import re +from collections.abc import Sequence +from dataclasses import dataclass +from functools import reduce +from typing import Final, Literal + +from tests.sdk_function_trace.profiler import FunctionTraceEvent + +Engine = Literal["python", "rust"] + + +@dataclass(frozen=True, slots=True) +class Step: + name: str + python: re.Pattern[str] | None + rust: str | None + + +def _step(name: str, python: str | None = None, rust: str | None = None) -> Step: + return Step(name, re.compile(python) if python is not None else None, rust) + + +_POST: Final = r"AsyncHTTPHandler\.post$|HTTPHandler\.post$" + +STEPS: Final[dict[str, tuple[Step, ...]]] = { + "ocr": ( + _step("ocr", r"ocr/main\.py:\d+ a?ocr$", "ocr"), + _step("prepare_ocr_call", r"ocr/main\.py:\d+ _prepare_ocr_request$", "prepare_ocr_call"), + _step("get_provider_ocr_config", r"ProviderConfigManager\.get_provider_ocr_config$", "ocr_provider_config"), + _step("supported_ocr_params", r"get_supported_ocr_params$", "supported_ocr_params"), + _step("map_ocr_params", r"(? tuple[str, ...]: + names: Final = tuple(event.function for event in events) + required: Final = tuple(step.name for step in STEPS[route] if getattr(step, engine) is not None) + missing: Final = tuple(f"missing {name}" for name in required if name not in names) + provider: Final = next(name for name in required if name.startswith("get_provider_")) + handler: Final = next(name for name in required if name.startswith("execute_")) + dispatch_only: Final = route == "audio_transcription" and engine == "python" + request: Final = next( + (name for name in required if name.startswith("transform_") and name.endswith("request")), handler + ) + response: Final = next( + (name for name in required if name.startswith("transform_") and name.endswith("response")), handler + ) + phases: Final = ( + (route, "map_transcription_params", provider, handler) + if dispatch_only + else (route, provider, request, "http_request", response) + ) + extra_edges: Final = ( + () + if dispatch_only + else ( + (handler, "http_request"), + *((name, request) for name in required if name.startswith(("map_", "supported_"))), + *((name, "http_request") for name in ("validate_environment", "complete_url") if name in required), + ) + ) + edges: Final = (*zip(phases, phases[1:]), *extra_edges) + return missing + tuple( + f"{before} must precede {after}" + for before, after in edges + if before in names and after in names and names.index(before) >= names.index(after) + ) + + +def _canonical_name(route: str, engine: Engine, function: str) -> str | None: + for step in STEPS[route]: + if engine == "python": + if step.python is not None and step.python.search(function): + return step.name + elif step.rust is not None and function == step.rust: + return step.name + return function if engine == "rust" else None + + +@dataclass(frozen=True, slots=True) +class _Projection: + shown: tuple[FunctionTraceEvent, ...] = () + stack: tuple[tuple[int, int], ...] = () + seen: frozenset[str] = frozenset() + + +def _project(route: str, engine: Engine, state: _Projection, event: FunctionTraceEvent) -> _Projection: + stack: Final = tuple(pair for pair in state.stack if event.depth > pair[0]) + name: Final = _canonical_name(route, engine, event.function) + if name is None or name in state.seen: + return _Projection(state.shown, stack, state.seen) + depth: Final = ( + next( + ( + kept.depth + 1 + for ancestor in event.ancestors + for kept in state.shown + if kept.function == _canonical_name(route, engine, ancestor) + ), + 0, + ) + if event.ancestors is not None + else stack[-1][1] + 1 + if stack + else 0 + ) + return _Projection( + state.shown + (FunctionTraceEvent(function=name, depth=depth),), + stack + ((event.depth, depth),), + state.seen | {name}, + ) + + +def pipeline_steps(route: str, engine: Engine, events: Sequence[FunctionTraceEvent]) -> tuple[FunctionTraceEvent, ...]: + projection: Final = reduce(lambda state, event: _project(route, engine, state, event), events, _Projection()) + return projection.shown diff --git a/tests/sdk_function_trace/table.py b/tests/sdk_function_trace/table.py new file mode 100644 index 00000000000..2124d7e3faf --- /dev/null +++ b/tests/sdk_function_trace/table.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from collections.abc import Iterator +from difflib import SequenceMatcher +from typing import Final + +from tests.sdk_function_trace.profiler import FunctionTraceEvent + + +def _aligned_rows( + python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...] +) -> Iterator[tuple[FunctionTraceEvent | None, FunctionTraceEvent | None]]: + matcher: Final = SequenceMatcher( + a=tuple(event.function for event in python), + b=tuple(event.function for event in rust), + autojunk=False, + ) + for tag, python_start, python_end, rust_start, rust_end in matcher.get_opcodes(): + if tag == "equal": + yield from zip(python[python_start:python_end], rust[rust_start:rust_end]) + else: + yield from ((event, None) for event in python[python_start:python_end]) + yield from ((None, event) for event in rust[rust_start:rust_end]) + + +def _label(event: FunctionTraceEvent | None) -> str: + return f"{' ' * event.depth}{event.function}" if event is not None else "" + + +def _status( + python: FunctionTraceEvent | None, + rust: FunctionTraceEvent | None, + python_names: frozenset[str], + rust_names: frozenset[str], +) -> tuple[str, str]: + if python is not None and rust is not None: + return "match", "\033[32m" + if python is not None: + return ("reordered", "\033[31m") if python.function in rust_names else ("python only", "\033[34m") + if rust is not None: + return ("reordered", "\033[31m") if rust.function in python_names else ("rust only", "\033[33m") + return "", "" + + +def format_trace_table( + python: tuple[FunctionTraceEvent, ...], + rust: tuple[FunctionTraceEvent, ...], + *, + colorize: bool, +) -> str: + python_header: Final = f"python ({len(python)} steps)" + rust_header: Final = f"rust ({len(rust)} steps)" + python_width: Final = max(len(python_header), *(len(_label(event)) for event in python), 0) + rust_width: Final = max(len(rust_header), *(len(_label(event)) for event in rust), 0) + python_names: Final = frozenset(event.function for event in python) + rust_names: Final = frozenset(event.function for event in rust) + border: Final = f"+-{'-' * python_width}-+-{'-' * rust_width}-+-------------+" + rows: Final = tuple( + f"{color}{line}\033[0m" if colorize else line + for left, right in _aligned_rows(python, rust) + for status, color in (_status(left, right, python_names, rust_names),) + for line in (f"| {_label(left):<{python_width}} | {_label(right):<{rust_width}} | {status:<11} |",) + ) + return "\n".join( + ( + border, + f"| {python_header:<{python_width}} | {rust_header:<{rust_width}} | {'comparison':<11} |", + border, + *rows, + border, + ) + ) diff --git a/tests/sdk_function_trace/test_mock_provider.py b/tests/sdk_function_trace/test_mock_provider.py new file mode 100644 index 00000000000..88d7d5392d0 --- /dev/null +++ b/tests/sdk_function_trace/test_mock_provider.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from contextlib import ExitStack +from typing import Final +from urllib.error import HTTPError +from urllib.request import Request, urlopen + +import pytest + +from tests.sdk_function_trace.mock_provider import MockProviderResponse, mock_provider + + +def test_mock_provider_preserves_error_response() -> None: + response: Final = MockProviderResponse(429, (("retry-after", "2"),), b'{"error":"rate limited"}') + with mock_provider(response) as api_base: + with pytest.raises(HTTPError) as error: + urlopen(Request(api_base, data=b"{}"), timeout=5) + with error.value as received: + assert received.code == 429 + assert received.headers["retry-after"] == "2" + assert received.read() == response.body + + +@pytest.mark.parametrize("request_count", [0, 2]) +def test_mock_provider_rejects_missing_or_duplicate_requests(request_count: int) -> None: + response: Final = MockProviderResponse(200, (), b"{}") + with ExitStack() as stack: + api_base: Final = stack.enter_context(mock_provider(response)) + for _ in range(request_count): + with urlopen(Request(api_base, data=b"{}"), timeout=5) as received: + assert received.read() == response.body + with pytest.raises(AssertionError, match=f"expected one provider request, received {request_count}"): + stack.close() diff --git a/tests/sdk_function_trace/test_profiler.py b/tests/sdk_function_trace/test_profiler.py new file mode 100644 index 00000000000..10a266fb1e8 --- /dev/null +++ b/tests/sdk_function_trace/test_profiler.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path +from types import FunctionType +from typing import Final, cast + +import pytest + +from tests.sdk_function_trace import ( + FunctionTraceEvent, + TraceScenario, + TraceStep, + assert_function_trace_parity, +) +from tests.sdk_function_trace.profiler import profile_python + + +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 profile_python((First.run,)) as outer: + with pytest.raises(RuntimeError, match="stop"): + with profile_python((Second.run,)): + raise RuntimeError("stop") + assert sys.getprofile() is outer + First.run() + + assert sys.getprofile() is previous + assert outer.events == [FunctionTraceEvent(function="run", depth=0)] + + +def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: + async def suspended() -> None: + await asyncio.sleep(0) + First.run() + await asyncio.sleep(0) + + with profile_python((suspended, First.run)) as profiler: + asyncio.run(suspended()) + + assert profiler.events == [ + FunctionTraceEvent(function="suspended", depth=0), + FunctionTraceEvent(function="run", depth=1), + ] + + +def test_source_profiler_records_real_frame_ancestry() -> None: + def outer() -> None: + First.run() + + with profile_python(source_root=Path(__file__).parent) as profiler: + outer() + Second.run() + + outer_event, first_event, second_event = ( + event for event in profiler.events if event.function.startswith("test_profiler.py:") + ) + assert first_event.ancestors is not None + assert outer_event.function in first_event.ancestors + assert second_event.ancestors is not None + assert outer_event.function not in second_event.ancestors + + +@pytest.mark.parametrize( + "rust_trace", + [ + (), + (FunctionTraceEvent(function="renamed", depth=0),), + (FunctionTraceEvent(function="run", depth=1),), + (FunctionTraceEvent(function="run", depth=0),) * 2, + ], + ids=["missing", "renamed", "wrong-depth", "extra-call"], +) +def test_harness_rejects_rust_function_trace_drift(rust_trace: tuple[FunctionTraceEvent, ...]) -> None: + with pytest.raises(AssertionError, match="Rust function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=First.run, + invoke_rust=lambda: rust_trace, + ) + ) + + +def test_harness_rejects_python_function_trace_drift() -> None: + with pytest.raises(AssertionError, match="Python function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=Second.run, + invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), + ) + ) + + +def test_harness_accepts_matching_traces() -> None: + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=First.run, + invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), + ) + ) + + +def test_harness_rejects_reordered_calls() -> None: + def begin() -> None: + return None + + def finish() -> None: + return None + + with pytest.raises(AssertionError, match="Rust function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=( + TraceStep(cast(FunctionType, begin), depth=0), + TraceStep(cast(FunctionType, finish), depth=0), + ), + invoke_python=lambda: (begin(), finish()), + invoke_rust=lambda: ( + FunctionTraceEvent(function="finish", depth=0), + FunctionTraceEvent(function="begin", depth=0), + ), + ) + ) diff --git a/tests/sdk_function_trace/test_runtime.py b/tests/sdk_function_trace/test_runtime.py new file mode 100644 index 00000000000..015cba55083 --- /dev/null +++ b/tests/sdk_function_trace/test_runtime.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.sdk_function_trace.runtime import ( + TraceFailed, + TraceSkipped, + attempt_trace, + run_trace, + trace_diff, +) +from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps + + +def test_sync_messages_records_the_known_python_limitation() -> None: + result: Final = attempt_trace("messages", engine="python", asynchronous=False) + + assert isinstance(result, TraceSkipped) + assert result.reason == "ValueError: anthropic_messages_handler is not implemented for sync calls" + + +def test_unexpected_call_failure_is_not_skipped() -> None: + result: Final = attempt_trace("unknown", engine="python", asynchronous=False) + + assert isinstance(result, TraceFailed) + assert result.reason == "ValueError: Unknown route: unknown" + + +@pytest.mark.parametrize( + ("route", "asynchronous"), + (("chat_completions", False), ("chat_completions", True), ("messages", True), ("ocr", False), ("ocr", True)), +) +def test_compiled_routes_match_python_steps(route: str, asynchronous: bool) -> None: + from litellm.rust_bridge import get_native_bridge + + if get_native_bridge() is None: + pytest.skip("build the native bridge to run executed route parity") + python: Final = pipeline_steps(route, "python", run_trace(route, engine="python", asynchronous=asynchronous)) + rust: Final = pipeline_steps(route, "rust", run_trace(route, engine="rust", asynchronous=asynchronous)) + + assert pipeline_issues(route, "python", python) == () + assert pipeline_issues(route, "rust", rust) == () + assert trace_diff(python, rust).matches + if route != "messages": + assert python == rust diff --git a/tests/sdk_function_trace/test_steps.py b/tests/sdk_function_trace/test_steps.py new file mode 100644 index 00000000000..b5432951187 --- /dev/null +++ b/tests/sdk_function_trace/test_steps.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.runtime import trace_diff +from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps + + +def test_python_ocr_projection_keeps_pipeline_and_drops_noise() -> None: + events: Final = ( + FunctionTraceEvent("utils.py:1747 client..wrapper_async", 0), + FunctionTraceEvent("ocr/main.py:331 aocr", 1), + FunctionTraceEvent("ocr/main.py:70 _prepare_ocr_request", 2), + FunctionTraceEvent("litellm_core_utils/get_llm_provider_logic.py:142 get_llm_provider", 3), + FunctionTraceEvent("utils.py:9303 ProviderConfigManager.get_provider_ocr_config", 3), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:72 MistralOCRConfig.map_ocr_params", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 5), + FunctionTraceEvent("llms/custom_httpx/llm_http_handler.py:1705 BaseLLMHTTPHandler.async_ocr", 2), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:94 MistralOCRConfig.validate_environment", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:124 MistralOCRConfig.get_complete_url", 4), + FunctionTraceEvent("llms/base_llm/ocr/transformation.py:209 BaseOCRConfig.async_transform_ocr_request", 5), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:149 MistralOCRConfig.transform_ocr_request", 6), + FunctionTraceEvent("llms/custom_httpx/http_handler.py:654 AsyncHTTPHandler.post", 6), + FunctionTraceEvent("llms/base_llm/ocr/transformation.py:255 BaseOCRConfig.async_transform_ocr_response", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:200 MistralOCRConfig.transform_ocr_response", 5), + FunctionTraceEvent("cost_calculator.py:1874 ocr_cost", 6), + ) + + assert pipeline_steps("ocr", "python", events) == ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("get_provider_ocr_config", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("map_ocr_params", 3), + FunctionTraceEvent("execute_ocr_provider_call", 1), + FunctionTraceEvent("validate_environment", 2), + FunctionTraceEvent("complete_url", 2), + FunctionTraceEvent("transform_ocr_request", 3), + FunctionTraceEvent("http_request", 3), + FunctionTraceEvent("transform_ocr_response", 2), + ) + + +def test_rust_ocr_projection_reuses_step_names_and_keeps_unknown_spans() -> None: + events: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("transform_ocr_request", 2), + FunctionTraceEvent("execute_ocr_provider_call", 2), + FunctionTraceEvent("transform_ocr_response", 3), + FunctionTraceEvent("new_uninstrumented_span", 3), + ) + + assert pipeline_steps("ocr", "rust", events) == ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("transform_ocr_request", 2), + FunctionTraceEvent("execute_ocr_provider_call", 2), + FunctionTraceEvent("transform_ocr_response", 3), + FunctionTraceEvent("new_uninstrumented_span", 3), + ) + + +def test_projection_resets_depth_on_thread_root() -> None: + events: Final = ( + FunctionTraceEvent("main.py:387 acompletion", 1), + FunctionTraceEvent("llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function", 2), + FunctionTraceEvent( + "llms/anthropic/experimental_pass_through/messages/handler.py:416 anthropic_messages_handler", 0 + ), + FunctionTraceEvent( + "llms/anthropic/experimental_pass_through/messages/transformation.py:575" + " AnthropicMessagesConfig.transform_anthropic_messages_request", + 4, + ), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("execute_chat_completions_provider_call", 1), + ) + assert pipeline_steps("messages", "python", events) == ( + FunctionTraceEvent("execute_messages_provider_call", 0), + FunctionTraceEvent("transform_request", 1), + ) + + +@pytest.mark.parametrize("function", ("completion", "completion_function", "acompletion_function")) +def test_chat_projection_includes_sync_and_async_handlers(function: str) -> None: + events: Final = (FunctionTraceEvent(f"llms/anthropic/chat/handler.py:100 AnthropicChatCompletion.{function}", 0),) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("execute_chat_completions_provider_call", 0), + ) + + +def test_trace_diff_reports_no_difference_for_identical_steps() -> None: + steps: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("transform_ocr_request", 1), + ) + + diff: Final = trace_diff(steps, steps) + + assert diff.python_only == () + assert diff.rust_only == () + assert diff.shared_order_matches + assert diff.matches + + +def test_trace_diff_reports_exclusive_steps_and_reordered_shared_steps() -> None: + python: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("supported_ocr_params", 1), + FunctionTraceEvent("map_ocr_params", 1), + FunctionTraceEvent("http_request", 2), + ) + rust: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("map_ocr_params", 1), + FunctionTraceEvent("supported_ocr_params", 2), + FunctionTraceEvent("transform_ocr_response", 2), + ) + + diff: Final = trace_diff(python, rust) + + assert diff.python_only == ("http_request",) + assert diff.rust_only == ("transform_ocr_response",) + assert not diff.shared_order_matches + assert not diff.matches + + +def test_trace_diff_does_not_claim_empty_or_disjoint_traces_match() -> None: + assert not trace_diff((), ()).shared_order_matches + assert not trace_diff((FunctionTraceEvent("ocr", 0),), (FunctionTraceEvent("messages", 0),)).shared_order_matches + + +def test_projection_uses_actual_ancestors_after_coroutine_resumption() -> None: + entrypoint: Final = "main.py:387 acompletion" + handler: Final = "llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function" + events: Final = ( + FunctionTraceEvent(entrypoint, 0, ()), + FunctionTraceEvent(handler, 1, (entrypoint,)), + FunctionTraceEvent("utils.py:100 unrelated_worker", 0, ()), + FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_response", 1, (handler,)), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("execute_chat_completions_provider_call", 1), + FunctionTraceEvent("transform_response", 2), + ) + + +def test_projection_does_not_nest_siblings_under_a_returned_config_lookup() -> None: + events: Final = ( + FunctionTraceEvent("main.py:387 completion", 0), + FunctionTraceEvent("utils.py:100 ProviderConfigManager.get_provider_chat_config", 1), + FunctionTraceEvent("utils.py:200 unrelated_helper", 1), + FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_request", 2), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("get_provider_chat_config", 1), + FunctionTraceEvent("transform_request", 1), + ) + + +CHAT_RUST_STEPS: Final = ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "execute_chat_completions_provider_call", + "validate_environment", + "transform_request", + "http_request", + "transform_response", +) + + +@pytest.mark.parametrize("missing", CHAT_RUST_STEPS) +def test_pipeline_check_rejects_missing_stages(missing: str) -> None: + steps: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS if name != missing) + + assert f"missing {missing}" in pipeline_issues("chat_completions", "rust", steps) + + +def test_pipeline_check_rejects_http_before_request_transformation() -> None: + steps: Final = tuple( + FunctionTraceEvent(name, 0) + for name in ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "execute_chat_completions_provider_call", + "validate_environment", + "http_request", + "transform_request", + "transform_response", + ) + ) + + assert "transform_request must precede http_request" in pipeline_issues("chat_completions", "rust", steps) + + +def test_step_parity_rejects_different_handler_boundaries_even_with_valid_stages() -> None: + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) + python: Final = tuple( + FunctionTraceEvent(name, 0) + for name in ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "validate_environment", + "transform_request", + "execute_chat_completions_provider_call", + "http_request", + "transform_response", + ) + ) + + assert not trace_diff(python, rust).shared_order_matches + assert not trace_diff(python, rust).matches + assert pipeline_issues("chat_completions", "python", python) == () + assert pipeline_issues("chat_completions", "rust", rust) == () + + +def test_step_parity_rejects_an_exclusive_helper_with_matching_shared_order() -> None: + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) + python: Final = (*rust, FunctionTraceEvent("unmatched_helper", 0)) + diff: Final = trace_diff(python, rust) + + assert diff.shared_order_matches + assert not diff.matches diff --git a/tests/sdk_function_trace/test_table.py b/tests/sdk_function_trace/test_table.py new file mode 100644 index 00000000000..c2341a391a9 --- /dev/null +++ b/tests/sdk_function_trace/test_table.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import re +from typing import Final + +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.table import format_trace_table + + +def test_table_aligns_matches_after_missing_steps_and_preserves_indentation() -> None: + python: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("python_helper", 1), + FunctionTraceEvent("http_request", 2), + ) + rust: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("rust_helper", 1), + FunctionTraceEvent("http_request", 1), + ) + output: Final = format_trace_table(python, rust, colorize=False) + rows: Final = tuple(line.split("|")[1:-1] for line in output.splitlines() if line.startswith("|")) + + assert tuple(tuple(cell.strip() for cell in row) for row in rows) == ( + ("python (3 steps)", "rust (3 steps)", "comparison"), + ("ocr", "ocr", "match"), + ("python_helper", "", "python only"), + ("", "rust_helper", "rust only"), + ("http_request", "http_request", "match"), + ) + assert rows[-1][0].startswith(" http_request") + assert rows[-1][1].startswith(" http_request") + assert len({len(line) for line in output.splitlines()}) == 1 + assert "\033[" not in output + + +def test_table_marks_reordered_calls_and_keeps_both_execution_orders() -> None: + python: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "map", "validate", "http")) + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "validate", "map", "http")) + output: Final = format_trace_table(python, rust, colorize=True) + plain: Final = re.sub(r"\033\[[0-9;]*m", "", output) + rows: Final = tuple(line.split("|")[1:-1] for line in plain.splitlines() if line.startswith("|"))[1:] + + assert tuple(row[0].strip() for row in rows if row[0].strip()) == tuple(event.function for event in python) + assert tuple(row[1].strip() for row in rows if row[1].strip()) == tuple(event.function for event in rust) + assert plain.count("reordered") == 2 + assert output.count("\033[31m") == 2 + assert "only" not in output + + +def test_table_colors_match_and_exclusive_rows_without_changing_alignment() -> None: + python: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("python_helper", 1)) + rust: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("rust_helper", 1)) + colored: Final = format_trace_table(python, rust, colorize=True) + + assert re.sub(r"\033\[[0-9;]*m", "", colored) == format_trace_table(python, rust, colorize=False) + assert next(line for line in colored.splitlines() if "match" in line).startswith("\033[32m") + assert next(line for line in colored.splitlines() if "python only" in line).startswith("\033[34m") + assert next(line for line in colored.splitlines() if "rust only" in line).startswith("\033[33m") + + +def test_table_handles_empty_traces() -> None: + output: Final = format_trace_table((), (), colorize=False) + + assert "python (0 steps)" in output + assert "rust (0 steps)" in output + assert "match" not in output