scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

refactor(python-bridge): split routes and add shared function tracing (#39031)

Browse source 
* refactor(python-bridge): split non-streaming bridge modules

* refactor(python-bridge): bring shared function tracing into route layer

* feat(dev): list Python route functions and call sites

* feat(dev): list Rust route functions and call sites

* docs(dev): record OCR parity gaps across Python and Rust

* feat(dev): list executed SDK calls with runtime tracing

* feat(dev): report Python vs Rust SDK pipeline steps in one CLI

* feat(dev): side-by-side pipeline step report in compare CLI

* fix(dev): drop invalid Final annotations in compare cell loop

* feat(dev): blue python-only and yellow rust-only steps in compare CLI

* feat(dev): vertical layout with section spacing in compare CLI

* fix(dev): validate SDK trace stages across sync and async routes

* refactor(rust): align SDK route call structure with Python

* refactor(python-bridge): share sync and async route call wrappers

* refactor(dev): split compare CLI into fixtures, runtime, and report modules

* fix(ci): run SDK trace tests and satisfy test lint
This commit is contained in:
yujonglee 2026-09-02 16:26:35 -07:00 • committed by GitHub
parent 3c6b0705b2
commit 198906495f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
66 changed files with 2993 additions and 713 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

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

40
litellm-rust/Cargo.lock generated
View file

@ -1392,6 +1392,12 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "libc"
version = "0.2.186"
@ -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"

2
litellm-rust/Cargo.toml
View file

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

1
litellm-rust/crates/ai-gateway/Cargo.toml
View file

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

7
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
View file

@ -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<Option<(&str, &str)>, Error> {
let Some(object) = document.as_object() else {
return Ok(None);

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

@ -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<Value, Error> {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn execute_ocr_provider_call(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
) -> Result<Value, Error> {
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()))?;

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

@ -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<ProviderOcrRequest, Error> {
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<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLifecycleHooks {
impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> 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<PreparedOcrRequest, ProviderOcrRequest, Value> 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<PreparedOcrRequest, ProviderOcrRequest, Value> 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,

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

@ -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<Value, Error> {
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
}

20
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
View file

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

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

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

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

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

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

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

6
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -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<Value, Error> {
@ -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();

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

@ -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<Value, Error> {
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
.await

2
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -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<ProviderAudioTranscriptionRequest, Error> {

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

@ -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<String, Value>) -> Map<String, Value> {
params
.iter()

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

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

10
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -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<ChatCompletionsResponse, Error> {
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.

5
litellm-rust/crates/core/src/chat_completions/mod.rs
View file

@ -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<ChatCompletionsResponse, Error> {
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

44
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -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<Vec<ChatMessage>, 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<ProviderChatCompletionsRequest, Error> {
) -> Result<ResolvedChatCompletionsRequest<'_>, 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<ProviderChatCompletionsRequest, Error> {
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,

10
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -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<ProviderChatCompletionsRequest, Error> {
prepare_provider_request(resolve_request(request)?)
}
fn request<'a>(
model: &'a str,

13
litellm-rust/crates/core/src/chat_completions/transformation.rs
View file

@ -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<String, Value>,
) -> Option<Unsupported> {
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<String, Value>,
) -> Option<Unsupported> {
@ -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"))

11
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -22,6 +22,17 @@ pub struct ChatCompletionsRequest<'a> {
pub timeout: Option<Duration>,
}
pub(super) struct ResolvedChatCompletionsRequest<'a> {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) messages: Vec<ChatMessage>,
pub(super) optional_params: Map<String, Value>,
pub(super) api_key: Option<&'a str>,
pub(super) api_base: Option<&'a str>,
pub(super) extra_headers: Option<Map<String, Value>>,
pub(super) timeout: Option<Duration>,
}
pub(super) struct ProviderChatCompletionsRequest {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,

7
litellm-rust/crates/core/src/http_utils.rs
View file

@ -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<reqwest::Response, reqwest::Error> {
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 {

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

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

17
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -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<AnthropicMessagesResponse, Error> {
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<reqwest::Response, Error> {
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();

6
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -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<AnthropicMessagesResponse, Error> {
execute_messages_provider_call(prepare_messages_call(request)?).await
execute_messages_provider_call(request).await
}
pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {
execute_messages_provider_stream(prepare_messages_call(request)?).await
execute_messages_provider_stream(request).await
}
#[cfg(test)]

63
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -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<ProviderMessagesRequest, Error> {
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<Map<String, Value>>,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Vec<(String, String)>, 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)
}

2
litellm-rust/crates/core/src/messages/transformation.rs
View file

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

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

@ -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<String, Value>) -> Map<String, Value> {
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<String>,
) -> Result<String, Error>;
#[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<String>,
) -> Result<Vec<(String, String)>, 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
}

14
litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs
View file

@ -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<String, Value>,
) -> Option<Unsupported> {
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,

1
litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs
View file

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

1
litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs
View file

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

3
litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs
View file

@ -46,10 +46,12 @@ fn optional_string<'a>(params: &'a Map<String, Value>, 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,

72
litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs
View file

@ -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<String, Value>) -> 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<String, Value>,
) -> Option<Unsupported> {
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(

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

@ -70,10 +70,12 @@ pub struct MistralOcrConfig;
pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
impl OcrProviderConfig for MistralOcrConfig {
#[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<String, Value>) -> Map<String, Va
MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn transform_ocr_request(
model: &str,
document: Value,
@ -169,6 +175,7 @@ pub fn transform_ocr_request(
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn transform_ocr_response(model: &str, response_json: Value) -> Result<OcrResponseData, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}

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

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

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

@ -0,0 +1 @@
pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

23
litellm-rust/crates/python-bridge/src/diagnostics.rs Normal file
View file

@ -0,0 +1,23 @@
use litellm_python_interop::release_count;
use pyo3::prelude::*;
use pyo3::types::PyDict;
#[pyfunction]
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
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(())
}

61
litellm-rust/crates/python-bridge/src/errors.rs Normal file
View file

@ -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::<RustBridgeDeclined>())?;
module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())
}

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

@ -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<T> {
Plain(T),
Traced {
response: T,
trace: Vec<FunctionTraceEvent>,
},
}
pub(crate) async fn trace_call<T, E>(
future: impl Future<Output = Result<T, E>>,
enabled: bool,
) -> Result<TraceResponse<T>, 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<Mutex<Vec<FunctionTraceEvent>>>,
}
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<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let depth = context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default();
self.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.push(FunctionTraceEvent {
function: attributes.metadata().name(),
depth,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[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,
},
]
);
}
}

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

@ -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<String, Value>>,
Map<String, Value>,
Option<Duration>,
);
fn messages_response_to_py(
py: Python<'_>,
response: AnthropicMessagesResponse,
) -> PyResult<Py<PyAny>> {
to_py(py, &response)
}
fn chat_completions_response_to_py(
py: Python<'_>,
response: ChatCompletionsResponse,
) -> PyResult<Py<PyAny>> {
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<Py<PyAny>>,
) -> PyResult<Map<String, Value>> {
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<f64>) -> Option<Duration> {
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<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
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<PyAny>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledOcrInputs> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
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<Map<String, Value>>, Option<Duration>);
fn marshal_messages_inputs(
py: Python<'_>,
body: Py<PyAny>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledMessagesInputs> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
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<String, Value>,
Option<Map<String, Value>>,
Option<Duration>,
);
fn marshal_chat_completions_inputs(
py: Python<'_>,
messages: Py<PyAny>,
optional_params: Option<Py<PyAny>>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledChatCompletionsInputs> {
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<PyAny>,
optional_params: Option<Py<PyAny>>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
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<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
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<Py<PyAny>> {
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::<RustBridgeDeclined>())?;
module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())?;
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::<ResponsesWebSocketConnection>()?;
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")
);
}
});
}
}

