diff --git a/litellm-rust/crates/core/src/call_lifecycle/dispatch.rs b/litellm-rust/crates/core/src/call_lifecycle/dispatch.rs new file mode 100644 index 00000000000..ed7bcd72ceb --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/dispatch.rs @@ -0,0 +1,239 @@ +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AsyncSuccessDelivery { + Skip, + Bookkeeping, + Background, + Deferred, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SuccessDispatch { + SyncBookkeeping, + SyncWorker, + Async(AsyncSuccessDelivery), +} + +pub trait SuccessFacts { + type Error; + + fn has_fallbacks(&self) -> Result; + fn callbacks_needed(&self, asynchronous: bool) -> Result; + fn defers_async_logging(&self) -> Result; +} + +pub fn success_dispatch( + asynchronous: bool, + internal: bool, + facts: &F, +) -> Result { + if !asynchronous { + return Ok(if facts.callbacks_needed(false)? { + SuccessDispatch::SyncWorker + } else { + SuccessDispatch::SyncBookkeeping + }); + } + let delivery = if internal || facts.has_fallbacks()? { + AsyncSuccessDelivery::Skip + } else if !facts.callbacks_needed(true)? { + AsyncSuccessDelivery::Bookkeeping + } else if facts.defers_async_logging()? { + AsyncSuccessDelivery::Deferred + } else { + AsyncSuccessDelivery::Background + }; + Ok(SuccessDispatch::Async(delivery)) +} + +pub fn failure_dispatch(asynchronous: bool, internal: bool, logger_available: bool) -> bool { + logger_available && !(asynchronous && internal) +} + +pub struct DeferredSuccess { + pending: bool, +} + +impl Default for DeferredSuccess { + fn default() -> Self { + Self { pending: true } + } +} + +impl DeferredSuccess { + pub fn resolve(&mut self, accepted: bool) -> bool { + std::mem::replace(&mut self.pending, false) && accepted + } +} + +#[cfg(test)] +mod tests { + use std::cell::RefCell; + + use super::*; + + struct Facts { + reads: RefCell>, + fallbacks: Result, + callbacks: Result, + deferred: Result, + } + + impl SuccessFacts for Facts { + type Error = &'static str; + + fn has_fallbacks(&self) -> Result { + self.reads.borrow_mut().push("fallbacks"); + self.fallbacks + } + + fn callbacks_needed(&self, asynchronous: bool) -> Result { + self.reads + .borrow_mut() + .push(if asynchronous { "async" } else { "sync" }); + self.callbacks + } + + fn defers_async_logging(&self) -> Result { + self.reads.borrow_mut().push("deferred"); + self.deferred + } + } + + #[test] + fn dispatch_reads_only_the_facts_needed_for_the_selected_delivery() { + use AsyncSuccessDelivery as Async; + for (asynchronous, internal, fallbacks, callbacks, deferred, expected, reads) in [ + ( + false, + false, + Err("unused"), + Ok(false), + Err("unused"), + SuccessDispatch::SyncBookkeeping, + vec!["sync"], + ), + ( + false, + true, + Err("unused"), + Ok(true), + Err("unused"), + SuccessDispatch::SyncWorker, + vec!["sync"], + ), + ( + true, + true, + Err("unused"), + Err("unused"), + Err("unused"), + SuccessDispatch::Async(Async::Skip), + vec![], + ), + ( + true, + false, + Ok(true), + Err("unused"), + Err("unused"), + SuccessDispatch::Async(Async::Skip), + vec!["fallbacks"], + ), + ( + true, + false, + Ok(false), + Ok(false), + Err("unused"), + SuccessDispatch::Async(Async::Bookkeeping), + vec!["fallbacks", "async"], + ), + ( + true, + false, + Ok(false), + Ok(true), + Ok(false), + SuccessDispatch::Async(Async::Background), + vec!["fallbacks", "async", "deferred"], + ), + ( + true, + false, + Ok(false), + Ok(true), + Ok(true), + SuccessDispatch::Async(Async::Deferred), + vec!["fallbacks", "async", "deferred"], + ), + ] { + let facts = Facts { + reads: RefCell::default(), + fallbacks, + callbacks, + deferred, + }; + assert_eq!( + success_dispatch(asynchronous, internal, &facts), + Ok(expected) + ); + assert_eq!(*facts.reads.borrow(), reads); + } + } + + #[test] + fn failed_observations_stop_dispatch_without_reading_later_facts() { + for (fallbacks, callbacks, deferred, reads) in [ + ( + Err("failure"), + Err("unused"), + Err("unused"), + vec!["fallbacks"], + ), + ( + Ok(false), + Err("failure"), + Err("unused"), + vec!["fallbacks", "async"], + ), + ( + Ok(false), + Ok(true), + Err("failure"), + vec!["fallbacks", "async", "deferred"], + ), + ] { + let facts = Facts { + reads: RefCell::default(), + fallbacks, + callbacks, + deferred, + }; + assert_eq!(success_dispatch(true, false, &facts), Err("failure")); + assert_eq!(*facts.reads.borrow(), reads); + } + } + + #[test] + fn failure_dispatch_preserves_sync_internal_calls_and_skips_async_internal_calls() { + for asynchronous in [false, true] { + for internal in [false, true] { + assert!(!failure_dispatch(asynchronous, internal, false)); + } + } + assert!(failure_dispatch(false, false, true)); + assert!(failure_dispatch(false, true, true)); + assert!(failure_dispatch(true, false, true)); + assert!(!failure_dispatch(true, true, true)); + } + + #[test] + fn deferred_success_can_be_accepted_or_rejected_only_once() { + for accepted in [false, true] { + let mut gate = DeferredSuccess::default(); + assert_eq!(gate.resolve(accepted), accepted); + assert!(!gate.resolve(true)); + assert!(!gate.resolve(false)); + } + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 49f3a6f0214..e85167ab14c 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -4,6 +4,7 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH}; use crate::Error; pub mod admission; +pub mod dispatch; pub mod host; #[cfg(test)] #[path = "../../tests/host_lifecycle.rs"] diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 359ad56c336..615827fe29d 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -39,17 +39,10 @@ pub enum Error { Http { status: u16, body: String }, #[error("upstream network error: {0}")] Network(String), - /// The provider was never reached: DNS, TCP, TLS or proxy setup failed - /// before any byte of the request went out. Nothing was billed, so a host - /// that keeps a reference implementation can serve the request itself. - /// A timeout is deliberately not this, since the provider may have received - /// and answered the request already. #[error("could not reach the provider: {0}")] Connect(String), #[error("routing error: {0}")] Routing(String), - /// The request is outside the surface this route covers in Rust. Hosts that - /// keep a reference implementation treat this as "fall back", not "fail". #[error("unsupported by the rust path: {0}")] Unsupported(&'static str), } diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 9262617156b..ed6f982a298 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -15,7 +15,7 @@ - Only core's typed, effect-free admission may return `Declined`; conversion errors and all post-admission failures are terminal - Admission cannot invoke hooks, acquire credentials, consume files/iterators, prepare requests or perform I/O - Disabled/unavailable native execution or an admission decline may select legacy once; callback exceptions never authorize fallback or replay -- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle in `src/lifecycle.rs` +- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle in `src/lifecycle/handle.rs` - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`; explicitly tagged `Await`/`Complete` preserve awaitable final values - Validate Created/Running/Suspended/Closed protocol states; core alone chooses lifecycle phases and result/error policy - Defer effectful setup/context reads/timestamps until start; unstarted-handle destruction releases inputs independently of Python `finally` diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index d25ae5a8130..430c180222a 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -1,43 +1,51 @@ -# CLAUDE.md +# Python bridge -Rules for `litellm-rust/crates/python-bridge`. +Follow `AGENTS.md` for the boundary invariants -## Responsibility +## Ownership -`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. -Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests, -maps domain errors to Python exceptions, and delegates generic conversion and -GIL handling to `litellm-python-interop`. +Core owns typed native state, effect-free admission, lifecycle sequencing, +provider preparation and I/O, normalization, and dispatch decisions -## Bridge Shape +This crate owns Python argument projection, retained Python references, public +response and exception construction, callback invocation, and host scheduling +Generic conversion and GIL utilities belong in `litellm-python-interop` -- Prefer one stable method per top-level LiteLLM route, for example - `messages(...)`, calling the matching `litellm-core` entrypoint. -- Do not add one exported PyO3 function per provider helper unless there is a - measured reason. -- Provider dispatch belongs in the `litellm-core` route module (e.g. - `litellm_core::messages`), not in this PyO3 crate. -- Python owns rollout state and fallback. Rust should return errors; Python - decides whether to raise or fall back. For a rust-only provider/route (no - Python reference), the Python side is a thin dispatch that calls Rust and - raises when the bridge is unavailable, with no fallback. -- Keep the Python interface minimal (well under 100 lines per route): it only - marshals inputs and calls Rust. Do not add per-route feature flags, and do - not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch - class under `litellm/llms///`. +Only core admission may authorize legacy fallback. Execution and conversion +failures are terminal, including authentication and connection failures -## Data Handling +## Structure -- OCR payloads can contain personal data and large base64 images. Do not log - payloads or provider responses. -- Avoid copying large payloads more than needed. The current JSON round-trip is - acceptable for the first scaffold, but future performance work should evaluate - direct PyO3 conversion before expanding Rust coverage to image-heavy paths. -- Do not expose raw Rust errors that include document contents or upstream - bodies. +Each route registers its lifecycle binding from `routes//lifecycle.rs` +Unimplemented routes use `unimplemented_lifecycle_route!`. Keep value bindings +in `value.rs` where needed, and add projection or callback modules when the +route requires them. Register route functions through `definition::add_function` +to reject duplicate exports -## Tests +`lifecycle/mod.rs` declares modules and exports the shared boundary types +`runner.rs` drives core calls, `state.rs` retains Python call state, and +`dispatch.rs` executes core-selected delivery and retains logging arguments +`handle.rs` owns the Created/Running/Suspended/Closed execution protocol +`bindings.rs` invokes Python integrations, and `preparation.rs` projects shared +preparation inputs. Runtime waiting and panic containment stay in `execution.rs` -- `cargo test --workspace` must compile this crate. -- Python tests must cover bridge disabled, bridge enabled, and module-missing - fallback behavior for every exposed route. +The Python coroutine driver lives in `litellm/rust_bridge/lifecycle.py` +Read Python state only at core-selected checkpoints. Preserve argument identity, +aliases, omitted values, and deliberate copies across suspension and callbacks + +## Data handling + +Project only consumed values at their reference read points. Keep native +provider state typed in core and preserve captured upload and request bytes +Do not log OCR documents or upstream bodies, or expose them through raw errors +Measure conversion and copy costs before optimizing large payloads + +## Verification + +`cargo test --workspace` must compile this crate. Cover disabled, enabled, and +unavailable execution, effect-free decline, terminal errors, callback delivery, +re-entry, GC, and cancellation. Validate the installed extension with a fresh +wheel and positive native execution evidence for lifecycle changes + +Keep `_native.pyi` consistent with the exported bindings, including the +Future-returning value bindings and coroutine-returning lifecycle bindings diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 701c6abb68c..e625b720044 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -6,7 +6,7 @@ pyo3::create_exception!( _native, RustBridgeDeclined, pyo3::exceptions::PyException, - "The route declined before calling the provider, so the host may retry on its own path." + "Core admission declined without effects, so the host may select its legacy path once." ); pyo3::create_exception!( @@ -28,33 +28,13 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { } } -/// Map a core error for a route whose host keeps a Python implementation. -/// -/// The distinction the host needs is whether the provider was already called. -/// Everything raised before the request goes out is safe for the host to retry -/// on its own path; anything after it is not, because the provider has already -/// done the work and billed for it. -pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::Unsupported(_) - | Error::Auth(_) - | Error::InvalidProvider(_) - | Error::InvalidRequest(_) - | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::MissingDocumentUrl - | Error::MissingApiKey { .. } - | Error::MissingAzureAiCredentials - | Error::MissingAzureDocumentIntelligenceCredentials - | Error::MissingReductoApiKey - | Error::Routing(_) - // Nothing reached the provider, so serving it on Python cannot double - // bill and is the only way the caller gets an answer at all. - | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), +pub(crate) fn execution_error_to_pyerr(error: Error) -> PyErr { + match error { Error::Http { status, body } => RustUpstreamError::new_err((status, body)), Error::Network(message) | Error::InvalidResponse(message) => { RustUpstreamError::new_err((0u16, message)) } + other => core_error_to_pyerr(other), } } @@ -63,3 +43,74 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn execution_failures_never_authorize_fallback() { + Python::initialize(); + Python::attach(|py| { + for error in [ + Error::Unsupported("unsupported option"), + Error::Auth("credential resolution failed".into()), + Error::InvalidProvider("unknown".into()), + Error::InvalidRequest("invalid request".into()), + Error::InvalidType { + expected: "object", + actual: "string", + }, + Error::MissingField("model"), + Error::MissingDocumentUrl, + Error::MissingApiKey { + provider: "anthropic", + }, + Error::MissingAzureAiCredentials, + Error::MissingAzureDocumentIntelligenceCredentials, + Error::MissingReductoApiKey, + Error::Routing("routing failed".into()), + Error::Connect("connection refused".into()), + ] { + let expected = core_error_to_pyerr(error.clone()); + let actual = execution_error_to_pyerr(error); + assert!(!actual.is_instance_of::(py)); + assert!(actual.get_type(py).is(expected.get_type(py))); + assert_eq!(actual.to_string(), expected.to_string()); + } + }); + } + + #[test] + fn upstream_failures_retain_the_status_and_message_contract() { + Python::initialize(); + Python::attach(|py| { + for (error, status, message) in [ + ( + Error::Http { + status: 429, + body: "rate limited".into(), + }, + 429u16, + "rate limited", + ), + ( + Error::Network("connection lost".into()), + 0, + "connection lost", + ), + ( + Error::InvalidResponse("invalid JSON".into()), + 0, + "invalid JSON", + ), + ] { + let actual = execution_error_to_pyerr(error); + assert!(actual.is_instance_of::(py)); + let args: (u16, String) = + actual.value(py).getattr("args").unwrap().extract().unwrap(); + assert_eq!(args, (status, message.to_string())); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs index 06b32b67fd5..586c23e0ad4 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -56,7 +56,7 @@ impl PythonLogger { pub(super) fn defer_success( &self, py: Python<'_>, - pending: Py, + pending: Py, ) -> PyResult<()> { self.object(py).setattr("_native_pending_logging", pending) } diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs b/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs new file mode 100644 index 00000000000..a5579d0f2ae --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs @@ -0,0 +1,179 @@ +use litellm_core::call_lifecycle::dispatch::{ + AsyncSuccessDelivery, DeferredSuccess, SuccessDispatch, SuccessFacts, failure_dispatch, + success_dispatch, +}; +use pyo3::exceptions::PyException; +use pyo3::prelude::*; + +use super::bindings::PythonLogger; +use super::state::PythonCallState; + +impl PythonCallState { + pub fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> { + match self.try_dispatch_success(py) { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, self.logger.as_ref().map(|logger| logger.object(py))); + Ok(()) + } + result => result, + } + } + + fn try_dispatch_success(&self, py: Python<'_>) -> PyResult<()> { + let logger = self.logger()?; + let pending = || PendingSuccess { + logger: logger.clone_ref(py), + response: self.response.as_ref().map(|value| value.clone_ref(py)), + start: self.start.clone_ref(py), + end: self.end.as_ref().map(|value| value.clone_ref(py)), + }; + let facts = PythonSuccessFacts { + py, + state: self, + logger, + }; + match success_dispatch(self.asynchronous, self.internal, &facts)? { + SuccessDispatch::SyncBookkeeping => { + logger.success_bookkeeping(py, &self.response, &self.start, &self.end, false) + } + SuccessDispatch::SyncWorker => pending().sync(py), + SuccessDispatch::Async(delivery) => { + match delivery { + AsyncSuccessDelivery::Skip => {} + AsyncSuccessDelivery::Bookkeeping => logger.success_bookkeeping( + py, + &self.response, + &self.start, + &self.end, + true, + )?, + AsyncSuccessDelivery::Background => pending().asynchronous(py)?, + AsyncSuccessDelivery::Deferred => { + logger.defer_success(py, Py::new(py, PendingLogging::new(pending()))?)? + } + } + logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } + } + } + + pub fn dispatch_failure( + &self, + py: Python<'_>, + asynchronous: bool, + ) -> PyResult>> { + if !failure_dispatch(self.asynchronous, self.internal, self.logger.is_some()) { + return Ok(None); + } + let Some(error) = &self.error else { + return Ok(None); + }; + self.logger()? + .failure(py, error, &self.start, &self.end, asynchronous) + } +} + +struct PythonSuccessFacts<'a, 'py> { + py: Python<'py>, + state: &'a PythonCallState, + logger: &'a PythonLogger, +} + +impl SuccessFacts for PythonSuccessFacts<'_, '_> { + type Error = PyErr; + + fn has_fallbacks(&self) -> PyResult { + Ok(self + .state + .kwargs + .bind(self.py) + .get_item("fallbacks")? + .is_some_and(|value| !value.is_none())) + } + + fn callbacks_needed(&self, asynchronous: bool) -> PyResult { + self.logger.callbacks_needed( + self.py, + if asynchronous { + "async_success" + } else { + "sync_success" + }, + ) + } + + fn defers_async_logging(&self) -> PyResult { + Ok(self.logger.defers_async_logging(self.py)) + } +} + +pub(super) struct PendingSuccess { + pub(super) logger: PythonLogger, + pub(super) response: Option>, + pub(super) start: Py, + pub(super) end: Option>, +} + +impl PendingSuccess { + fn sync(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .submit_success(py, &self.response, &self.start, &self.end) + } + + fn asynchronous(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .enqueue_success(py, &self.response, &self.start, &self.end) + } +} + +#[pyclass] +pub(super) struct PendingLogging { + pending: Option, + gate: DeferredSuccess, +} + +impl PendingLogging { + pub(super) fn new(pending: PendingSuccess) -> Self { + Self { + pending: Some(pending), + gate: DeferredSuccess::default(), + } + } + + fn take(&mut self, accepted: bool) -> (bool, Option) { + (self.gate.resolve(accepted), self.pending.take()) + } +} + +#[pymethods] +impl PendingLogging { + fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { + let (dispatch, pending) = slf.borrow_mut().take(success); + if let Some(pending) = pending + && dispatch + { + match pending.asynchronous(py) { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(pending.logger.object(py))); + } + result => return result, + } + } + Ok(()) + } + + fn __traverse__(&self, visit: pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { + if let Some(pending) = &self.pending { + pending.logger.traverse(&visit)?; + visit.call(&pending.response)?; + visit.call(&pending.start)?; + visit.call(&pending.end)?; + } + Ok(()) + } + + fn __clear__(slf: &Bound<'_, Self>) { + let (_, pending) = slf.borrow_mut().take(false); + drop(pending); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index 014564ae89d..7af600bb5ed 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -1,1175 +1,13 @@ -use std::sync::Arc; -use std::task::Poll; - -use futures_util::future::{AbortHandle, Abortable}; -#[cfg(test)] -use litellm_core::call_lifecycle::host::HostCallFuture; -use litellm_core::call_lifecycle::host::{ - HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, -}; -use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; -use pyo3::gc::{PyTraverseError, PyVisit}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; -use tokio::sync::Mutex; - -use crate::execution::{poll_async_value, run_async_value, run_sync_value}; - mod bindings; +mod dispatch; mod handle; mod preparation; +mod runner; +mod state; -use bindings::DeploymentHooks; pub(crate) use bindings::PythonLogger; -use handle::{Execution, ExecutionBody, ExecutionStep}; - -pub(crate) enum OperationClass { - Phase(HostPhase), - Route, -} - -pub(crate) trait PythonRoute: Send + Sync { - type Call: NativeCall + 'static; - - fn state(&self) -> &PythonCallState; - fn state_mut(&mut self) -> &mut PythonCallState; - fn classify(operation: &::Operation) -> OperationClass; - fn lifecycle_result() -> ::Result; - fn map_error(error: litellm_core::Error) -> PyErr; - fn invoke( - &mut self, - py: Python<'_>, - operation: ::Operation, - ) -> PyResult<::Result>; - fn cleanup(&mut self); - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; -} - -type NativeStep = NativeCallStep<::Operation, ::Complete>; -type NativeResult = Result, litellm_core::Error>; -type HostResumeStep = HostStep::Call>, Py>; - -struct NativeCallState { - call: C, - result: Option>, -} - -enum PendingOperation { - Native, - Host(HostPhase), -} - -struct PythonLifecycle { - route: R, - call: Option>>>, - pending: Option, - native_abort: Option, -} - -pub(crate) fn run_call( - py: Python<'_>, - call: R::Call, - route: R, -) -> PyResult> { - let asynchronous = route.state().asynchronous; - let mut lifecycle = PythonLifecycle { - route, - call: Some(Arc::new(Mutex::new(NativeCallState { call, result: None }))), - pending: None, - native_abort: None, - }; - if asynchronous { - let execution = Py::new(py, Execution::new(lifecycle))?; - return py - .import("litellm.rust_bridge.lifecycle")? - .getattr("drive")? - .call1((execution,)) - .map(Bound::unbind); - } - match lifecycle.resume(None)? { - ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Await(_) => Err(pyo3::exceptions::PyRuntimeError::new_err( - "sync call suspended", - )), - } -} - -pub(crate) fn missing_state() -> PyErr { - pyo3::exceptions::PyRuntimeError::new_err("missing native call state") -} - -impl PythonLifecycle { - fn resume_core( - &mut self, - py: Python<'_>, - result: Option::Result, HostFailure>>, - ) -> PyResult> { - let call = Arc::clone(self.call.as_ref().ok_or_else(missing_state)?); - let future = async move { - let mut call = call.lock().await; - let result = match result { - Some(Err(failure)) => call.call.interrupt(failure).await, - Some(Ok(result)) => call.call.resume(Some(result)).await, - None => call.call.resume(None).await, - }; - call.result = Some(result); - Ok(()) - }; - if self.route.state().asynchronous { - let mut future = Box::pin(future); - if let Poll::Ready(()) = poll_async_value(py, future.as_mut())? { - return Ok(HostStep::Ready(self.take_native_result()?)); - } - let (abort, registration) = AbortHandle::new_pair(); - self.native_abort = Some(abort); - self.pending = Some(PendingOperation::Native); - Ok(HostStep::Suspend( - run_async_value(py, async move { - Abortable::new(future, registration) - .await - .map_err(|_| PyRuntimeError::new_err("native execution closed"))? - })? - .unbind(), - )) - } else { - run_sync_value(py, future)?; - Ok(HostStep::Ready(self.take_native_result()?)) - } - } - - fn take_native_result(&self) -> PyResult> { - self.call - .as_ref() - .ok_or_else(missing_state)? - .try_lock() - .map_err(|_| missing_state())? - .result - .take() - .ok_or_else(missing_state)? - .map_err(R::map_error) - } - - fn host_failure( - &mut self, - py: Python<'_>, - error: PyErr, - phase: Option, - ) -> HostFailure { - let native = litellm_core::Error::InvalidRequest(error.to_string()); - let cancelled = !error.is_instance_of::(py); - let failure = if !cancelled { - HostFailure::Error(native) - } else { - HostFailure::Cancelled(native) - }; - let state = self.route.state_mut(); - if state.error.is_none() || (cancelled && phase != Some(HostPhase::DeploymentFailure)) { - state.retain_error(py, error); - } - if state.end.is_none() { - state.end = now(py).ok(); - } - failure - } - - fn drive( - &mut self, - py: Python<'_>, - result: Option>>, - ) -> PyResult { - let mut step = match (self.pending.take(), result) { - (None, None) => self.resume_core(py, None)?, - (Some(PendingOperation::Native), Some(result)) => match result { - Ok(_) => HostStep::Ready(self.take_native_result()?), - Err(error) => { - let failure = self.host_failure(py, error, None); - self.resume_core(py, Some(Err(failure)))? - } - }, - (Some(PendingOperation::Host(phase)), Some(result)) => { - let result = - result.and_then(|value| self.route.state_mut().accept(py, phase, value)); - let result = match result { - Ok(()) => Ok(R::lifecycle_result()), - Err(error) => Err(self.host_failure(py, error, Some(phase))), - }; - self.resume_core(py, Some(result))? - } - _ => return Err(missing_state()), - }; - loop { - let operation = match step { - HostStep::Suspend(awaitable) => return Ok(ExecutionStep::Await(awaitable)), - HostStep::Ready(NativeCallStep::Complete(_)) => { - return self - .route - .state_mut() - .response - .take() - .map(ExecutionStep::Return) - .ok_or_else(missing_state); - } - HostStep::Ready(NativeCallStep::Host(operation)) => operation, - }; - let phase = match R::classify(&operation) { - OperationClass::Phase(phase) => Some(phase), - OperationClass::Route => None, - }; - let result = match phase { - Some(phase) => match self.route.state_mut().invoke(py, phase) { - Ok(HostStep::Suspend(awaitable)) => { - self.pending = Some(PendingOperation::Host(phase)); - return Ok(ExecutionStep::Await(awaitable)); - } - Ok(HostStep::Ready(value)) => self - .route - .state_mut() - .accept(py, phase, value) - .map(|()| R::lifecycle_result()), - Err(error) => Err(error), - }, - None => self.route.invoke(py, operation), - }; - let result = match result { - Ok(result) => Ok(result), - Err(error) => Err(self.host_failure(py, error, phase)), - }; - step = self.resume_core(py, Some(result))?; - } - } -} - -impl ExecutionBody for PythonLifecycle { - fn resume(&mut self, result: Option>>) -> PyResult { - let result = Python::attach(|py| self.drive(py, result)); - match result { - Ok(ExecutionStep::Await(value)) => Ok(ExecutionStep::Await(value)), - result => result.map_err(|error| { - Python::attach(|py| { - self.route - .state_mut() - .error - .take() - .map(|value| PyErr::from_value(value.into_bound(py).into_any())) - .unwrap_or(error) - }) - }), - } - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.route.state().traverse(visit)?; - self.route.traverse(visit) - } -} - -impl PythonLifecycle { - fn clear(&mut self) { - if let Some(abort) = self.native_abort.take() { - abort.abort(); - } - if self.call.take().is_some() { - Python::attach(|py| self.route.state_mut().cleanup(py)); - self.route.cleanup(); - } - } -} - -impl Drop for PythonLifecycle { - fn drop(&mut self) { - self.clear(); - } -} - -pub(crate) struct PythonCallState { - pub args: Py, - pub kwargs: Py, - pub logger: Option, - pub start: Py, - pub end: Option>, - pub response: Option>, - pub error: Option>, - pub asynchronous: bool, - pub internal: bool, - pub call_type: &'static str, -} - -pub(crate) fn now(py: Python<'_>) -> PyResult> { - py.import("datetime")? - .getattr("datetime")? - .call_method0("now") - .map(Bound::unbind) -} - -impl PythonCallState { - fn invoke( - &mut self, - py: Python<'_>, - phase: HostPhase, - ) -> PyResult, Py>> { - match phase { - HostPhase::Setup => self.setup(py)?, - HostPhase::DeploymentPreCall => { - if !DeploymentHooks::needed(py)? { - return Ok(HostStep::Ready(self.kwargs.clone_ref(py).into_any())); - } - return Ok(HostStep::Suspend(DeploymentHooks::before_call( - py, - &self.kwargs, - self.call_type, - )?)); - } - HostPhase::Prepare => self.prepare(py)?, - HostPhase::DeploymentPostCall => { - if !DeploymentHooks::needed(py)? { - return self - .response - .as_ref() - .map(|value| HostStep::Ready(value.clone_ref(py))) - .ok_or_else(missing_state); - } - return Ok(HostStep::Suspend(DeploymentHooks::after_success( - py, - &self.kwargs, - &self.response, - self.call_type, - )?)); - } - HostPhase::Finalize => self.finalize(py)?, - HostPhase::Success => self.dispatch_success(py)?, - HostPhase::DeploymentFailure => { - if let Some(error) = &self.error - && DeploymentHooks::needed(py)? - { - return Ok(HostStep::Suspend(DeploymentHooks::after_failure( - py, - &self.kwargs, - error, - self.call_type, - )?)); - } - } - HostPhase::Failure | HostPhase::AsyncFailure => { - if let Some(awaitable) = - self.dispatch_failure(py, phase == HostPhase::AsyncFailure)? - { - return Ok(HostStep::Suspend(awaitable)); - } - } - HostPhase::Execute - | HostPhase::ConstructResponse - | HostPhase::MapFailure - | HostPhase::Complete => return Err(missing_state()), - } - Ok(HostStep::Ready(py.None())) - } - - fn accept(&mut self, py: Python<'_>, phase: HostPhase, value: Py) -> PyResult<()> { - match phase { - HostPhase::DeploymentPreCall => { - self.kwargs = value.into_bound(py).cast_into::()?.unbind() - } - HostPhase::DeploymentPostCall => self.response = Some(value), - _ => {} - } - Ok(()) - } - - pub fn new( - py: Python<'_>, - args: Py, - kwargs: Py, - asynchronous: bool, - call_type: &'static str, - ) -> PyResult { - Ok(Self { - args, - kwargs, - logger: None, - start: py.None(), - end: None, - response: None, - error: None, - asynchronous, - internal: false, - call_type, - }) - } - - pub fn logger(&self) -> PyResult<&PythonLogger> { - self.logger.as_ref().ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized") - }) - } - - pub fn setup(&mut self, py: Python<'_>) -> PyResult<()> { - self.start = now(py)?; - self.internal = bindings::is_internal_call(py)?; - let result = bindings::setup( - py, - self.call_type, - &self.args, - &self.kwargs, - &self.start, - self.asynchronous, - )?; - self.logger = Some(result.logger()?); - self.kwargs = result.kwargs()?; - Ok(()) - } - - pub fn prepare(&mut self, py: Python<'_>) -> PyResult<()> { - self.kwargs = preparation::prepare(py, self.kwargs.bind(py), self.logger()?)?.unbind(); - Ok(()) - } - - pub fn finalize(&self, py: Python<'_>) -> PyResult<()> { - bindings::finalize( - py, - &self.response, - self.logger()?, - &self.kwargs, - &self.start, - &self.end, - ) - } - - pub fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> { - match self.try_dispatch_success(py) { - Err(error) if error.is_instance_of::(py) => { - error.write_unraisable(py, self.logger.as_ref().map(|logger| logger.object(py))); - Ok(()) - } - result => result, - } - } - - fn try_dispatch_success(&self, py: Python<'_>) -> PyResult<()> { - let logger = self.logger()?; - let pending = || PendingSuccess { - logger: logger.clone_ref(py), - response: self.response.as_ref().map(|value| value.clone_ref(py)), - start: self.start.clone_ref(py), - end: self.end.as_ref().map(|value| value.clone_ref(py)), - }; - if !self.asynchronous { - if !logger.callbacks_needed(py, "sync_success")? { - return logger.success_bookkeeping( - py, - &self.response, - &self.start, - &self.end, - false, - ); - } - pending().sync(py) - } else { - if !self.internal - && self - .kwargs - .bind(py) - .get_item("fallbacks")? - .is_none_or(|value| value.is_none()) - { - if !logger.callbacks_needed(py, "async_success")? { - logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; - } else if logger.defers_async_logging(py) { - logger.defer_success( - py, - Py::new( - py, - PendingLogging { - pending: Some(pending()), - }, - )?, - )?; - } else { - pending().asynchronous(py)?; - } - } - logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) - } - } - - pub fn dispatch_failure( - &self, - py: Python<'_>, - asynchronous: bool, - ) -> PyResult>> { - if self.logger.is_none() || (self.asynchronous && self.internal) { - return Ok(None); - } - let Some(error) = &self.error else { - return Ok(None); - }; - self.logger()? - .failure(py, error, &self.start, &self.end, asynchronous) - } - - pub fn cleanup(&mut self, py: Python<'_>) { - if let Some(logger) = self.logger.take() - && let Err(error) = logger.restore_context(py) - { - error.write_unraisable(py, None); - } - } - - pub fn retain_error(&mut self, py: Python<'_>, error: PyErr) { - self.error = Some(error.into_value(py)); - } - - pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.args)?; - visit.call(&self.kwargs)?; - if let Some(logger) = &self.logger { - logger.traverse(visit)?; - } - visit.call(&self.start)?; - visit.call(&self.end)?; - visit.call(&self.response)?; - visit.call(&self.error) - } -} - -struct PendingSuccess { - logger: PythonLogger, - response: Option>, - start: Py, - end: Option>, -} - -impl PendingSuccess { - fn sync(&self, py: Python<'_>) -> PyResult<()> { - self.logger - .submit_success(py, &self.response, &self.start, &self.end) - } - - fn asynchronous(&self, py: Python<'_>) -> PyResult<()> { - self.logger - .enqueue_success(py, &self.response, &self.start, &self.end) - } -} - -#[pyclass] -struct PendingLogging { - pending: Option, -} - -#[pymethods] -impl PendingLogging { - fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { - let pending = slf.borrow_mut().pending.take(); - if let Some(pending) = pending - && success - { - match pending.asynchronous(py) { - Err(error) if error.is_instance_of::(py) => { - error.write_unraisable(py, Some(pending.logger.object(py))); - } - result => return result, - } - } - Ok(()) - } - - fn __traverse__(&self, visit: pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { - if let Some(pending) = &self.pending { - pending.logger.traverse(&visit)?; - visit.call(&pending.response)?; - visit.call(&pending.start)?; - visit.call(&pending.end)?; - } - Ok(()) - } - - fn __clear__(slf: &Bound<'_, Self>) { - let pending = slf.borrow_mut().pending.take(); - drop(pending); - } -} +pub(crate) use runner::{OperationClass, PythonRoute, run_call}; +pub(crate) use state::{PythonCallState, missing_state, now}; #[cfg(test)] -mod tests { - use super::*; - use pyo3::types::PyDict; - use std::sync::Mutex; - - static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); - - fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> { - py.import("litellm.litellm_core_utils.logging_worker")? - .setattr("GLOBAL_LOGGING_WORKER", worker) - } - - struct RetainingHost { - retained: Option>, - } - - impl ExecutionBody for RetainingHost { - fn resume(&mut self, _: Option>>) -> PyResult { - Python::attach(|py| Ok(ExecutionStep::Return(py.None()))) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.retained) - } - } - - #[pyfunction] - fn retaining_coroutine(py: Python<'_>, retained: Py) -> PyResult> { - Py::new( - py, - Execution::new(RetainingHost { - retained: Some(retained), - }), - ) - } - - struct AwaitBody(Option>); - - impl ExecutionBody for AwaitBody { - fn resume(&mut self, result: Option>>) -> PyResult { - match self.0.take() { - Some(awaitable) => Ok(ExecutionStep::Await(awaitable)), - None => result - .expect("selected await completed") - .map(ExecutionStep::Return), - } - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.0) - } - } - - #[pyfunction] - fn await_execution(awaitable: Py) -> Execution { - Execution::new(AwaitBody(Some(awaitable))) - } - - struct CallingBody(Py); - - impl ExecutionBody for CallingBody { - fn resume(&mut self, _: Option>>) -> PyResult { - Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return)) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.0) - } - } - - #[pyfunction] - fn calling_execution(callback: Py) -> Execution { - Execution::new(CallingBody(callback)) - } - - struct SyntheticCall(bool); - - impl NativeCall for SyntheticCall { - type Operation = (); - type Result = (); - type Complete = (); - - fn resume( - &mut self, - result: Option, - ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { - Box::pin(async move { - match (self.0, result) { - (false, None) => { - self.0 = true; - Ok(NativeCallStep::Host(())) - } - (true, Some(())) => Ok(NativeCallStep::Complete(())), - _ => Err(litellm_core::Error::InvalidRequest( - "invalid synthetic lifecycle state".into(), - )), - } - }) - } - - fn interrupt( - &mut self, - _: HostFailure, - ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { - Box::pin(async { Ok(NativeCallStep::Complete(())) }) - } - } - - struct SyntheticRoute(PythonCallState); - - impl PythonRoute for SyntheticRoute { - type Call = SyntheticCall; - - fn state(&self) -> &PythonCallState { - &self.0 - } - - fn state_mut(&mut self) -> &mut PythonCallState { - &mut self.0 - } - - fn classify(_: &()) -> OperationClass { - OperationClass::Route - } - - fn lifecycle_result() {} - - fn map_error(error: litellm_core::Error) -> PyErr { - crate::errors::core_error_to_pyerr(error) - } - - fn invoke(&mut self, py: Python<'_>, _: ()) -> PyResult<()> { - self.0.response = Some( - pyo3::types::PyString::new(py, "shared lifecycle") - .into_any() - .unbind(), - ); - Ok(()) - } - - fn cleanup(&mut self) {} - - fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { - Ok(()) - } - } - - #[test] - fn shared_runner_executes_a_non_ocr_adapter() { - Python::initialize(); - Python::attach(|py| { - let route = SyntheticRoute( - PythonCallState::new( - py, - PyTuple::empty(py).unbind(), - PyDict::new(py).unbind(), - false, - "synthetic", - ) - .unwrap(), - ); - let value: String = run_call(py, SyntheticCall(false), route) - .unwrap() - .extract(py) - .unwrap(); - assert_eq!(value, "shared lifecycle"); - }); - } - - #[test] - fn ready_native_lifecycle_completes_without_scheduling() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - Python::initialize(); - Python::attach(|py| { - let source = std::ffi::CString::new(include_str!( - "../../../../../litellm/rust_bridge/lifecycle.py" - )) - .unwrap(); - PyModule::from_code( - py, - &source, - pyo3::ffi::c_str!("lifecycle.py"), - pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), - ) - .unwrap(); - let route = SyntheticRoute( - PythonCallState::new( - py, - PyTuple::empty(py).unbind(), - PyDict::new(py).unbind(), - true, - "synthetic", - ) - .unwrap(), - ); - let coroutine = run_call(py, SyntheticCall(false), route).unwrap(); - let completed = coroutine - .call_method1(py, "send", (py.None(),)) - .unwrap_err(); - assert!(completed.is_instance_of::(py)); - assert_eq!( - completed - .value(py) - .getattr("value") - .unwrap() - .extract::() - .unwrap(), - "shared lifecycle", - ); - }); - } - - #[test] - fn python_driver_preserves_inline_await_and_native_ownership() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - Python::initialize(); - Python::attach(|py| { - py.import("asyncio").unwrap(); - let source = std::ffi::CString::new(include_str!( - "../../../../../litellm/rust_bridge/lifecycle.py" - )) - .unwrap(); - let module = PyModule::from_code( - py, - &source, - pyo3::ffi::c_str!("lifecycle.py"), - pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), - ) - .unwrap(); - let locals = PyDict::new(py); - locals - .set_item("drive", module.getattr("drive").unwrap()) - .unwrap(); - locals - .set_item( - "await_execution", - wrap_pyfunction!(await_execution, py).unwrap(), - ) - .unwrap(); - locals - .set_item( - "calling_execution", - wrap_pyfunction!(calling_execution, py).unwrap(), - ) - .unwrap(); - let probe = std::ffi::CString::new(include_str!("../../tests/lifecycle.py")).unwrap(); - py.run(&probe, Some(&locals), Some(&locals)).unwrap(); - }); - } - - struct ErrorBody(PythonCallState); - - impl ExecutionBody for ErrorBody { - fn resume(&mut self, _: Option>>) -> PyResult { - Python::attach(|py| { - Err(PyErr::from_value( - self.0.error.take().unwrap().into_bound(py).into_any(), - )) - }) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.0.traverse(visit) - } - } - - #[pyfunction] - fn error_execution(py: Python<'_>, error: Bound<'_, PyBaseException>) -> Execution { - let mut state = PythonCallState::new( - py, - PyTuple::empty(py).unbind(), - PyDict::new(py).unbind(), - true, - "test", - ) - .unwrap(); - state.retain_error(py, PyErr::from_value(error.into_any())); - Execution::new(ErrorBody(state)) - } - - #[test] - fn retained_exception_frames_and_duplicate_argument_edges_are_collectable() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - locals - .set_item( - "error_execution", - wrap_pyfunction!(error_execution, py).unwrap(), - ) - .unwrap(); - py.run( - pyo3::ffi::c_str!( - r#" -import gc -import weakref - -class Retained: - pass - -def cycle(): - retained = Retained() - try: - raise ValueError('retained traceback') - except ValueError as error: - retained.owner = error_execution(error) - return weakref.ref(retained) - -reference = cycle() -gc.collect() -assert reference() is None -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - }); - } - - fn state( - py: Python<'_>, - logger: Py, - response: Py, - asynchronous: bool, - ) -> PythonCallState { - PythonCallState { - args: PyTuple::empty(py).unbind(), - kwargs: PyDict::new(py).unbind(), - logger: Some(logger.extract(py).unwrap()), - start: py.None(), - end: Some(py.None()), - response: Some(response), - error: None, - asynchronous, - internal: false, - call_type: "test", - } - } - - #[test] - fn success_dispatch_reports_ordinary_failures_without_replacing_response() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!( - r#" -import sys - -response = object() -failure = ValueError('terminal diagnostic') -diagnostics = [] -old_hook = sys.unraisablehook -sys.unraisablehook = lambda event: diagnostics.append(event.exc_value) - -class Logger: - def handle_sync_success_callbacks_for_async_calls(self, *args): - raise failure - -logger = Logger() -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - let response = locals.get_item("response").unwrap().unwrap().unbind(); - let mut lifecycle_state = state( - py, - locals.get_item("logger").unwrap().unwrap().unbind(), - response.clone_ref(py), - true, - ); - lifecycle_state.internal = true; - lifecycle_state.dispatch_success(py).unwrap(); - assert!(lifecycle_state.response.as_ref().unwrap().is(&response)); - py.run( - pyo3::ffi::c_str!( - r#" -assert diagnostics == [failure] -sys.unraisablehook = old_hook -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - }); - } - - #[test] - fn retained_failure_preserves_exception_identity() { - Python::initialize(); - Python::attach(|py| { - let logger = PyDict::new(py).into_any().unbind(); - let response = py.None(); - let failure = pyo3::exceptions::PyValueError::new_err("identity"); - let failure_value = failure.value(py).clone().unbind(); - let mut lifecycle_state = state(py, logger, response, false); - lifecycle_state.retain_error(py, failure); - let retained = lifecycle_state.error.take().unwrap(); - assert!(retained.is(&failure_value)); - }); - } - - #[test] - fn deferred_release_uses_release_context_and_allows_reentry_once() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!( - r#" -import sys -import types -from contextvars import ContextVar - -litellm = types.ModuleType('litellm') -core_utils = types.ModuleType('litellm.litellm_core_utils') -logging_worker = types.ModuleType('litellm.litellm_core_utils.logging_worker') -litellm.litellm_core_utils = core_utils -core_utils.logging_worker = logging_worker -sys.modules['litellm'] = litellm -sys.modules['litellm.litellm_core_utils'] = core_utils -sys.modules['litellm.litellm_core_utils.logging_worker'] = logging_worker - -marker = ContextVar('marker', default='unset') -observed = [] - -class Coroutine: - def close(self): - observed.append('closed') - -class Worker: - def ensure_initialized_and_enqueue(self, coroutine): - observed.append(marker.get()) - pending.release(True) - coroutine.close() - -class Logger: - def async_success_handler(self, *args): - observed.append('created') - return Coroutine() - -worker = Worker() -logger = Logger() -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - install_logging_worker(py, &locals.get_item("worker").unwrap().unwrap()).unwrap(); - let pending = Py::new( - py, - PendingLogging { - pending: Some(PendingSuccess { - logger: locals - .get_item("logger") - .unwrap() - .unwrap() - .extract() - .unwrap(), - response: Some(py.None()), - start: py.None(), - end: Some(py.None()), - }), - }, - ) - .unwrap(); - locals.set_item("pending", &pending).unwrap(); - py.run( - pyo3::ffi::c_str!( - r#" -marker.set('release') -pending.release(True) -pending.release(True) -assert observed == ['created', 'release', 'closed'] -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - }); - } - - #[test] - fn deferred_logging_collects_cycles_through_typed_logger() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!("class Logger: pass\nlogger = Logger()"), - Some(&locals), - Some(&locals), - ) - .unwrap(); - let pending = Py::new( - py, - PendingLogging { - pending: Some(PendingSuccess { - logger: locals - .get_item("logger") - .unwrap() - .unwrap() - .extract() - .unwrap(), - response: None, - start: py.None(), - end: None, - }), - }, - ) - .unwrap(); - locals.set_item("pending", pending).unwrap(); - py.run( - pyo3::ffi::c_str!( - r#" -import gc -import weakref -logger.pending = pending -reference = weakref.ref(logger) -del logger, pending -gc.collect() -assert reference() is None -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - }); - } - - #[test] - fn coroutine_collects_cycles_retained_by_bridge_host() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - locals - .set_item( - "retaining_coroutine", - wrap_pyfunction!(retaining_coroutine, py).unwrap(), - ) - .unwrap(); - py.run( - pyo3::ffi::c_str!( - r#" -import gc -import weakref - -class Retained: - pass - -def cycle(): - retained = Retained() - coroutine = retaining_coroutine(retained) - retained.coroutine = coroutine - return weakref.ref(retained) - -retained_ref = cycle() -gc.collect() -assert retained_ref() is None -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - }); - } -} +mod tests; diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs b/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs new file mode 100644 index 00000000000..26760f743e5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/runner.rs @@ -0,0 +1,269 @@ +use std::sync::Arc; +use std::task::Poll; + +use futures_util::future::{AbortHandle, Abortable}; +use litellm_core::call_lifecycle::host::{ + HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, +}; +use pyo3::exceptions::{PyException, PyRuntimeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use tokio::sync::Mutex; + +use super::handle::{Execution, ExecutionBody, ExecutionStep}; +use super::state::{PythonCallState, missing_state, now}; +use crate::execution::{poll_async_value, run_async_value, run_sync_value}; + +pub(crate) enum OperationClass { + Phase(HostPhase), + Route, +} + +pub(crate) trait PythonRoute: Send + Sync { + type Call: NativeCall + 'static; + + fn state(&self) -> &PythonCallState; + fn state_mut(&mut self) -> &mut PythonCallState; + fn classify(operation: &::Operation) -> OperationClass; + fn lifecycle_result() -> ::Result; + fn map_error(error: litellm_core::Error) -> PyErr; + fn invoke( + &mut self, + py: Python<'_>, + operation: ::Operation, + ) -> PyResult<::Result>; + fn cleanup(&mut self); + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; +} + +type NativeStep = NativeCallStep<::Operation, ::Complete>; +type NativeResult = Result, litellm_core::Error>; +type HostResumeStep = HostStep::Call>, Py>; + +struct NativeCallState { + call: C, + result: Option>, +} + +enum PendingOperation { + Native, + Host(HostPhase), +} + +struct PythonLifecycle { + route: R, + call: Option>>>, + pending: Option, + native_abort: Option, +} + +pub(crate) fn run_call( + py: Python<'_>, + call: R::Call, + route: R, +) -> PyResult> { + let asynchronous = route.state().asynchronous; + let mut lifecycle = PythonLifecycle { + route, + call: Some(Arc::new(Mutex::new(NativeCallState { call, result: None }))), + pending: None, + native_abort: None, + }; + if asynchronous { + let execution = Py::new(py, Execution::new(lifecycle))?; + return py + .import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) + .map(Bound::unbind); + } + match lifecycle.resume(None)? { + ExecutionStep::Return(value) => Ok(value), + ExecutionStep::Await(_) => Err(pyo3::exceptions::PyRuntimeError::new_err( + "sync call suspended", + )), + } +} + +impl PythonLifecycle { + fn resume_core( + &mut self, + py: Python<'_>, + result: Option::Result, HostFailure>>, + ) -> PyResult> { + let call = Arc::clone(self.call.as_ref().ok_or_else(missing_state)?); + let future = async move { + let mut call = call.lock().await; + let result = match result { + Some(Err(failure)) => call.call.interrupt(failure).await, + Some(Ok(result)) => call.call.resume(Some(result)).await, + None => call.call.resume(None).await, + }; + call.result = Some(result); + Ok(()) + }; + if self.route.state().asynchronous { + let mut future = Box::pin(future); + if let Poll::Ready(()) = poll_async_value(py, future.as_mut())? { + return Ok(HostStep::Ready(self.take_native_result()?)); + } + let (abort, registration) = AbortHandle::new_pair(); + self.native_abort = Some(abort); + self.pending = Some(PendingOperation::Native); + Ok(HostStep::Suspend( + run_async_value(py, async move { + Abortable::new(future, registration) + .await + .map_err(|_| PyRuntimeError::new_err("native execution closed"))? + })? + .unbind(), + )) + } else { + run_sync_value(py, future)?; + Ok(HostStep::Ready(self.take_native_result()?)) + } + } + + fn take_native_result(&self) -> PyResult> { + self.call + .as_ref() + .ok_or_else(missing_state)? + .try_lock() + .map_err(|_| missing_state())? + .result + .take() + .ok_or_else(missing_state)? + .map_err(R::map_error) + } + + fn host_failure( + &mut self, + py: Python<'_>, + error: PyErr, + phase: Option, + ) -> HostFailure { + let native = litellm_core::Error::InvalidRequest(error.to_string()); + let cancelled = !error.is_instance_of::(py); + let failure = if !cancelled { + HostFailure::Error(native) + } else { + HostFailure::Cancelled(native) + }; + let state = self.route.state_mut(); + if state.error.is_none() || (cancelled && phase != Some(HostPhase::DeploymentFailure)) { + state.retain_error(py, error); + } + if state.end.is_none() { + state.end = now(py).ok(); + } + failure + } + + fn drive( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + let mut step = match (self.pending.take(), result) { + (None, None) => self.resume_core(py, None)?, + (Some(PendingOperation::Native), Some(result)) => match result { + Ok(_) => HostStep::Ready(self.take_native_result()?), + Err(error) => { + let failure = self.host_failure(py, error, None); + self.resume_core(py, Some(Err(failure)))? + } + }, + (Some(PendingOperation::Host(phase)), Some(result)) => { + let result = + result.and_then(|value| self.route.state_mut().accept(py, phase, value)); + let result = match result { + Ok(()) => Ok(R::lifecycle_result()), + Err(error) => Err(self.host_failure(py, error, Some(phase))), + }; + self.resume_core(py, Some(result))? + } + _ => return Err(missing_state()), + }; + loop { + let operation = match step { + HostStep::Suspend(awaitable) => return Ok(ExecutionStep::Await(awaitable)), + HostStep::Ready(NativeCallStep::Complete(_)) => { + return self + .route + .state_mut() + .response + .take() + .map(ExecutionStep::Return) + .ok_or_else(missing_state); + } + HostStep::Ready(NativeCallStep::Host(operation)) => operation, + }; + let phase = match R::classify(&operation) { + OperationClass::Phase(phase) => Some(phase), + OperationClass::Route => None, + }; + let result = match phase { + Some(phase) => match self.route.state_mut().invoke(py, phase) { + Ok(HostStep::Suspend(awaitable)) => { + self.pending = Some(PendingOperation::Host(phase)); + return Ok(ExecutionStep::Await(awaitable)); + } + Ok(HostStep::Ready(value)) => self + .route + .state_mut() + .accept(py, phase, value) + .map(|()| R::lifecycle_result()), + Err(error) => Err(error), + }, + None => self.route.invoke(py, operation), + }; + let result = match result { + Ok(result) => Ok(result), + Err(error) => Err(self.host_failure(py, error, phase)), + }; + step = self.resume_core(py, Some(result))?; + } + } +} + +impl ExecutionBody for PythonLifecycle { + fn resume(&mut self, result: Option>>) -> PyResult { + let result = Python::attach(|py| self.drive(py, result)); + match result { + Ok(ExecutionStep::Await(value)) => Ok(ExecutionStep::Await(value)), + result => result.map_err(|error| { + Python::attach(|py| { + self.route + .state_mut() + .error + .take() + .map(|value| PyErr::from_value(value.into_bound(py).into_any())) + .unwrap_or(error) + }) + }), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.route.state().traverse(visit)?; + self.route.traverse(visit) + } +} + +impl PythonLifecycle { + fn clear(&mut self) { + if let Some(abort) = self.native_abort.take() { + abort.abort(); + } + if self.call.take().is_some() { + Python::attach(|py| self.route.state_mut().cleanup(py)); + self.route.cleanup(); + } + } +} + +impl Drop for PythonLifecycle { + fn drop(&mut self) { + self.clear(); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/state.rs b/litellm-rust/crates/python-bridge/src/lifecycle/state.rs new file mode 100644 index 00000000000..3c163123962 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/state.rs @@ -0,0 +1,195 @@ +use litellm_core::call_lifecycle::host::{HostPhase, HostStep}; +use pyo3::exceptions::PyBaseException; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use super::bindings::{self, DeploymentHooks, PythonLogger}; +use super::preparation; + +pub(crate) fn missing_state() -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err("missing native call state") +} + +pub(crate) struct PythonCallState { + pub args: Py, + pub kwargs: Py, + pub logger: Option, + pub start: Py, + pub end: Option>, + pub response: Option>, + pub error: Option>, + pub asynchronous: bool, + pub internal: bool, + pub call_type: &'static str, +} + +pub(crate) fn now(py: Python<'_>) -> PyResult> { + py.import("datetime")? + .getattr("datetime")? + .call_method0("now") + .map(Bound::unbind) +} + +impl PythonCallState { + pub(super) fn invoke( + &mut self, + py: Python<'_>, + phase: HostPhase, + ) -> PyResult, Py>> { + match phase { + HostPhase::Setup => self.setup(py)?, + HostPhase::DeploymentPreCall => { + if !DeploymentHooks::needed(py)? { + return Ok(HostStep::Ready(self.kwargs.clone_ref(py).into_any())); + } + return Ok(HostStep::Suspend(DeploymentHooks::before_call( + py, + &self.kwargs, + self.call_type, + )?)); + } + HostPhase::Prepare => self.prepare(py)?, + HostPhase::DeploymentPostCall => { + if !DeploymentHooks::needed(py)? { + return self + .response + .as_ref() + .map(|value| HostStep::Ready(value.clone_ref(py))) + .ok_or_else(missing_state); + } + return Ok(HostStep::Suspend(DeploymentHooks::after_success( + py, + &self.kwargs, + &self.response, + self.call_type, + )?)); + } + HostPhase::Finalize => self.finalize(py)?, + HostPhase::Success => self.dispatch_success(py)?, + HostPhase::DeploymentFailure => { + if let Some(error) = &self.error + && DeploymentHooks::needed(py)? + { + return Ok(HostStep::Suspend(DeploymentHooks::after_failure( + py, + &self.kwargs, + error, + self.call_type, + )?)); + } + } + HostPhase::Failure | HostPhase::AsyncFailure => { + if let Some(awaitable) = + self.dispatch_failure(py, phase == HostPhase::AsyncFailure)? + { + return Ok(HostStep::Suspend(awaitable)); + } + } + HostPhase::Execute + | HostPhase::ConstructResponse + | HostPhase::MapFailure + | HostPhase::Complete => return Err(missing_state()), + } + Ok(HostStep::Ready(py.None())) + } + + pub(super) fn accept( + &mut self, + py: Python<'_>, + phase: HostPhase, + value: Py, + ) -> PyResult<()> { + match phase { + HostPhase::DeploymentPreCall => { + self.kwargs = value.into_bound(py).cast_into::()?.unbind() + } + HostPhase::DeploymentPostCall => self.response = Some(value), + _ => {} + } + Ok(()) + } + + pub fn new( + py: Python<'_>, + args: Py, + kwargs: Py, + asynchronous: bool, + call_type: &'static str, + ) -> PyResult { + Ok(Self { + args, + kwargs, + logger: None, + start: py.None(), + end: None, + response: None, + error: None, + asynchronous, + internal: false, + call_type, + }) + } + + pub fn logger(&self) -> PyResult<&PythonLogger> { + self.logger.as_ref().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized") + }) + } + + pub fn setup(&mut self, py: Python<'_>) -> PyResult<()> { + self.start = now(py)?; + self.internal = bindings::is_internal_call(py)?; + let result = bindings::setup( + py, + self.call_type, + &self.args, + &self.kwargs, + &self.start, + self.asynchronous, + )?; + self.logger = Some(result.logger()?); + self.kwargs = result.kwargs()?; + Ok(()) + } + + pub fn prepare(&mut self, py: Python<'_>) -> PyResult<()> { + self.kwargs = preparation::prepare(py, self.kwargs.bind(py), self.logger()?)?.unbind(); + Ok(()) + } + + pub fn finalize(&self, py: Python<'_>) -> PyResult<()> { + bindings::finalize( + py, + &self.response, + self.logger()?, + &self.kwargs, + &self.start, + &self.end, + ) + } + + pub fn cleanup(&mut self, py: Python<'_>) { + if let Some(logger) = self.logger.take() + && let Err(error) = logger.restore_context(py) + { + error.write_unraisable(py, None); + } + } + + pub fn retain_error(&mut self, py: Python<'_>, error: PyErr) { + self.error = Some(error.into_value(py)); + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.args)?; + visit.call(&self.kwargs)?; + if let Some(logger) = &self.logger { + logger.traverse(visit)?; + } + visit.call(&self.start)?; + visit.call(&self.end)?; + visit.call(&self.response)?; + visit.call(&self.error) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs b/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs new file mode 100644 index 00000000000..0b0ee88896d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/tests.rs @@ -0,0 +1,585 @@ +use litellm_core::call_lifecycle::host::{ + HostCall as NativeCall, HostCallFuture, HostCallStep as NativeCallStep, HostFailure, +}; +use pyo3::exceptions::PyBaseException; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +use super::dispatch::{PendingLogging, PendingSuccess}; +use super::handle::{Execution, ExecutionBody, ExecutionStep}; +use super::*; +use pyo3::types::PyDict; +use std::sync::Mutex; + +static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); + +fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> { + py.import("litellm.litellm_core_utils.logging_worker")? + .setattr("GLOBAL_LOGGING_WORKER", worker) +} + +struct RetainingHost { + retained: Option>, +} + +impl ExecutionBody for RetainingHost { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| Ok(ExecutionStep::Return(py.None()))) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.retained) + } +} + +#[pyfunction] +fn retaining_coroutine(py: Python<'_>, retained: Py) -> PyResult> { + Py::new( + py, + Execution::new(RetainingHost { + retained: Some(retained), + }), + ) +} + +struct AwaitBody(Option>); + +impl ExecutionBody for AwaitBody { + fn resume(&mut self, result: Option>>) -> PyResult { + match self.0.take() { + Some(awaitable) => Ok(ExecutionStep::Await(awaitable)), + None => result + .expect("selected await completed") + .map(ExecutionStep::Return), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } +} + +#[pyfunction] +fn await_execution(awaitable: Py) -> Execution { + Execution::new(AwaitBody(Some(awaitable))) +} + +struct CallingBody(Py); + +impl ExecutionBody for CallingBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } +} + +#[pyfunction] +fn calling_execution(callback: Py) -> Execution { + Execution::new(CallingBody(callback)) +} + +struct SyntheticCall(bool); + +impl NativeCall for SyntheticCall { + type Operation = (); + type Result = (); + type Complete = (); + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(async move { + match (self.0, result) { + (false, None) => { + self.0 = true; + Ok(NativeCallStep::Host(())) + } + (true, Some(())) => Ok(NativeCallStep::Complete(())), + _ => Err(litellm_core::Error::InvalidRequest( + "invalid synthetic lifecycle state".into(), + )), + } + }) + } + + fn interrupt(&mut self, _: HostFailure) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(async { Ok(NativeCallStep::Complete(())) }) + } +} + +struct SyntheticRoute(PythonCallState); + +impl PythonRoute for SyntheticRoute { + type Call = SyntheticCall; + + fn state(&self) -> &PythonCallState { + &self.0 + } + + fn state_mut(&mut self) -> &mut PythonCallState { + &mut self.0 + } + + fn classify(_: &()) -> OperationClass { + OperationClass::Route + } + + fn lifecycle_result() {} + + fn map_error(error: litellm_core::Error) -> PyErr { + crate::errors::core_error_to_pyerr(error) + } + + fn invoke(&mut self, py: Python<'_>, _: ()) -> PyResult<()> { + self.0.response = Some( + pyo3::types::PyString::new(py, "shared lifecycle") + .into_any() + .unbind(), + ); + Ok(()) + } + + fn cleanup(&mut self) {} + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } +} + +#[test] +fn shared_runner_executes_a_non_ocr_adapter() { + Python::initialize(); + Python::attach(|py| { + let route = SyntheticRoute( + PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + false, + "synthetic", + ) + .unwrap(), + ); + let value: String = run_call(py, SyntheticCall(false), route) + .unwrap() + .extract(py) + .unwrap(); + assert_eq!(value, "shared lifecycle"); + }); +} + +#[test] +fn ready_native_lifecycle_completes_without_scheduling() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap(); + let route = SyntheticRoute( + PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + true, + "synthetic", + ) + .unwrap(), + ); + let coroutine = run_call(py, SyntheticCall(false), route).unwrap(); + let completed = coroutine + .call_method1(py, "send", (py.None(),)) + .unwrap_err(); + assert!(completed.is_instance_of::(py)); + assert_eq!( + completed + .value(py) + .getattr("value") + .unwrap() + .extract::() + .unwrap(), + "shared lifecycle", + ); + }); +} + +#[test] +fn python_driver_preserves_inline_await_and_native_ownership() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + py.import("asyncio").unwrap(); + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + let module = PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap(); + let locals = PyDict::new(py); + locals + .set_item("drive", module.getattr("drive").unwrap()) + .unwrap(); + locals + .set_item( + "await_execution", + wrap_pyfunction!(await_execution, py).unwrap(), + ) + .unwrap(); + locals + .set_item( + "calling_execution", + wrap_pyfunction!(calling_execution, py).unwrap(), + ) + .unwrap(); + let probe = std::ffi::CString::new(include_str!("../../tests/lifecycle.py")).unwrap(); + py.run(&probe, Some(&locals), Some(&locals)).unwrap(); + }); +} + +struct ErrorBody(PythonCallState); + +impl ExecutionBody for ErrorBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| { + Err(PyErr::from_value( + self.0.error.take().unwrap().into_bound(py).into_any(), + )) + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.0.traverse(visit) + } +} + +#[pyfunction] +fn error_execution(py: Python<'_>, error: Bound<'_, PyBaseException>) -> Execution { + let mut state = PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + true, + "test", + ) + .unwrap(); + state.retain_error(py, PyErr::from_value(error.into_any())); + Execution::new(ErrorBody(state)) +} + +#[test] +fn retained_exception_frames_and_duplicate_argument_edges_are_collectable() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "error_execution", + wrap_pyfunction!(error_execution, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained: + pass + +def cycle(): + retained = Retained() + try: + raise ValueError('retained traceback') + except ValueError as error: + retained.owner = error_execution(error) + return weakref.ref(retained) + +reference = cycle() +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); +} + +fn state( + py: Python<'_>, + logger: Py, + response: Py, + asynchronous: bool, +) -> PythonCallState { + PythonCallState { + args: PyTuple::empty(py).unbind(), + kwargs: PyDict::new(py).unbind(), + logger: Some(logger.extract(py).unwrap()), + start: py.None(), + end: Some(py.None()), + response: Some(response), + error: None, + asynchronous, + internal: false, + call_type: "test", + } +} + +#[test] +fn success_dispatch_reports_ordinary_failures_without_replacing_response() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +import sys + +response = object() +failure = ValueError('terminal diagnostic') +diagnostics = [] +old_hook = sys.unraisablehook +sys.unraisablehook = lambda event: diagnostics.append(event.exc_value) + +class Logger: + def handle_sync_success_callbacks_for_async_calls(self, *args): + raise failure + +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let response = locals.get_item("response").unwrap().unwrap().unbind(); + let mut lifecycle_state = state( + py, + locals.get_item("logger").unwrap().unwrap().unbind(), + response.clone_ref(py), + true, + ); + lifecycle_state.internal = true; + lifecycle_state.dispatch_success(py).unwrap(); + assert!(lifecycle_state.response.as_ref().unwrap().is(&response)); + py.run( + pyo3::ffi::c_str!( + r#" +assert diagnostics == [failure] +sys.unraisablehook = old_hook +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); +} + +#[test] +fn retained_failure_preserves_exception_identity() { + Python::initialize(); + Python::attach(|py| { + let logger = PyDict::new(py).into_any().unbind(); + let response = py.None(); + let failure = pyo3::exceptions::PyValueError::new_err("identity"); + let failure_value = failure.value(py).clone().unbind(); + let mut lifecycle_state = state(py, logger, response, false); + lifecycle_state.retain_error(py, failure); + let retained = lifecycle_state.error.take().unwrap(); + assert!(retained.is(&failure_value)); + }); +} + +#[test] +fn deferred_release_uses_release_context_and_allows_reentry_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +import sys +import types +from contextvars import ContextVar + +litellm = types.ModuleType('litellm') +core_utils = types.ModuleType('litellm.litellm_core_utils') +logging_worker = types.ModuleType('litellm.litellm_core_utils.logging_worker') +litellm.litellm_core_utils = core_utils +core_utils.logging_worker = logging_worker +sys.modules['litellm'] = litellm +sys.modules['litellm.litellm_core_utils'] = core_utils +sys.modules['litellm.litellm_core_utils.logging_worker'] = logging_worker + +marker = ContextVar('marker', default='unset') +observed = [] + +class Coroutine: + def close(self): + observed.append('closed') + +class Worker: + def ensure_initialized_and_enqueue(self, coroutine): + observed.append(marker.get()) + pending.release(True) + coroutine.close() + +class Logger: + def async_success_handler(self, *args): + observed.append('created') + return Coroutine() + +worker = Worker() +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + install_logging_worker(py, &locals.get_item("worker").unwrap().unwrap()).unwrap(); + let pending = Py::new( + py, + PendingLogging::new(PendingSuccess { + logger: locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(), + response: Some(py.None()), + start: py.None(), + end: Some(py.None()), + }), + ) + .unwrap(); + locals.set_item("pending", &pending).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['created', 'release', 'closed'] +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); +} + +#[test] +fn deferred_logging_collects_cycles_through_typed_logger() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!("class Logger: pass\nlogger = Logger()"), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let pending = Py::new( + py, + PendingLogging::new(PendingSuccess { + logger: locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(), + response: None, + start: py.None(), + end: None, + }), + ) + .unwrap(); + locals.set_item("pending", pending).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref +logger.pending = pending +reference = weakref.ref(logger) +del logger, pending +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); +} + +#[test] +fn coroutine_collects_cycles_retained_by_bridge_host() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "retaining_coroutine", + wrap_pyfunction!(retaining_coroutine, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained: + pass + +def cycle(): + retained = Retained() + coroutine = retaining_coroutine(retained) + retained.coroutine = coroutine + return weakref.ref(retained) + +retained_ref = cycle() +gc.collect() +assert retained_ref() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs index e67bfa89cc7..867a879ab91 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs @@ -8,7 +8,7 @@ use litellm_core::chat_completions::{ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::chat_completions_error_to_pyerr; +use crate::errors::execution_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_array}; fn prepare_chat_completions( @@ -86,6 +86,6 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_chat_completions, - errors = chat_completions_error_to_pyerr, + errors = execution_error_to_pyerr, extra = [chat_completions_decline], } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index d43c2f88775..912807f2048 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -169,9 +169,12 @@ fn _ocr_upload_document( pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?; - module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?; - module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?; - module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?) + crate::routes::definition::add_function( + module, + wrap_pyfunction!(_ocr_upload_document, module)?, + )?; + crate::routes::definition::add_function(module, wrap_pyfunction!(_ocr_file_document, module)?)?; + crate::routes::definition::add_function(module, wrap_pyfunction!(_ocr_mime_type, module)?) } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 12d902a3544..e655290d8fb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -307,5 +307,5 @@ fn _ocr_lifecycle( } pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(_ocr_lifecycle, module)?) + crate::routes::definition::add_function(module, wrap_pyfunction!(_ocr_lifecycle, module)?) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 10fa40b65ea..e1d8a3ac947 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -17,3 +17,32 @@ pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { value::register_trace(module) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registration_preserves_existing_private_exports() { + Python::initialize(); + Python::attach(|py| { + for name in [ + "_ocr_lifecycle", + "_ocr_upload_document", + "_ocr_file_document", + "_ocr_mime_type", + ] { + let module = PyModule::new(py, "ocr").unwrap(); + let original = pyo3::types::PyDict::new(py); + module.add(name, &original).unwrap(); + let error = register(&module).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!( + error.value(py).to_string(), + format!("duplicate native route: {name}") + ); + assert!(module.getattr(name).unwrap().is(&original)); + } + }); + } +}