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

chore: merge litellm_internal_staging into litellm_fix_agent_mcp_grants

Browse source
This commit is contained in:
mateo-berri 2026-09-02 17:09:01 -07:00
commit f7ef2899ec
173 changed files with 13364 additions and 1304 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/test-rust.yml vendored
View file

@ -9,6 +9,7 @@ on:
- "rust-toolchain.toml"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
pull_request:
branches:
@ -23,6 +24,7 @@ on:
- "rust-toolchain.toml"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
permissions:
@ -121,3 +123,6 @@ jobs:
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl

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

14
basedpyright-code-budget.json
View file

@ -3,7 +3,7 @@
"limit": 14074
},
"reportArgumentType": {
"limit": 2216
"limit": 2215
},
"reportAssignmentType": {
"limit": 319
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5601
"limit": 5599
},
"reportMissingTypeArgument": {
"limit": 15306
"limit": 15290
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44364
"limit": 44360
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38350
"limit": 38332
},
"reportUnknownParameterType": {
"limit": 19625
"limit": 19623
},
"reportUnknownVariableType": {
"limit": 29877
"limit": 29858
},
"reportUnnecessaryCast": {
"limit": 111

62
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]]
@ -1442,13 +1450,18 @@ name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"criterion",
"futures-util",
"litellm-ai-gateway",
"litellm-core",
"litellm-python-interop",
"pyo3",
"pyo3-async-runtimes",
"serde",
"serde_json",
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1669,9 +1682,9 @@ dependencies = [
[[package]]
name = "pyo3"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"libc",
"once_cell",
@ -1697,18 +1710,18 @@ dependencies = [
[[package]]
name = "pyo3-build-config"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
dependencies = [
"libc",
"pyo3-build-config",
@ -1716,9 +1729,9 @@ dependencies = [
[[package]]
name = "pyo3-macros"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
@ -1728,9 +1741,9 @@ dependencies = [
[[package]]
name = "pyo3-macros-backend"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
dependencies = [
"heck",
"proc-macro2",
@ -2276,6 +2289,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 +2436,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 +2693,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"

4
litellm-rust/Cargo.toml
View file

@ -14,11 +14,13 @@ 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" }
axum = "0.7"
pyo3 = "0.29.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"

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.

46
litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs
View file

@ -1,46 +0,0 @@
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::error::Error;
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use serde_json::{Map, Value};
use std::collections::BTreeMap;
pub(super) fn audio_transcription_provider_config(
provider: &str,
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
match provider {
"bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
pub(super) fn string_headers(
headers: Option<Map<String, Value>>,
) -> Result<BTreeMap<String, String>, Error> {
headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
Error::InvalidRequest(format!(
"audio transcription extra_headers.{key} must be a string"
))
})
})
.collect()
}
pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
headers.keys().any(|key| key.eq_ignore_ascii_case(name))
}
pub(super) fn truncate_error_body(body: &str) -> String {
let truncated: String = body.chars().take(256).collect();
if truncated.chars().count() == body.chars().count() {
truncated
} else {
format!("{truncated}... (truncated)")
}
}

84
litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs
View file

@ -1,84 +0,0 @@
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::error::Error;
use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use serde_json::Value;
use std::time::SystemTime;
use super::common_utils::truncate_error_body;
use super::types::ProviderAudioTranscriptionRequest;
use crate::client::http_client;
pub(crate) async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> Result<Value, Error> {
let body = serde_json::to_vec(&request.body)
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
let mut request_builder = http_client().post(&request.url).body(body.clone());
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.into_json())
}
pub(crate) async fn sign_request(
request: &ProviderAudioTranscriptionRequest,
optional_params: &serde_json::Map<String, Value>,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let env_lookup = environment_lookup;
let auth = request
.config
.auth_strategy(&request.model, optional_params, &env_lookup)?;
let body = serde_json::to_vec(&request.body)
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
let mut headers = super::common_utils::string_headers(None)?;
headers.insert("Content-Type".to_string(), "application/json".to_string());
headers.extend(request.upstream_headers.iter().cloned());
match auth {
AudioTranscriptionAuth::Bearer => {}
AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
let credentials =
resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
.await?;
headers.extend(sign_bedrock_post(
&request.url,
&body,
&headers,
&region,
&credentials,
SystemTime::now(),
)?);
}
}
Ok(ProviderAudioTranscriptionRequest {
upstream_headers: headers.into_iter().collect(),
..request.clone()
})
}
pub(super) fn environment_lookup(key: &str) -> Option<String> {
std::env::var(key).ok()
}

82
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -1,13 +1,14 @@
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::audio_transcription::{
AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
prepare_audio_transcription_provider_call,
};
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
use super::handler::sign_request;
use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use super::types::PreparedAudioTranscriptionRequest;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
@ -88,46 +89,29 @@ impl AudioTranscriptionLifecycleHooks {
&self,
request: PreparedAudioTranscriptionRequest,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let config = audio_transcription_provider_config(&request.custom_llm_provider)
.ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?;
let env_lookup = super::handler::environment_lookup;
let headers = string_headers(request.extra_headers)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let body = config.transform_transcription_request(
&request.model,
request.audio,
filtered_params,
)?;
let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(
&upstream_headers
.iter()
.cloned()
.collect::<std::collections::BTreeMap<_, _>>(),
"authorization",
)
&& let Some(api_key) = request.api_key.as_deref()
{
upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
let provider_request = ProviderAudioTranscriptionRequest {
model: request.model,
config,
url,
body: body.body,
upstream_headers,
timeout: request.timeout,
};
let provider_request = self.run_during_call_guardrails(provider_request).await?;
sign_request(&provider_request, &request.optional_params).await
let PreparedAudioTranscriptionRequest {
model,
custom_llm_provider,
audio,
api_key,
api_base,
extra_headers,
optional_params,
timeout,
..
} = request;
let provider_request =
prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: Some(&custom_llm_provider),
extra_headers,
optional_params,
timeout,
})?;
self.run_during_call_guardrails(provider_request).await
}
async fn run_during_call_guardrails(
@ -142,10 +126,10 @@ impl AudioTranscriptionLifecycleHooks {
.run_during_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": "bedrock",
"url": request.url,
"body": request.body,
"model": request.model(),
"custom_llm_provider": request.custom_llm_provider(),
"url": request.url(),
"body": request.body(),
})),
)
.await
@ -158,7 +142,7 @@ impl AudioTranscriptionLifecycleHooks {
let body = data.remove("body").ok_or_else(|| {
Error::InvalidRequest("audio transcription guardrail removed body".to_string())
})?;
Ok(ProviderAudioTranscriptionRequest { body, ..request })
Ok(request.with_body(body))
}
fn logging_payload(

4
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs
View file

@ -1,16 +1,14 @@
use litellm_core::Error;
use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::AudioTranscriptionRequest;
use handler::execute_audio_transcription_provider_call;
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {

11
litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs
View file

@ -1,7 +1,6 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use serde_json::{Map, Value};
@ -46,13 +45,3 @@ impl CallLifecycleRequest for PreparedAudioTranscriptionRequest {
)
}
}
#[derive(Clone)]
pub(crate) struct ProviderAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

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 }

14
litellm-rust/crates/core/src/audio_transcription/client.rs Normal file
View file

@ -0,0 +1,14 @@
use std::sync::OnceLock;
use std::time::Duration;
use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;
pub(super) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS))
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
}

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

@ -0,0 +1,91 @@
use serde_json::Value;
use crate::error::Error;
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> {
let body = serde_json::to_vec(&request.body)
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
let headers = signed_headers(&request, &body).await?;
let mut request_builder = http_client().post(&request.url).body(body);
for (key, value) in headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.into_json())
}
#[cfg(feature = "bedrock-auth")]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
use crate::providers::bedrock::audio_transcription::aws_auth_config;
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());
};
let env_lookup = |key: &str| std::env::var(key).ok();
let credentials = resolve_credentials(
aws_auth_config(&request.optional_params, &env_lookup),
&env_lookup,
)
.await?;
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
let signature = sign_bedrock_post(
&request.url,
body,
&unsigned,
region,
&credentials,
SystemTime::now(),
)?;
Ok(unsigned.into_iter().chain(signature).collect())
}
#[cfg(not(feature = "bedrock-auth"))]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
_body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
match request.auth {
AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
}
}

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

@ -1,2 +1,21 @@
use crate::Error;
mod client;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
use serde_json::Value;
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
}
#[cfg(test)]
mod tests;

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

@ -0,0 +1,74 @@
use crate::error::Error;
use crate::http_utils::{has_header, string_headers};
#[cfg(feature = "bedrock-auth")]
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
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" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
}
let _ = provider;
None
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn prepare_audio_transcription_provider_call(
request: AudioTranscriptionRequest<'_>,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.or_else(|| {
request
.custom_llm_provider
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
})
})
.ok_or_else(|| {
Error::InvalidProvider(
"unable to resolve custom_llm_provider for audio transcription request".to_string(),
)
})?;
let model = provider_info.model.to_string();
let config = provider_config(provider_info.custom_llm_provider)
.ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers("audio transcription", request.extra_headers)?;
let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?;
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(&headers, "authorization")
&& let Some(api_key) = request.api_key
{
headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));
}
let url = config.complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
config,
url,
body: transformed.body,
upstream_headers: headers,
auth,
#[cfg(feature = "bedrock-auth")]
optional_params: request.optional_params,
timeout: request.timeout,
})
}

50
litellm-rust/crates/core/src/audio_transcription/tests.rs Normal file
View file

@ -0,0 +1,50 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use serde_json::{Map, json};
use super::audio_transcription;
use super::types::AudioTranscriptionRequest;
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("connection");
let mut request = Vec::new();
let mut buffer = [0_u8; 16_384];
let count = stream.read(&mut buffer).expect("request");
request.extend_from_slice(&buffer[..count]);
let request = String::from_utf8_lossy(&request);
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
assert!(request.contains("x-amz-date:"));
assert!(request.contains("\"bytes\":\"AQI=\""));
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
stream.write_all(response).expect("response");
});
let optional_params = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let api_base = format!("http://{address}");
let response = audio_transcription(AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
api_key: None,
api_base: Some(&api_base),
custom_llm_provider: Some("bedrock"),
extra_headers: None,
optional_params,
timeout: None,
})
.await
.expect("transcription");
assert_eq!(response, json!({"text": "hello"}));
server.join().expect("server");
}

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