63
litellm-rust/crates/python-bridge/src/marshal.rs Normal file
View file

@ -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<Py<PyAny>>,
) -> PyResult<Map<String, Value>> {
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<Py<PyAny>>,
) -> PyResult<Option<Map<String, Value>>> {
value
.map(|value| optional_object_to_map(py, name, Some(value)))
.transpose()
}
pub(crate) fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
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<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
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()
}

126
litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs Normal file
View file

@ -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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Map<String, Value>>,
optional_params: Map<String, Value>,
timeout: Option<Duration>,
}
#[allow(clippy::too_many_arguments)]
fn marshal_inputs(
py: Python<'_>,
model: String,
audio: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<TranscriptionInputs> {
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<Value, Error> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Bound<'_, PyAny>> {
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)?)
}

156
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs Normal file
View file

@ -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<String, Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Map<String, Value>>,
timeout: Option<Duration>,
}
#[allow(clippy::too_many_arguments)]
fn marshal_inputs(
py: Python<'_>,
model: String,
messages: Py<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<ChatCompletionsInputs> {
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<ChatCompletionsResponse, Error> {
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<PyAny>,
optional_params: Option<Py<PyAny>>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
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<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Bound<'_, PyAny>> {
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)?)
}

122
litellm-rust/crates/python-bridge/src/routes/messages.rs Normal file
View file

@ -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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Map<String, Value>>,
timeout: Option<Duration>,
}
#[allow(clippy::too_many_arguments)]
fn marshal_inputs(
py: Python<'_>,
model: String,
body: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MessagesInputs> {
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<AnthropicMessagesResponse, Error> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Bound<'_, PyAny>> {
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)?)
}

53
litellm-rust/crates/python-bridge/src/routes/mod.rs Normal file
View file

