diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9edbc31125e..dedd41b2efc 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3922,7 +3922,9 @@ version = "0.1.0" dependencies = [ "bytes", "futures-util", + "litellm-coroutine", "litellm-host", + "parking_lot", "pyo3", "pyo3-async-runtimes", "pythonize", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b94988eb8e1..f04eb115071 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -80,6 +80,7 @@ http = "1" google-cloud-auth = { version = "1.16.0", default-features = false } jsonwebtoken = { version = "11.1.0", default-features = false } hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] } +parking_lot = "0.12.5" proptest = "1.7.0" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index bd103b566d5..8d220b19def 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -6,10 +6,12 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-coroutine.workspace = true litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true +parking_lot.workspace = true serde.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } diff --git a/litellm-rust/crates/host-python/src/caller.rs b/litellm-rust/crates/host-python/src/caller.rs new file mode 100644 index 00000000000..3bae948c936 --- /dev/null +++ b/litellm-rust/crates/host-python/src/caller.rs @@ -0,0 +1,58 @@ +use std::ops::ControlFlow; + +use litellm_coroutine::{Abandoned, Co, Reply}; +use pyo3::prelude::*; + +use crate::Error; +use crate::handle::Suspension; + +pub(crate) type Answer = PyResult>; + +pub(crate) struct Yielded { + pub(crate) suspension: Suspension, + pub(crate) reply: Reply, +} + +/// Hands suspensions to the Python caller and receives its answers. +#[derive(Clone)] +pub(crate) struct PythonCaller { + co: Co, +} + +impl PythonCaller { + /// Connects the call body to its coroutine's yield handle. + pub(crate) fn new(co: Co) -> Self { + Self { co } + } + + /// Asks the caller to await Python work and return its value or exception. + pub(crate) async fn await_python(&self, awaitable: Py) -> Answer { + self.suspend(Suspension::Await(awaitable)).await + } + + /// Hands the caller a stream head and waits for demand or closure. + pub(crate) async fn open_stream(&self, head: Py) -> ControlFlow<()> { + self.stream_boundary(Suspension::Open(head)).await + } + + /// Delivers a chunk and waits for the caller to request another or close. + pub(crate) async fn yield_chunk(&self, chunk: Py) -> ControlFlow<()> { + self.stream_boundary(Suspension::Yield(chunk)).await + } + + /// Maps stream demand to continuation and any error to detachment. + async fn stream_boundary(&self, suspension: Suspension) -> ControlFlow<()> { + match self.suspend(suspension).await { + Ok(_) => ControlFlow::Continue(()), + Err(_) => ControlFlow::Break(()), + } + } + + /// Yields a suspension with its reply handle and waits for the answer. + async fn suspend(&self, suspension: Suspension) -> Answer { + match self.co.yield_(|reply| Yielded { suspension, reply }).await { + Ok(answer) => answer, + Err(Abandoned) => Err(Error::MissingState.into()), + } + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 7c89b6cc032..da38588088e 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -1,13 +1,16 @@ -use std::ops::ControlFlow; +use std::pin::pin; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; +use parking_lot::{Mutex, MutexGuard}; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::PyDict; +use litellm_coroutine::{Coroutine, CoroutineState}; use litellm_host::{ call::HostedCompletion, - interceptors::WireRequest, lifecycle::{CallEvent, ExecutionEvent, FailureOrigin, Timing, epoch_seconds}, machine::{HostFailure, Machine, MachineStep}, observation::ObservationSender, @@ -15,22 +18,22 @@ use litellm_host::{ }; use crate::PythonHostCalls; -use crate::handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle}; -use crate::hooks::{HookResume, HookStep, PythonCallEvent, PythonCallHooks}; +use crate::caller::{Answer, PythonCaller, Yielded}; +use crate::handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle, Suspension}; +use crate::hooks::{HookStep, PythonCallEvent, PythonCallHooks}; use crate::native::{NativeMachine, NativePoll}; use crate::{Error, InvokeError, PythonBinding}; type ProtocolOf = ::Protocol; type ErrorOf = as Protocol>::Error; type ResponseOf = as Protocol>::Response; -type NativeStep = MachineStep, HostedCompletion>>; -type NativeResult = Result, ErrorOf>; type Interruption = Option>>; type StartMachine = Box< dyn FnOnce(Python<'_>, &Bound<'_, PyDict>,