53
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -1,5 +1,56 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}
#[derive(Clone)]
pub struct ProviderAudioTranscriptionRequest {
pub(super) model: String,
pub(super) custom_llm_provider: String,
pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: AudioTranscriptionAuth,
#[cfg(feature = "bedrock-auth")]
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}
impl ProviderAudioTranscriptionRequest {
pub fn model(&self) -> &str {
&self.model
}
pub fn custom_llm_provider(&self) -> &str {
&self.custom_llm_provider
}
pub fn url(&self) -> &str {
&self.url
}
pub fn body(&self) -> &Value {
&self.body
}
pub fn with_body(self, body: Value) -> Self {
Self { body, ..self }
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AudioTranscriptionRequestData {

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,

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

@ -30,6 +30,8 @@ pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600;
/// Connect timeout for chat completions provider calls, in seconds.
pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600;
/// `object` field every non-streaming chat completion response carries.
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";

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

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

@ -16,16 +16,21 @@ extension-module = ["pyo3/extension-module"]
panic-test = []
[dependencies]
futures-util.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
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
tokio-tungstenite.workspace = true
[[bench]]
name = "serialization"

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

423
litellm-rust/crates/python-bridge/src/execution.rs Normal file
View file

@ -0,0 +1,423 @@
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::time::Duration;
use futures_util::FutureExt;
use litellm_core::error::Error;
use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(crate) fn run_sync<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
}
fn run_sync_on<T, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
let result = map_core_result(result, map_error)?;
Pythonized(result).into_pyobject(py).map(Bound::unbind)
}
pub(crate) fn run_async<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let result = catch_future_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
})
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),
Err(error) => Err(
std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
.map_err(panic_to_pyerr)?,
),
}
}
async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
AssertUnwindSafe(future)
.catch_unwind()
.await
.map_err(panic_to_pyerr)
}
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
let future = catch_future_panic(future);
tokio::pin!(future);
let signal_interval = Duration::from_millis(50);
let mut signal_checks =
time::interval_at(time::Instant::now() + signal_interval, signal_interval);
signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
result = &mut future => return result,
_ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::future::poll_fn;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, mpsc};
use std::task::Poll;
use std::thread;
use std::time::Instant;
use pyo3::panic::PanicException;
use pyo3::types::{PyDict, PyModule};
use serde::Serializer;
use tokio::runtime::Builder;
use super::*;
fn runtime_error(error: Error) -> PyErr {
PyRuntimeError::new_err(error.to_string())
}
fn panicking_error_mapper(_error: Error) -> PyErr {
panic!("error mapper panicked")
}
struct PanickingOutput;
static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
impl Serialize for PanickingOutput {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[pyfunction]
fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(py, async { Ok(PanickingOutput) }, runtime_error)
}
#[pyfunction]
fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(
py,
async {
ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
Ok(true)
},
runtime_error,
)
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
result
.expect("route should complete")
.bind(py)
.extract()
.expect("result should convert")
}
#[test]
fn sync_runner_polls_future_on_the_caller_thread() {
Python::initialize();
Python::attach(|py| {
let caller_thread = std::thread::current().id();
let result = run_sync(
py,
async move { Ok(std::thread::current().id() == caller_thread) },
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_releases_gil_while_waiting() {
Python::initialize();
Python::attach(|py| {
let result = run_sync(
py,
async {
let gil_acquired = tokio::time::timeout(
Duration::from_secs(2),
tokio::task::spawn_blocking(|| Python::attach(|_| true)),
)
.await;
Ok(matches!(gil_acquired, Ok(Ok(true))))
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_rejects_calls_from_a_tokio_context() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
let error = runtime.block_on(async {
Python::attach(|py| {
run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
.expect_err("sync route should reject a nested Tokio runtime")
})
});
assert_eq!(
error.to_string(),
"RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
);
}
#[test]
fn sync_runner_can_drive_a_current_thread_runtime() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
Python::attach(|py| {
let result = run_sync_on(
py,
&runtime,
async {
tokio::task::yield_now().await;
Ok(true)
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_maps_a_panicked_future() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
runtime_error,
)
.expect_err("panicked route should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: route future panicked");
});
}
#[test]
fn sync_runner_maps_a_panicked_error_mapper() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
async { Err(Error::InvalidRequest("invalid".to_string())) },
panicking_error_mapper,
)
.expect_err("panicked mapper should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: error mapper panicked");
});
}
#[test]
fn sync_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
#[test]
fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
Python::initialize();
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let callers: Vec<_> = (0..2)
.map(|_| {
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
Python::attach(|py| {
extract_bool(
py,
run_sync(
py,
async move {
Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
.await
.is_ok())
},
runtime_error,
),
)
})
})
})
.collect();
let results: Vec<_> = callers
.into_iter()
.map(|caller| caller.join().expect("caller should not panic"))
.collect();
assert_eq!(results, vec![true, true]);
}
#[test]
fn async_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
module
.add_function(
wrap_pyfunction!(async_serialization_panic, &module)
.expect("function should wrap"),
)
.expect("function should register");
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
try:
await runtime.async_serialization_panic()
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "serializer panicked"
else:
raise AssertionError("serializer panic was not raised")
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("serializer panic should reach the Python awaiter");
});
}
#[test]
fn async_result_delivery_does_not_stall_tokio_workers() {
Python::initialize();
ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
for function in [
wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
] {
module
.add_function(function)
.expect("function should register");
}
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
worker_count = runtime.runtime_worker_count()
awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
assert runtime.runtime_is_responsive(worker_count)
assert await asyncio.gather(*awaitables) == [True] * worker_count
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("result delivery should leave Tokio workers responsive");
});
}
}

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,
},
]
);
}
}

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

@ -1,142 +1,18 @@
use std::collections::HashMap;
use std::time::Duration;
mod constants;
mod diagnostics;
mod errors;
mod execution;
pub mod function_trace;
mod marshal;
mod routes;
use litellm_ai_gateway::io::audio_transcription::{
AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
};
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
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;
use serde_json::Value;
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 {
@ -151,16 +27,16 @@ impl ResponsesWebSocketConnection {
_cls: &Bound<'py, pyo3::types::PyType>,
py: Python<'py>,
url: String,
headers: Option<Py<PyAny>>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] headers: Option<Value>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(py, headers)?;
let headers = marshal_headers(headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner }))
Ok(ResponsesWebSocketConnection { inner })
})
}
@ -186,453 +62,126 @@ 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);
#[pymodule(gil_used = false)]
mod _native {
use pyo3::prelude::*;
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)),
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
super::errors::register(module)?;
super::routes::register(module)?;
module.add_class::<super::ResponsesWebSocketConnection>()?;
super::diagnostics::register(module)
}
}
#[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,
)?;
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::time::Duration;
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)?;
use futures_util::{SinkExt, StreamExt};
use pyo3::types::PyDict;
use tokio::net::TcpListener;
use tokio_tungstenite::{accept_async, tungstenite::Message};
Python::attach(|py| to_py(py, &value))
})
}
use super::*;
#[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,
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)),
#[test]
fn module_registration_preserves_the_public_surface() {
Python::initialize();
Python::attach(|py| {
let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
let expected = [
"RustBridgeDeclined",
"RustUpstreamError",
"ocr",
"aocr",
"transcription",
"atranscription",
"messages",
"amessages",
"chat_completions_decline",
"chat_completions",
"achat_completions",
"ResponsesWebSocketConnection",
"gil_stats",
];
let public_names: Vec<String> = module
.dict()
.keys()
.extract::<Vec<String>>()
.expect("module names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.collect();
assert_eq!(public_names, expected);
});
}
#[test]
fn responses_websocket_connection_round_trips_through_python() {
Python::initialize();
let runtime = pyo3_async_runtimes::tokio::get_runtime();
let listener = runtime
.block_on(TcpListener::bind("127.0.0.1:0"))
.expect("listener should bind");
let address = listener
.local_addr()
.expect("listener should have an address");
let server = runtime.spawn(async move {
let (stream, _) = listener.accept().await.expect("server should accept");
let mut socket = accept_async(stream)
.await
.expect("handshake should succeed");
let message = socket
.next()
.await
.expect("client should send a frame")
.expect("client frame should be valid");
assert_eq!(message, Message::Text("from-python".into()));
socket
.send(Message::Text("from-server".into()))
.await
.expect("server should reply");
assert!(matches!(socket.next().await, Some(Ok(Message::Close(_)))));
});
Python::attach(|py| {
let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
let locals = PyDict::new(py);
locals
.set_item("native", &module)
.expect("module should enter Python locals");
locals
.set_item("url", format!("ws://{address}"))
.expect("URL should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
connection = await native.ResponsesWebSocketConnection.connect(url)
assert type(connection) is native.ResponsesWebSocketConnection
await connection.send_text("from-python")
assert await connection.recv_text() == "from-server"
await connection.close()
assert await connection.recv_text() is None
asyncio.run(asyncio.wait_for(exercise(), timeout=5))
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("Python WebSocket methods should round trip");
});
runtime
.block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await })
.expect("server should finish")
.expect("server task should not panic");
}
}
#[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,
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))
})
}
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)?)?;
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(())
}

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