@ -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<T>(
py: Python<'_>,
call: impl Future<Output = Result<T, Error>> + Send,
trace: bool,
map_err: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
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<Output = Result<T, Error>> + Send + 'static,
trace: bool,
map_err: fn(Error) -> PyErr,
) -> PyResult<Bound<'py, PyAny>>
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))
})
}

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

@ -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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Map<String, Value>>,
optional_params: Map<String, Value>,
timeout: Option<Duration>,
}
#[allow(clippy::too_many_arguments)]
fn marshal_inputs(
py: Python<'_>,
model: String,
document: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<OcrInputs> {
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<Value, Error> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Py<PyAny>> {
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<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
trace: bool,
) -> PyResult<Bound<'_, PyAny>> {
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)?)
}

30
tests/sdk_function_trace/README.md Normal file
View file

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

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

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

45
tests/sdk_function_trace/compare.py Normal file
View file

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

200
tests/sdk_function_trace/fixtures.py Normal file
View file

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

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

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

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

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

59
tests/sdk_function_trace/ocr-comparison.md Normal file
View file

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

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

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

175
tests/sdk_function_trace/report.py Normal file
View file

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

26
tests/sdk_function_trace/route-comparison.md Normal file
View file

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

128
tests/sdk_function_trace/runtime.py Normal file
View file

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

181
tests/sdk_function_trace/steps.py Normal file
View file

@ -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"(?<!async_)map_ocr_params$", "map_ocr_params"),
_step("validate_environment", r"(?<!_)validate_environment$", "validate_environment"),
_step("complete_url", r"get_complete_url$", "complete_url"),
_step("transform_ocr_request", r"(?<!async_)transform_ocr_request$", "transform_ocr_request"),
_step("execute_ocr_provider_call", r"BaseLLMHTTPHandler\.(?:async_)?ocr$", "execute_ocr_provider_call"),
_step("http_request", _POST, "http_request"),
_step("transform_ocr_response", r"(?<!async_)transform_ocr_response$", "transform_ocr_response"),
),
"chat_completions": (
_step("chat_completions", r"main\.py:\d+ a?completion$", "chat_completions"),
_step(
"get_provider_chat_config",
r"ProviderConfigManager\.get_provider_chat_config$",
"chat_completions_provider_config",
),
_step("supported_openai_params", r"get_supported_openai_params$", "supported_openai_params"),
_step("validate_environment", r"(?<!_)validate_environment$", "validate_environment"),
_step("transform_request", r"(?<!async_)transform_request$", "transform_request"),
_step(
"execute_chat_completions_provider_call",
r"a?completion_function$|ChatCompletion\.completion$",
"execute_chat_completions_provider_call",
),
_step("http_request", _POST, "http_request"),
_step("transform_response", r"(?<!async_)transform_response$", "transform_response"),
),
"messages": (
_step("messages", r"anthropic_interface/messages/__init__\.py:\d+ a?create$", "messages"),
_step(
"get_provider_messages_config",
r"ProviderConfigManager\.get_provider_anthropic_messages_config$",
"messages_provider_config",
),
_step("validate_environment", r"validate_anthropic_messages_environment$", "validate_environment"),
_step("transform_request", r"(?<!async_)transform_anthropic_messages_request$", "transform_request"),
_step("complete_url", r"get_complete_url$", "complete_url"),
_step(
"execute_messages_provider_call",
r"(?:async_)?anthropic_messages_handler$",
"execute_messages_provider_call",
),
_step("http_request", _POST, "http_request"),
_step("transform_response", r"(?<!async_)transform_anthropic_messages_response$", "transform_response"),
),
"audio_transcription": (
_step("audio_transcription", r"main\.py:\d+ a?transcription$", "audio_transcription"),
_step("prepare_audio_transcription_provider_call", rust="prepare_audio_transcription_provider_call"),
_step("get_non_default_transcription_params", r"get_non_default_transcription_params$"),
_step("map_transcription_params", r"get_optional_params_transcription$", "map_transcription_params"),
_step(
"get_provider_transcription_config",
r"ProviderConfigManager\.get_provider_audio_transcription_config$",
"provider_config",
),
_step("supported_transcription_params", rust="supported_transcription_params"),
_step("transform_transcription_request", rust="transform_transcription_request"),
_step(
"execute_audio_transcription_provider_call",
r"BedrockAudioTranscriptionRustDispatch\.(?:async_)?audio_transcriptions$",
"execute_audio_transcription_provider_call",
),
_step("transform_transcription_response", rust="transform_transcription_response"),
_step("http_request", rust="http_request"),
),
}
def pipeline_issues(route: str, engine: Engine, events: Sequence[FunctionTraceEvent]) -> 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

72
tests/sdk_function_trace/table.py Normal file
View file

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

33
tests/sdk_function_trace/test_mock_provider.py Normal file
View file

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

167
tests/sdk_function_trace/test_profiler.py Normal file
View file

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

47
tests/sdk_function_trace/test_runtime.py Normal file
View file

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

244
tests/sdk_function_trace/test_steps.py Normal file
View file

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

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

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