::Request) -> PyResult + Send + Sync, >; +type Body = Coroutine>>; pub struct CallOptions { pub asynchronous: bool, @@ -48,67 +51,537 @@ impl CallOptions { } } -enum Stage { - Begin, - Call, - Streaming, - AfterSuccess, - Succeeded(Py), - Failed(Py), +fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { + !error.is_instance_of::(py) } -enum EventNext { - Started, - Emitted(Reply<()>), - Terminal, +struct Observation { + observers: Option, + observing: bool, + started_at: f64, + ended_at: Option, + terminal: Option, } -enum Pending { - Host, - Native, - Arguments(HookResume>), - Wire(HookResume>, Reply), - Response(HookResume>), - Event(HookResume, EventNext), - Consumer(Reply>), -} +impl Observation { + fn timing(&self) -> Timing { + Timing { + start_time: self.started_at, + end_time: self.ended_at.unwrap_or_else(epoch_seconds), + } + } -/// A route answer as the driver resumes on it: a Python exception interrupts the call as -/// raised, a native rejection resumes the machine with it. -fn answered(answer: Result<(), InvokeError>) -> PyResult> { - match answer { - Ok(()) => Ok(Ok(())), - Err(InvokeError::Native(error)) => Ok(Err(error)), - Err(InvokeError::Python(error)) => Err(error), + fn emit(&self, event: &PythonCallEvent<'_>) { + if let Some(observers) = &self.observers { + observers.emit(event.snapshot()); + } + } + + fn retain_terminal(&mut self, event: &PythonCallEvent<'_>) { + self.terminal = self.observers.as_ref().map(|_| event.snapshot()); + } + + fn finish(&mut self) { + if let Some(event) = self.terminal.take() + && let Some(observers) = &self.observers + { + observers.emit(event); + } + self.observing = false; } } -enum Next> { - Return(ExecutionStep), - Continue(NativePoll>), +/// What the running call shares with its handle: the binding and hooks the handle +/// traverses and closes, and the Python values the call retains across a suspension. +struct Parts { + binding: H, + hooks: L, + arguments: Option>, + interrupted: Option>, + selected: Option>, + observation: Observation, } -struct PythonDriver +type SharedParts = Arc>>; + +enum Setup { + Host(PyErr), + Decode(InvokeError), +} + +/// The route's public exception for a native failure. When classification itself fails, +/// that failure is raised with the native error's text as its `__context__`. +fn classified(py: Python<'_>, binding: &H, error: ErrorOf) -> PyErr { + let native = error.to_string(); + let classifier_error = match binding.map_error(py, error) { + Ok(failure) => return failure.into(), + Err(classifier_error) => classifier_error, + }; + classifier_error.set_context(py, Some(PyRuntimeError::new_err(native))); + classifier_error +} + +/// The call as one async body. Every hook step and host call is awaited in place on the +/// caller's task; the parts lock is never held across an await. +struct CallBody +where + H: PythonBinding, + M: Machine, +{ + caller: PythonCaller, + parts: SharedParts, + native: NativeMachine, + start: Option, M>>, + streaming: bool, +} + +impl CallBody where L: PythonCallHooks + 'static, - H: PythonBinding + PythonHostCalls, + H: PythonBinding + PythonHostCalls + 'static, M: Machine + 'static, M::Complete: Into>>, { - binding: H, - hooks: L, - native: NativeMachine, - start: Option, M>>, - closed: bool, - arguments: Option>, - started_at: f64, - ended_at: Option, - stage: Stage, - pending: Option>, - interrupted: Option>, - observers: Option, - observing: bool, - terminal_observation: Option, + fn parts(&self) -> MutexGuard<'_, Parts> { + self.parts.lock() + } + + async fn run(mut self) -> PyResult> { + match self.begin().await { + Ok(()) => self.execute().await, + Err(Setup::Host(error)) => self.failure(error, FailureOrigin::Host).await, + Err(Setup::Decode(InvokeError::Python(error))) => { + self.failure(error, FailureOrigin::Call).await + } + Err(Setup::Decode(InvokeError::Native(error))) => self.machine_failed(error).await, + } + } + + async fn begin(&mut self) -> Result<(), Setup>> { + let started_at = epoch_seconds(); + let started = { + let mut parts = self.parts(); + parts.observation.started_at = started_at; + parts.observation.observing = true; + let event = PythonCallEvent::Started { + start_time: started_at, + }; + parts.observation.emit(&event); + Python::attach(|py| parts.hooks.on_event(py, event)).map_err(Setup::Host)? + }; + self.settle(started).await.map_err(Setup::Host)?; + let arguments = self + .parts() + .arguments + .take() + .ok_or_else(|| Setup::Host(Error::MissingState.into()))?; + let prepare = { + let mut parts = self.parts(); + Python::attach(|py| parts.hooks.prepare_arguments(py, arguments, started_at)) + .map_err(Setup::Host)? + }; + let prepared = self.settle(prepare).await.map_err(Setup::Host)?; + let request = { + let mut parts = self.parts(); + let Parts { + binding, + hooks, + arguments, + .. + } = &mut *parts; + Python::attach(|py| { + hooks + .arguments_prepared(py, &prepared) + .map_err(Setup::Host)?; + let decoded = binding.decode_request(py, prepared.bind(py)); + *arguments = Some(prepared); + decoded.map_err(Setup::Decode) + })? + }; + let start = self + .start + .take() + .ok_or_else(|| Setup::Host(Error::MissingState.into()))?; + let machine = { + let parts = self.parts(); + let arguments = parts + .arguments + .as_ref() + .ok_or_else(|| Setup::Host(Error::MissingState.into()))?; + Python::attach(|py| start(py, arguments.bind(py), request)).map_err(Setup::Host)? + }; + self.native.start(machine); + Ok(()) + } + + /// Runs one hook step to its value, handing each awaitable it produces to the caller. + async fn settle(&self, step: HookStep) -> PyResult { + let mut step = step; + loop { + let (awaitable, resume) = match step { + HookStep::Ready(value) => return Ok(value), + HookStep::Await(awaitable, resume) => (awaitable, resume), + }; + let answer = self.caller.await_python(awaitable).await; + step = { + let mut parts = self.parts(); + Python::attach(|py| resume(&mut parts.hooks, py, answer))? + }; + } + } + + async fn execute(&mut self) -> PyResult> { + let mut interruption: Interruption = None; + loop { + let poll = Python::attach(|py| self.native.resume(py, interruption.take()))?; + let result = match poll { + NativePoll::Ready(result) => result, + NativePoll::Suspend(awaitable) => match self.caller.await_python(awaitable).await { + Ok(_) => self.native.take_result()?, + Err(error) if self.parts().interrupted.is_none() => { + interruption = Some(self.interrupt(error)); + continue; + } + Err(error) => return Err(error), + }, + }; + let op = match result { + Ok(MachineStep::Suspended(op)) => op, + Ok(MachineStep::Complete(completion)) => { + return self.completed(completion).await; + } + Err(error) => return self.machine_failed(error).await, + }; + interruption = match self.answer(op).await { + Ok(()) => None, + Err(InvokeError::Native(error)) => Some(HostFailure::Error(error)), + Err(InvokeError::Python(error)) => Some(self.interrupt(error)), + }; + } + } + + fn interrupt(&self, error: PyErr) -> HostFailure> { + Python::attach(|py| { + let cancelled = is_cancellation(py, &error); + let native = H::host_error(&error); + self.parts().interrupted = Some(error.into_value(py)); + if cancelled { + HostFailure::Cancelled(native) + } else { + HostFailure::Error(native) + } + }) + } + + /// Answers one machine op. A Python error interrupts the call as raised, a native + /// rejection resumes the machine with it. + async fn answer( + &mut self, + op: HostRequest>, + ) -> Result<(), InvokeError>> { + match op { + HostRequest::HostCall(call) => self.host_call(call).await, + HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + }) => { + let step = { + let mut parts = self.parts(); + Python::attach(|py| parts.hooks.before_provider_request(py, wire, &context))? + }; + let wire = self.settle(step).await?; + reply.send(*wire); + Ok(()) + } + HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { + self.notify(PythonCallEvent::Execution(ExecutionEvent::ResultReady { + facts, + })) + .await?; + reply.send(()); + Ok(()) + } + HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { + self.notify(PythonCallEvent::Execution( + ExecutionEvent::ProviderResponseReceived { raw: &raw }, + )) + .await?; + reply.send(()); + Ok(()) + } + HostRequest::Stream(StreamDelivery::Open(head, reply)) => { + self.streaming = true; + let head = { + let mut parts = self.parts(); + let Parts { binding, hooks, .. } = &mut *parts; + Python::attach(|py| { + let head = binding.encode_stream_head(py, head)?; + hooks.on_stream_open(py, &head)?; + Ok::<_, PyErr>(head) + })? + }; + reply.send(self.caller.open_stream(head).await); + Ok(()) + } + HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => { + let chunk = { + let mut parts = self.parts(); + let Parts { binding, hooks, .. } = &mut *parts; + Python::attach(|py| { + let chunk = binding.encode_chunk(py, chunk)?; + hooks.on_stream_chunk(py, &chunk)?; + Ok::<_, PyErr>(chunk) + })? + }; + reply.send(self.caller.yield_chunk(chunk).await); + Ok(()) + } + } + } + + async fn host_call( + &self, + call: as Protocol>::HostCall, + ) -> Result<(), InvokeError>> { + let mut pending = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.begin_host_call(py, call))? + }; + while let Some(awaitable) = pending { + let result = self.caller.await_python(awaitable).await; + pending = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.resume_host_call(py, result))? + }; + } + Ok(()) + } + + async fn notify(&self, event: PythonCallEvent<'_>) -> PyResult<()> { + let step = { + let mut parts = self.parts(); + parts.observation.emit(&event); + Python::attach(|py| parts.hooks.on_event(py, event))? + }; + self.settle(step).await + } + + async fn completed( + &mut self, + completion: HostedCompletion>, + ) -> PyResult> { + self.parts().observation.ended_at = Some(epoch_seconds()); + let response = match completion { + HostedCompletion::Complete(response) => response, + HostedCompletion::StreamEnded | HostedCompletion::Detached => { + let none = Python::attach(|py| py.None()); + return self.succeeded(none).await; + } + }; + let encoded = { + let mut parts = self.parts(); + Python::attach(|py| parts.binding.encode_response(py, response)) + }; + let public = match encoded { + Ok(public) => public, + Err(error) => return self.failure(error, FailureOrigin::Call).await, + }; + if self.streaming { + return self.succeeded(public).await; + } + let transform = { + let mut parts = self.parts(); + let timing = parts.observation.timing(); + Python::attach(|py| parts.hooks.transform_response(py, public, timing)) + }; + let step = match transform { + Ok(step) => step, + Err(error) => return self.failure(error, FailureOrigin::Host).await, + }; + match self.settle(step).await { + Ok(response) => self.succeeded(response).await, + Err(error) => self.failure(error, FailureOrigin::Host).await, + } + } + + async fn machine_failed(&mut self, error: ErrorOf) -> PyResult> { + let error = { + let mut parts = self.parts(); + parts.observation.ended_at.get_or_insert_with(epoch_seconds); + Python::attach(|py| match parts.interrupted.take() { + Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), + None => classified(py, &parts.binding, error), + }) + }; + self.failure(error, FailureOrigin::Call).await + } + + /// The terminal hooks may suspend. Until they finish, the selected response is held + /// only in the traversed parts, so a cycle through it stays collectable. + async fn succeeded(&mut self, response: Py) -> PyResult> { + let step = { + let mut parts = self.parts(); + let timing = parts.observation.timing(); + let event = PythonCallEvent::Succeeded { + timing, + response: &response, + }; + parts.observation.retain_terminal(&event); + let step = Python::attach(|py| parts.hooks.on_event(py, event))?; + parts.selected = Some(response); + step + }; + self.settle(step).await?; + let mut parts = self.parts(); + parts.observation.finish(); + parts + .selected + .take() + .ok_or_else(|| Error::MissingState.into()) + } + + /// As [`Self::succeeded`], with the selected exception held as its value and raised + /// again once the hooks finish. + async fn failure(&mut self, error: PyErr, origin: FailureOrigin) -> PyResult> { + let step = { + let mut parts = self.parts(); + parts.observation.ended_at.get_or_insert_with(epoch_seconds); + let timing = parts.observation.timing(); + Python::attach(|py| { + if is_cancellation(py, &error) { + return Err(error); + } + let event = PythonCallEvent::Failed { + timing, + origin, + error: &error, + }; + parts.observation.retain_terminal(&event); + let step = parts.hooks.on_event(py, event)?; + parts.selected = Some(error.into_value(py).into_any()); + Ok(step) + })? + }; + self.settle(step).await?; + let mut parts = self.parts(); + parts.observation.finish(); + let selected = parts.selected.take().ok_or(Error::MissingState)?; + Err(Python::attach(|py| { + PyErr::from_value(selected.into_bound(py)) + })) + } +} + +struct CallDriver +where + H: PythonBinding, + L: PythonCallHooks, +{ + body: Mutex, + reply: Option>, + parts: SharedParts, +} + +impl CallDriver +where + H: PythonBinding + 'static, + L: PythonCallHooks + 'static, +{ + /// A failure the body raised past its terminal hooks, or a cancellation, is the last + /// thing the observers hear about the call. + fn observe_host_failure(&self, error: &PyErr) { + let mut parts = self.parts.lock(); + if !parts.observation.observing { + return; + } + let timing = parts.observation.timing(); + Python::attach(|py| { + let event = if is_cancellation(py, error) { + PythonCallEvent::Cancelled { timing } + } else { + PythonCallEvent::Failed { + timing, + origin: FailureOrigin::Host, + error, + } + }; + parts.observation.emit(&event); + }); + parts.observation.observing = false; + parts.observation.terminal = None; + } +} + +impl ExecutionBody for CallDriver +where + H: PythonBinding + 'static, + L: PythonCallHooks + 'static, +{ + fn resume(&mut self, answer: Option) -> PyResult { + match (self.reply.take(), answer) { + (Some(reply), Some(answer)) => reply.send(answer), + (None, None) => {} + (Some(_), None) | (None, Some(_)) => return Err(Error::MissingState.into()), + } + let step = { + let body = self.body.get_mut(); + let mut resume = pin!(body.resume()); + match resume + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + { + Poll::Ready(step) => step, + Poll::Pending => return Err(Error::MissingState.into()), + } + }; + match step { + Ok(CoroutineState::Yielded(Yielded { suspension, reply })) => { + self.reply = Some(reply); + Ok(ExecutionStep::Suspended(suspension)) + } + Ok(CoroutineState::Complete(Ok(value))) => Ok(ExecutionStep::Return(value)), + Ok(CoroutineState::Complete(Err(error))) => { + self.observe_host_failure(&error); + Err(error) + } + Err(_) => Err(Error::MissingState.into()), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + let Some(parts) = self.parts.try_lock() else { + return Ok(()); + }; + parts.binding.traverse(visit)?; + parts.hooks.traverse(visit)?; + visit.call(&parts.arguments)?; + visit.call(&parts.interrupted)?; + visit.call(&parts.selected) + } +} + +impl Drop for CallDriver +where + H: PythonBinding, + L: PythonCallHooks, +{ + fn drop(&mut self) { + self.body.get_mut().cancel(); + self.reply = None; + let mut parts = self.parts.lock(); + Python::attach(|py| { + if parts.observation.observing { + let timing = parts.observation.timing(); + parts + .observation + .emit(&PythonCallEvent::Cancelled { timing }); + parts.observation.observing = false; + } + parts.hooks.close(py); + parts.binding.close(py); + }); + } } pub fn run_call( @@ -132,21 +605,38 @@ where M: Machine + 'static, M::Complete: Into>>, { - let mut driver = PythonDriver { + let parts = Arc::new(Mutex::new(Parts { binding, hooks, - native: NativeMachine::new(options.asynchronous), - start: Some(Box::new(start)), - closed: false, arguments: Some(arguments), - started_at: 0.0, - ended_at: None, - stage: Stage::Begin, - pending: None, interrupted: None, - observers: options.observers, - observing: false, - terminal_observation: None, + selected: None, + observation: Observation { + observers: options.observers, + observing: false, + started_at: 0.0, + ended_at: None, + terminal: None, + }, + })); + let body = Coroutine::new({ + let parts = Arc::clone(&parts); + let asynchronous = options.asynchronous; + move |co| { + CallBody { + caller: PythonCaller::new(co), + parts, + native: NativeMachine::new(asynchronous), + start: Some(Box::new(start)), + streaming: false, + } + .run() + } + }); + let mut driver = CallDriver { + body: Mutex::new(body), + reply: None, + parts, }; if options.asynchronous { return Execution::new(driver, options.lifecycle) @@ -155,547 +645,13 @@ where } match driver.resume(None)? { ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Open(head) => Execution::suspended(driver, options.lifecycle) - .into_sync_stream(py, head) - .map(Bound::unbind), - ExecutionStep::Await(_) | ExecutionStep::Yield(_) => Err(Error::SyncCallSuspended.into()), - } -} - -fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { - !error.is_instance_of::(py) -} - -impl PythonDriver -where - L: PythonCallHooks + 'static, - H: PythonBinding + PythonHostCalls, - M: Machine + 'static, - M::Complete: Into>>, -{ - fn timing(&self) -> Timing { - Timing { - start_time: self.started_at, - end_time: self.ended_at.unwrap_or_else(epoch_seconds), + ExecutionStep::Suspended(Suspension::Open(head)) => { + Execution::suspended(driver, options.lifecycle) + .into_sync_stream(py, head) + .map(Bound::unbind) } - } - - fn drive( - &mut self, - py: Python<'_>, - result: Option>>, - ) -> PyResult { - match (self.pending.take(), result) { - (None, None) => { - self.started_at = epoch_seconds(); - let started = PythonCallEvent::Started { - start_time: self.started_at, - }; - self.observing = true; - self.observe(&started); - match self.hooks.on_event(py, started) { - Ok(step) => self.on_event(py, step, EventNext::Started), - Err(error) => self.hook_failed(py, error), - } - } - (Some(Pending::Host), Some(result)) => { - match self.binding.resume_host_call(py, result) { - Ok(Some(awaitable)) => { - self.pending = Some(Pending::Host); - Ok(ExecutionStep::Await(awaitable)) - } - Ok(None) => self.resume_machine(py, None), - Err(InvokeError::Python(error)) => self.interrupt(py, error), - Err(InvokeError::Native(error)) => { - self.resume_machine(py, Some(HostFailure::Error(error))) - } - } - } - (Some(Pending::Native), Some(Ok(_))) => { - let result = self.native.take_result()?; - self.run_steps(py, NativePoll::Ready(result)) - } - (Some(Pending::Native), Some(Err(error))) => self.interrupt(py, error), - (Some(Pending::Consumer(reply)), Some(read)) => { - reply.send(if read.is_ok() { - ControlFlow::Continue(()) - } else { - ControlFlow::Break(()) - }); - self.resume_machine(py, None) - } - (Some(Pending::Arguments(resume)), Some(result)) => { - let step = resume(&mut self.hooks, py, result); - match step { - Ok(step) => self.on_arguments(py, step), - Err(error) => self.hook_failed(py, error), - } - } - (Some(Pending::Wire(resume, reply)), Some(result)) => { - let step = resume(&mut self.hooks, py, result); - match step { - Ok(step) => self.on_wire(py, step, reply), - Err(error) => self.hook_failed(py, error), - } - } - (Some(Pending::Response(resume)), Some(result)) => { - let step = resume(&mut self.hooks, py, result); - match step { - Ok(step) => self.on_response(py, step), - Err(error) => self.hook_failed(py, error), - } - } - (Some(Pending::Event(resume, next)), Some(result)) => { - let step = resume(&mut self.hooks, py, result); - match step { - Ok(step) => self.on_event(py, step, next), - Err(error) => self.hook_failed(py, error), - } - } - _ => Err(Error::MissingState.into()), - } - } - - fn on_arguments( - &mut self, - py: Python<'_>, - step: HookStep>, - ) -> PyResult { - match step { - HookStep::Await(awaitable, resume) => { - self.pending = Some(Pending::Arguments(resume)); - Ok(ExecutionStep::Await(awaitable)) - } - HookStep::Ready(arguments) => { - if let Err(error) = self.hooks.arguments_prepared(py, &arguments) { - return self.hook_failed(py, error); - } - let decoded = self.binding.decode_request(py, arguments.bind(py)); - self.arguments = Some(arguments); - let request = match decoded { - Ok(request) => request, - Err(InvokeError::Native(error)) => return self.machine_failed(py, error), - Err(InvokeError::Python(error)) => { - return self.failure(py, error, FailureOrigin::Call); - } - }; - let start = self.start.take().ok_or(Error::MissingState)?; - let arguments = self.arguments.as_ref().ok_or(Error::MissingState)?; - let machine = match start(py, arguments.bind(py), request) { - Ok(machine) => machine, - Err(error) => return self.failure(py, error, FailureOrigin::Host), - }; - self.native.start(machine); - self.stage = Stage::Call; - self.resume_machine(py, None) - } - } - } - - fn on_wire( - &mut self, - py: Python<'_>, - step: HookStep>, - reply: Reply, - ) -> PyResult { - match step { - HookStep::Await(awaitable, resume) => { - self.pending = Some(Pending::Wire(resume, reply)); - Ok(ExecutionStep::Await(awaitable)) - } - HookStep::Ready(wire) => { - reply.send(*wire); - self.resume_machine(py, None) - } - } - } - - fn on_response( - &mut self, - py: Python<'_>, - step: HookStep>, - ) -> PyResult { - match step { - HookStep::Await(awaitable, resume) => { - self.pending = Some(Pending::Response(resume)); - Ok(ExecutionStep::Await(awaitable)) - } - HookStep::Ready(response) => self.succeeded(py, response), - } - } - - fn on_event( - &mut self, - py: Python<'_>, - step: HookStep, - next: EventNext, - ) -> PyResult { - match step { - HookStep::Await(awaitable, resume) => { - self.pending = Some(Pending::Event(resume, next)); - Ok(ExecutionStep::Await(awaitable)) - } - HookStep::Ready(()) => match next { - EventNext::Started => self.begin(py), - EventNext::Emitted(reply) => { - reply.send(()); - self.resume_machine(py, None) - } - EventNext::Terminal => { - if let Some(event) = self.terminal_observation.take() - && let Some(observers) = &self.observers - { - observers.emit(event); - } - self.observing = false; - match &self.stage { - Stage::Succeeded(response) => { - Ok(ExecutionStep::Return(response.clone_ref(py))) - } - Stage::Failed(error) => { - Err(PyErr::from_value(error.bind(py).clone().into_any())) - } - _ => Err(Error::MissingState.into()), - } - } - }, - } - } - - fn begin(&mut self, py: Python<'_>) -> PyResult { - let arguments = self.arguments.take().ok_or(Error::MissingState)?; - match self.hooks.prepare_arguments(py, arguments, self.started_at) { - Ok(step) => self.on_arguments(py, step), - Err(error) => self.hook_failed(py, error), - } - } - - fn hook_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { - match self.stage { - Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), - Stage::Call | Stage::Streaming => self.interrupt(py, error), - Stage::Succeeded(_) | Stage::Failed(_) => Err(error), - } - } - - fn resume_machine( - &mut self, - py: Python<'_>, - interruption: Interruption, - ) -> PyResult { - let step = self.native.resume(py, interruption)?; - self.run_steps(py, step) - } - - fn run_steps( - &mut self, - py: Python<'_>, - mut step: NativePoll>, - ) -> PyResult { - loop { - let result = match step { - NativePoll::Suspend(awaitable) => { - self.pending = Some(Pending::Native); - return Ok(ExecutionStep::Await(awaitable)); - } - NativePoll::Ready(result) => result, - }; - step = match self.handle_native(py, result)? { - Next::Return(step) => return Ok(step), - Next::Continue(step) => step, - }; - } - } - - /// Answers one machine step: performs the op it asked for, or finishes the call. - fn handle_native(&mut self, py: Python<'_>, result: NativeResult) -> PyResult> { - let op = match result { - Ok(MachineStep::Suspended(op)) => op, - Ok(MachineStep::Complete(response)) => { - return self.completed(py, response).map(Next::Return); - } - Err(error) => return self.machine_failed(py, error).map(Next::Return), - }; - let answered = match op { - HostRequest::HostCall(op) => match self.binding.begin_host_call(py, op) { - Ok(Some(awaitable)) => { - self.pending = Some(Pending::Host); - return Ok(Next::Return(ExecutionStep::Await(awaitable))); - } - result => answered(result.map(|_| ())), - }, - HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { - wire, - context, - reply, - }) => match self.hooks.before_provider_request(py, wire, &context) { - Ok(HookStep::Ready(wire)) => { - reply.send(*wire); - Ok(Ok(())) - } - Ok(HookStep::Await(awaitable, resume)) => { - self.pending = Some(Pending::Wire(resume, reply)); - return Ok(Next::Return(ExecutionStep::Await(awaitable))); - } - Err(error) => Err(error), - }, - HostRequest::Stream(StreamDelivery::Open(head, reply)) => { - return self.opened(py, head, reply).map(Next::Return); - } - HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => { - return self.delivered(py, chunk, reply).map(Next::Return); - } - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { - let event = PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }); - self.observe(&event); - match self.hooks.on_event(py, event) { - Ok(HookStep::Ready(())) => { - reply.send(()); - Ok(Ok(())) - } - Ok(HookStep::Await(awaitable, resume)) => { - self.pending = Some(Pending::Event(resume, EventNext::Emitted(reply))); - return Ok(Next::Return(ExecutionStep::Await(awaitable))); - } - Err(error) => Err(error), - } - } - HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { - let event = PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { - raw: &raw, - }); - self.observe(&event); - match self.hooks.on_event(py, event) { - Ok(HookStep::Ready(())) => { - reply.send(()); - Ok(Ok(())) - } - Ok(HookStep::Await(awaitable, resume)) => { - self.pending = Some(Pending::Event(resume, EventNext::Emitted(reply))); - return Ok(Next::Return(ExecutionStep::Await(awaitable))); - } - Err(error) => Err(error), - } - } - }; - match answered { - Ok(Ok(())) => self.native.resume(py, None).map(Next::Continue), - Ok(Err(native)) => self - .native - .resume(py, Some(HostFailure::Error(native))) - .map(Next::Continue), - Err(error) => self.interrupt(py, error).map(Next::Return), - } - } - - fn opened( - &mut self, - py: Python<'_>, - head: as Protocol>::StreamHead, - reply: Reply>, - ) -> PyResult { - self.stage = Stage::Streaming; - let head = match self.binding.encode_stream_head(py, head) { - Ok(head) => head, - Err(error) => return self.interrupt(py, error), - }; - match self.hooks.on_stream_open(py, &head) { - Ok(()) => { - self.pending = Some(Pending::Consumer(reply)); - Ok(ExecutionStep::Open(head)) - } - Err(error) => self.interrupt(py, error), - } - } - - fn delivered( - &mut self, - py: Python<'_>, - chunk: as Protocol>::Chunk, - reply: Reply>, - ) -> PyResult { - let chunk = match self.binding.encode_chunk(py, chunk) { - Ok(chunk) => chunk, - Err(error) => return self.interrupt(py, error), - }; - match self.hooks.on_stream_chunk(py, &chunk) { - Ok(()) => { - self.pending = Some(Pending::Consumer(reply)); - Ok(ExecutionStep::Yield(chunk)) - } - Err(error) => self.interrupt(py, error), - } - } - - fn interrupt(&mut self, py: Python<'_>, error: PyErr) -> PyResult { - let cancelled = is_cancellation(py, &error); - let native = H::host_error(&error); - self.interrupted = Some(error.into_value(py)); - let failure = if cancelled { - HostFailure::Cancelled(native) - } else { - HostFailure::Error(native) - }; - self.resume_machine(py, Some(failure)) - } - - fn completed( - &mut self, - py: Python<'_>, - response: HostedCompletion>, - ) -> PyResult { - self.ended_at = Some(epoch_seconds()); - let response = match response { - HostedCompletion::Complete(response) => response, - HostedCompletion::StreamEnded | HostedCompletion::Detached => { - return self.succeeded(py, py.None()); - } - }; - let public = match self.binding.encode_response(py, response) { - Ok(public) => public, - Err(error) => return self.failure(py, error, FailureOrigin::Call), - }; - if let Stage::Streaming = self.stage { - return self.succeeded(py, public); - } - self.stage = Stage::AfterSuccess; - match self.hooks.transform_response(py, public, self.timing()) { - Ok(step) => self.on_response(py, step), - Err(error) => self.failure(py, error, FailureOrigin::Host), - } - } - - fn machine_failed(&mut self, py: Python<'_>, error: ErrorOf) -> PyResult { - self.ended_at.get_or_insert_with(epoch_seconds); - let error = match self.interrupted.take() { - Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), - None => self.classified(py, error), - }; - self.failure(py, error, FailureOrigin::Call) - } - - /// The route's public exception for a native failure. When classification itself - /// fails, that failure is raised with the native error's text as its `__context__`. - fn classified(&self, py: Python<'_>, error: ErrorOf) -> PyErr { - let native = error.to_string(); - let classifier_error = match self.binding.map_error(py, error) { - Ok(failure) => return failure.into(), - Err(classifier_error) => classifier_error, - }; - classifier_error.set_context(py, Some(PyRuntimeError::new_err(native))); - classifier_error - } - - fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { - let event = PythonCallEvent::Succeeded { - timing: self.timing(), - response: &response, - }; - self.terminal_observation = self.observers.as_ref().map(|_| event.snapshot()); - let step = self.hooks.on_event(py, event)?; - self.stage = Stage::Succeeded(response); - self.on_event(py, step, EventNext::Terminal) - } - - fn failure( - &mut self, - py: Python<'_>, - error: PyErr, - origin: FailureOrigin, - ) -> PyResult { - self.ended_at.get_or_insert_with(epoch_seconds); - if is_cancellation(py, &error) { - return Err(error); - } - let event = PythonCallEvent::Failed { - timing: self.timing(), - origin, - error: &error, - }; - self.terminal_observation = self.observers.as_ref().map(|_| event.snapshot()); - let step = self.hooks.on_event(py, event)?; - self.stage = Stage::Failed(error.into_value(py)); - self.on_event(py, step, EventNext::Terminal) - } - - fn observe(&self, event: &PythonCallEvent<'_>) { - if let Some(observers) = &self.observers { - observers.emit(event.snapshot()); - } - } - - fn clear(&mut self) { - if !self.closed { - self.closed = true; - if self.observing { - self.observe(&PythonCallEvent::Cancelled { - timing: self.timing(), - }); - self.observing = false; - } - self.native.close(); - self.start = None; - Python::attach(|py| { - self.hooks.close(py); - self.binding.close(py); - }); + ExecutionStep::Suspended(Suspension::Await(_) | Suspension::Yield(_)) => { + Err(Error::SyncCallSuspended.into()) } } } - -impl ExecutionBody for PythonDriver -where - L: PythonCallHooks + 'static, - H: PythonBinding + PythonHostCalls, - M: Machine + 'static, - M::Complete: Into>>, -{ - fn resume(&mut self, result: Option>>) -> PyResult { - Python::attach(|py| { - let outcome = self.drive(py, result); - if let Err(error) = &outcome - && self.observing - { - let event = if is_cancellation(py, error) { - PythonCallEvent::Cancelled { - timing: self.timing(), - } - } else { - PythonCallEvent::Failed { - timing: self.timing(), - origin: FailureOrigin::Host, - error, - } - }; - self.observe(&event); - self.observing = false; - self.terminal_observation = None; - } - outcome - }) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.binding.traverse(visit)?; - self.hooks.traverse(visit)?; - visit.call(&self.arguments)?; - visit.call(&self.interrupted)?; - match &self.stage { - Stage::Succeeded(response) => visit.call(response), - Stage::Failed(error) => visit.call(error), - _ => Ok(()), - } - } -} - -impl Drop for PythonDriver -where - L: PythonCallHooks + 'static, - H: PythonBinding + PythonHostCalls, - M: Machine + 'static, - M::Complete: Into>>, -{ - fn drop(&mut self) { - self.clear(); - } -} diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index 134c088583e..a9f03dc9724 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -7,8 +7,9 @@ use pyo3::prelude::*; pub type PythonLifecycle = for<'py> fn(Python<'py>) -> PyResult>; -pub enum ExecutionStep { - Return(Py), +/// What a suspended execution waits on: an awaitable the caller awaits, or a stream +/// boundary the caller's stream reads past. +pub enum Suspension { Await(Py), /// The call streams: the caller gets a stream over this execution carrying this head, /// and the execution stays suspended until the stream asks for a chunk. @@ -16,6 +17,11 @@ pub enum ExecutionStep { Yield(Py), } +pub enum ExecutionStep { + Suspended(Suspension), + Return(Py), +} + pub trait ExecutionBody: Send + Sync { fn resume(&mut self, result: Option>>) -> PyResult; fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; @@ -94,9 +100,9 @@ impl Execution { let outcome = catch_unwind(AssertUnwindSafe(|| { let step = body.resume(result)?; let (tag, value, suspended) = match step { - ExecutionStep::Await(value) => ("Await", value, true), - ExecutionStep::Open(head) => ("Open", head, true), - ExecutionStep::Yield(value) => ("Yield", value, true), + ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true), + ExecutionStep::Suspended(Suspension::Open(head)) => ("Open", head, true), + ExecutionStep::Suspended(Suspension::Yield(value)) => ("Yield", value, true), ExecutionStep::Return(value) => ("Complete", value, false), }; let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind(); diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index a83480be4a3..580e6674b95 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -5,6 +5,7 @@ mod argument; mod binding; +mod caller; mod conversion_cache; mod driver; mod error; @@ -27,7 +28,7 @@ pub use error::{Error, InvokeError}; pub use file_reader::{FileContent, PythonFileReader, py_bytes}; pub use fork_gate::RuntimeAlreadyStarted; pub use gil::{PythonContext, attach_blocking, release_count, release_gil}; -pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle}; +pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle, Suspension}; pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime}; pub use marshal::{ Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py, diff --git a/litellm-rust/crates/host-python/tests/driver.rs b/litellm-rust/crates/host-python/tests/driver.rs index fc4a8ee726c..522e4b4ae8c 100644 --- a/litellm-rust/crates/host-python/tests/driver.rs +++ b/litellm-rust/crates/host-python/tests/driver.rs @@ -9,12 +9,8 @@ use litellm_host::{ lifecycle::{CallEvent, ExecutionEvent, FailureOrigin}, machine::CallMachine, observation::observation_channel, - protocol::Reply, -}; -use litellm_host_python::{ - CallOptions, Error as HostError, HookChain, InvokeError, PythonBinding, PythonHostCalls, - PythonOwned, PythonRuntime, run_call, }; +use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call}; use pyo3::{ exceptions::{ PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError, @@ -26,9 +22,9 @@ use pyo3::{ }; use rstest::rstest; use support::{ - Classified, Error, HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, - call_options, failing_machine, hooks, interpreter, lifecycle_binding, open_stream, read_all, - run_composed, run_scripted, settled, success_machine, + HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options, + failing_machine, hooks, interpreter, lifecycle_binding, open_stream, read_all, run_composed, + run_scripted, settled, success_machine, }; fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec) { @@ -261,7 +257,6 @@ fn startup_failures_release_the_factory_without_constructing_a_machine( self.0.store(true, Ordering::SeqCst); } } - Python::attach(|py| { let constructed = Arc::new(AtomicBool::new(false)); let did_construct = constructed.clone(); @@ -404,131 +399,103 @@ fn a_stream_carries_its_head_before_the_first_chunk( }); } +struct Classification { + provider_fails: bool, + op: OpScript, + classifier_fails: bool, + message: &'static str, + native_context: Option<&'static str>, + steps: &'static [&'static str], +} + #[rstest] -fn a_native_failure_is_classified_once_and_reported_classified( +#[case::provider_failure(Classification { + provider_fails: true, + op: OpScript::Answer, + classifier_fails: false, + message: "classified: provider exploded", + native_context: None, + steps: &["project", "classify:provider exploded"], +})] +#[case::host_op_rejection(Classification { + provider_fails: false, + op: OpScript::RejectNatively, + classifier_fails: false, + message: "classified: op rejected", + native_context: None, + steps: &["project", "op:sign", "classify:op rejected"], +})] +#[case::host_op_python_exception(Classification { + provider_fails: false, + op: OpScript::RaisePython, + classifier_fails: false, + message: "op failed", + native_context: None, + steps: &["project", "op:sign"], +})] +#[case::failing_classifier(Classification { + provider_fails: true, + op: OpScript::Answer, + classifier_fails: true, + message: "classifier failed", + native_context: Some("provider exploded"), + steps: &["project", "classify:provider exploded"], +})] +fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is( #[from(interpreter)] _python: Interpreter, + #[case] case: Classification, #[values(false, true)] asynchronous: bool, ) { Python::attach(|py| { - let (result, log) = run_scripted( - py, - failing_machine(), - OpScript::Answer, - HookScript::Plain, - asynchronous, - ); - let error = result.unwrap_err(); - assert!(error.is_instance_of::(py)); - assert_eq!(error.value(py).to_string(), "classified: provider exploded"); - assert_eq!( - log, - [ - "started", - "begin", - "project", - "classify:provider exploded", - "failed:Call:classified: provider exploded", - "adapter.close", - "host.close", - ] - ); - }); -} - -#[rstest] -fn a_native_rejection_from_a_host_operation_is_classified_once( - #[from(interpreter)] _python: Interpreter, -) { - Python::attach(|py| { - let (result, log) = run_scripted( - py, - success_machine(), - OpScript::RejectNatively, - HookScript::Plain, - false, - ); - assert_eq!( - result.unwrap_err().value(py).to_string(), - "classified: op rejected" - ); - assert_eq!( - log, - [ - "started", - "begin", - "project", - "op:sign", - "classify:op rejected", - "failed:Call:classified: op rejected", - "adapter.close", - "host.close", - ] - ); - }); -} - -#[rstest] -fn a_python_exception_from_a_host_operation_is_reported_as_raised( - #[from(interpreter)] _python: Interpreter, -) { - Python::attach(|py| { - let (result, log) = run_scripted( - py, - success_machine(), - OpScript::RaisePython, - HookScript::Plain, - false, - ); - let error = result.unwrap_err(); - assert!(error.is_instance_of::(py)); - assert_eq!(error.value(py).to_string(), "op failed"); - assert_eq!( - log, - [ - "started", - "begin", - "project", - "op:sign", - "failed:Call:op failed", - "adapter.close", - "host.close", - ] - ); - }); -} - -#[rstest] -fn a_failing_classifier_surfaces_with_the_native_error_as_context( - #[from(interpreter)] _python: Interpreter, -) { - Python::attach(|py| { - let log_handle = Log::default(); + let Classification { + provider_fails, + op, + classifier_fails, + message, + native_context, + steps, + } = case; + let log = Log::default(); + let binding = if classifier_fails { + binding(&log, op).with_failing_classifier() + } else { + binding(&log, op) + }; + let machine = if provider_fails { + failing_machine() + } else { + success_machine() + }; let (result, log) = run_composed( py, - failing_machine(), - binding(&log_handle, OpScript::Answer).with_failing_classifier(), - hooks(&log_handle, HookScript::Plain), + machine, + binding, + hooks(&log, HookScript::Plain), std::convert::identity, - call_options(false), + call_options(asynchronous), ); let error = result.unwrap_err(); - assert!(error.is_instance_of::(py)); - assert_eq!(error.value(py).to_string(), "classifier failed"); - let context = error.context(py).unwrap(); - assert!(context.is_instance_of::(py)); - assert_eq!(context.value(py).to_string(), "provider exploded"); - assert_eq!( - log, - [ - "started", - "begin", - "project", - "classify:provider exploded", - "failed:Call:classifier failed", - "adapter.close", - "host.close", - ] - ); + assert_eq!(error.value(py).to_string(), message); + match native_context { + Some(context) => { + assert!(error.is_instance_of::(py)); + let chained = error.context(py).unwrap(); + assert!(chained.is_instance_of::(py)); + assert_eq!(chained.value(py).to_string(), context); + } + None => assert!(error.is_instance_of::(py)), + } + let expected: Vec = ["started", "begin"] + .iter() + .chain(steps) + .map(|step| step.to_string()) + .chain([ + format!("failed:Call:{message}"), + "adapter.close".into(), + "host.close".into(), + ]) + .collect(); + assert_eq!(log, expected); }); } @@ -738,25 +705,6 @@ fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote( }); } -#[rstest] -fn the_adapters_finalized_response_is_what_the_call_returns_and_reports( - #[from(interpreter)] _python: Interpreter, - #[values(false, true)] asynchronous: bool, -) { - Python::attach(|py| { - let (result, log) = run_scripted( - py, - success_machine(), - OpScript::Answer, - HookScript::ReplaceResponse, - asynchronous, - ); - assert_eq!(result.unwrap().extract::(py).unwrap(), "replaced"); - assert!(log.contains(&"succeeded:replaced".to_string())); - assert!(!log.contains(&"succeeded:project:1|rewritten".to_string())); - }); -} - #[rstest] fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding( #[from(interpreter)] _python: Interpreter, @@ -788,79 +736,23 @@ fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding( } #[rstest] -fn cancellation_ends_the_call_without_terminal_dispatch(#[from(interpreter)] _python: Interpreter) { +fn cancellation_ends_the_call_without_terminal_dispatch( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { Python::attach(|py| { - struct Cancelling(Log); - impl PythonBinding for Cancelling { - type Protocol = Synthetic; - type Failure = Classified; - fn decode_request( - &mut self, - _: Python<'_>, - _: &Bound<'_, PyDict>, - ) -> Result> { - self.0.push("project"); - Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into()) - } - - fn encode_stream_head( - &mut self, - _: Python<'_>, - head: std::convert::Infallible, - ) -> PyResult> { - match head {} - } - fn encode_chunk( - &mut self, - _: Python<'_>, - chunk: std::convert::Infallible, - ) -> PyResult> { - match chunk {} - } - fn encode_response(&mut self, _: Python<'_>, _: String) -> PyResult> { - Err(HostError::MissingState.into()) - } - fn map_error(&self, _: Python<'_>, error: Error) -> PyResult { - self.0.push("classify"); - Ok(Classified(error.0)) - } - fn host_error(error: &PyErr) -> Error { - Error(error.to_string()) - } - } - - impl PythonHostCalls for Cancelling { - fn handle_host_call( - &mut self, - _: Python<'_>, - _: (&'static str, Reply), - ) -> Result<(), InvokeError> { - Err(PyErr::from(HostError::MissingState).into()) - } - } - - impl PythonOwned for Cancelling { - fn close(&mut self, _: Python<'_>) {} - fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { - Ok(()) - } - } - let log = Log::default(); - let host = Cancelling(log.clone()); - let adapter = hooks(&log, HookScript::Plain); - let error = run_call( + let (result, log) = run_scripted( py, - |_, _, request| Ok(success_machine()(request)), - host, - adapter, - PyDict::new(py).unbind(), - call_options(false), - ) - .unwrap_err(); + success_machine(), + OpScript::CancelRequest, + HookScript::Plain, + asynchronous, + ); + let error = result.unwrap_err(); assert!(!error.is_instance_of::(py)); assert_eq!( - log.entries(), - ["started", "begin", "project", "adapter.close"] + log, + ["started", "begin", "project", "adapter.close", "host.close"] ); }); } diff --git a/litellm-rust/crates/host-python/tests/handle.rs b/litellm-rust/crates/host-python/tests/handle.rs index 37b27300973..916850f4b84 100644 --- a/litellm-rust/crates/host-python/tests/handle.rs +++ b/litellm-rust/crates/host-python/tests/handle.rs @@ -1,42 +1,22 @@ mod support; -use litellm_host_python::{Execution, ExecutionBody, ExecutionStep}; +use litellm_host::{hooks::CallHooks, lifecycle::Timing}; +use litellm_host_python::{ + Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime, + Suspension, run_call, +}; use pyo3::{ exceptions::PyBaseException, gc::{PyTraverseError, PyVisit}, prelude::*, - types::PyDict, + types::{PyCFunction, PyDict}, }; use rstest::rstest; -use support::{Interpreter, install_lifecycle_module, interpreter, lifecycle_binding}; +use support::{ + Interpreter, Log, OpScript, binding, call_options, install_lifecycle_module, interpreter, + lifecycle_binding, success_machine, +}; -#[rstest] -fn python_driver_preserves_inline_await_and_native_ownership( - #[from(interpreter)] _python: Interpreter, -) { - Python::attach(|py| { - py.import("asyncio").unwrap(); - let module = install_lifecycle_module(py); - 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!("lifecycle.py")).unwrap(); - py.run(&probe, Some(&locals), Some(&locals)).unwrap(); - }); -} struct RetainingHost { retained: Option>, } @@ -69,7 +49,7 @@ struct AwaitBody(Option>); impl ExecutionBody for AwaitBody { fn resume(&mut self, result: Option>>) -> PyResult { match self.0.take() { - Some(awaitable) => Ok(ExecutionStep::Await(awaitable)), + Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))), None => result .expect("selected await completed") .map(ExecutionStep::Return), @@ -124,6 +104,116 @@ fn error_execution(error: Bound<'_, PyBaseException>) -> Execution { Execution::new(ErrorBody(Some(error.unbind())), lifecycle_binding) } +/// Where a paused call keeps the test's object: in its hooks, or selected as the call's +/// response or exception, with the call paused in the hook that is told about it. +#[derive(Clone, Copy, Debug)] +enum Holder { + Hooks, + Response, + Exception, +} + +struct PausingHooks { + retained: Option>, + pause: Option>, + holder: Holder, +} + +impl CallHooks for PausingHooks { + fn transform_response( + &mut self, + py: Python<'_>, + response: Py, + _: Timing, + ) -> PyResult>> { + match self.holder { + Holder::Hooks => Ok(HookStep::Ready(response)), + Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())), + Holder::Exception => Err(PyErr::from_value( + self.retained.take().unwrap().into_bound(py), + )), + } + } + + fn on_event( + &mut self, + _: Python<'_>, + event: PythonCallEvent<'_>, + ) -> PyResult> { + let pauses_here = match self.holder { + Holder::Hooks => true, + Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }), + Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }), + }; + Ok(match self.pause.take_if(|_| pauses_here) { + Some(pause) => { + HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(()))) + } + None => HookStep::Ready(()), + }) + } +} + +impl PythonOwned for PausingHooks { + fn close(&mut self, _: Python<'_>) { + self.retained = None; + self.pause = None; + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.retained)?; + visit.call(&self.pause) + } +} + +fn paused_call( + py: Python<'_>, + retained: Py, + pause: Py, + holder: Holder, +) -> PyResult> { + let machine = success_machine(); + run_call( + py, + move |_, _, request| Ok(machine(request)), + binding(&Log::default(), OpScript::Answer), + PausingHooks { + retained: Some(retained), + pause: Some(pause), + holder, + }, + PyDict::new(py).unbind(), + call_options(true), + ) +} + +#[rstest] +fn every_suspension_is_awaited_inline_on_the_callers_task( + #[from(interpreter)] _python: Interpreter, +) { + Python::attach(|py| { + let module = install_lifecycle_module(py); + 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!("lifecycle.py")).unwrap(); + py.run(&probe, Some(&locals), Some(&locals)).unwrap(); + }); +} + #[rstest] fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) { Python::attach(|py| { @@ -199,3 +289,51 @@ assert retained_ref() is None .unwrap(); }); } + +#[rstest] +#[case::held_by_the_hooks(Holder::Hooks)] +#[case::selected_as_the_response(Holder::Response)] +#[case::selected_as_the_exception(Holder::Exception)] +fn cycles_through_a_call_suspended_in_a_hook_are_collectable( + #[from(interpreter)] _python: Interpreter, + #[case] holder: Holder, +) { + Python::attach(|py| { + let locals = PyDict::new(py); + let call = PyCFunction::new_closure(py, None, None, move |args, _| { + let (retained, pause) = args.extract::<(Py, Py)>()?; + paused_call(args.py(), retained, pause, holder) + }) + .unwrap(); + locals.set_item("paused_call", call).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained(Exception): + pass + +class Pause: + def __await__(self): + yield + +def cycle(): + retained = Retained() + coroutine = paused_call(retained, Pause()) + coroutine.send(None) + retained.coroutine = coroutine + return weakref.ref(retained) + +reference = cycle() +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); +} diff --git a/litellm-rust/crates/host-python/tests/hook_chain.rs b/litellm-rust/crates/host-python/tests/hook_chain.rs index ca28bce5d17..1943cdfcb2f 100644 --- a/litellm-rust/crates/host-python/tests/hook_chain.rs +++ b/litellm-rust/crates/host-python/tests/hook_chain.rs @@ -504,7 +504,9 @@ impl litellm_host_python::ExecutionBody for Preparing { } HookStep::Await(awaitable, resume) => { self.resume = Some(resume); - Ok(litellm_host_python::ExecutionStep::Await(awaitable)) + Ok(litellm_host_python::ExecutionStep::Suspended( + litellm_host_python::Suspension::Await(awaitable), + )) } } })