@ -0,0 +1,104 @@
use std::collections::HashMap;
use std::time::Duration;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::{Map, Value};
pub(crate) struct RouteOptions {
pub(crate) model: String,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) custom_llm_provider: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) timeout: Option<Duration>,
}
pub(crate) struct RouteOptionsInputs {
pub(crate) model: String,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) custom_llm_provider: Option<String>,
pub(crate) extra_headers: Option<Value>,
pub(crate) timeout_seconds: Option<f64>,
}
impl RouteOptions {
pub(crate) fn from_python(inputs: RouteOptionsInputs) -> PyResult<Self> {
Ok(Self {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: optional_object("extra_headers", inputs.extra_headers)?,
timeout: optional_timeout(inputs.timeout_seconds),
})
}
}
pub(crate) fn required_value(
name: &'static str,
value: Value,
expected: fn(&Value) -> bool,
expected_name: &'static str,
) -> PyResult<Value> {
if expected(&value) {
return Ok(value);
}
Err(PyValueError::new_err(format!(
"{name} must be a {expected_name}"
)))
}
pub(crate) fn object_or_empty(
name: &'static str,
value: Option<Value>,
) -> PyResult<Map<String, Value>> {
match value {
Some(value) => object(name, value),
None => Ok(Map::new()),
}
}
fn optional_object(
name: &'static str,
value: Option<Value>,
) -> PyResult<Option<Map<String, Value>>> {
value.map(|value| object(name, value)).transpose()
}
fn object(name: &'static str, value: Value) -> PyResult<Map<String, Value>> {
match value {
Value::Object(map) => Ok(map),
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
}
}
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(headers: Option<Value>) -> PyResult<HashMap<String, String>> {
let value = match headers {
Some(headers) => headers,
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()
}

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

@ -0,0 +1,71 @@
use litellm_core::Error;
use std::future::Future;
use litellm_core::audio_transcription::{
AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
};
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
fn prepare_transcription(
inputs: AudioTranscriptionInputs,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let audio = inputs.audio;
let options = RouteOptions::from_python(RouteOptionsInputs {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: inputs.extra_headers,
timeout_seconds: inputs.timeout_seconds,
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
Ok(async move {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
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
})
}
bridge_route! {
sync = transcription,
asynchronous = atranscription,
inputs = AudioTranscriptionInputs,
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
audio: Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_transcription,
errors = core_error_to_pyerr,
}

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

@ -0,0 +1,91 @@
use litellm_core::Error;
use std::future::Future;
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
use litellm_core::chat_completions::{
chat_completions as run_chat_completions, chat_completions_decline_reason,
};
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::chat_completions_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value};
fn prepare_chat_completions(
inputs: ChatCompletionsInputs,
) -> PyResult<impl Future<Output = Result<ChatCompletionsResponse, Error>> + Send + 'static> {
let messages = required_value("messages", inputs.messages, Value::is_array, "list")?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
let options = RouteOptions::from_python(RouteOptionsInputs {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: inputs.extra_headers,
timeout_seconds: inputs.timeout_seconds,
})?;
Ok(async move {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
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
})
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
fn chat_completions_decline(
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
let optional_params = object_or_empty("optional_params", optional_params)?;
Ok(chat_completions_decline_reason(
&model,
custom_llm_provider.as_deref(),
messages,
&optional_params,
)
.map(str::to_string))
}
bridge_route! {
sync = chat_completions,
asynchronous = achat_completions,
inputs = ChatCompletionsInputs,
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
messages: Value,
},
optional = {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_chat_completions,
errors = chat_completions_error_to_pyerr,
extra = [chat_completions_decline],
}

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

@ -0,0 +1,429 @@
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use pyo3::types::PyCFunction;
macro_rules! bridge_route {
(
sync = $sync_name:ident,
asynchronous = $async_name:ident,
inputs = $inputs:ident,
required = { $($(#[$required_attr:meta])* $required_name:ident: $required_type:ty),+ $(,)? },
optional = { $($(#[$optional_attr:meta])* $optional_name:ident: $optional_type:ty),* $(,)? },
prepare = $prepare:path,
errors = $map_error:path
$(, extra = [$($extra:ident),* $(,)?])?
$(,)?
) => {
struct $inputs {
$($required_name: $required_type,)*
$($optional_name: $optional_type),*
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
trace: bool,
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_sync(
py,
$crate::function_trace::trace_call(future, trace),
$map_error,
)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
trace: bool,
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_async(
py,
$crate::function_trace::trace_call(future, trace),
$map_error,
)
}
pub(super) fn register(
module: &pyo3::Bound<'_, pyo3::types::PyModule>,
) -> pyo3::PyResult<()> {
$($($crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($extra, module)?)?;)*)?
$crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($sync_name, module)?)?;
$crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?;
Ok(())
}
};
}
pub(super) fn add_function(
module: &Bound<'_, PyModule>,
function: Bound<'_, PyCFunction>,
) -> PyResult<()> {
let name: String = function.getattr("__name__")?.extract()?;
if module.hasattr(&name)? {
return Err(PyRuntimeError::new_err(format!(
"duplicate native route: {name}"
)));
}
module.add_function(function)
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::sync::atomic::{AtomicBool, Ordering};
use litellm_core::error::Error;
use pyo3::exceptions::PyLookupError;
use pyo3::types::{PyDict, PyList};
use super::*;
mod synthetic {
use std::future::{Future, pending};
use super::*;
static FUTURE_DROPPED: AtomicBool = AtomicBool::new(false);
struct DropGuard;
impl Drop for DropGuard {
fn drop(&mut self) {
FUTURE_DROPPED.store(true, Ordering::SeqCst);
}
}
#[pyfunction]
fn future_dropped() -> bool {
FUTURE_DROPPED.load(Ordering::SeqCst)
}
bridge_route! {
sync = echo,
asynchronous = aecho,
inputs = EchoInputs,
required = { value: String },
optional = {},
prepare = prepare_echo,
errors = map_error,
extra = [future_dropped],
}
fn prepare_echo(
inputs: EchoInputs,
) -> PyResult<impl Future<Output = Result<String, Error>> + Send + 'static> {
FUTURE_DROPPED.store(false, Ordering::SeqCst);
let drop_guard = (inputs.value == "pending").then_some(DropGuard);
Ok(async move {
let _drop_guard = drop_guard;
tokio::task::yield_now().await;
match inputs.value.as_str() {
"error" => Err(Error::InvalidRequest("synthetic error".to_string())),
"map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
"panic" => panic!("synthetic panic"),
"pending" => {
pending::<()>().await;
unreachable!()
}
_ => Ok(inputs.value),
}
})
}
fn map_error(error: Error) -> PyErr {
if matches!(&error, Error::InvalidRequest(message) if message == "panic in mapper") {
panic!("synthetic mapper panic")
}
PyLookupError::new_err(error.to_string())
}
}
#[test]
fn sync_and_async_route_signatures_match_the_python_contract() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "routes").expect("module should be created");
crate::routes::register(&module).expect("routes should register");
let routes = [
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
),
(
"transcription",
"atranscription",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
),
(
"messages",
"amessages",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
),
(
"chat_completions",
"achat_completions",
"(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
),
];
for (sync_name, async_name, expected) in routes {
let sync_signature: String = module
.getattr(sync_name)
.and_then(|function| function.getattr("__text_signature__"))
.and_then(|signature| signature.extract())
.expect("sync signature should be available");
let async_signature: String = module
.getattr(async_name)
.and_then(|function| function.getattr("__text_signature__"))
.and_then(|signature| signature.extract())
.expect("async signature should be available");
assert_eq!(sync_signature, expected);
assert_eq!(async_signature, expected);
}
});
}
#[test]
fn sync_and_async_routes_apply_the_same_input_validation() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "routes").expect("module should be created");
crate::routes::register(&module).expect("routes should register");
let invalid_messages = PyDict::new(py);
let sync_chat_error = module
.getattr("chat_completions")
.and_then(|function| function.call1(("model", &invalid_messages)))
.expect_err("sync chat should reject a non-list messages value");
let async_chat_error = module
.getattr("achat_completions")
.and_then(|function| function.call1(("model", &invalid_messages)))
.expect_err("async chat should reject a non-list messages value");
assert_eq!(
sync_chat_error.to_string(),
"ValueError: messages must be a list"
);
assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string());
let invalid_body = PyList::empty(py);
let sync_messages_error = module
.getattr("messages")
.and_then(|function| function.call1(("model", &invalid_body)))
.expect_err("sync Messages should reject a non-dict body");
let async_messages_error = module
.getattr("amessages")
.and_then(|function| function.call1(("model", &invalid_body)))
.expect_err("async Messages should reject a non-dict body");
assert_eq!(
sync_messages_error.to_string(),
"ValueError: body must be a dict"
);
assert_eq!(
async_messages_error.to_string(),
sync_messages_error.to_string()
);
let invalid_headers = PyList::empty(py);
let kwargs = PyDict::new(py);
kwargs
.set_item("extra_headers", &invalid_headers)
.expect("kwargs should accept extra_headers");
let document = PyDict::new(py);
for (sync_name, async_name) in [("ocr", "aocr"), ("transcription", "atranscription")] {
let sync_error = module
.getattr(sync_name)
.and_then(|function| function.call(("model", &document), Some(&kwargs)))
.expect_err("sync route should reject non-dict extra_headers");
let async_error = module
.getattr(async_name)
.and_then(|function| function.call(("model", &document), Some(&kwargs)))
.expect_err("async route should reject non-dict extra_headers");
assert_eq!(
sync_error.to_string(),
"ValueError: extra_headers must be a dict"
);
assert_eq!(async_error.to_string(), sync_error.to_string());
}
});
}
#[test]
fn route_input_validation_preserves_left_to_right_order() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "routes").expect("module should be created");
crate::routes::register(&module).expect("routes should register");
let invalid = PyList::empty(py);
let chat_kwargs = PyDict::new(py);
chat_kwargs
.set_item("optional_params", &invalid)
.expect("kwargs should accept optional_params");
chat_kwargs
.set_item("extra_headers", &invalid)
.expect("kwargs should accept extra_headers");
let invalid_messages = PyDict::new(py);
let error = module
.getattr("chat_completions")
.and_then(|function| {
function.call(("model", &invalid_messages), Some(&chat_kwargs))
})
.expect_err("messages should be validated first");
assert_eq!(error.to_string(), "ValueError: messages must be a list");
let valid_messages = PyList::empty(py);
let error = module
.getattr("chat_completions")
.and_then(|function| function.call(("model", &valid_messages), Some(&chat_kwargs)))
.expect_err("optional_params should be validated before headers");
assert_eq!(
error.to_string(),
"ValueError: optional_params must be a dict"
);
let headers_kwargs = PyDict::new(py);
headers_kwargs
.set_item("extra_headers", &invalid)
.expect("kwargs should accept extra_headers");
let invalid_body = PyList::empty(py);
let error = module
.getattr("messages")
.and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs)))
.expect_err("body should be validated before headers");
assert_eq!(error.to_string(), "ValueError: body must be a dict");
let invalid_payload =
PyModule::new(py, "invalid_payload").expect("invalid payload should be created");
for name in ["ocr", "transcription"] {
let error = module
.getattr(name)
.and_then(|function| {
function.call(("model", &invalid_payload), Some(&headers_kwargs))
})
.expect_err("payload should be validated before headers");
assert!(!error.to_string().contains("extra_headers"));
}
});
}
#[test]
fn generated_routes_execute_sync_and_async_contracts() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "synthetic").expect("module should be created");
synthetic::register(&module).expect("routes should register");
let sync_value: String = module
.getattr("echo")
.and_then(|function| function.call1(("sync",)))
.and_then(|value| value.extract())
.expect("sync route should return its value");
assert_eq!(sync_value, "sync");
let sync_error = module
.getattr("echo")
.and_then(|function| function.call1(("error",)))
.expect_err("sync route should map its error");
assert!(sync_error.is_instance_of::<PyLookupError>(py));
assert_eq!(
sync_error.to_string(),
"LookupError: invalid request: synthetic error"
);
let locals = PyDict::new(py);
locals
.set_item("routes", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
assert await routes.aecho("async") == "async"
try:
await routes.aecho("error")
except LookupError as error:
assert str(error) == "invalid request: synthetic error"
else:
raise AssertionError("mapped error was not raised")
try:
await routes.aecho("panic")
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "synthetic panic"
else:
raise AssertionError("panic was not raised")
try:
await routes.aecho("map_panic")
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "synthetic mapper panic"
else:
raise AssertionError("mapper panic was not raised")
task = asyncio.ensure_future(routes.aecho("pending"))
await asyncio.sleep(0)
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
else:
raise AssertionError("cancelled route completed")
for _ in range(100):
if routes.future_dropped():
break
await asyncio.sleep(0.001)
assert routes.future_dropped()
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("async route contract should hold");
});
}
#[test]
fn route_registration_rejects_duplicate_python_names() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "synthetic").expect("module should be created");
synthetic::register(&module).expect("first registration should succeed");
let error = synthetic::register(&module)
.expect_err("duplicate registration should be rejected");
assert_eq!(
error.to_string(),
"RuntimeError: duplicate native route: future_dropped"
);
});
}
}

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

@ -0,0 +1,65 @@
use litellm_core::Error;
use litellm_core::messages::messages as run_messages;
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
use pyo3::prelude::*;
use serde_json::Value;
use std::future::Future;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value};
fn prepare_messages(
inputs: MessagesInputs,
) -> PyResult<impl Future<Output = Result<AnthropicMessagesResponse, Error>> + Send + 'static> {
let body = required_value("body", inputs.body, Value::is_object, "dict")?;
let options = RouteOptions::from_python(RouteOptionsInputs {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: inputs.extra_headers,
timeout_seconds: inputs.timeout_seconds,
})?;
Ok(async move {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
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
})
}
bridge_route! {
sync = messages,
asynchronous = amessages,
inputs = MessagesInputs,
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
body: Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_messages,
errors = core_error_to_pyerr,
}

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

@ -0,0 +1,16 @@
use pyo3::prelude::*;
#[macro_use]
mod definition;
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)
}

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

