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

refactor(python-bridge): declare sync and async routes once

Browse source
This commit is contained in:
Yujong Lee 2026-09-01 21:06:20 -07:00 • committed by GitHub
parent d69da842b8
commit d549376af3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 728 additions and 440 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,25 +1,79 @@
use std::collections::HashMap;
use std::time::Duration;
use litellm_python_interop::from_py;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::{Map, Value};
pub(crate) fn optional_object_to_map(
py: Python<'_>,
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: Option<Py<PyAny>>,
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) => match from_py(value.bind(py))? {
Value::Object(map) => Ok(map),
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
},
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 {

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

@ -1,95 +1,71 @@
use litellm_core::Error;
use std::future::Future;
use litellm_core::audio_transcription::{
AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
};
use litellm_python_interop::from_py;
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
use super::runtime::{run_async, run_sync};
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)?;
#[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);
run_sync(
py,
async move {
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
},
core_error_to_pyerr,
)
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
})
}
#[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);
run_async(
py,
async move {
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
},
core_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(transcription, module)?)?;
module.add_function(wrap_pyfunction!(atranscription, module)?)
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,
}

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

@ -1,64 +1,62 @@
use std::time::Duration;
use litellm_core::Error;
use std::future::Future;
use litellm_core::chat_completions::types::ChatCompletionsRequest;
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_python_interop::from_py;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use serde_json::Value;
use crate::errors::chat_completions_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value};
use super::runtime::{run_async, run_sync};
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,
})?;
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),
))
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
})
}
/// 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>>,
#[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 messages = from_py(messages.bind(py))?;
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let optional_params = object_or_empty("optional_params", optional_params)?;
Ok(chat_completions_decline_reason(
&model,
custom_llm_provider.as_deref(),
@ -68,90 +66,26 @@ fn chat_completions_decline(
.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,
)?;
run_sync(
py,
async move {
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
},
chat_completions_error_to_pyerr,
)
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[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,
)?;
run_async(
py,
async move {
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
},
chat_completions_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?;
module.add_function(wrap_pyfunction!(chat_completions, module)?)?;
module.add_function(wrap_pyfunction!(achat_completions, module)?)
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],
}

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

@ -0,0 +1,419 @@
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),*))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type),*
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::routes::runtime::run_sync(py, future, $map_error)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type),*
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::routes::runtime::run_async(py, future, $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)",
),
(
"transcription",
"atranscription",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"messages",
"amessages",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
(
"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)",
),
];
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"
);
});
}
}

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

@ -1,105 +1,65 @@
use std::time::Duration;
use litellm_core::Error;
use litellm_core::messages::messages as run_messages;
use litellm_core::messages::types::MessagesRequest;
use litellm_python_interop::from_py;
use pyo3::exceptions::PyValueError;
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
use pyo3::prelude::*;
use serde_json::{Map, Value};
use serde_json::Value;
use std::future::Future;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value};
use super::runtime::{run_async, run_sync};
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,
})?;
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)))
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
})
}
#[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)?;
run_sync(
py,
async move {
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
},
core_error_to_pyerr,
)
}
#[pyfunction]
#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[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)?;
run_async(
py,
async move {
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
},
core_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(messages, module)?)?;
module.add_function(wrap_pyfunction!(amessages, module)?)
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,
}

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

@ -1,10 +1,13 @@
use pyo3::prelude::*;
#[macro_use]
mod definition;
mod runtime;
mod audio_transcription;
mod chat_completions;
mod messages;
mod ocr;
mod runtime;
pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
ocr::register(module)?;

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

@ -1,131 +1,73 @@
use std::time::Duration;
use litellm_core::Error;
use std::future::Future;
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_python_interop::from_py;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
use crate::marshal::{optional_object_to_map, optional_timeout};
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
use super::runtime::{run_async, run_sync};
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)?;
type MarshaledOcrInputs = (
Value,
Option<Map<String, Value>>,
Map<String, Value>,
Option<Duration>,
);
fn marshal_inputs(
py: Python<'_>,
document: Py<PyAny>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledOcrInputs> {
let document = from_py(document.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
};
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let timeout = optional_timeout(timeout_seconds);
Ok((document, extra_headers, optional_params, timeout))
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
})
}
#[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,
)?;
run_sync(
py,
async move {
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
},
core_error_to_pyerr,
)
}
#[pyfunction]
#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[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,
)?;
run_async(
py,
async move {
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
},
core_error_to_pyerr,
)
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(ocr, module)?)?;
module.add_function(wrap_pyfunction!(aocr, module)?)
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,
}