diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 1c780b2baf9..a14e4b55d82 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -1,33 +1,77 @@ 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, + pub(crate) api_base: Option, + pub(crate) custom_llm_provider: Option, + pub(crate) extra_headers: Option>, + pub(crate) timeout: Option, +} + +pub(crate) struct RouteOptionsInputs { + pub(crate) model: String, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) custom_llm_provider: Option, + pub(crate) extra_headers: Option, + pub(crate) timeout_seconds: Option, +} + +impl RouteOptions { + pub(crate) fn from_python(inputs: RouteOptionsInputs) -> PyResult { + 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>, + value: Value, + expected: fn(&Value) -> bool, + expected_name: &'static str, +) -> PyResult { + 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, ) -> 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"))), - }, + Some(value) => object(name, value), None => Ok(Map::new()), } } -pub(crate) fn optional_object( - py: Python<'_>, +fn optional_object( name: &'static str, - value: Option>, + value: Option, ) -> PyResult>> { - value - .map(|value| optional_object_to_map(py, name, Some(value))) - .transpose() + value.map(|value| object(name, value)).transpose() +} + +fn object(name: &'static str, value: Value) -> PyResult> { + 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) -> Option { diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 1c925ed8592..10b86132be7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -1,126 +1,71 @@ -use std::time::Duration; +use litellm_core::Error; +use std::future::Future; use litellm_core::audio_transcription::{ AudioTranscriptionRequest, audio_transcription as run_audio_transcription, }; -use litellm_core::error::Error; -use litellm_python_interop::from_py; use pyo3::prelude::*; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; -struct TranscriptionInputs { - model: String, - audio: Value, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Map, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - Ok(TranscriptionInputs { - model, - audio: from_py(audio.bind(py))?, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - optional_params: optional_object_to_map(py, "optional_params", optional_params)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: TranscriptionInputs) -> Result { - run_audio_transcription(AudioTranscriptionRequest { - model: &inputs.model, - audio: inputs.audio, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_transcription( + inputs: AudioTranscriptionInputs, +) -> PyResult> + 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, - optional_params: inputs.optional_params, - timeout: inputs.timeout, + 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 }) - .await } -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn transcription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - audio, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn atranscription( - py: Python<'_>, - model: String, - audio: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - audio, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(transcription, module)?)?; - module.add_function(wrap_pyfunction!(atranscription, module)?) +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, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + timeout_seconds: Option, + }, + prepare = prepare_transcription, + errors = core_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 4dc446f9985..68b7762cb10 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -1,87 +1,62 @@ -use std::time::Duration; +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 litellm_core::error::Error; -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::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; -struct ChatCompletionsInputs { - model: String, - messages: Value, - optional_params: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let messages: Value = from_py(messages.bind(py))?; - if !messages.is_array() { - return Err(PyValueError::new_err("messages must be a list")); - } - Ok(ChatCompletionsInputs { - model, - messages, - optional_params: optional_object_to_map(py, "optional_params", optional_params)?, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: ChatCompletionsInputs) -> Result { - run_chat_completions(ChatCompletionsRequest { - model: &inputs.model, - messages: inputs.messages, - optional_params: inputs.optional_params, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_chat_completions( + inputs: ChatCompletionsInputs, +) -> PyResult> + 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: inputs.timeout, + 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 }) - .await } -/// The decline reason for this request, or `None` when the Rust path accepts -/// it. Resolves no credentials and performs no I/O, so a host can ask before -/// committing to either path. #[pyfunction] #[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))] fn chat_completions_decline( - py: Python<'_>, model: String, - messages: Py, - optional_params: Option>, + #[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value, + #[pyo3(from_py_with = litellm_python_interop::from_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)?; + let optional_params = object_or_empty("optional_params", optional_params)?; Ok(chat_completions_decline_reason( &model, custom_llm_provider.as_deref(), @@ -91,74 +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, trace=false))] -#[allow(clippy::too_many_arguments)] -fn chat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - messages, - optional_params, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_sync( - py, - trace_call(call(inputs), trace), - chat_completions_error_to_pyerr, - ) -} - -#[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn achat_completions( - py: Python<'_>, - model: String, - messages: Py, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - messages, - optional_params, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_async( - py, - trace_call(call(inputs), trace), - chat_completions_error_to_pyerr, - ) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?; - module.add_function(wrap_pyfunction!(chat_completions, module)?)?; - module.add_function(wrap_pyfunction!(achat_completions, module)?) +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, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + timeout_seconds: Option, + }, + prepare = prepare_chat_completions, + errors = chat_completions_error_to_pyerr, + extra = [chat_completions_decline], } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs new file mode 100644 index 00000000000..21a7fd5a766 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -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> { + 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> { + 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> + 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::(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" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index b91f1567e06..2bb64a7a763 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -1,122 +1,65 @@ -use std::time::Duration; - -use litellm_core::error::Error; +use litellm_core::Error; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; -use litellm_python_interop::from_py; -use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use serde_json::{Map, Value}; +use serde_json::Value; +use std::future::Future; use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; -struct MessagesInputs { - model: String, - body: Value, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - let body: Value = from_py(body.bind(py))?; - if !body.is_object() { - return Err(PyValueError::new_err("body must be a dict")); - } - Ok(MessagesInputs { - model, - body, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: MessagesInputs) -> Result { - run_messages(MessagesRequest { - model: &inputs.model, - body: inputs.body, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_messages( + inputs: MessagesInputs, +) -> PyResult> + 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: inputs.timeout, + 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 }) - .await } -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn messages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - body, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn amessages( - py: Python<'_>, - model: String, - body: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - body, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout_seconds, - )?; - run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(messages, module)?)?; - module.add_function(wrap_pyfunction!(amessages, module)?) +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, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + timeout_seconds: Option, + }, + prepare = prepare_messages, + errors = core_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index a2eb8355767..bf611c26d44 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -1,5 +1,8 @@ use pyo3::prelude::*; +#[macro_use] +mod definition; + mod audio_transcription; mod chat_completions; mod messages; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index bed63a871b5..5588c400972 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -1,128 +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_core::error::Error; -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::execution::{run_async, run_sync}; -use crate::function_trace::trace_call; -use crate::marshal::{optional_object, optional_object_to_map, optional_timeout}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; -struct OcrInputs { - model: String, - document: Value, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Map, - timeout: Option, -} - -#[allow(clippy::too_many_arguments)] -fn marshal_inputs( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - Ok(OcrInputs { - model, - document: from_py(document.bind(py))?, - api_key, - api_base, - custom_llm_provider, - extra_headers: optional_object(py, "extra_headers", extra_headers)?, - optional_params: optional_object_to_map(py, "optional_params", optional_params)?, - timeout: optional_timeout(timeout_seconds), - }) -} - -async fn call(inputs: OcrInputs) -> Result { - run_ocr(OcrRequest { - model: &inputs.model, - document: inputs.document, - api_key: inputs.api_key.as_deref(), - api_base: inputs.api_base.as_deref(), - custom_llm_provider: inputs.custom_llm_provider.as_deref(), +fn prepare_ocr( + inputs: OcrInputs, +) -> PyResult> + 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, - optional_params: inputs.optional_params, - timeout: inputs.timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, + 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 }) - .await } -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn ocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_sync(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=false))] -#[allow(clippy::too_many_arguments)] -fn aocr( - py: Python<'_>, - model: String, - document: Py, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, - trace: bool, -) -> PyResult> { - let inputs = marshal_inputs( - py, - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - timeout_seconds, - )?; - run_async(py, trace_call(call(inputs), trace), core_error_to_pyerr) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(ocr, module)?)?; - module.add_function(wrap_pyfunction!(aocr, module)?) +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, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + timeout_seconds: Option, + }, + prepare = prepare_ocr, + errors = core_error_to_pyerr, } diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs new file mode 100644 index 00000000000..87a0c3e0104 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/runtime.rs @@ -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( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_on( + py, + pyo3_async_runtimes::tokio::get_runtime(), + future, + map_error, + ) +} + +fn run_sync_on( + py: Python<'_>, + runtime: &Runtime, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + 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( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + 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(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { + 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(future: F) -> PyResult> +where + F: Future>, +{ + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(panic_to_pyerr) +} + +async fn wait_for_sync_result(future: F) -> PyResult> +where + F: Future>, +{ + 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(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[pyfunction] + fn async_serialization_panic(py: Python<'_>) -> PyResult> { + run_async(py, async { Ok(PanickingOutput) }, runtime_error) + } + + #[pyfunction] + fn async_runtime_probe(py: Python<'_>) -> PyResult> { + 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>) -> 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::(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::( + py, + poll_fn(|_| -> Poll> { panic!("route future panicked") }), + runtime_error, + ) + .expect_err("panicked route should become a Python exception"); + + assert!(error.is_instance_of::(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::( + 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::(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::(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"); + }); + } +}