@ -0,0 +1,73 @@
use litellm_core::Error;
use std::future::Future;
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
fn prepare_ocr(
inputs: OcrInputs,
) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
let document = inputs.document;
let options = RouteOptions::from_python(RouteOptionsInputs {
model: inputs.model,
api_key: inputs.api_key,
api_base: inputs.api_base,
custom_llm_provider: inputs.custom_llm_provider,
extra_headers: inputs.extra_headers,
timeout_seconds: inputs.timeout_seconds,
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
Ok(async move {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
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
})
}
bridge_route! {
sync = ocr,
asynchronous = aocr,
inputs = OcrInputs,
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
document: Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,
errors = core_error_to_pyerr,
}

423
litellm-rust/crates/python-bridge/src/routes/runtime.rs Normal file
View file

@ -0,0 +1,423 @@
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::time::Duration;
use futures_util::FutureExt;
use litellm_core::error::Error;
use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(super) fn run_sync<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
}
fn run_sync_on<T, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
let result = map_core_result(result, map_error)?;
Pythonized(result).into_pyobject(py).map(Bound::unbind)
}
pub(super) fn run_async<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let result = catch_route_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
})
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),
Err(error) => Err(
std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
.map_err(panic_to_pyerr)?,
),
}
}
async fn catch_route_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
AssertUnwindSafe(future)
.catch_unwind()
.await
.map_err(panic_to_pyerr)
}
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
let future = catch_route_panic(future);
tokio::pin!(future);
let signal_interval = Duration::from_millis(50);
let mut signal_checks =
time::interval_at(time::Instant::now() + signal_interval, signal_interval);
signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
result = &mut future => return result,
_ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::future::poll_fn;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, mpsc};
use std::task::Poll;
use std::thread;
use std::time::Instant;
use pyo3::panic::PanicException;
use pyo3::types::{PyDict, PyModule};
use serde::Serializer;
use tokio::runtime::Builder;
use super::*;
fn runtime_error(error: Error) -> PyErr {
PyRuntimeError::new_err(error.to_string())
}
fn panicking_error_mapper(_error: Error) -> PyErr {
panic!("error mapper panicked")
}
struct PanickingOutput;
static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
impl Serialize for PanickingOutput {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[pyfunction]
fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(py, async { Ok(PanickingOutput) }, runtime_error)
}
#[pyfunction]
fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(
py,
async {
ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
Ok(true)
},
runtime_error,
)
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
result
.expect("route should complete")
.bind(py)
.extract()
.expect("result should convert")
}
#[test]
fn sync_runner_polls_future_on_the_caller_thread() {
Python::initialize();
Python::attach(|py| {
let caller_thread = std::thread::current().id();
let result = run_sync(
py,
async move { Ok(std::thread::current().id() == caller_thread) },
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_releases_gil_while_waiting() {
Python::initialize();
Python::attach(|py| {
let result = run_sync(
py,
async {
let gil_acquired = tokio::time::timeout(
Duration::from_secs(2),
tokio::task::spawn_blocking(|| Python::attach(|_| true)),
)
.await;
Ok(matches!(gil_acquired, Ok(Ok(true))))
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_rejects_calls_from_a_tokio_context() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
let error = runtime.block_on(async {
Python::attach(|py| {
run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
.expect_err("sync route should reject a nested Tokio runtime")
})
});
assert_eq!(
error.to_string(),
"RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
);
}
#[test]
fn sync_runner_can_drive_a_current_thread_runtime() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
Python::attach(|py| {
let result = run_sync_on(
py,
&runtime,
async {
tokio::task::yield_now().await;
Ok(true)
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_maps_a_panicked_future() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
runtime_error,
)
.expect_err("panicked route should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: route future panicked");
});
}
#[test]
fn sync_runner_maps_a_panicked_error_mapper() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
async { Err(Error::InvalidRequest("invalid".to_string())) },
panicking_error_mapper,
)
.expect_err("panicked mapper should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: error mapper panicked");
});
}
#[test]
fn sync_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
#[test]
fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
Python::initialize();
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let callers: Vec<_> = (0..2)
.map(|_| {
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
Python::attach(|py| {
extract_bool(
py,
run_sync(
py,
async move {
Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
.await
.is_ok())
},
runtime_error,
),
)
})
})
})
.collect();
let results: Vec<_> = callers
.into_iter()
.map(|caller| caller.join().expect("caller should not panic"))
.collect();
assert_eq!(results, vec![true, true]);
}
#[test]
fn async_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
module
.add_function(
wrap_pyfunction!(async_serialization_panic, &module)
.expect("function should wrap"),
)
.expect("function should register");
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
try:
await runtime.async_serialization_panic()
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "serializer panicked"
else:
raise AssertionError("serializer panic was not raised")
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("serializer panic should reach the Python awaiter");
});
}
#[test]
fn async_result_delivery_does_not_stall_tokio_workers() {
Python::initialize();
ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
for function in [
wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
] {
module
.add_function(function)
.expect("function should register");
}
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
worker_count = runtime.runtime_worker_count()
awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
assert runtime.runtime_is_responsive(worker_count)
assert await asyncio.gather(*awaitables) == [True] * worker_count
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("result delivery should leave Tokio workers responsive");
});
}
}

2
litellm-rust/crates/python-interop/src/lib.rs
View file

@ -2,4 +2,4 @@ mod gil;
mod marshal;
pub use gil::{release_count, release_gil};
pub use marshal::{from_py, to_py};
pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py};

72
litellm-rust/crates/python-interop/src/marshal.rs
View file

