From 42b8cd80beb350e3bfdb85a47d00234e88315974 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 8 Sep 2026 08:10:45 -0700 Subject: [PATCH] wip --- .../crates/python-bridge/src/driver.rs | 236 +++++++++++++----- .../src/routes/chat_completions.rs | 168 ++++--------- .../python-bridge/src/routes/messages.rs | 183 +++----------- .../crates/python-bridge/src/routes/ocr.rs | 192 ++++---------- litellm/rust_bridge/_lifecycle.py | 154 +++++++++++- litellm/rust_bridge/chat_completions.py | 141 +++++++++++ litellm/rust_bridge/messages.py | 154 +++++++++++- litellm/rust_bridge/ocr.py | 147 ++++++++++- .../test_prometheus_caller_identity.py | 10 +- ...prometheus_custom_metadata_label_counts.py | 9 +- .../integrations/test_prometheus_labels.py | 19 +- .../test_prometheus_rate_limit_labels.py | 27 +- .../test_prometheus_service_tier_label.py | 15 +- .../rust_bridge/test_lifecycle.py | 91 +++++++ tests/test_litellm_rust/recording_server.py | 11 +- 15 files changed, 1016 insertions(+), 541 deletions(-) create mode 100644 tests/test_litellm/rust_bridge/test_lifecycle.py diff --git a/litellm-rust/crates/python-bridge/src/driver.rs b/litellm-rust/crates/python-bridge/src/driver.rs index e1e35845850..f1db0ddb8fd 100644 --- a/litellm-rust/crates/python-bridge/src/driver.rs +++ b/litellm-rust/crates/python-bridge/src/driver.rs @@ -1,83 +1,197 @@ -use std::ffi::CString; - +use litellm_core::lifecycle::program::Operation; +use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; -use crate::errors::RustBridgeDriverError; +pub(crate) const ADDITIONAL_ARGS: &str = "additional_args"; +pub(crate) const API_BASE: &str = "api_base"; +pub(crate) const API_KEY: &str = "api_key"; +pub(crate) const COMPLETE_INPUT_DICT: &str = "complete_input_dict"; +pub(crate) const HEADERS: &str = "headers"; +pub(crate) const INPUT: &str = "input"; -const DRIVE: &str = r#" -def drive_sync(arguments): - host = Host(arguments, False) - while host.machine.complete() is None: - try: - _invoke(host.machine, host) - except Exception as error: - host.advance(1, error) - except BaseException as error: - host.advance(2, error) - else: - host.advance(0) - return host.result() +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct OperationBinding { + method: &'static str, + awaiting: bool, +} -async def drive_async(arguments): - host = Host(arguments, True) - while host.machine.complete() is None: - try: - awaiting, value = _invoke(host.machine, host) - if awaiting: - await value - except Exception as error: - host.advance(1, error) - except BaseException as error: - host.advance(2, error) - else: - host.advance(0) - return host.result() -"#; - -pub(crate) fn compile<'py>( - py: Python<'py>, +fn operation_binding( + operation: Operation, + asynchronous: bool, + supports_pre_call: bool, route: &str, - host: &str, -) -> PyResult> { - let source = CString::new(format!("{host}\n{DRIVE}")) - .map_err(|_| RustBridgeDriverError::new_err("driver source contains a null byte"))?; - let filename = CString::new(format!("{route}_driver.py")) - .map_err(|_| RustBridgeDriverError::new_err("driver route name contains a null byte"))?; - let module_name = CString::new(format!("_{route}_driver")) - .map_err(|_| RustBridgeDriverError::new_err("driver route name contains a null byte"))?; - PyModule::from_code(py, &source, &filename, &module_name) +) -> PyResult { + let binding = match operation { + Operation::Setup => OperationBinding { + method: "setup", + awaiting: false, + }, + Operation::DeploymentPre => OperationBinding { + method: "deployment_pre", + awaiting: true, + }, + Operation::Prepare => OperationBinding { + method: "prepare", + awaiting: false, + }, + Operation::PreCall if supports_pre_call => OperationBinding { + method: "pre_call", + awaiting: false, + }, + Operation::PreCall => { + return Err(PyRuntimeError::new_err(format!( + "{route} lifecycle selected an unsupported pre-call operation" + ))); + } + Operation::Send if asynchronous => OperationBinding { + method: "send", + awaiting: true, + }, + Operation::Send => OperationBinding { + method: "send_sync", + awaiting: false, + }, + Operation::DeploymentSuccess => OperationBinding { + method: "deployment_success", + awaiting: true, + }, + Operation::DeploymentFailure => OperationBinding { + method: "deployment_failure", + awaiting: true, + }, + Operation::SyncSuccess => OperationBinding { + method: "sync_success", + awaiting: false, + }, + Operation::AsyncSuccess => OperationBinding { + method: "async_success", + awaiting: false, + }, + Operation::SyncSuccessIfNeeded => OperationBinding { + method: "sync_success_if_needed", + awaiting: false, + }, + Operation::SyncFailure => OperationBinding { + method: "sync_failure", + awaiting: false, + }, + Operation::AsyncFailure => OperationBinding { + method: "async_failure", + awaiting: true, + }, + Operation::Restore => OperationBinding { + method: "restore", + awaiting: false, + }, + Operation::Complete(_) => { + return Err(PyRuntimeError::new_err(format!( + "{route} lifecycle is complete" + ))); + } + }; + Ok(binding) +} + +pub(crate) fn invoke( + py: Python<'_>, + operation: Operation, + asynchronous: bool, + supports_pre_call: bool, + route: &str, + host: Py, +) -> PyResult<(bool, Py)> { + let binding = operation_binding(operation, asynchronous, supports_pre_call, route)?; + Ok(( + binding.awaiting, + host.getattr(py, binding.method)?.call0(py)?, + )) } #[cfg(test)] mod tests { + use litellm_core::lifecycle::Outcome; + use super::*; #[test] - fn null_byte_in_route_raises_driver_error() { + fn operation_bindings_cover_the_lifecycle_contract() { Python::initialize(); - Python::attach(|py| { - let error = compile(py, "invalid\0route", "class Host: pass") - .expect_err("route names containing null bytes should fail"); - - assert!(error.is_instance_of::(py)); - assert_eq!( - error.to_string(), - "RustBridgeDriverError: driver route name contains a null byte" - ); + Python::attach(|_| { + let cases = [ + (Operation::Setup, false, false, "setup", false), + ( + Operation::DeploymentPre, + false, + false, + "deployment_pre", + true, + ), + (Operation::Prepare, false, false, "prepare", false), + (Operation::PreCall, false, true, "pre_call", false), + (Operation::Send, false, false, "send_sync", false), + (Operation::Send, true, false, "send", true), + ( + Operation::DeploymentSuccess, + false, + false, + "deployment_success", + true, + ), + ( + Operation::DeploymentFailure, + false, + false, + "deployment_failure", + true, + ), + (Operation::SyncSuccess, false, false, "sync_success", false), + ( + Operation::AsyncSuccess, + false, + false, + "async_success", + false, + ), + ( + Operation::SyncSuccessIfNeeded, + false, + false, + "sync_success_if_needed", + false, + ), + (Operation::SyncFailure, false, false, "sync_failure", false), + (Operation::AsyncFailure, false, false, "async_failure", true), + (Operation::Restore, false, false, "restore", false), + ]; + for (operation, asynchronous, pre_call, method, awaiting) in cases { + assert_eq!( + operation_binding(operation, asynchronous, pre_call, "test").unwrap(), + OperationBinding { method, awaiting }, + ); + } }); } #[test] - fn null_byte_in_source_raises_driver_error() { + fn invalid_operations_raise_route_specific_errors() { Python::initialize(); - Python::attach(|py| { - let error = compile(py, "test", "class Host:\0 pass") - .expect_err("driver source containing null bytes should fail"); - - assert!(error.is_instance_of::(py)); + Python::attach(|_| { + let pre_call = operation_binding(Operation::PreCall, false, false, "messages") + .expect_err("unsupported pre-call should fail"); assert_eq!( - error.to_string(), - "RustBridgeDriverError: driver source contains a null byte" + pre_call.to_string(), + "RuntimeError: messages lifecycle selected an unsupported pre-call operation" + ); + let complete = operation_binding( + Operation::Complete(Outcome::Success), + false, + false, + "messages", + ) + .expect_err("complete lifecycle should fail"); + assert_eq!( + complete.to_string(), + "RuntimeError: messages lifecycle is complete" ); }); } 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 6478d3f553a..9f40b8562e5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -17,6 +17,7 @@ use pyo3::sync::PyOnceLock; use pyo3::types::PyDict; use serde_json::{Map, Value}; +use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT}; use crate::errors::{RustBridgeDeclined, chat_completions_error_to_pyerr, core_error_to_pyerr}; use crate::marshal::optional_timeout; @@ -154,36 +155,15 @@ fn invoke( let machine = machine.borrow(py); (machine.machine.operation(), machine.asynchronous) }; - let (method, awaiting) = match operation { - Operation::Setup => ("setup", false), - Operation::DeploymentPre => ("deployment_pre", true), - Operation::Prepare => ("prepare", false), - Operation::PreCall => { - return Err(PyRuntimeError::new_err( - "chat completions lifecycle selected an unsupported pre-call operation", - )); - } - Operation::Send if asynchronous => ("send", true), - Operation::Send => ("send_sync", false), - Operation::DeploymentSuccess => ("deployment_success", true), - Operation::DeploymentFailure => ("deployment_failure", true), - Operation::SyncSuccess => ("sync_success", false), - Operation::AsyncSuccess => ("async_success", false), - Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false), - Operation::SyncFailure => ("sync_failure", false), - Operation::AsyncFailure => ("async_failure", true), - Operation::Restore => ("restore", false), - Operation::Complete(_) => { - return Err(PyRuntimeError::new_err( - "chat completions lifecycle is complete", - )); - } - }; - Ok((awaiting, host.getattr(py, method)?.call0(py)?)) + crate::driver::invoke(py, operation, asynchronous, false, "chat completions", host) } #[pyfunction] -fn prepare(py: Python<'_>, arguments: Py) -> PyResult> { +fn prepare( + py: Python<'_>, + arguments: Py, + logging: Py, +) -> PyResult> { let bag = arguments.bind(py); let admission = admission(bag)?; let api_key = scalar(bag, "api_key")?; @@ -195,10 +175,6 @@ fn prepare(py: Python<'_>, arguments: Py) -> PyResult()) .transpose()?, )?; - let logging = bag - .get_item("litellm_logging_obj")? - .filter(|value| !value.is_none()) - .ok_or_else(|| PyRuntimeError::new_err("chat completions logging was not initialized"))?; let complete_input = PyDict::new(py); complete_input.set_item("model", &admission.model)?; complete_input.set_item("messages", bag.get_item("messages")?)?; @@ -206,14 +182,16 @@ fn prepare(py: Python<'_>, arguments: Py) -> PyResult) -> PyResult<()> { #[pyfunction] fn chat_completions(py: Python<'_>, arguments: Py) -> PyResult> { validate_arguments(arguments.bind(py))?; - driver(py)?.getattr("drive_sync")?.call1((arguments,)) + runner(py)? + .getattr("_drive_sync")? + .call1((arguments, bindings(py)?)) } #[pyfunction] fn achat_completions(py: Python<'_>, arguments: Py) -> PyResult> { validate_arguments(arguments.bind(py))?; - driver(py)?.getattr("drive_async")?.call1((arguments,)) + runner(py)? + .getattr("_drive_async")? + .call1((arguments, bindings(py)?)) } #[pyfunction] @@ -376,94 +358,32 @@ fn chat_completions_decline( .map(str::to_string) } -fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { - static DRIVER: PyOnceLock> = PyOnceLock::new(); - if let Some(module) = DRIVER.get(py) { +fn runner(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static RUNNER: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = RUNNER.get(py) { return Ok(module.bind(py)); } - let module = crate::driver::compile(py, "chat_completions", HOST)?; - module.add("_Lifecycle", py.get_type::())?; - module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?; - module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; - module.add("_send", wrap_pyfunction!(send, &module)?)?; - module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?; - module.add( - "_terminal_record", - wrap_pyfunction!(terminal_record, &module)?, - )?; - Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) + let module = py.import("litellm.rust_bridge.chat_completions")?; + Ok(RUNNER.get_or_init(py, || module.unbind()).bind(py)) } -const HOST: &str = r#" -from datetime import datetime -from litellm import utils -from litellm.types.utils import CallTypes -from litellm.rust_bridge.chat_completions import build_model_response, initialize_logging, invoke_terminal - -class Host: - def __init__(self, arguments, asynchronous): - self.machine = _Lifecycle(arguments, asynchronous, utils.is_internal_call.get()) - self.arguments = arguments - self.current = arguments - self.asynchronous = asynchronous - self.logger = arguments.get('litellm_logging_obj') - self.state = None - self.response = None - self.error = None - self.start = datetime.now() - self.end = None - - def setup(self): - self.logger = initialize_logging(self.arguments, self.asynchronous) - self.arguments['litellm_logging_obj'] = self.logger - - async def deployment_pre(self): - modified = await utils.async_pre_call_deployment_hook(self.current, 'acompletion') - if modified is not None: - self.current = modified - self.current['litellm_logging_obj'] = self.logger - - def prepare(self): self.state = _prepare(self.current) - - def send_sync(self): - self.response = build_model_response(_send_sync(self.state), self.arguments['model_response']) - self.end = datetime.now() - - async def send(self): - self.response = build_model_response(await _send(self.state), self.arguments['model_response']) - self.end = datetime.now() - - async def deployment_success(self): - self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.acompletion) - - async def deployment_failure(self): - await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'acompletion') - - def terminal(self, action, value): - record = _terminal_record(self.state) if self.state is not None else None - return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, record, value, self.start, self.end) - - def sync_success(self): return self.terminal('sync_success', self.response) - def async_success(self): return self.terminal('async_success', self.response) - def sync_success_if_needed(self): return self.terminal('sync_success_if_needed', self.response) - def sync_failure(self): return self.terminal('sync_failure', self.error) - def async_failure(self): return self.terminal('async_failure', self.error) - def restore(self): utils._restore_correlation_context_if_supported(self.logger) - - def advance(self, outcome, error=None): - if error is not None and self.end is None: - self.end = datetime.now() - if self.logger is None: - self.logger = self.arguments.get('litellm_logging_obj') - replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None) - if replace: - self.error = error - - def result(self): - if self.machine.complete(): - return self.response - raise self.error -"#; +fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static BINDINGS: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = BINDINGS.get(py) { + return Ok(module.bind(py)); + } + let module = PyModule::new(py, "_chat_completions_bindings")?; + module.add("Lifecycle", py.get_type::())?; + module.add("invoke", wrap_pyfunction!(invoke, &module)?)?; + module.add("prepare", wrap_pyfunction!(prepare, &module)?)?; + module.add("send", wrap_pyfunction!(send, &module)?)?; + module.add("send_sync", wrap_pyfunction!(send_sync, &module)?)?; + module.add( + "terminal_record", + wrap_pyfunction!(terminal_record, &module)?, + )?; + Ok(BINDINGS.get_or_init(py, || module.unbind()).bind(py)) +} pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { crate::routes::definition::add_function(module, wrap_pyfunction!(chat_completions, module)?)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index a6adfd66f9f..479c77aaed4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -11,6 +11,7 @@ use pyo3::sync::PyOnceLock; use pyo3::types::PyDict; use serde_json::{Map, Value}; +use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT}; use crate::errors::{RustUpstreamError, core_error_to_pyerr, messages_provider_error_to_pyerr}; use crate::marshal::optional_timeout; @@ -145,34 +146,15 @@ fn invoke( let machine = machine.borrow(py); (machine.machine.operation(), machine.asynchronous) }; - let (method, awaiting) = match operation { - Operation::Setup => ("setup", false), - Operation::DeploymentPre => ("deployment_pre", true), - Operation::Prepare => ("prepare", false), - Operation::PreCall => { - return Err(PyRuntimeError::new_err( - "messages lifecycle selected an unsupported pre-call operation", - )); - } - Operation::Send if asynchronous => ("send", true), - Operation::Send => ("send_sync", false), - Operation::DeploymentSuccess => ("deployment_success", true), - Operation::DeploymentFailure => ("deployment_failure", true), - Operation::SyncSuccess => ("sync_success", false), - Operation::AsyncSuccess => ("async_success", false), - Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false), - Operation::SyncFailure => ("sync_failure", false), - Operation::AsyncFailure => ("async_failure", true), - Operation::Restore => ("restore", false), - Operation::Complete(_) => { - return Err(PyRuntimeError::new_err("messages lifecycle is complete")); - } - }; - Ok((awaiting, host.getattr(py, method)?.call0(py)?)) + crate::driver::invoke(py, operation, asynchronous, false, "messages", host) } #[pyfunction] -fn prepare(py: Python<'_>, arguments: Py) -> PyResult> { +fn prepare( + py: Python<'_>, + arguments: Py, + logging: Py, +) -> PyResult> { let bag = arguments.bind(py); let request = decode_request(py, bag)?; let prepared = py @@ -185,23 +167,18 @@ fn prepare(py: Python<'_>, arguments: Py) -> PyResult> for (name, value) in &prepared.upstream_headers { headers.set_item(name, value)?; } - let logging = bag - .get_item("litellm_logging_obj")? - .filter(|value| !value.is_none()) - .map(|value| value.unbind()) - .ok_or_else(|| PyRuntimeError::new_err("messages logging was not initialized"))?; let additional = PyDict::new(py); - additional.set_item(pyo3::intern!(py, "complete_input_dict"), &body)?; - additional.set_item(pyo3::intern!(py, "api_base"), &prepared.url)?; - additional.set_item(pyo3::intern!(py, "headers"), &headers)?; + additional.set_item(COMPLETE_INPUT_DICT, &body)?; + additional.set_item(API_BASE, &prepared.url)?; + additional.set_item(HEADERS, &headers)?; let kwargs = PyDict::new(py); let serialized = py.import("json")?.call_method1("dumps", (&body,))?; let message = PyDict::new(py); message.set_item("role", "user")?; message.set_item("content", serialized)?; - kwargs.set_item("input", vec![message])?; - kwargs.set_item("api_key", "")?; - kwargs.set_item("additional_args", additional)?; + kwargs.set_item(INPUT, vec![message])?; + kwargs.set_item(API_KEY, "")?; + kwargs.set_item(ADDITIONAL_ARGS, additional)?; logging .bind(py) .call_method(pyo3::intern!(py, "pre_call"), (), Some(&kwargs))?; @@ -296,123 +273,45 @@ fn committed_failure() -> PyResult<()> { #[pyfunction] fn messages(py: Python<'_>, arguments: Py) -> PyResult> { validate_arguments(arguments.bind(py))?; - driver(py)?.getattr("drive_sync")?.call1((arguments,)) + runner(py)? + .getattr("_drive_sync")? + .call1((arguments, bindings(py)?)) } #[pyfunction] fn amessages(py: Python<'_>, arguments: Py) -> PyResult> { validate_arguments(arguments.bind(py))?; - driver(py)?.getattr("drive_async")?.call1((arguments,)) + runner(py)? + .getattr("_drive_async")? + .call1((arguments, bindings(py)?)) } -fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { - static DRIVER: PyOnceLock> = PyOnceLock::new(); - if let Some(module) = DRIVER.get(py) { +fn runner(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static RUNNER: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = RUNNER.get(py) { return Ok(module.bind(py)); } - let module = crate::driver::compile(py, "messages", HOST)?; - module.add("_Lifecycle", py.get_type::())?; - module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?; - module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; - module.add("_send", wrap_pyfunction!(send, &module)?)?; - module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?; - module.add( - "_committed_failure", - wrap_pyfunction!(committed_failure, &module)?, - )?; - Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) + let module = py.import("litellm.rust_bridge.messages")?; + Ok(RUNNER.get_or_init(py, || module.unbind()).bind(py)) } -const HOST: &str = r#" -from datetime import datetime -from litellm import utils -from litellm.types.utils import CallTypes -from litellm.rust_bridge.messages import initialize_logging, invoke_terminal, retain_stream_response - -class Host: - def __init__(self, arguments, asynchronous): - self.machine = _Lifecycle(asynchronous, utils.is_internal_call.get()) - self.arguments = arguments - self.current = arguments - self.asynchronous = asynchronous - self.logger = arguments.get('litellm_logging_obj') - self.lifecycle_owned = self.logger is None - self.state = None - self.response = None - self.error = None - self.start = datetime.now() - self.end = None - self.streaming = False - - def setup(self): - self.logger = initialize_logging(self.arguments, self.asynchronous) - self.arguments['litellm_logging_obj'] = self.logger - self.streaming = self.logger.stream is True - - async def deployment_pre(self): - if not self.lifecycle_owned: - return - modified = await utils.async_pre_call_deployment_hook(self.current, 'anthropic_messages') - if modified is not None: - self.current = modified - self.current['litellm_logging_obj'] = self.logger - - def prepare(self): - self.state = _prepare(self.current) - - def send_sync(self): - self.response = _send_sync(self.state) - self.end = datetime.now() - - async def send(self): - self.response = await _send(self.state) - self.end = datetime.now() - - async def deployment_success(self): - if self.lifecycle_owned: - self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages) - if self.streaming: - self.response = retain_stream_response( - self.response, - (self.arguments, self.current, self.state), - self.logger, - self.start, - ) - - async def deployment_failure(self): - if self.lifecycle_owned: - await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'anthropic_messages') - - def terminal(self, action, value): - if self.streaming or not self.lifecycle_owned: - return None - return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, None, value, self.start, self.end) - - def sync_success(self): return self.terminal('sync_success', self.response) - def async_success(self): return self.terminal('async_success', self.response) - def sync_success_if_needed(self): return self.terminal('sync_success_if_needed', self.response) - def sync_failure(self): return self.terminal('sync_failure', self.error) - def async_failure(self): return self.terminal('async_failure', self.error) - def restore(self): - if not self.streaming and self.lifecycle_owned: - utils._restore_correlation_context_if_supported(self.logger) - - def advance(self, outcome, error=None): - if error is not None and self.end is None: - self.end = datetime.now() - if self.logger is None: - self.logger = self.arguments.get('litellm_logging_obj') - replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None) - if replace: - self.error = error - - def result(self): - if self.machine.complete(): - return self.response - if self.machine.failed_after_provider_response(): - _committed_failure() - raise self.error -"#; +fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static BINDINGS: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = BINDINGS.get(py) { + return Ok(module.bind(py)); + } + let module = PyModule::new(py, "_messages_bindings")?; + module.add("Lifecycle", py.get_type::())?; + module.add("invoke", wrap_pyfunction!(invoke, &module)?)?; + module.add("prepare", wrap_pyfunction!(prepare, &module)?)?; + module.add("send", wrap_pyfunction!(send, &module)?)?; + module.add("send_sync", wrap_pyfunction!(send_sync, &module)?)?; + module.add( + "committed_failure", + wrap_pyfunction!(committed_failure, &module)?, + )?; + Ok(BINDINGS.get_or_init(py, || module.unbind()).bind(py)) +} pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { crate::routes::definition::add_function(module, wrap_pyfunction!(messages, module)?)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index cfd2c1bad9a..01840ce51cd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -19,6 +19,7 @@ use pyo3::pyclass::{PyTraverseError, PyVisit}; use pyo3::sync::PyOnceLock; use pyo3::types::PyDict; +use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT}; use crate::errors::core_error_to_pyerr; use litellm_python_interop::{run_async_value, run_sync_value}; @@ -213,15 +214,13 @@ impl OcrLifecycle { fn new( py: Python<'_>, arguments: &Bound<'_, PyDict>, + logger: Option<&Bound<'_, PyAny>>, asynchronous: bool, internal_call: bool, ) -> PyResult { let request = decode_request(py, arguments)?; - let logger = arguments - .get_item("litellm_logging_obj")? - .filter(|value| !value.is_none()); let identity = |name: &str| -> PyResult> { - if let Some(logger) = &logger { + if let Some(logger) = logger { match logger.getattr(name) { Ok(value) => { if let Ok(value) = value.extract::() { @@ -313,30 +312,16 @@ fn invoke( let machine = machine.borrow(py); (machine.machine.operation(), machine.asynchronous) }; - let (method, awaiting) = match operation { - Operation::Setup => ("setup", false), - Operation::DeploymentPre => ("deployment_pre", true), - Operation::Prepare => ("prepare", false), - Operation::PreCall => ("pre_call", false), - Operation::Send if asynchronous => ("send", true), - Operation::Send => ("send_sync", false), - Operation::DeploymentSuccess => ("deployment_success", true), - Operation::DeploymentFailure => ("deployment_failure", true), - Operation::SyncSuccess => ("sync_success", false), - Operation::AsyncSuccess => ("async_success", false), - Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false), - Operation::SyncFailure => ("sync_failure", false), - Operation::AsyncFailure => ("async_failure", true), - Operation::Restore => ("restore", false), - Operation::Complete(_) => return Err(PyRuntimeError::new_err("OCR lifecycle is complete")), - }; - let value = host.getattr(py, method)?.call0(py)?; - Ok((awaiting, value)) + crate::driver::invoke(py, operation, asynchronous, true, "OCR", host) } #[pyfunction] -#[pyo3(signature = (arguments, asynchronous=false))] -fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResult> { +fn prepare( + py: Python<'_>, + arguments: Py, + logging: Py, + asynchronous: bool, +) -> PyResult> { let bag = arguments.bind(py); let request = decode_request(py, bag)?; let model = request.model.clone(); @@ -380,10 +365,6 @@ fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResul for (name, value) in &draft_headers { headers.set_item(name, value)?; } - let logging = py - .import("litellm.rust_bridge.ocr")? - .getattr("initialize_logging")? - .call1((bag, asynchronous))?; let litellm_params = PyDict::new(py); litellm_params.set_item("litellm_call_id", bag.get_item("litellm_call_id")?)?; litellm_params.set_item( @@ -396,17 +377,18 @@ fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResul update.set_item("optional_params", optional_params)?; update.set_item("litellm_params", litellm_params)?; update.set_item("custom_llm_provider", endpoint.custom_llm_provider())?; - logging.call_method("update_from_kwargs", (), Some(&update))?; + logging + .bind(py) + .call_method("update_from_kwargs", (), Some(&update))?; let additional_args = PyDict::new(py); - additional_args.set_item("complete_input_dict", &body)?; - additional_args.set_item(pyo3::intern!(py, "api_base"), endpoint.url())?; - additional_args.set_item(pyo3::intern!(py, "headers"), &headers)?; + additional_args.set_item(COMPLETE_INPUT_DICT, &body)?; + additional_args.set_item(API_BASE, endpoint.url())?; + additional_args.set_item(HEADERS, &headers)?; let pre_call = PyDict::new(py); - pre_call.set_item("input", "OCR document processing")?; - pre_call.set_item("api_key", bag.get_item("api_key")?)?; - pre_call.set_item("additional_args", additional_args)?; - let logging = logging.unbind(); + pre_call.set_item(INPUT, "OCR document processing")?; + pre_call.set_item(API_KEY, bag.get_item("api_key")?)?; + pre_call.set_item(ADDITIONAL_ARGS, additional_args)?; Py::new( py, OcrState { @@ -593,127 +575,45 @@ fn terminal_record(py: Python<'_>, state: Py) -> PyResult> { #[pyfunction] fn ocr(py: Python<'_>, arguments: Py) -> PyResult> { - driver(py)?.getattr("drive_sync")?.call1((arguments,)) + runner(py)? + .getattr("_drive_sync")? + .call1((arguments, bindings(py)?)) } #[pyfunction] fn aocr(py: Python<'_>, arguments: Py) -> PyResult> { - driver(py)?.getattr("drive_async")?.call1((arguments,)) + runner(py)? + .getattr("_drive_async")? + .call1((arguments, bindings(py)?)) } -// Compilation can re-enter through audit hooks; publish only a finished module. -fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { - static DRIVER: PyOnceLock> = PyOnceLock::new(); - if let Some(module) = DRIVER.get(py) { +fn runner(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static RUNNER: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = RUNNER.get(py) { return Ok(module.bind(py)); } - let module = crate::driver::compile( - py, - "ocr", - "from datetime import datetime -from litellm import utils -from litellm.types.utils import CallTypes -from litellm.rust_bridge.ocr import initialize_logging, invoke_terminal + let module = py.import("litellm.rust_bridge.ocr")?; + Ok(RUNNER.get_or_init(py, || module.unbind()).bind(py)) +} -class Host: - def __init__(self, arguments, asynchronous): - self.machine = _Lifecycle(arguments, asynchronous, utils.is_internal_call.get()) - self.arguments = arguments - self.current = arguments - self.asynchronous = asynchronous - self.logger = arguments.get('litellm_logging_obj') - self.state = None - self.response = None - self.error = None - self.start = datetime.now() - self.end = None - - def setup(self): - call_id, trace_id = self.machine.identity() - self.arguments['litellm_call_id'] = call_id - self.arguments['litellm_trace_id'] = trace_id - self.logger = initialize_logging(self.arguments, self.asynchronous) - self.arguments['litellm_logging_obj'] = self.logger - - async def deployment_pre(self): - modified = await utils.async_pre_call_deployment_hook(self.current, 'aocr') - if modified is not None: - self.current = modified - self.current['litellm_logging_obj'] = self.logger - call_id, trace_id = self.machine.identity() - self.current['litellm_call_id'] = call_id - self.current['litellm_trace_id'] = trace_id - - def prepare(self): - self.state = _prepare(self.current, self.asynchronous) - - def pre_call(self): - _pre_call(self.state) - - def send_sync(self): - self.response = _send_sync(self.state) - self.end = datetime.now() - - async def send(self): - self.response = _finish(await _send(self.state)) - self.end = datetime.now() - - async def deployment_success(self): - self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aocr) - - async def deployment_failure(self): - await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'aocr') - - def terminal(self, action, value): - record = _terminal_record(self.state) if self.state is not None else None - return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, record, value, self.start, self.end) - - def sync_success(self): - return self.terminal('sync_success', self.response) - - def async_success(self): - return self.terminal('async_success', self.response) - - def sync_success_if_needed(self): - return self.terminal('sync_success_if_needed', self.response) - - def sync_failure(self): - return self.terminal('sync_failure', self.error) - - def async_failure(self): - return self.terminal('async_failure', self.error) - - def restore(self): - utils._restore_correlation_context_if_supported(self.logger) - - def advance(self, outcome, error=None): - if error is not None and self.end is None: - self.end = datetime.now() - if self.logger is None: - self.logger = self.arguments.get('litellm_logging_obj') - replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None) - if replace: - self.error = error - - def result(self): - if self.machine.complete(): - return self.response - raise self.error - -", - )?; - module.add("_Lifecycle", py.get_type::())?; - module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?; - module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; - module.add("_pre_call", wrap_pyfunction!(pre_call, &module)?)?; - module.add("_send", wrap_pyfunction!(send, &module)?)?; - module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?; - module.add("_finish", wrap_pyfunction!(finish, &module)?)?; +fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static BINDINGS: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = BINDINGS.get(py) { + return Ok(module.bind(py)); + } + let module = PyModule::new(py, "_ocr_bindings")?; + module.add("Lifecycle", py.get_type::())?; + module.add("invoke", wrap_pyfunction!(invoke, &module)?)?; + module.add("prepare", wrap_pyfunction!(prepare, &module)?)?; + module.add("pre_call", wrap_pyfunction!(pre_call, &module)?)?; + module.add("send", wrap_pyfunction!(send, &module)?)?; + module.add("send_sync", wrap_pyfunction!(send_sync, &module)?)?; + module.add("finish", wrap_pyfunction!(finish, &module)?)?; module.add( - "_terminal_record", + "terminal_record", wrap_pyfunction!(terminal_record, &module)?, )?; - Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) + Ok(BINDINGS.get_or_init(py, || module.unbind()).bind(py)) } pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm/rust_bridge/_lifecycle.py b/litellm/rust_bridge/_lifecycle.py index edfd0496866..20632d99e8d 100644 --- a/litellm/rust_bridge/_lifecycle.py +++ b/litellm/rust_bridge/_lifecycle.py @@ -1,7 +1,157 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime +from enum import IntEnum +from typing import TYPE_CHECKING, Final, Literal, Protocol + +from pydantic import TypeAdapter + +if TYPE_CHECKING: + from litellm.types.utils import CallTypes + + +LOGGING_OBJECT_KEY: Final = "litellm_logging_obj" +FALLBACKS_KEY: Final = "fallbacks" +TerminalAction = Literal[ + "sync_success", + "async_success", + "sync_success_if_needed", + "sync_failure", + "async_failure", +] +_OPTIONAL_ARGUMENTS_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter( + dict[str, object] | None +) + + +class NativeOutcome(IntEnum): + SUCCESS = 0 + FAILURE = 1 + ABORT = 2 + + +class NativeLifecycle(Protocol): + def complete(self) -> bool | None: ... + + def advance(self, outcome: int, logger_available: bool, has_fallbacks: bool) -> bool: ... + + +class NativeLifecycleBindings(Protocol): + invoke: Callable[[NativeLifecycle, object], tuple[bool, object]] + + +class LifecycleHost(Protocol): + @property + def machine(self) -> NativeLifecycle: ... + + def invoke(self) -> tuple[bool, object]: ... + + def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None: ... + + def result(self) -> object: ... + + +class MutableLifecycleHost(LifecycleHost, Protocol): + arguments: dict[str, object] + current: dict[str, object] + logger: object | None + response: object + error: BaseException | None + end: datetime | None + + +def advance_host(host: MutableLifecycleHost, outcome: NativeOutcome, error: BaseException | None) -> None: + if error is not None and host.end is None: + host.end = datetime.now() + if host.logger is None: + host.logger = host.arguments.get(LOGGING_OBJECT_KEY) + replace: Final = host.machine.advance( + outcome, + host.logger is not None, + host.current.get(FALLBACKS_KEY) is not None, + ) + if replace: + host.error = error + + +def host_result(host: MutableLifecycleHost) -> object: + if host.machine.complete(): + return host.response + if host.error is None: + raise RuntimeError("native lifecycle failed without an error") + raise host.error + + +async def deployment_pre(arguments: dict[str, object], call_type: str) -> dict[str, object]: + from litellm import utils + + modified: Final = _OPTIONAL_ARGUMENTS_ADAPTER.validate_python( + await utils.async_pre_call_deployment_hook(arguments, call_type) + ) + return arguments if modified is None else modified + + +async def deployment_success(arguments: dict[str, object], response: object, call_type: CallTypes) -> object: + from litellm import utils + + updated: object = await utils.async_post_call_success_deployment_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook annotations expose an unknown return + arguments, response, call_type + ) + return updated + + +async def deployment_failure(arguments: dict[str, object], error: BaseException | None, call_type: str) -> None: + from litellm import utils + + if not isinstance(error, Exception): + raise RuntimeError("native lifecycle failure did not retain an exception") + await utils.async_post_call_failure_deployment_hook(arguments, error, call_type) + + +def restore_correlation_context(logger: object | None) -> None: + from litellm import utils + + utils._restore_correlation_context_if_supported(logger) # pyright: ignore[reportPrivateUsage] # lifecycle cleanup has no public wrapper + + +def _invoke_sync(host: LifecycleHost) -> None: + host.invoke() + + +async def _invoke_async(host: LifecycleHost) -> None: + awaiting, value = host.invoke() + if not awaiting: + return + if not isinstance(value, Awaitable): + raise TypeError("native lifecycle operation did not return an awaitable") + await value + + +def drive_sync(host: LifecycleHost) -> object: + while host.machine.complete() is None: + try: + _invoke_sync(host) + except Exception as error: + host.advance(NativeOutcome.FAILURE, error) + except BaseException as error: + host.advance(NativeOutcome.ABORT, error) + else: + host.advance(NativeOutcome.SUCCESS) + return host.result() + + +async def drive_async(host: LifecycleHost) -> object: + while host.machine.complete() is None: + try: + await _invoke_async(host) + except Exception as error: + host.advance(NativeOutcome.FAILURE, error) + except BaseException as error: + host.advance(NativeOutcome.ABORT, error) + else: + host.advance(NativeOutcome.SUCCESS) + return host.result() def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: str) -> object: @@ -11,7 +161,7 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: def invoke_terminal( - action: str, + action: TerminalAction, roots: object, logger: object, record: Mapping[str, object] | None, diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index 5a7e84b2162..3cbadcaebfe 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -16,6 +16,7 @@ import inspect import json from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from typing import ( TYPE_CHECKING, Final, @@ -32,6 +33,22 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge._lifecycle import ( + LOGGING_OBJECT_KEY, + NativeLifecycle, + NativeLifecycleBindings, + NativeOutcome, + TerminalAction, + advance_host, + deployment_failure, + deployment_pre, + deployment_success, + drive_async, + drive_sync, + host_result, + invoke_terminal, + restore_correlation_context, +) from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -579,3 +596,127 @@ def _arguments( "timeout_seconds": timeout_to_seconds(timeout), "logging_api_key": logging_api_key if logging_api_key is not None else api_key or "", } + + +class _ChatCompletionsBindings(NativeLifecycleBindings, Protocol): + Lifecycle: Callable[[dict[str, object], bool, bool], NativeLifecycle] + prepare: Callable[[dict[str, object], object], object] + send: Callable[[object], Awaitable[Mapping[str, object]]] + send_sync: Callable[[object], Mapping[str, object]] + terminal_record: Callable[[object], Mapping[str, object]] + + +class _ChatCompletionsHost: + def __init__( + self, + arguments: dict[str, object], + asynchronous: bool, + bindings: _ChatCompletionsBindings, + ) -> None: + from litellm import utils + + self.bindings: _ChatCompletionsBindings = bindings + self.machine: NativeLifecycle = bindings.Lifecycle(arguments, asynchronous, utils.is_internal_call.get()) + self.arguments: dict[str, object] = arguments + self.current: dict[str, object] = arguments + self.asynchronous: bool = asynchronous + self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY) + self.state: object | None = None + self.response: object = None + self.error: BaseException | None = None + self.start: datetime = datetime.now() + self.end: datetime | None = None + + def invoke(self) -> tuple[bool, object]: + return self.bindings.invoke(self.machine, self) + + def setup(self) -> None: + self.logger = initialize_logging(self.arguments, self.asynchronous) + self.arguments[LOGGING_OBJECT_KEY] = self.logger + + async def deployment_pre(self) -> None: + self.current = await deployment_pre(self.current, "acompletion") + self.current[LOGGING_OBJECT_KEY] = self.logger + + def prepare(self) -> None: + if self.logger is None: + raise RuntimeError("chat completions logging was not initialized") + self.state = self.bindings.prepare(self.current, self.logger) + + def send_sync(self) -> None: + model_response: Final = self.arguments["model_response"] + if not isinstance(model_response, ModelResponse): + raise TypeError("chat completions model_response must be a ModelResponse") + self.response = build_model_response(self.bindings.send_sync(self.state), model_response) + self.end = datetime.now() + + async def send(self) -> None: + model_response: Final = self.arguments["model_response"] + if not isinstance(model_response, ModelResponse): + raise TypeError("chat completions model_response must be a ModelResponse") + self.response = build_model_response(await self.bindings.send(self.state), model_response) + self.end = datetime.now() + + async def deployment_success(self) -> None: + from litellm.types.utils import CallTypes + + self.response = await deployment_success(self.current, self.response, CallTypes.acompletion) + + async def deployment_failure(self) -> None: + await deployment_failure(self.current, self.error, "acompletion") + + def terminal(self, action: TerminalAction, value: object) -> object: + if self.logger is None or self.end is None: + raise RuntimeError("chat completions terminal state was not initialized") + record: Final = self.bindings.terminal_record(self.state) if self.state is not None else None + return invoke_terminal( + action, + (self.arguments, self.current, self.state), + self.logger, + record, + value, + self.start, + self.end, + ) + + def sync_success(self) -> object: + return self.terminal("sync_success", self.response) + + def async_success(self) -> object: + return self.terminal("async_success", self.response) + + def sync_success_if_needed(self) -> object: + return self.terminal("sync_success_if_needed", self.response) + + def sync_failure(self) -> object: + return self.terminal("sync_failure", self.error) + + def async_failure(self) -> object: + return self.terminal("async_failure", self.error) + + def restore(self) -> None: + restore_correlation_context(self.logger) + + def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None: + advance_host(self, outcome, error) + + def result(self) -> object: + return host_result(self) + + +def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension + arguments: dict[str, object], bindings: _ChatCompletionsBindings +) -> ModelResponse: + result: Final = drive_sync(_ChatCompletionsHost(arguments, False, bindings)) + if not isinstance(result, ModelResponse): + raise TypeError("native chat completions driver returned an invalid response") + return result + + +async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension + arguments: dict[str, object], bindings: _ChatCompletionsBindings +) -> ModelResponse: + result: Final = await drive_async(_ChatCompletionsHost(arguments, True, bindings)) + if not isinstance(result, ModelResponse): + raise TypeError("native chat completions driver returned an invalid response") + return result diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index a115acade13..1d3d2694c2f 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -1,16 +1,31 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import Awaitable, Callable from dataclasses import dataclass from datetime import datetime, timezone from typing import Final, Protocol, cast import httpx +from litellm.rust_bridge._lifecycle import ( + LOGGING_OBJECT_KEY, + NativeLifecycle, + NativeLifecycleBindings, + NativeOutcome, + TerminalAction, + advance_host, + deployment_failure, + deployment_pre, + deployment_success, + drive_async, + drive_sync, + host_result, + invoke_terminal, + restore_correlation_context, +) from litellm.rust_bridge._lifecycle import ( initialize_logging as initialize_lifecycle_logging, ) -from litellm.rust_bridge._lifecycle import invoke_terminal from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -206,6 +221,141 @@ async def amessages( ) +class _MessagesLifecycle(NativeLifecycle, Protocol): + def failed_after_provider_response(self) -> bool: ... + + +class _MessagesBindings(NativeLifecycleBindings, Protocol): + Lifecycle: Callable[[bool, bool], _MessagesLifecycle] + prepare: Callable[[dict[str, object], object], object] + send: Callable[[object], Awaitable[AnthropicMessagesResponse]] + send_sync: Callable[[object], AnthropicMessagesResponse] + committed_failure: Callable[[], None] + + +class _MessagesHost: + def __init__(self, arguments: dict[str, object], asynchronous: bool, bindings: _MessagesBindings) -> None: + from litellm import utils + + self.bindings: _MessagesBindings = bindings + self.machine: _MessagesLifecycle = bindings.Lifecycle(asynchronous, utils.is_internal_call.get()) + self.arguments: dict[str, object] = arguments + self.current: dict[str, object] = arguments + self.asynchronous: bool = asynchronous + self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY) + self.lifecycle_owned: bool = self.logger is None + self.state: object | None = None + self.response: object = None + self.error: BaseException | None = None + self.start: datetime = datetime.now() + self.end: datetime | None = None + self.streaming: bool = False + + def invoke(self) -> tuple[bool, object]: + return self.bindings.invoke(self.machine, self) + + def setup(self) -> None: + self.logger = initialize_logging(self.arguments, self.asynchronous) + self.arguments[LOGGING_OBJECT_KEY] = self.logger + self.streaming = getattr(self.logger, "stream", False) is True + + async def deployment_pre(self) -> None: + if not self.lifecycle_owned: + return + self.current = await deployment_pre(self.current, "anthropic_messages") + self.current[LOGGING_OBJECT_KEY] = self.logger + + def prepare(self) -> None: + if self.logger is None: + raise RuntimeError("messages logging was not initialized") + self.state = self.bindings.prepare(self.current, self.logger) + + def send_sync(self) -> None: + self.response = self.bindings.send_sync(self.state) + self.end = datetime.now() + + async def send(self) -> None: + self.response = await self.bindings.send(self.state) + self.end = datetime.now() + + async def deployment_success(self) -> None: + from litellm.types.utils import CallTypes + + if self.lifecycle_owned: + self.response = await deployment_success(self.current, self.response, CallTypes.aanthropic_messages) + if not self.streaming: + return + if self.logger is None: + raise RuntimeError("messages logging was not initialized") + self.response = retain_stream_response( + cast(AnthropicMessagesResponse, self.response), + (self.arguments, self.current, self.state), + cast(_MessagesLogging, self.logger), + self.start, + ) + + async def deployment_failure(self) -> None: + if not self.lifecycle_owned: + return + await deployment_failure(self.current, self.error, "anthropic_messages") + + def terminal(self, action: TerminalAction, value: object) -> object: + if self.streaming or not self.lifecycle_owned: + return None + if self.logger is None or self.end is None: + raise RuntimeError("messages terminal state was not initialized") + return invoke_terminal( + action, + (self.arguments, self.current, self.state), + self.logger, + None, + value, + self.start, + self.end, + ) + + def sync_success(self) -> object: + return self.terminal("sync_success", self.response) + + def async_success(self) -> object: + return self.terminal("async_success", self.response) + + def sync_success_if_needed(self) -> object: + return self.terminal("sync_success_if_needed", self.response) + + def sync_failure(self) -> object: + return self.terminal("sync_failure", self.error) + + def async_failure(self) -> object: + return self.terminal("async_failure", self.error) + + def restore(self) -> None: + if not self.streaming and self.lifecycle_owned: + restore_correlation_context(self.logger) + + def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None: + advance_host(self, outcome, error) + + def result(self) -> object: + if self.machine.complete(): + return self.response + if self.machine.failed_after_provider_response(): + self.bindings.committed_failure() + return host_result(self) + + +def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension + arguments: dict[str, object], bindings: _MessagesBindings +) -> AnthropicMessagesResponse: + return cast(AnthropicMessagesResponse, drive_sync(_MessagesHost(arguments, False, bindings))) + + +async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension + arguments: dict[str, object], bindings: _MessagesBindings +) -> AnthropicMessagesResponse: + return cast(AnthropicMessagesResponse, await drive_async(_MessagesHost(arguments, True, bindings))) + + __all__ = ( "amessages", "initialize_logging", diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 904d3e510ad..e5ef4e73e3d 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -3,13 +3,28 @@ from __future__ import annotations import traceback -from collections.abc import Awaitable, Mapping +from collections.abc import Awaitable, Callable, Mapping from contextvars import copy_context from datetime import datetime from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables from uuid import uuid4 from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge._lifecycle import ( + LOGGING_OBJECT_KEY, + NativeLifecycle, + NativeLifecycleBindings, + NativeOutcome, + TerminalAction, + advance_host, + deployment_failure, + deployment_pre, + deployment_success, + drive_async, + drive_sync, + host_result, + restore_correlation_context, +) from litellm.rust_bridge.bindings import NativeBinding @@ -63,7 +78,7 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.litellm_logging import Logging, set_callbacks - supplied: Final = arguments.get("litellm_logging_obj") + supplied: Final = arguments.get(LOGGING_OBJECT_KEY) if supplied is not None: return supplied callbacks: Final = tuple( # cast-ok: callback registry accepts heterogeneous legacy callback objects @@ -185,12 +200,12 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: cb for cb in dict.fromkeys(logger.dynamic_input_callbacks or ()) if cb not in litellm.input_callback ] arguments["litellm_call_id"] = call_id - arguments["litellm_logging_obj"] = logger + arguments[LOGGING_OBJECT_KEY] = logger return logger def invoke_terminal( - action: str, + action: TerminalAction, roots: object, logger: object, record: Mapping[str, object] | None, @@ -250,3 +265,127 @@ def invoke_terminal( logging.failure_handler(exception, trace, start_time, end_time) return None return logging.async_failure_handler(exception, trace, start_time, end_time) + + +class _OcrLifecycle(NativeLifecycle, Protocol): + def identity(self) -> tuple[str, str | None]: ... + + +class _OcrBindings(NativeLifecycleBindings, Protocol): + Lifecycle: Callable[[dict[str, object], object | None, bool, bool], _OcrLifecycle] + prepare: Callable[[dict[str, object], object, bool], object] + pre_call: Callable[[object], None] + send: Callable[[object], Awaitable[dict[str, object]]] + send_sync: Callable[[object], OCRResponse] + finish: Callable[[dict[str, object]], OCRResponse] + terminal_record: Callable[[object], Mapping[str, object]] + + +class _OcrHost: + def __init__(self, arguments: dict[str, object], asynchronous: bool, bindings: _OcrBindings) -> None: + from litellm import utils + + self.bindings: _OcrBindings = bindings + self.arguments: dict[str, object] = arguments + self.current: dict[str, object] = arguments + self.asynchronous: bool = asynchronous + self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY) + self.machine: _OcrLifecycle = bindings.Lifecycle( + arguments, self.logger, asynchronous, utils.is_internal_call.get() + ) + self.state: object | None = None + self.response: object = None + self.error: BaseException | None = None + self.start: datetime = datetime.now() + self.end: datetime | None = None + + def invoke(self) -> tuple[bool, object]: + return self.bindings.invoke(self.machine, self) + + def setup(self) -> None: + call_id, trace_id = self.machine.identity() + self.arguments["litellm_call_id"] = call_id + self.arguments["litellm_trace_id"] = trace_id + self.logger = initialize_logging(self.arguments, self.asynchronous) + self.arguments[LOGGING_OBJECT_KEY] = self.logger + + async def deployment_pre(self) -> None: + self.current = await deployment_pre(self.current, "aocr") + self.current[LOGGING_OBJECT_KEY] = self.logger + call_id, trace_id = self.machine.identity() + self.current["litellm_call_id"] = call_id + self.current["litellm_trace_id"] = trace_id + + def prepare(self) -> None: + if self.logger is None: + raise RuntimeError("OCR logging was not initialized") + self.state = self.bindings.prepare(self.current, self.logger, self.asynchronous) + + def pre_call(self) -> None: + self.bindings.pre_call(self.state) + + def send_sync(self) -> None: + self.response = self.bindings.send_sync(self.state) + self.end = datetime.now() + + async def send(self) -> None: + self.response = self.bindings.finish(await self.bindings.send(self.state)) + self.end = datetime.now() + + async def deployment_success(self) -> None: + from litellm.types.utils import CallTypes + + self.response = await deployment_success(self.current, self.response, CallTypes.aocr) + + async def deployment_failure(self) -> None: + await deployment_failure(self.current, self.error, "aocr") + + def terminal(self, action: TerminalAction, value: object) -> object: + if self.logger is None or self.end is None: + raise RuntimeError("OCR terminal state was not initialized") + record: Final = self.bindings.terminal_record(self.state) if self.state is not None else None + return invoke_terminal( + action, + (self.arguments, self.current, self.state), + self.logger, + record, + value, + self.start, + self.end, + ) + + def sync_success(self) -> object: + return self.terminal("sync_success", self.response) + + def async_success(self) -> object: + return self.terminal("async_success", self.response) + + def sync_success_if_needed(self) -> object: + return self.terminal("sync_success_if_needed", self.response) + + def sync_failure(self) -> object: + return self.terminal("sync_failure", self.error) + + def async_failure(self) -> object: + return self.terminal("async_failure", self.error) + + def restore(self) -> None: + restore_correlation_context(self.logger) + + def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None: + advance_host(self, outcome, error) + + def result(self) -> object: + return host_result(self) + + +def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension + arguments: dict[str, object], bindings: _OcrBindings +) -> OCRResponse: + return cast(OCRResponse, drive_sync(_OcrHost(arguments, False, bindings))) + + +async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension + arguments: dict[str, object], bindings: _OcrBindings +) -> OCRResponse: + return cast(OCRResponse, await drive_async(_OcrHost(arguments, True, bindings))) diff --git a/tests/test_litellm/integrations/test_prometheus_caller_identity.py b/tests/test_litellm/integrations/test_prometheus_caller_identity.py index abe54cc4d99..f9204567cec 100644 --- a/tests/test_litellm/integrations/test_prometheus_caller_identity.py +++ b/tests/test_litellm/integrations/test_prometheus_caller_identity.py @@ -24,6 +24,7 @@ from litellm.types.integrations.prometheus import ( validate_prometheus_deployment_and_latency_caller_identity, ) from litellm.types.utils import StandardLoggingPayload +from tests._prometheus_helpers import clear_prometheus_registry TARGET_METRICS: Final[tuple[DEFINED_PROMETHEUS_METRICS, ...]] = cast( tuple[DEFINED_PROMETHEUS_METRICS, ...], @@ -32,14 +33,9 @@ TARGET_METRICS: Final[tuple[DEFINED_PROMETHEUS_METRICS, ...]] = cast( IDENTITY_MODES: Final = ("api_key_alias", "user_email", "both") -def _clear_prometheus_registry() -> None: - for collector in list(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] - REGISTRY.unregister(collector) - - @pytest.fixture(autouse=True) def reset_prometheus_settings(monkeypatch: pytest.MonkeyPatch): - _clear_prometheus_registry() + clear_prometheus_registry() monkeypatch.setattr(litellm, "prometheus_deployment_and_latency_caller_identity", "api_key_alias") monkeypatch.setattr(litellm, "prometheus_metrics_config", None) monkeypatch.setattr(litellm, "prometheus_exclude_metrics", None) @@ -47,7 +43,7 @@ def reset_prometheus_settings(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", []) monkeypatch.setattr(litellm, "custom_prometheus_tags", []) yield - _clear_prometheus_registry() + clear_prometheus_registry() def _expected_identity_labels(baseline: list[str], mode: str) -> list[str]: diff --git a/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py b/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py index 99eb5abb7b5..02a907fb2f4 100644 --- a/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py +++ b/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py @@ -6,12 +6,7 @@ from prometheus_client import REGISTRY import litellm from litellm.integrations.prometheus import PrometheusLogger - - -def _clear_prometheus_registry() -> None: - collectors = list(REGISTRY._collector_to_names.keys()) - for collector in collectors: - REGISTRY.unregister(collector) +from tests._prometheus_helpers import clear_prometheus_registry def _create_prometheus_logger_with_custom_labels(monkeypatch: pytest.MonkeyPatch): @@ -20,7 +15,7 @@ def _create_prometheus_logger_with_custom_labels(monkeypatch: pytest.MonkeyPatch "custom_prometheus_metadata_labels", ["metadata.department", "metadata.environment"], ) - _clear_prometheus_registry() + clear_prometheus_registry() return PrometheusLogger() diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 859cdd30c11..f0c2bb9a408 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -8,16 +8,7 @@ from litellm.types.integrations.prometheus import ( PrometheusMetricLabels, UserAPIKeyLabelNames, ) - - -def _clear_prometheus_registry() -> None: - from prometheus_client import REGISTRY - - for collector in list(REGISTRY._collector_to_names.keys()): - try: - REGISTRY.unregister(collector) - except Exception: - pass +from tests._prometheus_helpers import clear_prometheus_registry def _collected_samples(metric_name: str): @@ -661,7 +652,7 @@ async def test_success_hook_emits_api_provider_value_on_token_metric(): "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, } - _clear_prometheus_registry() + clear_prometheus_registry() try: logger = PrometheusLogger() now = datetime.datetime.now() @@ -682,7 +673,7 @@ async def test_success_hook_emits_api_provider_value_on_token_metric(): f"{[s.labels.get('api_provider') for s in samples]}" ) finally: - _clear_prometheus_registry() + clear_prometheus_registry() @pytest.mark.asyncio @@ -698,7 +689,7 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import UserAPIKeyAuth - _clear_prometheus_registry() + clear_prometheus_registry() try: logger = PrometheusLogger() await logger.async_post_call_failure_hook( @@ -713,7 +704,7 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() f"{[s.labels.get('api_provider') for s in samples]}" ) finally: - _clear_prometheus_registry() + clear_prometheus_registry() if __name__ == "__main__": diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index bf1d68c7714..cc5cb99e886 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -31,6 +31,7 @@ from litellm.types.integrations.prometheus import ( UserAPIKeyLabelNames, UserAPIKeyLabelValues, ) +from tests._prometheus_helpers import clear_prometheus_registry # --------------------------------------------------------------------------- @@ -482,16 +483,6 @@ KEY_AND_TEAM_RATE_LIMIT_METRICS = ( ) -def _clear_prometheus_registry() -> None: - from prometheus_client import REGISTRY - - for collector in list(REGISTRY._collector_to_names.keys()): - try: - REGISTRY.unregister(collector) - except Exception: - pass - - def _collected_samples(metric_name: str) -> dict[tuple[tuple[str, str], ...], float]: from prometheus_client import REGISTRY @@ -569,7 +560,7 @@ async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_head and the window consumption as ``limit - remaining`` for each key / team dimension, split by ``rate_limit_type``. """ - _clear_prometheus_registry() + clear_prometheus_registry() try: await _run_success_event( { @@ -624,7 +615,7 @@ async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_head team_tokens: 40, } finally: - _clear_prometheus_registry() + clear_prometheus_registry() @pytest.mark.asyncio @@ -634,7 +625,7 @@ async def test_should_emit_only_the_dimensions_the_limiter_enforced(): key/requests headers, so no tokens series and no team series may appear (a phantom 0 or sys.maxsize series would misreport an unlimited dimension). """ - _clear_prometheus_registry() + clear_prometheus_registry() try: await _run_success_event( { @@ -653,7 +644,7 @@ async def test_should_emit_only_the_dimensions_the_limiter_enforced(): assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {} assert _collected_samples("litellm_team_rate_limit_used_metric") == {} finally: - _clear_prometheus_registry() + clear_prometheus_registry() @pytest.mark.asyncio @@ -664,7 +655,7 @@ async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_ requests. The old allowed/used samples must disappear instead of keeping a limit that no longer exists on the scrape. """ - _clear_prometheus_registry() + clear_prometheus_registry() try: logger = PrometheusLogger() await _run_success_event( @@ -698,7 +689,7 @@ async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_ assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {team_requests: 50} assert _collected_samples("litellm_team_rate_limit_used_metric") == {team_requests: 4} finally: - _clear_prometheus_registry() + clear_prometheus_registry() @pytest.mark.asyncio @@ -715,11 +706,11 @@ async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_ async def test_should_emit_no_key_or_team_rate_limit_series_without_a_complete_int_pair( additional_headers, ): - _clear_prometheus_registry() + clear_prometheus_registry() try: await _run_success_event(additional_headers) for metric_name in KEY_AND_TEAM_RATE_LIMIT_METRICS: assert _collected_samples(metric_name) == {}, metric_name finally: - _clear_prometheus_registry() + clear_prometheus_registry() diff --git a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py b/tests/test_litellm/integrations/test_prometheus_service_tier_label.py index b2212c4ff41..2f0e5cca402 100644 --- a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py +++ b/tests/test_litellm/integrations/test_prometheus_service_tier_label.py @@ -23,6 +23,7 @@ from litellm.types.integrations.prometheus import ( UserAPIKeyLabelNames, UserAPIKeyLabelValues, ) +from tests._prometheus_helpers import clear_prometheus_registry SERVICE_TIER_METRICS = [ "litellm_llm_api_latency_metric", @@ -32,16 +33,6 @@ SERVICE_TIER_METRICS = [ ] -def _clear_prometheus_registry() -> None: - from prometheus_client import REGISTRY - - for collector in list(REGISTRY._collector_to_names.keys()): - try: - REGISTRY.unregister(collector) - except Exception: - pass - - def _collected_samples(metric_name: str): from prometheus_client import REGISTRY @@ -211,7 +202,7 @@ async def test_success_event_emits_service_tier_on_latency_and_spend_metrics(): "end_time": now, } - _clear_prometheus_registry() + clear_prometheus_registry() try: logger = PrometheusLogger() await logger.async_log_success_event(kwargs, None, now, now) @@ -229,7 +220,7 @@ async def test_success_event_emits_service_tier_on_latency_and_spend_metrics(): f"{sorted({sample.labels.get('service_tier') for sample in samples})}" ) finally: - _clear_prometheus_registry() + clear_prometheus_registry() def test_allowlist_covers_every_modeled_service_tier(): diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/test_litellm/rust_bridge/test_lifecycle.py new file mode 100644 index 00000000000..eb9ad25a680 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_lifecycle.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Final + +import pytest + +from litellm.rust_bridge._lifecycle import NativeOutcome, drive_async, drive_sync + + +class _Machine: + def __init__(self, replace: bool = False) -> None: + self.outcome: NativeOutcome | None = None + self.replace: Final = replace + + def complete(self) -> bool | None: + if self.outcome is None: + return None + return self.outcome is NativeOutcome.SUCCESS + + def advance(self, outcome: int, logger_available: bool, has_fallbacks: bool) -> bool: + self.outcome = NativeOutcome(outcome) + return self.replace + + +class _Host: + def __init__(self, invoke: Callable[[], tuple[bool, object]], replace: bool = False) -> None: + self.machine: Final = _Machine(replace) + self._invoke: Final = invoke + self.error: BaseException | None = None + + def invoke(self) -> tuple[bool, object]: + return self._invoke() + + def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None: + if self.machine.advance(outcome, True, False): + self.error = error + + def result(self) -> object: + if self.machine.complete(): + return "complete" + if self.error is None: + raise RuntimeError("missing test error") + raise self.error + + +def test_drive_sync_returns_terminal_result() -> None: + host: Final = _Host(lambda: (False, None)) + + assert drive_sync(host) == "complete" + assert host.machine.outcome is NativeOutcome.SUCCESS + + +def test_drive_sync_replaces_an_ordinary_failure() -> None: + failure: Final = ValueError("failed") + host: Final = _Host(lambda: (_ for _ in ()).throw(failure), replace=True) + + with pytest.raises(ValueError, match="failed") as raised: + drive_sync(host) + + assert raised.value is failure + assert host.machine.outcome is NativeOutcome.FAILURE + + +def test_drive_sync_classifies_base_exception_as_abort() -> None: + class Abort(BaseException): + pass + + failure: Final = Abort("aborted") + host: Final = _Host(lambda: (_ for _ in ()).throw(failure), replace=True) + + with pytest.raises(Abort, match="aborted"): + drive_sync(host) + + assert host.machine.outcome is NativeOutcome.ABORT + + +@pytest.mark.asyncio +async def test_drive_async_awaits_the_selected_operation() -> None: + completed: Final = asyncio.Event() + + async def operation() -> None: + await asyncio.sleep(0) + completed.set() + + host: Final = _Host(lambda: (True, operation())) + + assert await drive_async(host) == "complete" + assert completed.is_set() + assert host.machine.outcome is NativeOutcome.SUCCESS diff --git a/tests/test_litellm_rust/recording_server.py b/tests/test_litellm_rust/recording_server.py index 0b8dc8f563a..7ea20406cac 100644 --- a/tests/test_litellm_rust/recording_server.py +++ b/tests/test_litellm_rust/recording_server.py @@ -3,6 +3,7 @@ import copy import threading import time from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final @@ -44,8 +45,8 @@ class RecordingServer: self.responses.append(response) -@pytest.fixture -def recording_server() -> Iterator[RecordingServer]: +@contextmanager +def recording_service() -> Iterator[RecordingServer]: requests: list[RecordedRequest] = [] responses: list[ResponseSpec] = [] @@ -105,3 +106,9 @@ def recording_server() -> Iterator[RecordingServer]: if recording_server.expected_requests is not None: assert len(recording_server.requests) == recording_server.expected_requests assert recording_server.responses == [] + + +@pytest.fixture +def recording_server() -> Iterator[RecordingServer]: + with recording_service() as server: + yield server