From 67b7ba37506fd73e52fc53341444738f1b3f5b2f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 15:36:42 -0700 Subject: [PATCH] refactor(python-bridge): split non-streaming bridge modules --- .../crates/python-bridge/src/diagnostics.rs | 23 + .../crates/python-bridge/src/errors.rs | 61 ++ litellm-rust/crates/python-bridge/src/lib.rs | 629 ++---------------- .../crates/python-bridge/src/marshal.rs | 53 ++ .../src/routes/audio_transcription.rs | 100 +++ .../src/routes/chat_completions.rs | 163 +++++ .../python-bridge/src/routes/messages.rs | 109 +++ .../crates/python-bridge/src/routes/mod.rs | 13 + .../crates/python-bridge/src/routes/ocr.rs | 128 ++++ 9 files changed, 698 insertions(+), 581 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/diagnostics.rs create mode 100644 litellm-rust/crates/python-bridge/src/errors.rs create mode 100644 litellm-rust/crates/python-bridge/src/marshal.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/chat_completions.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/ocr.rs diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs new file mode 100644 index 00000000000..cc153a89b8f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -0,0 +1,23 @@ +use litellm_python_interop::release_count; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +#[pyfunction] +fn gil_stats(py: Python<'_>) -> PyResult> { + let stats = PyDict::new(py); + stats.set_item("releases", release_count())?; + Ok(stats.into_any().unbind()) +} + +#[cfg(feature = "panic-test")] +#[pyfunction] +fn _panic_for_test() { + panic!("intentional PyO3 panic smoke test"); +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + #[cfg(feature = "panic-test")] + module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; + Ok(()) +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs new file mode 100644 index 00000000000..ec68cf7bfa3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -0,0 +1,61 @@ +use litellm_core::error::CoreError; +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: CoreError) -> PyErr { + match err { + CoreError::Auth(message) => PyValueError::new_err(message), + CoreError::InvalidProvider(_) + | CoreError::InvalidRequest(_) + | CoreError::InvalidType { .. } + | CoreError::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: CoreError) -> PyErr { + match err { + CoreError::Unsupported(_) + | CoreError::Auth(_) + | CoreError::InvalidProvider(_) + | CoreError::InvalidRequest(_) + | CoreError::InvalidType { .. } + | CoreError::MissingField(_) + | CoreError::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. + | CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + CoreError::Http { status, body } => { + RustUpstreamError::new_err((status, format!("{status}: {body}"))) + } + CoreError::Network(message) | CoreError::InvalidResponse(message) => { + RustUpstreamError::new_err((0u16, message)) + } + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + let py = module.py(); + module.add("RustBridgeDeclined", py.get_type::())?; + module.add("RustUpstreamError", py.get_type::()) +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 746e0770f9b..f628c987220 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,142 +1,14 @@ -use std::collections::HashMap; -use std::time::Duration; +mod diagnostics; +mod errors; +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::CoreError; -use litellm_core::messages::messages as run_messages; -use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; -use litellm_python_interop::{from_py, release_count, release_gil, to_py}; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyAny, PyDict}; -use serde_json::{Map, Value}; +use pyo3::types::PyAny; -pyo3::create_exception!( - _native, - RustBridgeDeclined, - pyo3::exceptions::PyException, - "The route declined before calling the provider, so the host may retry on its own path." -); - -pyo3::create_exception!( - _native, - RustUpstreamError, - pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." -); - -type MarshaledOcrInputs = ( - Value, - Option>, - Map, - Option, -); - -fn messages_response_to_py( - py: Python<'_>, - response: AnthropicMessagesResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn chat_completions_response_to_py( - py: Python<'_>, - response: ChatCompletionsResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn core_error_to_pyerr(err: CoreError) -> PyErr { - match err { - CoreError::Auth(message) => PyValueError::new_err(message), - CoreError::InvalidProvider(_) - | CoreError::InvalidRequest(_) - | CoreError::InvalidType { .. } - | CoreError::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: CoreError) -> PyErr { - match err { - CoreError::Unsupported(_) - | CoreError::Auth(_) - | CoreError::InvalidProvider(_) - | CoreError::InvalidRequest(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) - | CoreError::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. - | CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - CoreError::Http { status, body } => { - RustUpstreamError::new_err((status, format!("{status}: {body}"))) - } - CoreError::Network(message) | CoreError::InvalidResponse(message) => { - RustUpstreamError::new_err((0u16, message)) - } - } -} - -fn optional_object_to_map( - py: Python<'_>, - name: &'static str, - value: Option>, -) -> PyResult> { - match value { - Some(value) => match from_py(value.bind(py))? { - Value::Object(map) => Ok(map), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), - }, - None => Ok(Map::new()), - } -} - -fn optional_timeout(timeout_seconds: Option) -> Option { - timeout_seconds.and_then(|secs| { - if secs.is_finite() && secs > 0.0 { - Some(Duration::from_secs_f64(secs)) - } else { - None - } - }) -} - -fn marshal_headers( - py: Python<'_>, - headers: Option>, -) -> PyResult> { - let value = match headers { - Some(headers) => from_py(headers.bind(py))?, - None => Value::Object(Map::new()), - }; - let Value::Object(headers) = value else { - return Err(PyValueError::new_err("headers must be a dict")); - }; - headers - .into_iter() - .map(|(name, value)| { - value - .as_str() - .map(|value| (name, value.to_string())) - .ok_or_else(|| PyValueError::new_err("header values must be strings")) - }) - .collect() -} +use crate::errors::core_error_to_pyerr; +use crate::marshal::{marshal_headers, optional_timeout}; #[pyclass] struct ResponsesWebSocketConnection { @@ -186,453 +58,48 @@ impl ResponsesWebSocketConnection { } } -fn marshal_inputs( - py: Python<'_>, - document: Py, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - let document = from_py(document.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - - Ok((document, extra_headers, optional_params, timeout)) -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn ocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (document, extra_headers, optional_params, timeout) = marshal_inputs( - py, - document, - extra_headers, - optional_params, - timeout_seconds, - )?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - })) - }); - - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn aocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (document, extra_headers, optional_params, timeout) = marshal_inputs( - py, - document, - extra_headers, - optional_params, - timeout_seconds, - )?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }) - .await - .map_err(core_error_to_pyerr)?; - - Python::attach(|py| to_py(py, &value)) - }) -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn transcription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( - AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }, - )) - }); - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn atranscription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_audio_transcription(AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - 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>, Option); - -fn marshal_messages_inputs( - py: Python<'_>, - body: Py, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let body: Value = from_py(body.bind(py))?; - if !body.is_object() { - return Err(PyValueError::new_err("body must be a dict")); - } - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - Ok((body, extra_headers, optional_timeout(timeout_seconds))) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn messages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (body, extra_headers, timeout) = - marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { - model: &model, - body, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - })) - }); - - match result { - Ok(response) => messages_response_to_py(py, response), - Err(err) => Err(core_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn amessages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (body, extra_headers, timeout) = - marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = run_messages(MessagesRequest { - model: &model, - body, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await - .map_err(core_error_to_pyerr)?; - - Python::attach(|py| messages_response_to_py(py, response)) - }) -} - -type MarshaledChatCompletionsInputs = ( - Value, - Map, - Option>, - Option, -); - -fn marshal_chat_completions_inputs( - py: Python<'_>, - messages: Py, - optional_params: Option>, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let messages: Value = from_py(messages.bind(py))?; - if !messages.is_array() { - return Err(PyValueError::new_err("messages must be a list")); - } - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - Ok(( - messages, - optional_params, - extra_headers, - optional_timeout(timeout_seconds), - )) -} - -/// The decline reason for this request, or `None` when the Rust path accepts -/// it. Resolves no credentials and performs no I/O, so a host can ask before -/// committing to either path. -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] -fn chat_completions_decline( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - custom_llm_provider: Option, -) -> PyResult> { - let messages = from_py(messages.bind(py))?; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - Ok(chat_completions_decline_reason( - &model, - custom_llm_provider.as_deref(), - messages, - &optional_params, - ) - .map(str::to_string)) -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn chat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( - py, - messages, - optional_params, - extra_headers, - timeout_seconds, - )?; - - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( - ChatCompletionsRequest { - model: &model, - messages, - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }, - )) - }); - - match result { - Ok(response) => chat_completions_response_to_py(py, response), - Err(err) => Err(chat_completions_error_to_pyerr(err)), - } -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[allow(clippy::too_many_arguments)] -fn achat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( - py, - messages, - optional_params, - extra_headers, - timeout_seconds, - )?; - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let response = run_chat_completions(ChatCompletionsRequest { - model: &model, - messages, - optional_params, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await - .map_err(chat_completions_error_to_pyerr)?; - - Python::attach(|py| chat_completions_response_to_py(py, response)) - }) -} - -#[pyfunction] -fn gil_stats(py: Python<'_>) -> PyResult> { - let stats = PyDict::new(py); - stats.set_item("releases", release_count())?; - Ok(stats.into_any().unbind()) -} - -#[cfg(feature = "panic-test")] -#[pyfunction] -fn _panic_for_test() { - panic!("intentional PyO3 panic smoke test"); -} - #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { - let py = module.py(); - module.add_function(wrap_pyfunction!(ocr, module)?)?; - module.add_function(wrap_pyfunction!(aocr, module)?)?; - module.add_function(wrap_pyfunction!(transcription, module)?)?; - module.add_function(wrap_pyfunction!(atranscription, module)?)?; - module.add_function(wrap_pyfunction!(messages, module)?)?; - module.add_function(wrap_pyfunction!(amessages, module)?)?; - module.add("RustBridgeDeclined", py.get_type::())?; - module.add("RustUpstreamError", py.get_type::())?; - module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; - module.add_function(wrap_pyfunction!(chat_completions, module)?)?; - module.add_function(wrap_pyfunction!(achat_completions, module)?)?; + errors::register(module)?; + routes::register(module)?; module.add_class::()?; - module.add_function(wrap_pyfunction!(gil_stats, module)?)?; - #[cfg(feature = "panic-test")] - module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; - Ok(()) + diagnostics::register(module) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn module_registration_preserves_the_public_surface() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "_native").expect("module should be created"); + _native(&module).expect("module should register"); + + let expected = [ + "RustBridgeDeclined", + "RustUpstreamError", + "ocr", + "aocr", + "transcription", + "atranscription", + "messages", + "amessages", + "chat_completions_decline", + "chat_completions", + "achat_completions", + "ResponsesWebSocketConnection", + "gil_stats", + ]; + + for name in expected { + assert!( + module + .hasattr(name) + .expect("attribute lookup should succeed") + ); + } + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..6c070d2a9ee --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,53 @@ +use std::collections::HashMap; +use std::time::Duration; + +use litellm_python_interop::from_py; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +pub(crate) fn optional_object_to_map( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult> { + match value { + Some(value) => match from_py(value.bind(py))? { + Value::Object(map) => Ok(map), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + }, + None => Ok(Map::new()), + } +} + +pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { + timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }) +} + +pub(crate) fn marshal_headers( + py: Python<'_>, + headers: Option>, +) -> PyResult> { + let value = match headers { + Some(headers) => from_py(headers.bind(py))?, + None => Value::Object(Map::new()), + }; + let Value::Object(headers) = value else { + return Err(PyValueError::new_err("headers must be a dict")); + }; + headers + .into_iter() + .map(|(name, value)| { + value + .as_str() + .map(|value| (name, value.to_string())) + .ok_or_else(|| PyValueError::new_err("header values must be strings")) + }) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs new file mode 100644 index 00000000000..ec0dfc75df2 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -0,0 +1,100 @@ +use litellm_ai_gateway::io::audio_transcription::{ + AudioTranscriptionRequest, audio_transcription as run_audio_transcription, +}; +use litellm_python_interop::{from_py, release_gil, to_py}; +use pyo3::prelude::*; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object_to_map, optional_timeout}; + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn transcription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let audio = from_py(audio.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + let result = release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( + AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }, + )) + }); + match result { + Ok(value) => to_py(py, &value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn atranscription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let audio = from_py(audio.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = run_audio_transcription(AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + 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)) + }) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(transcription, module)?)?; + module.add_function(wrap_pyfunction!(atranscription, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs new file mode 100644 index 00000000000..b353765e7cc --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -0,0 +1,163 @@ +use std::time::Duration; + +use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; +use litellm_core::chat_completions::{ + chat_completions as run_chat_completions, chat_completions_decline_reason, +}; +use litellm_python_interop::{from_py, release_gil, to_py}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::chat_completions_error_to_pyerr; +use crate::marshal::{optional_object_to_map, optional_timeout}; + +fn chat_completions_response_to_py( + py: Python<'_>, + response: ChatCompletionsResponse, +) -> PyResult> { + to_py(py, &response) +} + +type MarshaledChatCompletionsInputs = ( + Value, + Map, + Option>, + Option, +); + +fn marshal_chat_completions_inputs( + py: Python<'_>, + messages: Py, + optional_params: Option>, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let messages: Value = from_py(messages.bind(py))?; + if !messages.is_array() { + return Err(PyValueError::new_err("messages must be a list")); + } + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + Ok(( + messages, + optional_params, + extra_headers, + optional_timeout(timeout_seconds), + )) +} + +/// The decline reason for this request, or `None` when the Rust path accepts +/// it. Resolves no credentials and performs no I/O, so a host can ask before +/// committing to either path. +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] +fn chat_completions_decline( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + custom_llm_provider: Option, +) -> PyResult> { + let messages = from_py(messages.bind(py))?; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + Ok(chat_completions_decline_reason( + &model, + custom_llm_provider.as_deref(), + messages, + &optional_params, + ) + .map(str::to_string)) +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn chat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( + py, + messages, + optional_params, + extra_headers, + timeout_seconds, + )?; + + let result = release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( + ChatCompletionsRequest { + model: &model, + messages, + optional_params, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }, + )) + }); + + match result { + Ok(response) => chat_completions_response_to_py(py, response), + Err(err) => Err(chat_completions_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn achat_completions( + py: Python<'_>, + model: String, + messages: Py, + optional_params: Option>, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs( + py, + messages, + optional_params, + extra_headers, + timeout_seconds, + )?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let response = run_chat_completions(ChatCompletionsRequest { + model: &model, + messages, + optional_params, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await + .map_err(chat_completions_error_to_pyerr)?; + + Python::attach(|py| chat_completions_response_to_py(py, response)) + }) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; + module.add_function(wrap_pyfunction!(chat_completions, module)?)?; + module.add_function(wrap_pyfunction!(achat_completions, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs new file mode 100644 index 00000000000..ac361bb0e59 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -0,0 +1,109 @@ +use std::time::Duration; + +use litellm_core::messages::messages as run_messages; +use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_python_interop::{from_py, release_gil, to_py}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object_to_map, optional_timeout}; + +fn messages_response_to_py( + py: Python<'_>, + response: AnthropicMessagesResponse, +) -> PyResult> { + to_py(py, &response) +} + +type MarshaledMessagesInputs = (Value, Option>, Option); + +fn marshal_messages_inputs( + py: Python<'_>, + body: Py, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult { + let body: Value = from_py(body.bind(py))?; + if !body.is_object() { + return Err(PyValueError::new_err("body must be a dict")); + } + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + Ok((body, extra_headers, optional_timeout(timeout_seconds))) +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn messages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (body, extra_headers, timeout) = + marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; + + let result = release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { + model: &model, + body, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + })) + }); + + match result { + Ok(response) => messages_response_to_py(py, response), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn amessages( + py: Python<'_>, + model: String, + body: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (body, extra_headers, timeout) = + marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let response = run_messages(MessagesRequest { + model: &model, + body, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await + .map_err(core_error_to_pyerr)?; + + Python::attach(|py| messages_response_to_py(py, response)) + }) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(messages, module)?)?; + module.add_function(wrap_pyfunction!(amessages, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs new file mode 100644 index 00000000000..a2eb8355767 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -0,0 +1,13 @@ +use pyo3::prelude::*; + +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) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs new file mode 100644 index 00000000000..047bf245d7c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -0,0 +1,128 @@ +use std::time::Duration; + +use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; +use litellm_python_interop::{from_py, release_gil, to_py}; +use pyo3::prelude::*; +use serde_json::{Map, Value}; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{optional_object_to_map, optional_timeout}; + +type MarshaledOcrInputs = ( + Value, + Option>, + Map, + Option, +); + +fn marshal_inputs( + py: Python<'_>, + document: Py, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + let document = from_py(document.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + + Ok((document, extra_headers, optional_params, timeout)) +} + +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn ocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (document, extra_headers, optional_params, timeout) = marshal_inputs( + py, + document, + extra_headers, + optional_params, + timeout_seconds, + )?; + + let result = release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { + model: &model, + document, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + })) + }); + + match result { + Ok(value) => to_py(py, &value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn aocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let (document, extra_headers, optional_params, timeout) = marshal_inputs( + py, + document, + extra_headers, + optional_params, + timeout_seconds, + )?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = run_ocr(OcrRequest { + model: &model, + document, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await + .map_err(core_error_to_pyerr)?; + + Python::attach(|py| to_py(py, &value)) + }) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(aocr, module)?) +}