@ -1,4 +1,8 @@
use std::any::Any;
use std::panic::{AssertUnwindSafe, catch_unwind};
use pyo3::exceptions::PyValueError;
use pyo3::panic::PanicException;
use pyo3::prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;
@ -18,3 +22,71 @@ where
.map(Bound::unbind)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
pub struct Pythonized<T>(pub T);
impl<'py, T> IntoPyObject<'py> for Pythonized<T>
where
T: Serialize,
{
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
.map_err(panic_to_pyerr)?
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
pub fn panic_to_pyerr(payload: Box<dyn Any + Send>) -> PyErr {
let message = payload
.downcast_ref::<String>()
.map(String::as_str)
.or_else(|| payload.downcast_ref::<&str>().copied())
.unwrap_or("panic from Rust code");
PanicException::new_err(message.to_string())
}
#[cfg(test)]
mod tests {
use serde::Serializer;
use super::*;
struct PanickingSerializer;
impl Serialize for PanickingSerializer {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[test]
fn pythonized_converts_on_the_attached_thread() {
Python::initialize();
Python::attach(|py| {
let value: Vec<i32> = Pythonized(vec![1, 2, 3])
.into_pyobject(py)
.and_then(|value| value.extract())
.expect("value should convert");
assert_eq!(value, vec![1, 2, 3]);
});
}
#[test]
fn pythonized_maps_serializer_panics_to_a_base_exception() {
Python::initialize();
Python::attach(|py| {
let error = Pythonized(PanickingSerializer)
.into_pyobject(py)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
}

2
litellm/__init__.py
View file

@ -1417,7 +1417,7 @@ from .skills.main import (
)
from .containers.main import *
from .ocr.main import *
from .rust_bridge.ocr import use_litellm_rust
from .rust_bridge import use_litellm_rust
from .rag.main import *
from .sandbox.main import *
from .search.main import *

9
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -1955,11 +1955,10 @@ Model Info:
if not thresholds_enabled and not anomalies_enabled:
return
if prisma_client is None:
from litellm.proxy.proxy_server import prisma_client as global_prisma_client
from litellm.proxy.proxy_server import prisma_client as global_prisma_client
prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client
if prisma_client is None:
client: Final = prisma_client if prisma_client is not None else global_prisma_client
if client is None:
return
from litellm.integrations.SlackAlerting.user_spend_alerts import (
@ -1970,7 +1969,7 @@ Model Info:
try:
today: Final = datetime.datetime.now(datetime.timezone.utc).date()
rows: Final = await fetch_user_spend_rows(
prisma_client=prisma_client,
prisma_client=client,
today=today,
baseline_days=self.alerting_args.spend_anomaly_baseline_days,
)

1
litellm/integrations/websearch_interception/handler.py
View file

@ -419,7 +419,6 @@ class WebSearchInterceptionLogger(CustomLogger):
if call_type in (CallTypes.responses, CallTypes.aresponses):
return self._convert_responses_tools(kwargs=kwargs, tools=tools)
# Check if any tool is a web search tool (native or already LiteLLM standard)
has_websearch: Final = any(is_web_search_tool(t) for t in tools)
if not has_websearch:

3
litellm/llms/anthropic/chat/transformation.py
View file

@ -1565,6 +1565,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
optional_params.pop("thinking", None)
else:
optional_params["thinking"] = value
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider
)
elif param == "reasoning_effort":
# Accept both string ("low") and dict ({"effort": "low",
# "summary": "concise"}). The Responses->Chat parser keeps the

52
litellm/llms/anthropic/common_utils.py
View file

@ -13,7 +13,12 @@ import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME
from litellm.constants import (
DEFAULT_MODEL_CREATED_AT_TIME,
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_file_ids_from_messages,
)
@ -534,6 +539,51 @@ class AnthropicModelInfo(BaseLLMModelInfo):
)
optional_params.pop("thinking", None)
@staticmethod
def translate_legacy_thinking_for_adaptive_model(
model: str,
optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers
custom_llm_provider: str,
) -> None:
"""Translate legacy ``thinking.type=enabled`` to adaptive for the
adaptive-thinking models that reject it (4.7+ and the 5 families).
Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the
legacy shape natively, so it is forwarded verbatim and the caller's
``budget_tokens`` cap keeps applying. Caller-provided
``output_config.effort`` is never overridden.
"""
if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
return
if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
return
thinking: Final = optional_params.get("thinking")
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
return
effort: Final = AnthropicModelInfo._legacy_budget_to_effort(
model=model,
budget_tokens=int(thinking.get("budget_tokens") or 0),
custom_llm_provider=custom_llm_provider,
)
existing_output_config: Final = optional_params.get("output_config")
optional_params["thinking"] = {"type": "adaptive"}
optional_params["output_config"] = {
"effort": effort,
**(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})),
}
@staticmethod
def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str:
if budget_tokens >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and (
AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider)
):
return "xhigh"
if budget_tokens >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET:
return "high"
if budget_tokens >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET:
return "medium"
return "low"
def is_effort_used(
self,
optional_params: dict | None,

47
litellm/llms/anthropic/experimental_pass_through/messages/transformation.py
View file

@ -3,11 +3,6 @@ from typing import Any, Final
import httpx
from litellm.constants import (
DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
)
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import verbose_logger
@ -400,46 +395,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
existing_output_config.setdefault("effort", mapped_effort)
optional_params["output_config"] = existing_output_config
@staticmethod
def _translate_legacy_thinking_for_adaptive_model(
model: str, optional_params: dict, custom_llm_provider: str
) -> None:
"""Translate legacy ``thinking.type=enabled`` to adaptive for the
adaptive-thinking models that reject it (4.7+ and the 5 families).
Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the
legacy shape natively, so it is forwarded verbatim and the caller's
``budget_tokens`` cap keeps applying. Caller-provided
``output_config.effort`` is never overridden.
"""
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
return
if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
return
thinking: Final = optional_params.get("thinking")
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
return
budget: Final = int(thinking.get("budget_tokens") or 0)
if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and (
AnthropicConfig._supports_effort_level(model, "xhigh", custom_llm_provider)
):
effort = "xhigh"
elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET:
effort = "high"
elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET:
effort = "medium"
else:
effort = "low"
optional_params["thinking"] = {"type": "adaptive"}
existing_output_config = optional_params.get("output_config")
if not isinstance(existing_output_config, dict):
existing_output_config = {}
existing_output_config.setdefault("effort", effort)
optional_params["output_config"] = existing_output_config
@staticmethod
def _translate_adaptive_effort_for_non_adaptive_model(
model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str
@ -606,7 +561,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
custom_llm_provider=self._resolved_provider,
)
self._translate_legacy_thinking_for_adaptive_model(
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
model=model,
optional_params=anthropic_messages_optional_request_params,
custom_llm_provider=self._resolved_provider,

7
litellm/llms/azure_ai/anthropic/transformation.py
View file

@ -17,13 +17,14 @@ def _promote_extra_body_to_optional_params(optional_params: dict) -> None:
``output_config`` get auto-routed into ``extra_body`` by
``add_provider_specific_params_to_optional_params``. For the Azure→Anthropic
route those keys must reach the request body and be validated, so promote
them. ``setdefault`` keeps explicit top-level values authoritative.
them. The caller's values overwrite mapped top-level duplicates, matching
the native ``anthropic`` provider, where the same passthrough lands on
top-level ``optional_params`` after mapping.
"""
extra_body: Final = optional_params.get("extra_body")
if not isinstance(extra_body, dict) or not extra_body:
return
for k, v in extra_body.items():
optional_params.setdefault(k, v)
optional_params.update(extra_body)
optional_params.pop("extra_body", None)

4
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -943,6 +943,9 @@ class AmazonConverseConfig(BaseConfig):
litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model)
else:
optional_params["thinking"] = value
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
model=model, optional_params=optional_params, custom_llm_provider="bedrock"
)
elif param == "reasoning_effort" and isinstance(value, str):
self._handle_reasoning_effort_parameter(
model=model, reasoning_effort=value, optional_params=optional_params
@ -1334,6 +1337,7 @@ class AmazonConverseConfig(BaseConfig):
)
additional_request_params.pop("parallel_tool_calls", None)
additional_request_params.pop("client_metadata", None)
# Only set the topK value in for models that support it
additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params))

4
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -107,6 +107,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
# Restore original model name
model = original_model
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
model=original_model, optional_params=optional_params, custom_llm_provider="bedrock"
)
# The stub model hides the original model from the parent's forced-tool-use backstop
response_format_tool_choice: Final = optional_params.get("tool_choice")
if (

20
litellm/llms/bedrock/common_utils.py
View file

@ -748,6 +748,15 @@ def strip_bedrock_throughput_suffix(model: str) -> str:
MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages"
_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1")
def _mantle_api_base_from_env() -> str | None:
env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE")
if env_base is None:
return None
base: Final = env_base.rstrip("/")
return next((base[: -len(suffix)] for suffix in _MANTLE_OPENAI_BASE_SUFFIXES if base.endswith(suffix)), base)
def build_mantle_messages_url(
@ -758,12 +767,15 @@ def build_mantle_messages_url(
"""Build the bedrock-mantle Anthropic /messages URL.
Honors an explicit endpoint override (``api_base``, then
``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle
endpoints are reachable; otherwise falls back to the public regional host.
``aws_bedrock_runtime_endpoint``, then ``BEDROCK_MANTLE_API_BASE``) so
private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise
falls back to the public regional host.
The mantle messages path is appended unless the override already carries it,
so callers can pass either the host or the full messages URL.
so callers can pass either the host or the full messages URL. The env var is
shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to
carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first.
"""
override: Final = api_base or aws_bedrock_runtime_endpoint
override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env()
if override:
base: Final = override.rstrip("/")
if base.endswith(MANTLE_MESSAGES_PATH):

17
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1,6 +1,5 @@
import asyncio
import json
import os
import ssl
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
@ -159,7 +158,11 @@ def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
litellm_params: GenericLiteLLMParams,
) -> bool:
return custom_llm_provider == "openai" and litellm_params.get("rust") is True
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return custom_llm_provider == "openai" and rust_enabled(request_override=request_override)
from .http_handler import get_shared_realtime_ssl_context
@ -2364,10 +2367,6 @@ class BaseLLMHTTPHandler:
"anthropic_messages",
)
@staticmethod
def _rust_env_enabled() -> bool:
return os.getenv("LITELLM_RUST", "").strip().lower() in {"1", "true", "yes", "on"}
@staticmethod
async def _maybe_rust_anthropic_messages(
*,
@ -2383,7 +2382,11 @@ class BaseLLMHTTPHandler:
) -> AnthropicMessagesResponse | None:
if custom_llm_provider not in ("azure_ai", "anthropic"):
return None
if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled():
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
if not rust_enabled(request_override=request_override):
return None
if has_agentic_hook:
return None

4
litellm/llms/databricks/chat/transformation.py
View file

@ -330,6 +330,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
) -> dict:
is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params)
mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
if "claude" in model:
AnthropicConfig.translate_legacy_thinking_for_adaptive_model(
model=model, optional_params=mapped_params, custom_llm_provider="databricks"
)
if "tools" in mapped_params:
mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"])
if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens:

28
litellm/llms/gigachat/authenticator.py
View file

@ -8,6 +8,7 @@ Based on official GigaChat SDK authentication flow.
import time
import uuid
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import httpx
@ -32,8 +33,8 @@ GIGACHAT_SCOPE: Final = "GIGACHAT_API_PERS"
# Token expiry buffer in milliseconds (refresh token 60s before expiry)
TOKEN_EXPIRY_BUFFER_MS: Final = 60000
# Cache for access tokens
_token_cache: Final = InMemoryCache()
_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
class GigaChatAuthError(BaseLLMException):
@ -80,10 +81,9 @@ def get_access_token(
Raises:
GigaChatAuthError: If authentication fails
"""
if not litellm_params:
litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default
params: Final = litellm_params or _NO_LITELLM_PARAMS
access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
if access_token:
return access_token
@ -94,24 +94,20 @@ def get_access_token(
message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.",
)
effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope()
effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url()
effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope()
effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url()
# Check cache
cache_key: Final = f"gigachat_token:{effective_credentials[:16]}"
cached: Final = _token_cache.get_cache(cache_key)
if cached:
_token, _expires_at = cached
# Check if token is still valid (with buffer)
if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS:
verbose_logger.debug("Using cached GigaChat access token")
return _token
# Request new token
new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str
if new_expires_at:
# Cache token
ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000)
if ttl_seconds > 0:
_token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds)
@ -126,10 +122,9 @@ async def get_access_token_async(
litellm_params: Mapping[str, object] | None = None,
) -> str:
"""Async version of get_access_token."""
if not litellm_params:
litellm_params = {} # mutable-ok: empty dict default; rebind-ok: provide default
params: Final = litellm_params or _NO_LITELLM_PARAMS
access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
if access_token:
return access_token
@ -140,10 +135,9 @@ async def get_access_token_async(
message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.",
)
effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope()
effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url()
effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope()
effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url()
# Check cache
cache_key: Final = f"gigachat_token:{effective_credentials[:16]}"
cached: Final = _token_cache.get_cache(cache_key)
if cached:
@ -152,11 +146,9 @@ async def get_access_token_async(
verbose_logger.debug("Using cached GigaChat access token")
return _token
# Request new token
new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url) # pyright: ignore[reportArgumentType] # credential keys may be broader than str
if new_expires_at:
# Cache token
ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000)
if ttl_seconds > 0:
_token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds)

1
litellm/llms/gigachat/chat/streaming.py
View file

@ -52,7 +52,6 @@ class GigaChatModelResponseIterator:
tool_use: ChatCompletionToolCallChunk | None = None # rebind-ok: conditionally assigned on function_call
finish_reason: str | None = chunk_finish_reason
# Handle function_call in stream
raw_function_call: Final = delta.get("function_call")
if chunk_finish_reason == "function_call" and isinstance(raw_function_call, Mapping) and raw_function_call:
func_call: Final[Mapping[str, object]] = raw_function_call

44
litellm/llms/gigachat/chat/transformation.py
View file

@ -10,6 +10,7 @@ import json
import time
import uuid
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -34,6 +35,9 @@ else:
LiteLLMLoggingObj = Any
_EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({})
def is_valid_json(value: str) -> bool:
"""Checks whether the value passed is a valid serialized JSON string"""
try:
@ -111,11 +115,9 @@ class GigaChatConfig(BaseConfig):
"""
Set up headers with OAuth token.
"""
# Get access token
credentials: Final = api_key or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY")
access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params)
# Store credentials for image uploads
self._current_credentials = credentials
self._current_api_base = api_base
@ -208,18 +210,18 @@ class GigaChatConfig(BaseConfig):
def _convert_tools_to_functions(self, tools: Sequence) -> Sequence[dict]:
"""Convert OpenAI tools format to GigaChat functions format."""
functions: Final[list[dict]] = [] # mutable-ok: accumulator for building functions list
for tool in tools:
if isinstance(tool, dict) and tool.get("type") == "function":
func = tool.get("function", {})
functions.append(
{
"name": func.get("name", ""),
"description": func.get("description", ""),
"parameters": func.get("parameters", {}),
}
)
return functions
return [
{
"name": function.get("name", ""),
"description": function.get("description", ""),
"parameters": function.get("parameters", {}),
}
for function in (
tool.get("function", _EMPTY_FUNCTION)
for tool in tools
if isinstance(tool, dict) and tool.get("type") == "function"
)
]
def _map_tool_choice(self, tool_choice: str | Mapping[str, object]) -> str | Mapping[str, object] | None:
"""
@ -299,7 +301,6 @@ class GigaChatConfig(BaseConfig):
if part.get("type") == "text":
texts.append(part.get("text", ""))
elif part.get("type") == "image_url":
# Extract image URL and upload to GigaChat
image_url: object = part.get("image_url", {})
upload_url: str
if isinstance(image_url, str):
@ -322,16 +323,13 @@ class GigaChatConfig(BaseConfig):
headers: Mapping[str, object],
) -> dict: # mutable-ok: request payload sent to httpx
"""Transform OpenAI request to GigaChat format."""
# Transform messages
giga_messages: Final = self._transform_messages(messages)
# Build request
request_data: Final[dict[str, object]] = {
"model": model.replace("gigachat/", ""),
"messages": giga_messages,
}
# Add optional params
for key in [
"temperature",
"top_p",
@ -343,7 +341,6 @@ class GigaChatConfig(BaseConfig):
if key in optional_params:
request_data[key] = optional_params[key]
# Add functions if present
if "functions" in optional_params:
request_data["functions"] = optional_params["functions"]
if "function_call" in optional_params:
@ -358,10 +355,8 @@ class GigaChatConfig(BaseConfig):
for i, msg in enumerate(messages):
message = dict(msg)
# Remove unsupported fields
message.pop("name", None)
# Transform roles
role = message.get("role", "user")
if role == "developer":
message["role"] = "system"
@ -374,18 +369,15 @@ class GigaChatConfig(BaseConfig):
if not isinstance(content, str) or not is_valid_json(content):
message["content"] = json.dumps(content, ensure_ascii=False)
# Handle None content
if message.get("content") is None:
message["content"] = ""
# Handle list content (multimodal) - extract text and images
content = message.get("content")
if isinstance(content, list):
message["content"], attachments = self._transform_list_content(content)
if attachments:
message["attachments"] = attachments
# Transform tool_calls to function_call
tool_calls = message.get("tool_calls")
if tool_calls and isinstance(tool_calls, list) and len(tool_calls) > 0:
tool_call = tool_calls[0]
@ -436,13 +428,11 @@ class GigaChatConfig(BaseConfig):
message_data = choice.get("message", {})
finish_reason = choice.get("finish_reason", "stop")
# Transform function_call to tool_calls or content
if finish_reason == "function_call" and message_data.get("function_call"):
func_call = message_data["function_call"]
args = func_call.get("arguments", {})
if is_structured_output:
# Convert to content for structured output
if isinstance(args, dict):
content = json.dumps(args, ensure_ascii=False)
else:
@ -452,7 +442,6 @@ class GigaChatConfig(BaseConfig):
message_data.pop("functions_state_id", None)
finish_reason = "stop"
else:
# Convert to tool_calls format
if isinstance(args, dict):
args = json.dumps(args, ensure_ascii=False)
message_data["tool_calls"] = [
@ -468,7 +457,6 @@ class GigaChatConfig(BaseConfig):
message_data.pop("function_call", None)
finish_reason = "tool_calls"
# Clean up GigaChat-specific fields
message_data.pop("functions_state_id", None)
choices.append(

14
litellm/llms/gigachat/embedding/transformation.py
View file

@ -112,18 +112,10 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig):
"input": ["text1", "text2", ...]
}
"""
# Normalize input to list
if isinstance(input, str):
input_list: list = [input] # rebind-ok: locally scoped conversion
else:
input_list = input
# Remove gigachat/ prefix from model if present
model = model.removeprefix("gigachat/") # rebind-ok: parameter reassignment for normalization
normalized_input: Final = [input] if isinstance(input, str) else input # mutable-ok: preserve list API
return {
"model": model,
"input": input_list,
"model": model.removeprefix("gigachat/"),
"input": normalized_input,
}
def transform_embedding_response(

2
litellm/llms/gigachat/passthrough/transformation.py
View file

@ -60,7 +60,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig):
"""
Set up headers with OAuth token.
"""
# Get access token
access_token: Final = get_access_token(credentials=api_key, litellm_params=litellm_params)
headers["Authorization"] = f"Bearer {access_token}" # rebind-ok: mutating for OAuth setup
@ -82,7 +81,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig):
from litellm.types.utils import LlmProviders, ModelResponse
from litellm.utils import ProviderConfigManager
# cost tracking only for completions and embeddings
if "completions" in endpoint:
provider_chat_config: Final = ProviderConfigManager.get_provider_chat_config(
provider=LlmProviders(custom_llm_provider),

1
litellm/llms/gigachat/utils.py
View file

@ -4,7 +4,6 @@ from typing import Final
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
# GigaChat API endpoint
GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1"

4
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py
View file

@ -177,6 +177,10 @@ class VertexAIAnthropicConfig(AnthropicConfig):
# Restore original model name for any other processing
model = original_model
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
model=original_model, optional_params=optional_params, custom_llm_provider="vertex_ai"
)
return optional_params
def transform_response(

82
litellm/model_prices_and_context_window_backup.json
View file

@ -33093,6 +33093,88 @@
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true
},
"meta/muse-spark-1.3": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 1.25e-06,
"litellm_provider": "meta",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.25e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.0025,
"search_context_size_low": 0.0025,
"search_context_size_medium": 0.0025
},
"source": "https://ai.developer.meta.com/docs/pricing-rate-limits",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses",
"/v1/messages"
],
"supported_modalities": [
"text",
"image",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_minimal_reasoning_effort": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true
},
"meta/muse-spark-1.3-contributor": {
"cache_read_input_token_cost": 2e-09,
"input_cost_per_token": 1e-07,
"litellm_provider": "meta",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2e-07,
"search_context_cost_per_query": {
"search_context_size_high": 0.0025,
"search_context_size_low": 0.0025,
"search_context_size_medium": 0.0025
},
"source": "https://ai.developer.meta.com/docs/pricing-rate-limits",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses",
"/v1/messages"
],
"supported_modalities": [
"text",
"image",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_minimal_reasoning_effort": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true
},
"meta_llama/Llama-3.3-70B-Instruct": {
"litellm_provider": "meta_llama",
"max_input_tokens": 128000,

10
litellm/ocr/main.py
View file

@ -194,6 +194,12 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
raw_request_override: Final = prepared_request.litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
@ -422,7 +428,7 @@ async def aocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
from litellm.secret_managers.main import get_secret_str
rust_response: Final = await _run_rust_aocr(
@ -694,7 +700,7 @@ def ocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
from litellm.secret_managers.main import get_secret_str
rust_response: Final = _run_rust_ocr(

3
litellm/passthrough/main.py
View file

@ -113,10 +113,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]):
)
)
# Compliant: Save a strong reference to prevent GC
self._background_tasks.add(task)
# Remove the task from the set when it finishes to avoid memory leaks
task.add_done_callback(self._background_tasks.discard)
except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging
verbose_logger.exception(
@ -578,7 +576,6 @@ def llm_passthrough_route(
else:
return response
except Exception as e:
# provider_config is guaranteed non-None here due to the earlier guard
assert provider_config is not None
raise base_llm_http_handler._handle_error(
e=e,

13
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -92,6 +92,9 @@ def _connection_error_message(exc: BaseException) -> str:
if MCP_AVAILABLE:
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.client import MCPClient
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
global_mcp_server_manager,
@ -876,7 +879,6 @@ if MCP_AVAILABLE:
return (), classify_list_exception(e)
return tools_result, ServerListOk(tool_count=len(tools_result))
# Query all servers the user has access to
queried_servers: Final = tuple(
server
for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids)
@ -1141,6 +1143,11 @@ if MCP_AVAILABLE:
scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None
return client_id, client_secret, scopes
async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None:
with anyio.move_on_after(deadline):
return await client.list_tools(raise_on_error=True)
return None
async def _execute_with_mcp_client(
request: NewMCPServerRequest,
operation: Callable[..., Awaitable[Mapping[str, object]]],
@ -1422,9 +1429,7 @@ if MCP_AVAILABLE:
getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT,
MCP_TOOL_LISTING_TIMEOUT,
)
list_tools_result = None # rebind-ok: set inside the timeout scope below
with anyio.move_on_after(listing_deadline):
list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above
list_tools_result: Final = await _list_tools_within(client, listing_deadline)
if list_tools_result is None:
verbose_logger.warning(
"MCP tools/list preview timed out after %s seconds while paginating upstream tools",

35
litellm/proxy/auth/ip_address_utils.py
View file

@ -6,8 +6,11 @@ External callers (public IPs) only see servers with available_on_public_internet
"""
import ipaddress
import os
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Final
from urllib.parse import urlparse
from fastapi import Request
from pydantic import TypeAdapter, ValidationError
@ -137,7 +140,7 @@ class IPAddressUtils:
@staticmethod
def is_request_from_trusted_proxy(
request: Request,
general_settings: dict[str, Any] | None = None,
general_settings: Mapping[str, Any] | None = None,
) -> bool:
"""
Return True if X-Forwarded-* headers on this request should be trusted.
@ -190,6 +193,36 @@ class IPAddressUtils:
trusted_networks: Final = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges)
return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks)
@staticmethod
def is_request_https(
request: Request,
general_settings: Mapping[str, Any] | None = None,
) -> bool:
"""
Whether this request's PUBLIC-facing origin is HTTPS, for deciding
whether a cookie set on the response should be marked ``Secure``.
litellm only sees a plain-HTTP hop whenever TLS terminates at a
reverse proxy, so ``request.url.scheme`` alone cannot answer this in
that deployment shape. Resolved from the first trusted signal:
1. ``PROXY_BASE_URL`` (operator-declared public origin).
2. ``X-Forwarded-Proto``, only when the request's direct peer is a
configured trusted proxy -- see ``is_request_from_trusted_proxy``.
An untrusted caller cannot spoof this header to strip Secure.
3. The request's own literal scheme (direct TLS termination, or no
reverse proxy in front of litellm).
"""
configured_base_url: Final = os.environ.get("PROXY_BASE_URL", "").strip()
if configured_base_url:
return urlparse(configured_base_url).scheme == "https"
if IPAddressUtils.is_request_from_trusted_proxy(request, general_settings=general_settings):
forwarded_proto: Final = request.headers.get("X-Forwarded-Proto")
if forwarded_proto:
return forwarded_proto.split(",")[0].strip().lower() == "https"
return request.url.scheme == "https"
@staticmethod
def extract_client_ip_from_xff_hops(
xff_header: str,

16
litellm/proxy/guardrails/guardrail_hooks/alice/alice.py
View file

@ -8,6 +8,7 @@
import json
import os
from collections.abc import Mapping
from itertools import islice
from typing import (
TYPE_CHECKING,
Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml
@ -341,19 +342,18 @@ def _json_safe(
if depth >= _MAX_DEPTH or id(value) in seen:
return None
nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately
nested: Final = seen | frozenset((id(value),))
if isinstance(value, dict):
out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is
for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view
if isinstance(key, str) and key not in strip_keys:
out[key] = _json_safe(item, depth + 1, nested, strip_keys)
return out
return {
key: _json_safe(item, depth + 1, nested, strip_keys)
for key, item in islice(value.items(), _MAX_ITEMS)
if isinstance(key, str) and key not in strip_keys
}
if isinstance(value, (list, tuple, set, frozenset)):
return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use
_json_safe(item, depth + 1, nested, strip_keys)
for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view
_json_safe(item, depth + 1, nested, strip_keys) for item in islice(value, _MAX_ITEMS)
]
dump: Final = getattr(value, "model_dump", None)

3
litellm/proxy/management_endpoints/sso/saml_sso.py
View file

@ -36,6 +36,7 @@ from pydantic import ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
from litellm.proxy.utils import get_custom_url
@ -131,7 +132,7 @@ class SAMLAuthHandler:
@staticmethod
def _is_https(request: Request) -> bool:
return SAMLAuthHandler._base_url(request).startswith("https")
return IPAddressUtils.is_request_https(request)
@staticmethod
def _acs_url(request: Request) -> str:

39
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -92,6 +92,7 @@ from litellm.proxy.auth.auth_utils import (
has_user_setup_sso,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.admin_ui_utils import (
admin_ui_disabled,
@ -1118,7 +1119,7 @@ async def google_login(
request=request,
)
if sso_redirect is not None:
_persist_return_to_cookie(sso_redirect, return_to)
_persist_return_to_cookie(sso_redirect, return_to, request)
return sso_redirect
from fastapi.responses import HTMLResponse
@ -1138,7 +1139,7 @@ async def google_login(
# helper the SSO branch uses, so /login can resume the connect flow instead of dead-ending at the
# dashboard. One implementation → the two sign-in branches cannot diverge (and the login form always
# renders, since the helper never raises on a bad return_to).
_persist_return_to_cookie(form_response, return_to)
_persist_return_to_cookie(form_response, return_to, request)
return form_response
@ -2741,6 +2742,7 @@ async def _sso_return_to_redirect(
jwt_token: str,
redis_usage_cache,
user_api_key_cache,
request: Request,
) -> RedirectResponse | None:
"""Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard.
@ -2759,7 +2761,7 @@ async def _sso_return_to_redirect(
if _is_same_origin_return_path(return_to):
redirect_response = RedirectResponse(url=return_to, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token)
set_session_token_cookie(redirect_response, request, jwt_token)
redirect_response.delete_cookie("litellm_cp_return_to")
return redirect_response
@ -2782,7 +2784,25 @@ async def _sso_return_to_redirect(
return None
def _persist_return_to_cookie(response: Response, return_to: str | None) -> None:
def set_session_token_cookie(response: Response, request: Request, jwt_token: str) -> None:
"""Set the ``token`` session cookie shared by every sign-in path.
Not HttpOnly: the dashboard reads this cookie via ``document.cookie`` to
populate its own Authorization headers (see
``ui/litellm-dashboard/src/utils/cookieUtils.ts``), so marking it
HttpOnly would break login. Secure is still required whenever the public
origin is HTTPS, resolved the same trust-aware way as every other
litellm cookie."""
response.set_cookie(
key="token",
value=jwt_token,
secure=IPAddressUtils.is_request_https(request),
httponly=False,
samesite="lax",
)
def _persist_return_to_cookie(response: Response, return_to: str | None, request: Request) -> None:
"""Best-effort: persist a SAFE ``return_to`` on ``response`` as the one-shot ``litellm_cp_return_to``
cookie so ANY sign-in path — SSO / Okta / generic OR the username/password form — can resume there
afterwards. THIS is the single source of truth, called by every sign-in branch so they cannot
@ -2803,6 +2823,7 @@ def _persist_return_to_cookie(response: Response, return_to: str | None) -> None
max_age=600,
httponly=True,
samesite="lax",
secure=IPAddressUtils.is_request_https(request),
)
@ -3079,8 +3100,11 @@ class SSOAuthenticationHandler:
# incoming request is HTTP (local dev). Without
# ``Secure`` the cookie is sent over plain HTTP,
# letting a network observer read and replay the
# state value and bypass this protection.
secure_flag: Final = request is None or request.url.scheme == "https"
# state value and bypass this protection. Trust-aware:
# honors PROXY_BASE_URL / a trusted reverse proxy's
# X-Forwarded-Proto instead of only the literal scheme
# litellm sees on the wire.
secure_flag: Final = request is None or IPAddressUtils.is_request_https(request)
redirect_response.set_cookie(
key="litellm_oauth_state",
value=state_value,
@ -3628,6 +3652,7 @@ class SSOAuthenticationHandler:
jwt_token=jwt_token,
redis_usage_cache=redis_usage_cache,
user_api_key_cache=user_api_key_cache,
request=request,
)
if return_to_redirect is not None:
return return_to_redirect
@ -3636,7 +3661,7 @@ class SSOAuthenticationHandler:
litellm_dashboard_ui += "?login=success"
verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui)
redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token)
set_session_token_cookie(redirect_response, request, jwt_token)
return redirect_response
@staticmethod

4
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -1731,7 +1731,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict:
def get_vertex_pass_through_handler(
call_type: Literal["discovery", "aiplatform"], # noqa: UP037
call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here
) -> BaseVertexAIPassThroughHandler:
if call_type == "discovery":
return VertexAIDiscoveryPassThroughHandler()
@ -2961,7 +2961,6 @@ async def handle_gigachat_passthrough_router_model(
"""
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
# Detect streaming based on request body
is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown]
data: dict[str, Any] = await _read_request_body(
@ -2997,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model(
data["json"] = request_body
data["custom_llm_provider"] = "gigachat"
# Remove sensitive keys from data
keys: Final = [ # mutable-ok: list of keys to remove from data
"gigachat_auth_url",
"gigachat_access_token",

79
litellm/proxy/proxy_server.py
View file

@ -4744,6 +4744,63 @@ class ProxyConfig:
)
return coordination_redis_cache
@staticmethod
async def _init_coordination_redis_env_fallback(litellm_settings: Mapping[str, object]) -> RedisCache | None:
"""
Last-resort coordination Redis, tried after an explicit
`general_settings.coordination_redis` block and `litellm_settings.cache`
have both had a chance to resolve one. Without this, a deployment that
only exports REDIS_HOST/REDIS_PORT (no cache block, no coordination_redis
block) gets NO cross-pod coordination at all: spend counters, budget-window
enforcement, and the reset_spend cache-eviction broadcast all silently stay
per-pod local, so a key reset on one pod never clears another pod's stale
enforcement.
Unlike the explicit block and cache-backend paths (a deliberate opt-in, so a
bad connection target or a malformed REDIS_CLUSTER_NODES/REDIS_SENTINEL_NODES
value should fail loudly), this one is inferred from bare env vars that may be
set for an unrelated reason -- e.g. a REDIS_HOST left over from a different
job/service, or a REDIS_CLUSTER_NODES value nothing here ever asked to be
parsed. Wrongly guessing "coordination available" must not turn a previously
harmless in-memory-only proxy into one that fails to boot or raises on every
cache write, so a malformed value or a failed/slow ping are both treated the
same as no REDIS_* vars at all.
"""
try:
env_coordination_redis_cache: Final = _build_redis_usage_cache_from_environment()
except Exception as e: # noqa: BLE001 # a malformed inferred Redis env var must not block startup
verbose_proxy_logger.warning(
"coordination_redis: could not build a Redis client from REDIS_* environment variables "
"(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis "
"explicitly to require it.",
e,
)
return None
if env_coordination_redis_cache is None:
return None
try:
reachable: Final = await asyncio.wait_for(env_coordination_redis_cache.ping(), timeout=2.0)
except Exception as e: # noqa: BLE001 # an unreachable inferred Redis must not block startup or writes
verbose_proxy_logger.warning(
"coordination_redis: REDIS_* environment variables named a Redis that is not reachable "
"(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis "
"explicitly to require it.",
e,
)
return None
if not reachable:
return None
_attach_redis_usage_cache(
env_coordination_redis_cache,
enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True,
)
verbose_proxy_logger.info(
"coordination_redis: using a standalone Redis built from REDIS_* "
"environment variables for usage tracking, rate limiting, and "
"cross-pod coordination."
)
return env_coordination_redis_cache
def _init_cache(
self,
cache_params: dict,
@ -5412,6 +5469,13 @@ class ProxyConfig:
reset_audit_log_callback_cache()
_in_memory_loggers[:] = [cb for cb in _in_memory_loggers if not isinstance(cb, S3V2Logger)]
if redis_usage_cache is None:
env_coordination_redis_cache: Final = await self._init_coordination_redis_env_fallback(
litellm_settings=litellm_settings
)
if env_coordination_redis_cache is not None:
_set_redis_usage_cache(env_coordination_redis_cache)
## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging
general_settings = config.get("general_settings", {})
if general_settings is None:
@ -15265,7 +15329,10 @@ async def login(request: Request):
# authorize round-trip), mirroring the SSO callback; otherwise land on the dashboard. Gated by
# _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the
# one-shot cookie is cleared after use.
from litellm.proxy.management_endpoints.ui_sso import _sso_return_to_redirect
from litellm.proxy.management_endpoints.ui_sso import (
_sso_return_to_redirect,
set_session_token_cookie,
)
# Resume through the SAME resumer the SSO callback uses, rather than a second, narrower arm.
# _persist_return_to_cookie stores both shapes it accepts (a relative same-origin path AND a
@ -15282,6 +15349,7 @@ async def login(request: Request):
jwt_token=jwt_token,
redis_usage_cache=redis_usage_cache,
user_api_key_cache=user_api_key_cache,
request=request,
)
except Exception: # noqa: BLE001 # resuming must NEVER block a completed sign-in
# The symmetric half of _persist_return_to_cookie's "never raises" contract. The resumer
@ -15296,7 +15364,7 @@ async def login(request: Request):
# Create redirect response with cookie
redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token)
set_session_token_cookie(redirect_response, request, jwt_token)
if cp_return_to:
redirect_response.delete_cookie(key="litellm_cp_return_to")
return redirect_response
@ -15306,6 +15374,7 @@ async def login(request: Request):
async def login_v2(request: Request):
global premium_user, general_settings, master_key
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt
from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie
from litellm.proxy.utils import get_custom_url
try:
@ -15340,7 +15409,7 @@ async def login_v2(request: Request):
content={"redirect_url": litellm_dashboard_ui, "token": jwt_token},
status_code=status.HTTP_200_OK,
)
json_response.set_cookie(key="token", value=jwt_token)
set_session_token_cookie(json_response, request, jwt_token)
return json_response
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v2(): Exception occurred - %s", e)
@ -15440,6 +15509,8 @@ async def login_v3(request: Request):
@router.post("/v3/login/exchange", include_in_schema=False) # exchange single-use opaque code for JWT
async def login_v3_exchange(request: Request):
from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie
try:
if not general_settings.get("control_plane_url"):
raise ProxyException(
@ -15486,7 +15557,7 @@ async def login_v3_exchange(request: Request):
},
status_code=status.HTTP_200_OK,
)
json_response.set_cookie(key="token", value=cached_data["token"])
set_session_token_cookie(json_response, request, cached_data["token"])
return json_response
except ProxyException:
raise

30
litellm/router_strategy/complexity_router/README.md
View file

@ -68,6 +68,36 @@ still resolve to a deployment in `model_list`; this configuration does not creat
- abc
```
### Heuristic v2
Set `classifier_type: heuristic_v2` to classify with the bundled calibrated
success-probability model instead of the hand-written weighted scorer
```yaml
model_list:
- model_name: smart-router
litellm_params:
model: auto_router/complexity_router
complexity_router_config:
classifier_type: heuristic_v2
tiers:
SIMPLE: luna
MEDIUM: terra
COMPLEX: sol
REASONING: sol-ultra
```
No classifier model call or per-model training data is required. The classifier
uses global tier quality, request-type quality, and similar-request cohorts from
the bundled UltraFeedback artifact. It estimates success at every tier, enforces
monotonic probabilities, and returns the first tier meeting the trained 0.75
threshold. The existing complexity-router tier pool then selects and dispatches
a model from that tier
Spend logs record `routing_decision.cause: heuristic_v2`, the detected request
type, and all four predicted probabilities. Existing `classifier_type: heuristic`
configurations keep the original weighted scorer unchanged
### Renaming the tiers
`tier_labels` puts your own vocabulary on the four tiers:

4069
litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json Normal file
View file

File diff suppressed because it is too large Load diff

31
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -33,6 +33,11 @@ from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal
from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
from litellm.router_strategy.complexity_router.tier_predictor import (
TierSuccessPredictor,
resolve_tier_artifact,
)
from litellm.types.utils import (
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
ModelResponse,
@ -790,6 +795,7 @@ class ClassificationOutcome(NamedTuple):
signals: tuple[str, ...]
cause: Literal[
"heuristic_scorer",
"heuristic_v2",
"reasoning_override",
"llm_classifier",
"heuristic_first_short_circuit",
@ -978,6 +984,11 @@ class ComplexityRouter(CustomLogger):
if llm_classifier_configured
else None
)
self._tier_success_predictor: TierSuccessPredictor | None = (
TierSuccessPredictor(resolve_tier_artifact(self.config.heuristic_v2_artifact))
if self.config.classifier_type == "heuristic_v2"
else None
)
verbose_router_logger.debug("ComplexityRouter initialized for %s with tiers: %s", model_name, self.config.tiers)
@ -1350,6 +1361,8 @@ class ComplexityRouter(CustomLogger):
custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and
default_model. The outcome's `cause` reports which path actually ran.
"""
if self.config.classifier_type == "heuristic_v2":
return self._classify_with_heuristic_v2(prompt)
if self.config.classifier_type == "custom":
return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages)
if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None:
@ -1359,6 +1372,24 @@ class ComplexityRouter(CustomLogger):
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
def _classify_with_heuristic_v2(self, prompt: str) -> ClassificationOutcome:
predictor: Final = self._tier_success_predictor
if predictor is None:
raise ValueError("heuristic v2 predictor is not configured")
request_type: Final = classify_prompt(prompt)
prediction: Final = predictor.predict(prompt, request_type)
tier: Final = TIER_SEVERITY_ORDER[prediction.required_tier - 1]
probability_signals: Final = tuple(
f"tier-probability:{candidate.value.lower()}={prediction.probabilities[index]:.6f}"
for index, candidate in enumerate(TIER_SEVERITY_ORDER, start=1)
)
return ClassificationOutcome(
tier=tier,
score=None,
signals=(f"request-type:{request_type.value}", *probability_signals),
cause="heuristic_v2",
)
async def _classify_heuristic_first(
self,
prompt: str,

21
litellm/router_strategy/complexity_router/config.py
View file

@ -14,6 +14,8 @@ from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_seriali
from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
from .tier_predictor import TrainedTierArtifact
class ComplexityTier(str, Enum):
"""Complexity tiers for routing decisions."""
@ -625,12 +627,19 @@ class ComplexityRouterConfig(BaseModel):
)
# Classifier strategy
classifier_type: Literal["heuristic", "llm", "custom", "heuristic_first"] = Field(
classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first"] = Field(
default="heuristic",
description=(
"Classification strategy: local regex/keyword scoring, an LLM call, a custom classifier "
"plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier "
"when the local scorer does not confidently land a cheap tier"
"Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, "
"an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays "
"for the LLM classifier when the local scorer does not confidently land a cheap tier"
),
)
heuristic_v2_artifact: TrainedTierArtifact | Literal["ultrafeedback"] = Field(
default="ultrafeedback",
description=(
"Success-probability artifact used by classifier_type 'heuristic_v2'. The bundled "
"UltraFeedback artifact is selected by default; an inline trained artifact may replace it"
),
)
classifier_llm_config: ClassifierLLMConfig | None = Field(
@ -1248,10 +1257,10 @@ class ComplexityRouterConfig(BaseModel):
)
if duplicated:
raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}")
if self.classifier_type in ("heuristic", "heuristic_first"):
if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first"):
raise ValueError(
"tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only "
"produces the built-in tiers"
"produces the four built-in tiers, as does heuristic_v2"
)
conflicts: Final = self._tier_definition_conflicts()
if conflicts:

156
litellm/router_strategy/complexity_router/tier_predictor.py Normal file
View file

@ -0,0 +1,156 @@
from __future__ import annotations
import re
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, Field, model_validator
from litellm.types.router import RequestType
class TierGlobalStatistic(BaseModel):
tier: int = Field(ge=1, le=4)
successes: float = Field(ge=0.0)
observations: float = Field(gt=0.0)
@model_validator(mode="after")
def _successes_do_not_exceed_observations(self) -> TierGlobalStatistic:
if self.successes > self.observations:
raise ValueError("successes cannot exceed observations")
return self
class TierDomainStatistic(TierGlobalStatistic):
request_type: RequestType
class TierCohortStatistic(TierGlobalStatistic):
cohort: str = Field(min_length=1)
class TierDataset(BaseModel):
name: str = Field(min_length=1)
url: str = Field(min_length=1)
license: str = Field(min_length=1)
rows: int = Field(gt=0)
success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1)
class TrainedTierArtifact(BaseModel):
schema_version: Literal[1] = 1
global_statistics: tuple[TierGlobalStatistic, ...]
domain_statistics: tuple[TierDomainStatistic, ...] = ()
cohort_statistics: tuple[TierCohortStatistic, ...] = ()
domain_prior_mass: float = Field(default=200.0, gt=0.0)
cohort_prior_mass: float = Field(default=20.0, gt=0.0)
routing_threshold: float = Field(default=0.75, ge=0.0, le=1.0)
datasets: tuple[TierDataset, ...] = ()
success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1)
split_method: str = Field(default="sha256(prompt): 70% train, 15% validation, 15% test", min_length=1)
@model_validator(mode="after")
def _statistics_are_unique(self) -> TrainedTierArtifact:
global_tiers: Final = tuple(stat.tier for stat in self.global_statistics)
if frozenset(global_tiers) != frozenset((1, 2, 3, 4)) or len(global_tiers) != 4:
raise ValueError("global statistics must contain each tier exactly once")
domain_keys: Final = tuple((stat.request_type, stat.tier) for stat in self.domain_statistics)
if len(domain_keys) != len(frozenset(domain_keys)):
raise ValueError("domain statistics must contain unique request_type and tier pairs")
cohort_keys: Final = tuple((stat.cohort, stat.tier) for stat in self.cohort_statistics)
if len(cohort_keys) != len(frozenset(cohort_keys)):
raise ValueError("cohort statistics must contain unique cohort and tier pairs")
return self
_CODE_PATTERN: Final = re.compile(
r"```|\b(def|class|function|python|javascript|typescript|sql|code)\b",
re.IGNORECASE,
)
_MATH_PATTERN: Final = re.compile(
r"\b(solve|calculate|equation|probability|theorem|proof|integral)\b|[$=]",
re.IGNORECASE,
)
_MULTIPLE_CHOICE_PATTERN: Final = re.compile(r"(?:^|\s)[A-D][.)]\s")
_TIERS: Final = (1, 2, 3, 4)
_BUILTIN_ARTIFACTS: Final = MappingProxyType({"ultrafeedback": "ultrafeedback_tiers.json"})
def resolve_tier_artifact(artifact: TrainedTierArtifact | str) -> TrainedTierArtifact:
if isinstance(artifact, TrainedTierArtifact):
return artifact
filename: Final = _BUILTIN_ARTIFACTS.get(artifact)
if filename is None:
raise ValueError(f"unknown complexity router tier artifact: {artifact}")
path: Final = Path(__file__).with_name("artifacts") / filename
return TrainedTierArtifact.model_validate_json(path.read_text())
def similarity_cohort(prompt: str, request_type: RequestType) -> str:
length: Final = len(prompt)
length_bucket: Final = (
"short" if length < 200 else "medium" if length < 800 else "long" if length < 2000 else "very_long"
)
code: Final = int(bool(_CODE_PATTERN.search(prompt)))
math: Final = int(bool(_MATH_PATTERN.search(prompt)))
multiple_choice: Final = int(bool(_MULTIPLE_CHOICE_PATTERN.search(prompt)))
non_ascii: Final = int(sum(ord(character) > 127 for character in prompt) / max(1, length) > 0.1)
return f"{request_type.value}|{length_bucket}|code={code}|math={math}|mc={multiple_choice}|intl={non_ascii}"
@dataclass(frozen=True, slots=True)
class TierPrediction:
probabilities: Mapping[int, float]
required_tier: int
class TierSuccessPredictor:
def __init__(self, artifact: TrainedTierArtifact) -> None:
self._artifact = artifact
self._global: Mapping[int, TierGlobalStatistic] = MappingProxyType(
{stat.tier: stat for stat in artifact.global_statistics}
)
self._domain: Mapping[tuple[RequestType, int], TierDomainStatistic] = MappingProxyType(
{(stat.request_type, stat.tier): stat for stat in artifact.domain_statistics}
)
self._cohort: Mapping[tuple[str, int], TierCohortStatistic] = MappingProxyType(
{(stat.cohort, stat.tier): stat for stat in artifact.cohort_statistics}
)
@property
def routing_threshold(self) -> float:
return self._artifact.routing_threshold
def predict(self, prompt: str, request_type: RequestType) -> TierPrediction:
cohort: Final = similarity_cohort(prompt, request_type)
raw: Final = tuple(self._probability(tier, request_type, cohort) for tier in _TIERS)
monotonic: Final = tuple(max(raw[:index]) for index in range(1, len(raw) + 1))
probabilities: Final[Mapping[int, float]] = MappingProxyType(
{int(tier): probability for tier, probability in zip(_TIERS, monotonic)}
)
required_tier: Final = next(
(tier for tier in _TIERS if probabilities[tier] >= self._artifact.routing_threshold),
4,
)
return TierPrediction(probabilities=probabilities, required_tier=required_tier)
def _probability(self, tier: int, request_type: RequestType, cohort: str) -> float:
global_stat: Final = self._global[tier]
global_mean: Final = (global_stat.successes + 1.0) / (global_stat.observations + 2.0)
domain_stat: Final = self._domain.get((request_type, tier))
domain_mean: Final = self._posterior_mean(domain_stat, self._artifact.domain_prior_mass, global_mean)
cohort_stat: Final = self._cohort.get((cohort, tier))
return self._posterior_mean(cohort_stat, self._artifact.cohort_prior_mass, domain_mean)
@staticmethod
def _posterior_mean(
statistic: TierGlobalStatistic | None,
prior_mass: float,
prior_mean: float,
) -> float:
if statistic is None:
return prior_mean
return (statistic.successes + prior_mass * prior_mean) / (statistic.observations + prior_mass)

2
litellm/rust_bridge/__init__.py
View file

@ -1,9 +1,9 @@
"""LiteLLM Rust bridge package."""
from litellm.rust_bridge.configuration import use_litellm_rust
from litellm.rust_bridge.loader import (
get_native_bridge,
native_bridge_available,
)
from litellm.rust_bridge.ocr import use_litellm_rust
__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"]

49
litellm/rust_bridge/bindings.py Normal file
View file

@ -0,0 +1,49 @@
from __future__ import annotations
from collections.abc import Callable
from typing import Final, Generic, TypeVar
from litellm.rust_bridge.loader import get_native_bridge
BindingT = TypeVar("BindingT")
class _Unset:
pass
_UNSET: Final = _Unset()
class NativeBinding(Generic[BindingT]):
"""Resolve one native attribute with an explicit, resettable test override."""
def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None:
self._attribute: Final = attribute
self._validate: Final = validate
self._override: BindingT | None | _Unset = _UNSET
def load(self) -> BindingT | None:
if not isinstance(self._override, _Unset):
return self._override
native: Final = get_native_bridge()
if native is None:
return None
return self._validate(getattr(native, self._attribute, None))
def override(self, value: BindingT | None) -> None:
self._override = value
def reset(self) -> None:
self._override = _UNSET
def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None:
native: Final = get_native_bridge()
if native is None:
return None
declined: Final = getattr(native, "RustBridgeDeclined", None)
upstream: Final = getattr(native, "RustUpstreamError", None)
if not isinstance(declined, type) or not isinstance(upstream, type):
return None
return declined, upstream

Some files were not shown because too many files have changed in this diff Show more