From 327bdcaa7f978719df403e4c8a48dea8fd6bb672 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 14 Oct 2026 17:19:05 -0700 Subject: [PATCH] test(host-python): move driver tests out of the inline module (#46325) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/host-python/src/driver.rs | 1840 ----------------- .../crates/host-python/tests/driver.rs | 866 ++++++++ .../crates/host-python/tests/handle.rs | 201 ++ .../crates/host-python/tests/support/mod.rs | 656 ++++++ 4 files changed, 1723 insertions(+), 1840 deletions(-) create mode 100644 litellm-rust/crates/host-python/tests/driver.rs create mode 100644 litellm-rust/crates/host-python/tests/handle.rs create mode 100644 litellm-rust/crates/host-python/tests/support/mod.rs diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index c70f2a5f1a0..a1450ed6308 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -701,1843 +701,3 @@ where self.clear(); } } - -#[cfg(test)] -mod tests { - use std::sync::{Arc, Mutex}; - - use litellm_host::{ - interceptors::{RawResponse, RequestContext}, - machine::{CallMachine, MachineFault}, - }; - use pyo3::exceptions::{PyBaseException, PyValueError}; - use pyo3::types::PyDict; - - use super::*; - use crate::{PythonOwned, PythonRuntime}; - use litellm_host::hooks::CallHooks; - use litellm_host::interceptors::Interceptors; - - static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); - - fn lifecycle_binding(py: Python<'_>) -> PyResult> { - py.import("host_test_lifecycle") - } - - fn call_options(asynchronous: bool) -> CallOptions { - CallOptions::new(asynchronous, lifecycle_binding) - } - - fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> { - 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!("host_test_lifecycle"), - ) - .unwrap() - } - - #[derive(Clone, Debug, PartialEq, Eq)] - struct Error(String); - - impl std::fmt::Display for Error { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str(&self.0) - } - } - - impl From for Error { - fn from(fault: MachineFault) -> Self { - Self(format!("{fault:?}")) - } - } - - struct Synthetic; - - impl Protocol for Synthetic { - type Response = String; - type Error = Error; - type Request = String; - type HostCall = (&'static str, Reply); - type Chunk = std::convert::Infallible; - type StreamHead = std::convert::Infallible; - } - - fn wire() -> WireRequest { - WireRequest { - url: "https://example.invalid".into(), - headers: Vec::new(), - body: serde_json::json!({}), - } - } - - fn context() -> RequestContext { - RequestContext { - model: "model".into(), - custom_llm_provider: "provider".into(), - optional_params: serde_json::json!({}), - secret_fields: Vec::new(), - api_key: None, - } - } - - #[derive(Default)] - struct Log(Arc>>); - - impl Log { - fn push(&self, entry: impl Into) { - self.0.lock().unwrap().push(entry.into()); - } - - fn entries(&self) -> Vec { - self.0.lock().unwrap().clone() - } - } - - #[derive(Clone, Copy)] - enum OpScript { - Answer, - RaisePython, - RejectNatively, - RejectRequestNatively, - RaiseRequestPython, - AwaitAnswer, - AwaitFailure, - } - - struct SyntheticBinding { - log: Log, - op: OpScript, - classifier_fails: bool, - pending_reply: Option>, - } - - /// The fake route's public exception, kept as a value so a test sees what `classify` - /// produced before the driver raises it. - #[derive(Debug, PartialEq, Eq)] - struct Classified(String); - - impl From for PyErr { - fn from(classified: Classified) -> Self { - PyValueError::new_err(format!("classified: {}", classified.0)) - } - } - - impl SyntheticBinding { - fn answer(&self, value: impl FnOnce() -> String) -> Result> { - match self.op { - OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()), - OpScript::RaisePython | OpScript::RaiseRequestPython => { - Err(PyValueError::new_err("op failed").into()) - } - OpScript::RejectNatively | OpScript::RejectRequestNatively => { - Err(InvokeError::Native(Error("op rejected".into()))) - } - } - } - } - - impl PythonBinding for SyntheticBinding { - type Protocol = Synthetic; - type Failure = Classified; - - fn decode_request( - &mut self, - _: Python<'_>, - arguments: &Bound<'_, PyDict>, - ) -> Result> { - self.log.push("project"); - match self.op { - OpScript::RejectRequestNatively | OpScript::RaiseRequestPython => { - self.answer(String::new) - } - _ => Ok(format!("project:{}", arguments.len())), - } - } - - 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, py: Python<'_>, response: String) -> PyResult> { - self.log.push("complete"); - Ok(pyo3::types::PyString::new(py, &response) - .into_any() - .unbind()) - } - - fn map_error(&self, _: Python<'_>, error: Error) -> PyResult { - self.log.push(format!("classify:{error}")); - if self.classifier_fails { - return Err(pyo3::exceptions::PyTypeError::new_err("classifier failed")); - } - Ok(Classified(error.0)) - } - - fn host_error(error: &PyErr) -> Error { - Error(error.to_string()) - } - } - - impl PythonHostCalls for SyntheticBinding { - fn handle_host_call( - &mut self, - _: Python<'_>, - (op, reply): (&'static str, Reply), - ) -> Result<(), InvokeError> { - self.log.push(format!("op:{op}")); - self.answer(|| op.to_string()) - .map(|answer| reply.send(answer)) - } - - fn begin_host_call( - &mut self, - py: Python<'_>, - (op, reply): (&'static str, Reply), - ) -> Result>, InvokeError> { - if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) { - return self.handle_host_call(py, (op, reply)).map(|()| None); - } - self.pending_reply = Some(reply); - let module = PyModule::from_code( - py, - pyo3::ffi::c_str!( - "async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n" - ), - pyo3::ffi::c_str!("host_op.py"), - pyo3::ffi::c_str!("host_op"), - )?; - Ok(Some( - module - .getattr("answer")? - .call1((matches!(self.op, OpScript::AwaitFailure),))? - .unbind(), - )) - } - - fn resume_host_call( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> Result>, InvokeError> { - let answer = result?.extract::(py)?; - self.pending_reply.take().unwrap().send(answer); - Ok(None) - } - } - - impl PythonOwned for SyntheticBinding { - fn close(&mut self, _: Python<'_>) { - self.log.push("host.close"); - } - fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { - Ok(()) - } - } - - #[derive(Clone, Copy)] - enum HookScript { - Plain, - RewriteArguments, - ObserveArguments, - FailBegin, - ReplaceResponse, - FailAfterSuccess, - CancelTerminal, - FailTerminal, - } - - struct SyntheticHooks { - log: Log, - script: HookScript, - } - - impl CallHooks for SyntheticHooks { - fn prepare_arguments( - &mut self, - _: Python<'_>, - arguments: Py, - _: f64, - ) -> PyResult>> { - self.log.push("begin"); - if matches!(self.script, HookScript::FailBegin) { - return Err(PyValueError::new_err("begin failed")); - } - if matches!(self.script, HookScript::RewriteArguments) { - let prepared = Python::attach(|py| -> PyResult> { - let prepared = arguments.bind(py).copy()?; - prepared.set_item("prepared", "hook")?; - prepared.set_item("api_key", "hook-key")?; - Ok(prepared.unbind()) - })?; - return Ok(HookStep::Ready(prepared)); - } - Ok(HookStep::Ready(arguments)) - } - - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - if matches!(self.script, HookScript::ObserveArguments) { - let value: String = arguments - .bind(py) - .get_item("prepared")? - .unwrap() - .extract()?; - self.log.push(format!("adopted:{value}")); - } - Ok(()) - } - - fn before_provider_request( - &mut self, - _: Python<'_>, - wire: Box, - _: &RequestContext, - ) -> PyResult>> { - self.log.push("before_provider_request"); - Ok(HookStep::Ready(Box::new(WireRequest { - url: "rewritten".into(), - ..*wire - }))) - } - - fn transform_response( - &mut self, - py: Python<'_>, - response: Py, - _: Timing, - ) -> PyResult>> { - self.log.push("after_success"); - match self.script { - HookScript::ReplaceResponse => Ok(HookStep::Ready( - "replaced".into_pyobject(py)?.into_any().unbind(), - )), - HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")), - HookScript::Plain - | HookScript::ObserveArguments - | HookScript::RewriteArguments - | HookScript::FailBegin - | HookScript::CancelTerminal - | HookScript::FailTerminal => Ok(HookStep::Ready(response)), - } - } - - fn on_event( - &mut self, - py: Python<'_>, - event: PythonCallEvent<'_>, - ) -> PyResult> { - if matches!(self.script, HookScript::CancelTerminal) - && matches!( - event, - PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. } - ) - { - return Err(pyo3::exceptions::asyncio::CancelledError::new_err( - "callback cancelled", - )); - } - if matches!(self.script, HookScript::FailTerminal) - && matches!( - event, - PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. } - ) - { - return Err(PyValueError::new_err("callback failed")); - } - self.log.push(match event { - PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => { - "cache_hit".into() - } - PythonCallEvent::Started { .. } => "started".into(), - PythonCallEvent::Cancelled { .. } => "cancelled".into(), - PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - format!("response:{}", raw.body) - } - PythonCallEvent::Succeeded { response, .. } => { - format!("succeeded:{}", response.bind(py)) - } - PythonCallEvent::Failed { origin, error, .. } => { - format!("failed:{origin:?}:{}", error.value(py)) - } - }); - Ok(HookStep::Ready(())) - } - - fn on_stream_open(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { - self.log.push("opened"); - Ok(()) - } - - fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { - self.log.push("delivered"); - Ok(()) - } - } - - impl PythonOwned for SyntheticHooks { - fn close(&mut self, _: Python<'_>) { - self.log.push("adapter.close"); - } - fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { - Ok(()) - } - } - - fn run_scripted( - py: Python<'_>, - machine: impl FnOnce(String) -> CallMachine + Send + Sync + 'static, - op: OpScript, - script: HookScript, - asynchronous: bool, - ) -> (PyResult>, Vec) { - run_hosted( - py, - machine, - SyntheticBinding { - log: Log::default(), - op, - classifier_fails: false, - pending_reply: None, - }, - script, - asynchronous, - ) - } - - fn run_hosted( - py: Python<'_>, - machine: impl FnOnce(String) -> CallMachine + Send + Sync + 'static, - host: SyntheticBinding, - script: HookScript, - asynchronous: bool, - ) -> (PyResult>, Vec) { - run_composed( - py, - machine, - host, - script, - std::convert::identity, - call_options(asynchronous), - ) - } - - #[rstest::rstest] - #[case::success(OpScript::AwaitAnswer)] - #[case::failure(OpScript::AwaitFailure)] - fn asynchronous_host_operations_resume_the_same_machine(#[case] op: OpScript) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - Python::initialize(); - Python::attach(|py| { - install_lifecycle_module(py); - let (result, log) = run_scripted( - py, - |_| { - CallMachine::::new(None, |host| { - Box::pin(async move { host.services.call(|reply| ("read", reply)).await }) - }) - }, - op, - HookScript::Plain, - true, - ); - match op { - OpScript::AwaitAnswer => { - assert_eq!(result.unwrap().extract::(py).unwrap(), "awaited") - } - OpScript::AwaitFailure => { - assert!( - result - .unwrap_err() - .is_instance_of::(py) - ); - assert_eq!( - log.iter() - .filter(|entry| entry.starts_with("failed:")) - .count(), - 1 - ); - assert!(!log.iter().any(|entry| entry.starts_with("succeeded:"))); - } - _ => unreachable!(), - } - }); - } - - #[rstest::rstest] - #[case::synchronous(false)] - #[case::asynchronous(true)] - fn composed_hooks_share_prepared_arguments_and_one_execution(#[case] asynchronous: bool) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let log = Log::default(); - let hook = |script| SyntheticHooks { - log: Log(log.0.clone()), - script, - }; - let hooks = crate::HookChain::new() - .with(hook(HookScript::ObserveArguments)) - .with(hook(HookScript::RewriteArguments)) - .with(hook(HookScript::ReplaceResponse)); - let result = run_call( - py, - move |_, _, request| Ok(success_machine()(request)), - SyntheticBinding { - log: Log(log.0.clone()), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - hooks, - PyDict::new(py).unbind(), - call_options(asynchronous), - ) - .unwrap(); - let response = if asynchronous { - assert!(log.entries().is_empty()); - let completion = result.call_method1(py, "send", (py.None(),)).unwrap_err(); - assert!(completion.is_instance_of::(py)); - completion.value(py).getattr("value").unwrap().unbind() - } else { - result - }; - assert_eq!(response.extract::(py).unwrap(), "replaced"); - let entries = log.entries(); - assert!(entries.iter().any(|entry| entry == "adopted:hook")); - assert_eq!( - entries - .iter() - .filter(|entry| entry.as_str() == "project") - .count(), - 1 - ); - assert_eq!( - entries - .iter() - .filter(|entry| entry.as_str() == "op:sign") - .count(), - 1 - ); - assert_eq!( - entries - .iter() - .filter(|entry| entry.as_str() == "succeeded:replaced") - .count(), - 3 - ); - assert!(!entries.iter().any(|entry| entry.starts_with("failed:"))); - }); - } - - fn run_composed( - py: Python<'_>, - machine: impl FnOnce(String) -> CallMachine + Send + Sync + 'static, - host: SyntheticBinding, - script: HookScript, - compose: impl FnOnce(SyntheticHooks) -> L, - options: CallOptions, - ) -> (PyResult>, Vec) { - let asynchronous = options.asynchronous; - let log = Log(host.log.0.clone()); - let adapter = SyntheticHooks { - log: Log(log.0.clone()), - script, - }; - let arguments = PyDict::new(py); - arguments.set_item("model", "m").unwrap(); - let result = run_call( - py, - move |_, _, request| Ok(machine(request)), - host, - compose(adapter), - arguments.unbind(), - options, - ); - let result = if asynchronous { - result.and_then(|coroutine| { - let completed = coroutine - .call_method1(py, "send", (py.None(),)) - .unwrap_err(); - if !completed.is_instance_of::(py) { - return Err(completed); - } - completed.value(py).getattr("value").map(Bound::unbind) - }) - } else { - result - }; - (result, log.entries()) - } - - /// Answers to projection, to the route op and to `before_provider_request` all reach the - /// response, so a driver that misroutes a reply changes what the call returns. - fn success_machine() -> impl FnOnce(String) -> CallMachine + Send + Sync { - move |projected| { - CallMachine::new(None, move |host| { - Box::pin(async move { - let signed = host.services.call(|reply| ("sign", reply)).await?; - let wire = host - .interceptors - .before_provider_request(wire(), context()) - .await?; - host.interceptors - .after_provider_response(RawResponse { body: "raw".into() }) - .await?; - Ok(format!("{projected}|{signed}|{}", wire.url)) - }) - }) - } - } - - #[rstest::rstest] - #[case::draining(4, false)] - #[case::full(1, false)] - #[case::closed(4, true)] - fn observation_delivery_preserves_callback_order_and_response( - #[case] capacity: usize, - #[case] closed: bool, - #[values(false, true)] asynchronous: bool, - ) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let (expected, expected_log) = run_scripted( - py, - success_machine(), - OpScript::Answer, - HookScript::ReplaceResponse, - asynchronous, - ); - let (sender, mut receiver) = litellm_host::observation::observation_channel( - std::num::NonZeroUsize::new(capacity).unwrap(), - ); - if closed { - receiver.close(); - } - let (actual, actual_log) = run_composed( - py, - success_machine(), - SyntheticBinding { - log: Log::default(), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - HookScript::ReplaceResponse, - std::convert::identity, - CallOptions { - asynchronous, - lifecycle: lifecycle_binding, - observers: Some(sender.clone()), - }, - ); - assert_eq!( - actual.unwrap().extract::(py).unwrap(), - expected.unwrap().extract::(py).unwrap() - ); - assert_eq!(actual_log, expected_log); - let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect(); - assert_eq!(events.len() as u64 + sender.dropped_events(), 3); - if !closed && capacity >= 3 { - let [ - CallEvent::Started { start_time }, - CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }), - CallEvent::Succeeded { timing, .. }, - ] = events.as_slice() - else { - panic!("expected start, provider response and success: {events:?}"); - }; - assert_eq!(raw.body, "raw"); - assert_eq!(*start_time, timing.start_time); - assert!(timing.end_time >= timing.start_time); - } - }); - } - - #[rstest::rstest] - #[case::preparation(HookScript::FailBegin, false, 1)] - #[case::transformation(HookScript::FailAfterSuccess, false, 1)] - #[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)] - #[case::terminal_failure(HookScript::FailTerminal, false, 0)] - fn observation_reports_the_final_outcome_without_replaying_callbacks( - #[case] script: HookScript, - #[case] cancelled: bool, - #[case] failure_callbacks: usize, - #[values(false, true)] asynchronous: bool, - ) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let (sender, mut receiver) = litellm_host::observation::observation_channel( - std::num::NonZeroUsize::new(4).unwrap(), - ); - let (result, log) = run_composed( - py, - success_machine(), - SyntheticBinding { - log: Log::default(), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - script, - std::convert::identity, - CallOptions { - asynchronous, - lifecycle: lifecycle_binding, - observers: Some(sender), - }, - ); - let error = result.unwrap_err(); - assert_eq!( - error.is_instance_of::(py), - cancelled - ); - let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect(); - assert!(matches!(events.first(), Some(CallEvent::Started { .. }))); - assert_eq!( - matches!(events.last(), Some(CallEvent::Cancelled { .. })), - cancelled - ); - assert_eq!( - matches!( - events.last(), - Some(CallEvent::Failed { - origin: FailureOrigin::Host, - .. - }) - ), - !cancelled - ); - assert_eq!( - events - .iter() - .filter(|event| matches!( - event, - CallEvent::Failed { .. } - | CallEvent::Succeeded { .. } - | CallEvent::Cancelled { .. } - )) - .count(), - 1 - ); - assert_eq!( - log.iter() - .filter(|entry| entry.starts_with("failed:")) - .count(), - failure_callbacks - ); - assert!( - !log.iter() - .any(|entry| entry.starts_with("succeeded:") || entry == "cancelled") - ); - }); - } - - #[rstest::rstest] - #[case::preparation(HookScript::FailBegin, OpScript::Answer)] - #[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)] - #[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)] - fn startup_failures_release_the_factory_without_constructing_a_machine( - #[case] script: HookScript, - #[case] op: OpScript, - #[values(false, true)] asynchronous: bool, - ) { - use std::sync::atomic::{AtomicBool, Ordering}; - struct Release(Arc); - impl Drop for Release { - fn drop(&mut self) { - self.0.store(true, Ordering::SeqCst); - } - } - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let constructed = Arc::new(AtomicBool::new(false)); - let did_construct = constructed.clone(); - let released = Arc::new(AtomicBool::new(false)); - let release = Release(released.clone()); - let (result, log) = run_scripted( - py, - move |request| { - let _release = release; - did_construct.store(true, Ordering::SeqCst); - success_machine()(request) - }, - op, - script, - asynchronous, - ); - assert!(result.is_err()); - assert!(!constructed.load(Ordering::SeqCst)); - assert!(released.load(Ordering::SeqCst)); - assert_eq!( - log.iter() - .filter(|event| event.starts_with("failed:")) - .count(), - 1 - ); - assert_eq!( - log.iter().filter(|event| *event == "adapter.close").count(), - 1 - ); - assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1); - }); - } - - #[rstest::rstest] - fn success_runs_every_step_in_order_and_returns_the_public_response() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - for asynchronous in [false, true] { - let (result, log) = run_scripted( - py, - success_machine(), - OpScript::Answer, - HookScript::Plain, - asynchronous, - ); - assert_eq!( - result.unwrap().extract::(py).unwrap(), - "project:1|sign|rewritten" - ); - assert_eq!( - log, - [ - "started", - "begin", - "project", - "op:sign", - "before_provider_request", - "response:raw", - "complete", - "after_success", - "succeeded:project:1|sign|rewritten", - "adapter.close", - "host.close", - ] - ); - } - }); - } - - struct Streaming; - - impl Protocol for Streaming { - type Response = (); - type Error = Error; - type Request = (); - type HostCall = std::convert::Infallible; - type Chunk = &'static str; - type StreamHead = Vec<(&'static str, &'static str)>; - } - - struct StreamingBinding; - - impl PythonBinding for StreamingBinding { - type Protocol = Streaming; - type Failure = Classified; - - fn decode_request( - &mut self, - _: Python<'_>, - _: &Bound<'_, PyDict>, - ) -> Result<(), InvokeError> { - Ok(()) - } - - fn encode_stream_head( - &mut self, - py: Python<'_>, - head: Vec<(&'static str, &'static str)>, - ) -> PyResult> { - let headers = PyDict::new(py); - for (name, value) in head { - headers.set_item(name, value)?; - } - let hidden = PyDict::new(py); - hidden.set_item("additional_headers", headers)?; - Ok(hidden.into_any().unbind()) - } - - fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult> { - Ok(pyo3::types::PyString::new(py, chunk).into_any().unbind()) - } - - fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult> { - Ok(py.None()) - } - - fn map_error(&self, _: Python<'_>, error: Error) -> PyResult { - Ok(Classified(error.0)) - } - - fn host_error(error: &PyErr) -> Error { - Error(error.to_string()) - } - } - - impl PythonHostCalls for StreamingBinding { - fn handle_host_call( - &mut self, - _: Python<'_>, - op: std::convert::Infallible, - ) -> Result<(), InvokeError> { - match op {} - } - } - - impl PythonOwned for StreamingBinding { - fn close(&mut self, _: Python<'_>) {} - fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { - Ok(()) - } - } - - fn streaming_machine() - -> impl FnOnce(()) -> litellm_host::call::HostedMachine + Send + Sync { - move |()| { - litellm_host::call::hosted_call((), None, |(), _, _, _observations| async { - Ok(litellm_host::call::CallOutput::Stream { - head: vec![("request-id", "req_1")], - chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])), - }) - }) - } - } - - #[rstest::rstest] - #[case::sync(false)] - #[case::asynchronous(true)] - fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks( - #[case] asynchronous: bool, - ) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let log = Log::default(); - let (sender, mut receiver) = litellm_host::observation::observation_channel( - std::num::NonZeroUsize::new(3).unwrap(), - ); - let handed = run_call( - py, - |_, _, request| Ok(streaming_machine()(request)), - StreamingBinding, - SyntheticHooks { - log: Log(log.0.clone()), - script: HookScript::Plain, - }, - PyDict::new(py).unbind(), - CallOptions { - asynchronous, - lifecycle: lifecycle_binding, - observers: Some(sender), - }, - ) - .unwrap(); - let stream = if asynchronous { - let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err(); - stop.value(py).getattr("value").unwrap() - } else { - handed.into_bound(py) - }; - if asynchronous { - for method in ["__anext__", "aclose", "aclose"] { - let stop = stream - .call_method0(method) - .unwrap() - .call_method1("send", (py.None(),)) - .unwrap_err(); - assert!(stop.is_instance_of::(py)); - } - } else { - assert_eq!( - stream - .call_method0("__next__") - .unwrap() - .extract::() - .unwrap(), - "first" - ); - stream.call_method0("close").unwrap(); - stream.call_method0("close").unwrap(); - } - let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect(); - assert!(matches!( - events.as_slice(), - [CallEvent::Started { .. }, CallEvent::Succeeded { .. }] - )); - assert_eq!( - log.entries(), - [ - "started", - "begin", - "opened", - "delivered", - "succeeded:None", - "adapter.close" - ] - ); - }); - } - - /// Drives a `Stream` (async) or `SyncStream` to completion from a sync test. - fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec { - if !asynchronous { - return stream - .try_iter() - .unwrap() - .map(|chunk| chunk.unwrap().extract().unwrap()) - .collect(); - } - std::iter::from_fn(|| { - let stop = stream - .call_method0("__anext__") - .unwrap() - .call_method1("send", (py.None(),)) - .unwrap_err(); - if stop.is_instance_of::(py) { - return None; - } - assert!(stop.is_instance_of::(py)); - Some(stop.value(py).getattr("value").unwrap().extract().unwrap()) - }) - .collect() - } - - #[rstest::rstest] - fn a_stream_carries_its_head_before_the_first_chunk() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - for asynchronous in [false, true] { - let log = Log::default(); - let adapter = SyntheticHooks { - log: Log(log.0.clone()), - script: HookScript::Plain, - }; - let handed = run_call( - py, - |_, _, request| Ok(streaming_machine()(request)), - StreamingBinding, - adapter, - PyDict::new(py).unbind(), - call_options(asynchronous), - ) - .unwrap(); - let stream = if asynchronous { - let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err(); - stop.value(py).getattr("value").unwrap() - } else { - handed.into_bound(py) - }; - let hidden: std::collections::HashMap< - String, - std::collections::HashMap, - > = stream.getattr("head").unwrap().extract().unwrap(); - assert_eq!( - hidden["additional_headers"], - std::collections::HashMap::from([( - "request-id".to_string(), - "req_1".to_string() - )]) - ); - assert_eq!(log.entries(), ["started", "begin", "opened"]); - assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]); - } - }); - } - - fn failing_machine() -> impl FnOnce(String) -> CallMachine + Send + Sync { - move |_| { - CallMachine::new(None, |_| { - Box::pin(async move { Err(Error("provider exploded".into())) }) - }) - } - } - - #[rstest::rstest] - fn a_native_failure_is_classified_once_and_reported_classified() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - for asynchronous in [false, true] { - 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::rstest] - fn a_native_rejection_from_a_host_operation_is_classified_once() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - 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::rstest] - fn a_python_exception_from_a_host_operation_is_reported_as_raised() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - 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::rstest] - fn a_failing_classifier_surfaces_with_the_native_error_as_context() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - let (result, log) = run_hosted( - py, - failing_machine(), - SyntheticBinding { - log: Log::default(), - op: OpScript::Answer, - classifier_fails: true, - pending_reply: None, - }, - HookScript::Plain, - false, - ); - 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", - ] - ); - }); - } - - #[rstest::rstest] - fn begin_failures_are_host_failures_without_provider_mapping() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - let (result, log) = run_scripted( - py, - success_machine(), - OpScript::Answer, - HookScript::FailBegin, - false, - ); - let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "begin failed"); - assert_eq!( - log, - [ - "started", - "begin", - "failed:Host:begin failed", - "adapter.close", - "host.close" - ] - ); - }); - } - - enum ArgumentPolicy { - Inherit, - Reject(Py), - } - - impl CallHooks for ArgumentPolicy { - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - match self { - Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"), - Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())), - } - } - } - - impl PythonOwned for ArgumentPolicy { - fn close(&mut self, _: Python<'_>) {} - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - match self { - Self::Inherit => Ok(()), - Self::Reject(error) => visit.call(error), - } - } - } - - #[rstest::rstest] - #[case::sync_success(false, false)] - #[case::async_success(true, false)] - #[case::sync_failure(false, true)] - #[case::async_failure(true, true)] - fn resource_setup_uses_prepared_arguments_and_failures_are_terminal( - #[case] asynchronous: bool, - #[case] fail_setup: bool, - ) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let log = Log::default(); - let setup_log = Log(log.0.clone()); - let error = PyValueError::new_err("setup failed").into_value(py); - let setup_error = error.clone_ref(py); - let arguments = PyDict::new(py); - arguments.set_item("api_key", "original").unwrap(); - let result = run_call( - py, - move |py, prepared, request| { - let source: String = prepared.get_item("prepared")?.unwrap().extract()?; - let key: String = prepared.get_item("api_key")?.unwrap().extract()?; - setup_log.push(format!("setup:{source}:{key}")); - if fail_setup { - return Err(PyErr::from_value(setup_error.into_bound(py).into_any())); - } - Ok(success_machine()(request)) - }, - SyntheticBinding { - log: Log(log.0.clone()), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - crate::HookChain::new() - .with(SyntheticHooks { - log: Log(log.0.clone()), - script: HookScript::RewriteArguments, - }) - .with(ArgumentPolicy::Inherit), - arguments.clone().unbind(), - call_options(asynchronous), - ); - let settled = if asynchronous { - assert!(log.entries().is_empty()); - let completion = result - .unwrap() - .call_method1(py, "send", (py.None(),)) - .unwrap_err(); - if completion.is_instance_of::(py) { - completion.value(py).getattr("value").map(Bound::unbind) - } else { - Err(completion) - } - } else { - result - }; - assert_eq!( - arguments - .get_item("api_key") - .unwrap() - .unwrap() - .extract::() - .unwrap(), - "original" - ); - assert_eq!( - &log.entries()[..4], - ["started", "begin", "project", "setup:hook:inherited"] - ); - if fail_setup { - assert!(settled.unwrap_err().value(py).is(error.bind(py))); - assert_eq!( - log.entries(), - [ - "started", - "begin", - "project", - "setup:hook:inherited", - "failed:Host:setup failed", - "adapter.close", - "host.close", - ] - ); - } else { - assert!(settled.is_ok()); - assert!( - log.entries() - .iter() - .any(|entry| entry.starts_with("succeeded:")) - ); - } - }); - } - - #[rstest::rstest] - fn closing_an_unstarted_call_never_initializes_resources() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let log = Log::default(); - let setup_log = Log(log.0.clone()); - let pending = run_call( - py, - move |_, _, request| { - setup_log.push("setup"); - Ok(success_machine()(request)) - }, - SyntheticBinding { - log: Log(log.0.clone()), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - SyntheticHooks { - log: Log(log.0.clone()), - script: HookScript::Plain, - }, - PyDict::new(py).unbind(), - call_options(true), - ) - .unwrap(); - assert!(log.entries().is_empty()); - pending.call_method0(py, "close").unwrap(); - drop(pending); - assert_eq!(log.entries(), ["adapter.close", "host.close"]); - }); - } - - #[rstest::rstest] - #[case::synchronous(false)] - #[case::asynchronous(true)] - fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts( - #[case] asynchronous: bool, - ) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let raised = PyValueError::new_err("over budget").into_value(py); - let (result, log) = run_composed( - py, - success_machine(), - SyntheticBinding { - log: Log::default(), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - HookScript::Plain, - |hooks| { - crate::HookChain::new() - .with(hooks) - .with(ArgumentPolicy::Reject(raised.clone_ref(py))) - }, - call_options(asynchronous), - ); - assert!(result.unwrap_err().value(py).is(raised.bind(py))); - assert_eq!( - log, - [ - "started", - "begin", - "failed:Host:over budget", - "adapter.close", - "host.close" - ] - ); - }); - } - - #[rstest::rstest] - #[case::synchronous(false)] - #[case::asynchronous(true)] - fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote( - #[case] asynchronous: bool, - ) { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - let (result, _) = run_composed( - py, - success_machine(), - SyntheticBinding { - log: Log::default(), - op: OpScript::Answer, - classifier_fails: false, - pending_reply: None, - }, - HookScript::Plain, - |hooks| { - crate::HookChain::new() - .with(hooks) - .with(ArgumentPolicy::Inherit) - }, - call_options(asynchronous), - ); - assert_eq!( - result.unwrap().extract::(py).unwrap(), - "project:2|sign|rewritten" - ); - }); - } - - #[rstest::rstest] - fn the_adapters_finalized_response_is_what_the_call_returns_and_reports() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - for asynchronous in [false, true] { - 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::rstest] - fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - Python::attach(|py| { - install_lifecycle_module(py); - for asynchronous in [false, true] { - let (result, log) = run_scripted( - py, - success_machine(), - OpScript::Answer, - HookScript::FailAfterSuccess, - asynchronous, - ); - let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "after_success failed"); - assert_eq!( - &log[log.len() - 4..], - [ - "after_success", - "failed:Host:after_success failed", - "adapter.close", - "host.close" - ] - ); - assert!(!log.iter().any(|entry| entry.starts_with("succeeded"))); - } - }); - } - - #[rstest::rstest] - fn cancellation_ends_the_call_without_terminal_dispatch() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - 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(missing_state()) - } - 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(missing_state().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(log.0.clone())); - let adapter = SyntheticHooks { - log: Log(log.0.clone()), - script: HookScript::Plain, - }; - let error = run_call( - py, - |_, _, request| Ok(success_machine()(request)), - host, - adapter, - PyDict::new(py).unbind(), - call_options(false), - ) - .unwrap_err(); - assert!(!error.is_instance_of::(py)); - assert_eq!( - log.entries(), - ["started", "begin", "project", "adapter.close"] - ); - }); - } - - #[rstest::rstest] - fn python_driver_preserves_inline_await_and_native_ownership() { - let _guard = PYTHON_GLOBALS - .lock() - .unwrap_or_else(|error| error.into_inner()); - crate::initialize_python(); - 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!("../tests/lifecycle.py")).unwrap(); - py.run(&probe, Some(&locals), Some(&locals)).unwrap(); - }); - } - 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), - }, - lifecycle_binding, - ), - ) - } - - 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)), lifecycle_binding) - } - - 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), lifecycle_binding) - } - - struct ErrorBody(Option>); - - impl ExecutionBody for ErrorBody { - fn resume(&mut self, _: Option>>) -> PyResult { - Python::attach(|py| { - Err(PyErr::from_value( - self.0.take().unwrap().into_bound(py).into_any(), - )) - }) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.0) - } - } - - #[pyfunction] - fn error_execution(error: Bound<'_, PyBaseException>) -> Execution { - Execution::new(ErrorBody(Some(error.unbind())), lifecycle_binding) - } - - #[rstest::rstest] - fn retained_exception_frames_are_collectable() { - crate::initialize_python(); - 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(); - }); - } - - #[rstest::rstest] - fn coroutine_collects_cycles_retained_by_bridge_host() { - crate::initialize_python(); - 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/host-python/tests/driver.rs b/litellm-rust/crates/host-python/tests/driver.rs new file mode 100644 index 00000000000..51a3653c2e2 --- /dev/null +++ b/litellm-rust/crates/host-python/tests/driver.rs @@ -0,0 +1,866 @@ +mod support; + +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use litellm_host::{ + hooks::CallHooks, + lifecycle::{CallEvent, ExecutionEvent, FailureOrigin}, + machine::CallMachine, + observation::observation_channel, + protocol::Reply, +}; +use litellm_host_python::{ + CallOptions, HookChain, InvokeError, PythonBinding, PythonHostCalls, PythonOwned, + PythonRuntime, missing_state, run_call, +}; +use pyo3::{ + exceptions::{ + PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError, + PyValueError, asyncio::CancelledError, + }, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; +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, +}; + +fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec) { + let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap()); + let options = CallOptions { + asynchronous: false, + lifecycle: lifecycle_binding, + observers: Some(sender), + }; + (options, move || { + std::iter::from_fn(|| receiver.try_recv().ok()).collect() + }) +} + +#[rstest] +#[case::success(OpScript::AwaitAnswer)] +#[case::failure(OpScript::AwaitFailure)] +fn asynchronous_host_operations_resume_the_same_machine( + #[from(interpreter)] _python: Interpreter, + #[case] op: OpScript, +) { + Python::attach(|py| { + let machine: MachineFactory = Box::new(|_| { + CallMachine::::new(None, |host| { + Box::pin(async move { host.services.call(|reply| ("read", reply)).await }) + }) + }); + let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true); + match op { + OpScript::AwaitAnswer => { + assert_eq!(result.unwrap().extract::(py).unwrap(), "awaited") + } + OpScript::AwaitFailure => { + assert!(result.unwrap_err().is_instance_of::(py)); + assert_eq!( + log.iter() + .filter(|entry| entry.starts_with("failed:")) + .count(), + 1 + ); + assert!(!log.iter().any(|entry| entry.starts_with("succeeded:"))); + } + _ => unreachable!(), + } + }); +} + +#[rstest] +fn composed_hooks_share_prepared_arguments_and_one_execution( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let log = Log::default(); + let chain = HookChain::new() + .with(hooks(&log, HookScript::ObserveArguments)) + .with(hooks(&log, HookScript::RewriteArguments)) + .with(hooks(&log, HookScript::ReplaceResponse)); + let result = run_call( + py, + move |_, _, request| Ok(success_machine()(request)), + binding(&log, OpScript::Answer), + chain, + PyDict::new(py).unbind(), + call_options(asynchronous), + ); + if asynchronous { + assert!(log.entries().is_empty()); + } + let response = settled(py, result, asynchronous).unwrap(); + assert_eq!(response.extract::(py).unwrap(), "replaced"); + let entries = log.entries(); + assert!(entries.iter().any(|entry| entry == "adopted:hook")); + assert_eq!( + entries.iter().filter(|entry| *entry == "project").count(), + 1 + ); + assert_eq!( + entries.iter().filter(|entry| *entry == "op:sign").count(), + 1 + ); + assert_eq!( + entries + .iter() + .filter(|entry| *entry == "succeeded:replaced") + .count(), + 3 + ); + assert!(!entries.iter().any(|entry| entry.starts_with("failed:"))); + }); +} + +#[rstest] +#[case::draining(4, false)] +#[case::full(1, false)] +#[case::closed(4, true)] +fn observation_delivery_preserves_callback_order_and_response( + #[from(interpreter)] _python: Interpreter, + #[case] capacity: usize, + #[case] closed: bool, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let (expected, expected_log) = run_scripted( + py, + success_machine(), + OpScript::Answer, + HookScript::ReplaceResponse, + asynchronous, + ); + let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap()); + if closed { + receiver.close(); + } + let log = Log::default(); + let (actual, actual_log) = run_composed( + py, + success_machine(), + binding(&log, OpScript::Answer), + hooks(&log, HookScript::ReplaceResponse), + std::convert::identity, + CallOptions { + asynchronous, + lifecycle: lifecycle_binding, + observers: Some(sender.clone()), + }, + ); + assert_eq!( + actual.unwrap().extract::(py).unwrap(), + expected.unwrap().extract::(py).unwrap() + ); + assert_eq!(actual_log, expected_log); + let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect(); + assert_eq!(events.len() as u64 + sender.dropped_events(), 3); + if !closed && capacity >= 3 { + let [ + CallEvent::Started { start_time }, + CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }), + CallEvent::Succeeded { timing, .. }, + ] = events.as_slice() + else { + panic!("expected start, provider response and success: {events:?}"); + }; + assert_eq!(raw.body, "raw"); + assert_eq!(*start_time, timing.start_time); + assert!(timing.end_time >= timing.start_time); + } + }); +} + +#[rstest] +#[case::preparation(HookScript::FailBegin, false, 1)] +#[case::transformation(HookScript::FailAfterSuccess, false, 1)] +#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)] +#[case::terminal_failure(HookScript::FailTerminal, false, 0)] +fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks( + #[from(interpreter)] _python: Interpreter, + #[case] script: HookScript, + #[case] cancelled: bool, + #[case] failure_callbacks: usize, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let (options, mut events) = observed(4); + let log = Log::default(); + let (result, log) = run_composed( + py, + success_machine(), + binding(&log, OpScript::Answer), + hooks(&log, script), + std::convert::identity, + CallOptions { + asynchronous, + ..options + }, + ); + let error = result.unwrap_err(); + assert_eq!(error.is_instance_of::(py), cancelled); + let events = events(); + assert!(matches!(events.first(), Some(CallEvent::Started { .. }))); + assert_eq!( + matches!(events.last(), Some(CallEvent::Cancelled { .. })), + cancelled + ); + assert_eq!( + matches!( + events.last(), + Some(CallEvent::Failed { + origin: FailureOrigin::Host, + .. + }) + ), + !cancelled + ); + let terminal = events.iter().filter(|event| { + matches!( + event, + CallEvent::Failed { .. } + | CallEvent::Succeeded { .. } + | CallEvent::Cancelled { .. } + ) + }); + assert_eq!(terminal.count(), 1); + assert_eq!( + log.iter() + .filter(|entry| entry.starts_with("failed:")) + .count(), + failure_callbacks + ); + assert!( + !log.iter() + .any(|entry| entry.starts_with("succeeded:") || entry == "cancelled") + ); + }); +} + +#[rstest] +#[case::preparation(HookScript::FailBegin, OpScript::Answer)] +#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)] +#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)] +fn startup_failures_release_the_factory_without_constructing_a_machine( + #[from(interpreter)] _python: Interpreter, + #[case] script: HookScript, + #[case] op: OpScript, + #[values(false, true)] asynchronous: bool, +) { + struct Release(Arc); + impl Drop for Release { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + Python::attach(|py| { + let constructed = Arc::new(AtomicBool::new(false)); + let did_construct = constructed.clone(); + let released = Arc::new(AtomicBool::new(false)); + let release = Release(released.clone()); + let machine: MachineFactory = Box::new(move |request| { + let _release = release; + did_construct.store(true, Ordering::SeqCst); + success_machine()(request) + }); + let (result, log) = run_scripted(py, machine, op, script, asynchronous); + assert!(result.is_err()); + assert!(!constructed.load(Ordering::SeqCst)); + assert!(released.load(Ordering::SeqCst)); + assert_eq!( + log.iter() + .filter(|event| event.starts_with("failed:")) + .count(), + 1 + ); + assert_eq!( + log.iter().filter(|event| *event == "adapter.close").count(), + 1 + ); + assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1); + }); +} + +#[rstest] +fn success_runs_every_step_in_order_and_returns_the_public_response( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::Answer, + HookScript::Plain, + asynchronous, + ); + assert_eq!( + result.unwrap().extract::(py).unwrap(), + "project:1|sign|rewritten" + ); + assert_eq!( + log, + [ + "started", + "begin", + "project", + "op:sign", + "before_provider_request", + "response:raw", + "complete", + "after_success", + "succeeded:project:1|sign|rewritten", + "adapter.close", + "host.close", + ] + ); + }); +} + +#[rstest] +fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let log = Log::default(); + let (options, mut events) = observed(3); + let stream = open_stream( + py, + hooks(&log, HookScript::Plain), + CallOptions { + asynchronous, + ..options + }, + ); + if asynchronous { + for method in ["__anext__", "aclose", "aclose"] { + let stop = stream + .call_method0(method) + .unwrap() + .call_method1("send", (py.None(),)) + .unwrap_err(); + assert!(stop.is_instance_of::(py)); + } + } else { + assert_eq!( + stream + .call_method0("__next__") + .unwrap() + .extract::() + .unwrap(), + "first" + ); + stream.call_method0("close").unwrap(); + stream.call_method0("close").unwrap(); + } + assert!(matches!( + events().as_slice(), + [CallEvent::Started { .. }, CallEvent::Succeeded { .. }] + )); + assert_eq!( + log.entries(), + [ + "started", + "begin", + "opened", + "delivered", + "succeeded:None", + "adapter.close" + ] + ); + }); +} + +#[rstest] +fn a_stream_carries_its_head_before_the_first_chunk( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let log = Log::default(); + let stream = open_stream( + py, + hooks(&log, HookScript::Plain), + call_options(asynchronous), + ); + let hidden: std::collections::HashMap> = + stream.getattr("head").unwrap().extract().unwrap(); + assert_eq!( + hidden["additional_headers"], + std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())]) + ); + assert_eq!(log.entries(), ["started", "begin", "opened"]); + assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]); + }); +} + +#[rstest] +fn a_native_failure_is_classified_once_and_reported_classified( + #[from(interpreter)] _python: Interpreter, + #[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 (result, log) = run_composed( + py, + failing_machine(), + binding(&log_handle, OpScript::Answer).with_failing_classifier(), + hooks(&log_handle, HookScript::Plain), + std::convert::identity, + call_options(false), + ); + 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", + ] + ); + }); +} + +#[rstest] +fn begin_failures_are_host_failures_without_provider_mapping( + #[from(interpreter)] _python: Interpreter, +) { + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::Answer, + HookScript::FailBegin, + false, + ); + assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed"); + assert_eq!( + log, + [ + "started", + "begin", + "failed:Host:begin failed", + "adapter.close", + "host.close" + ] + ); + }); +} + +enum ArgumentPolicy { + Inherit, + Reject(Py), +} + +impl CallHooks for ArgumentPolicy { + fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { + match self { + Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"), + Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())), + } + } +} + +impl PythonOwned for ArgumentPolicy { + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + match self { + Self::Inherit => Ok(()), + Self::Reject(error) => visit.call(error), + } + } +} + +#[rstest] +fn resource_setup_uses_prepared_arguments_and_failures_are_terminal( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, + #[values(false, true)] fail_setup: bool, +) { + Python::attach(|py| { + let log = Log::default(); + let setup_log = log.clone(); + let error = PyValueError::new_err("setup failed").into_value(py); + let setup_error = error.clone_ref(py); + let arguments = PyDict::new(py); + arguments.set_item("api_key", "original").unwrap(); + let result = run_call( + py, + move |py, prepared, request| { + let source: String = prepared.get_item("prepared")?.unwrap().extract()?; + let key: String = prepared.get_item("api_key")?.unwrap().extract()?; + setup_log.push(format!("setup:{source}:{key}")); + if fail_setup { + return Err(PyErr::from_value(setup_error.into_bound(py).into_any())); + } + Ok(success_machine()(request)) + }, + binding(&log, OpScript::Answer), + HookChain::new() + .with(hooks(&log, HookScript::RewriteArguments)) + .with(ArgumentPolicy::Inherit), + arguments.clone().unbind(), + call_options(asynchronous), + ); + if asynchronous { + assert!(log.entries().is_empty()); + } + let settled = settled(py, result, asynchronous); + assert_eq!( + arguments + .get_item("api_key") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "original" + ); + assert_eq!( + &log.entries()[..4], + ["started", "begin", "project", "setup:hook:inherited"] + ); + if fail_setup { + assert!(settled.unwrap_err().value(py).is(error.bind(py))); + assert_eq!( + log.entries(), + [ + "started", + "begin", + "project", + "setup:hook:inherited", + "failed:Host:setup failed", + "adapter.close", + "host.close", + ] + ); + } else { + assert!(settled.is_ok()); + assert!( + log.entries() + .iter() + .any(|entry| entry.starts_with("succeeded:")) + ); + } + }); +} + +#[rstest] +fn closing_an_unstarted_call_never_initializes_resources( + #[from(interpreter)] _python: Interpreter, +) { + Python::attach(|py| { + let log = Log::default(); + let setup_log = log.clone(); + let pending = run_call( + py, + move |_, _, request| { + setup_log.push("setup"); + Ok(success_machine()(request)) + }, + binding(&log, OpScript::Answer), + hooks(&log, HookScript::Plain), + PyDict::new(py).unbind(), + call_options(true), + ) + .unwrap(); + assert!(log.entries().is_empty()); + pending.call_method0(py, "close").unwrap(); + drop(pending); + assert_eq!(log.entries(), ["adapter.close", "host.close"]); + }); +} + +#[rstest] +fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let raised = PyValueError::new_err("over budget").into_value(py); + let log = Log::default(); + let (result, log) = run_composed( + py, + success_machine(), + binding(&log, OpScript::Answer), + hooks(&log, HookScript::Plain), + |hooks| { + HookChain::new() + .with(hooks) + .with(ArgumentPolicy::Reject(raised.clone_ref(py))) + }, + call_options(asynchronous), + ); + assert!(result.unwrap_err().value(py).is(raised.bind(py))); + assert_eq!( + log, + [ + "started", + "begin", + "failed:Host:over budget", + "adapter.close", + "host.close" + ] + ); + }); +} + +#[rstest] +fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let log = Log::default(); + let (result, _) = run_composed( + py, + success_machine(), + binding(&log, OpScript::Answer), + hooks(&log, HookScript::Plain), + |hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit), + call_options(asynchronous), + ); + assert_eq!( + result.unwrap().extract::(py).unwrap(), + "project:2|sign|rewritten" + ); + }); +} + +#[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, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::Answer, + HookScript::FailAfterSuccess, + asynchronous, + ); + assert_eq!( + result.unwrap_err().value(py).to_string(), + "after_success failed" + ); + assert_eq!( + &log[log.len() - 4..], + [ + "after_success", + "failed:Host:after_success failed", + "adapter.close", + "host.close" + ] + ); + assert!(!log.iter().any(|entry| entry.starts_with("succeeded"))); + }); +} + +#[rstest] +fn cancellation_ends_the_call_without_terminal_dispatch(#[from(interpreter)] _python: Interpreter) { + 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(missing_state()) + } + 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(missing_state().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( + py, + |_, _, request| Ok(success_machine()(request)), + host, + adapter, + PyDict::new(py).unbind(), + call_options(false), + ) + .unwrap_err(); + assert!(!error.is_instance_of::(py)); + assert_eq!( + log.entries(), + ["started", "begin", "project", "adapter.close"] + ); + }); +} diff --git a/litellm-rust/crates/host-python/tests/handle.rs b/litellm-rust/crates/host-python/tests/handle.rs new file mode 100644 index 00000000000..37b27300973 --- /dev/null +++ b/litellm-rust/crates/host-python/tests/handle.rs @@ -0,0 +1,201 @@ +mod support; + +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep}; +use pyo3::{ + exceptions::PyBaseException, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; +use rstest::rstest; +use support::{Interpreter, install_lifecycle_module, interpreter, lifecycle_binding}; + +#[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>, +} + +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), + }, + lifecycle_binding, + ), + ) +} + +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)), lifecycle_binding) +} + +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), lifecycle_binding) +} + +struct ErrorBody(Option>); + +impl ExecutionBody for ErrorBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| { + Err(PyErr::from_value( + self.0.take().unwrap().into_bound(py).into_any(), + )) + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } +} + +#[pyfunction] +fn error_execution(error: Bound<'_, PyBaseException>) -> Execution { + Execution::new(ErrorBody(Some(error.unbind())), lifecycle_binding) +} + +#[rstest] +fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) { + 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(); + }); +} + +#[rstest] +fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) { + 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/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs new file mode 100644 index 00000000000..c2899c16878 --- /dev/null +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -0,0 +1,656 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use litellm_host::{ + call::{CallOutput, HostedMachine, hosted_call}, + hooks::CallHooks, + interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}, + lifecycle::{ExecutionEvent, Timing}, + machine::{CallMachine, MachineFault}, + protocol::{Protocol, Reply}, +}; +use litellm_host_python::{ + CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks, + PythonHostCalls, PythonOwned, PythonRuntime, run_call, +}; +use pyo3::{ + exceptions::{PyTypeError, PyValueError, asyncio::CancelledError}, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyDict, PyString}, +}; +use rstest::fixture; + +static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); + +/// One initialized interpreter with the real lifecycle protocol installed, held for the +/// length of a test so tests sharing the interpreter run one at a time. +pub struct Interpreter(MutexGuard<'static, ()>); + +#[fixture] +pub fn interpreter() -> Interpreter { + let guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(PoisonError::into_inner); + Python::initialize(); + Python::attach(|py| { + py.import("asyncio").unwrap(); + install_lifecycle_module(py); + }); + Interpreter(guard) +} + +pub fn lifecycle_binding(py: Python<'_>) -> PyResult> { + py.import("host_test_lifecycle") +} + +pub fn call_options(asynchronous: bool) -> CallOptions { + CallOptions::new(asynchronous, lifecycle_binding) +} + +pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> { + 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!("host_test_lifecycle"), + ) + .unwrap() +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Error(pub String); + +impl std::fmt::Display for Error { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } +} + +impl From for Error { + fn from(fault: MachineFault) -> Self { + Self(format!("{fault:?}")) + } +} + +pub struct Synthetic; + +impl Protocol for Synthetic { + type Response = String; + type Error = Error; + type Request = String; + type HostCall = (&'static str, Reply); + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; +} + +pub fn wire() -> WireRequest { + WireRequest { + url: "https://example.invalid".into(), + headers: Vec::new(), + body: serde_json::json!({}), + } +} + +pub fn context() -> RequestContext { + RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params: serde_json::json!({}), + secret_fields: Vec::new(), + api_key: None, + } +} + +#[derive(Clone, Default)] +pub struct Log(Arc>>); + +impl Log { + pub fn push(&self, entry: impl Into) { + self.0.lock().unwrap().push(entry.into()); + } + + pub fn entries(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +#[derive(Clone, Copy, Debug)] +pub enum OpScript { + Answer, + RaisePython, + RejectNatively, + RejectRequestNatively, + RaiseRequestPython, + CancelRequest, + AwaitAnswer, + AwaitFailure, +} + +/// The fake route's public exception, kept as a value so a test sees what `classify` +/// produced before the driver raises it. +#[derive(Debug, PartialEq, Eq)] +pub struct Classified(pub String); + +impl From for PyErr { + fn from(classified: Classified) -> Self { + PyValueError::new_err(format!("classified: {}", classified.0)) + } +} + +pub struct SyntheticBinding { + log: Log, + op: OpScript, + classifier_fails: bool, + pending_reply: Option>, +} + +pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding { + SyntheticBinding { + log: log.clone(), + op, + classifier_fails: false, + pending_reply: None, + } +} + +impl SyntheticBinding { + pub fn with_failing_classifier(self) -> Self { + Self { + classifier_fails: true, + ..self + } + } + + fn answer(&self, value: impl FnOnce() -> String) -> Result> { + match self.op { + OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()), + OpScript::RaisePython | OpScript::RaiseRequestPython => { + Err(PyValueError::new_err("op failed").into()) + } + OpScript::RejectNatively | OpScript::RejectRequestNatively => { + Err(InvokeError::Native(Error("op rejected".into()))) + } + OpScript::CancelRequest => Err(CancelledError::new_err(()).into()), + } + } +} + +impl PythonBinding for SyntheticBinding { + type Protocol = Synthetic; + type Failure = Classified; + + fn decode_request( + &mut self, + _: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> Result> { + self.log.push("project"); + match self.op { + OpScript::RejectRequestNatively + | OpScript::RaiseRequestPython + | OpScript::CancelRequest => self.answer(String::new), + _ => Ok(format!("project:{}", arguments.len())), + } + } + + 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, py: Python<'_>, response: String) -> PyResult> { + self.log.push("complete"); + Ok(PyString::new(py, &response).into_any().unbind()) + } + + fn map_error(&self, _: Python<'_>, error: Error) -> PyResult { + self.log.push(format!("classify:{error}")); + if self.classifier_fails { + return Err(PyTypeError::new_err("classifier failed")); + } + Ok(Classified(error.0)) + } + + fn host_error(error: &PyErr) -> Error { + Error(error.to_string()) + } +} + +impl PythonHostCalls for SyntheticBinding { + fn handle_host_call( + &mut self, + _: Python<'_>, + (op, reply): (&'static str, Reply), + ) -> Result<(), InvokeError> { + self.log.push(format!("op:{op}")); + self.answer(|| op.to_string()) + .map(|answer| reply.send(answer)) + } + + fn begin_host_call( + &mut self, + py: Python<'_>, + (op, reply): (&'static str, Reply), + ) -> Result>, InvokeError> { + if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) { + return self.handle_host_call(py, (op, reply)).map(|()| None); + } + self.pending_reply = Some(reply); + let module = PyModule::from_code( + py, + pyo3::ffi::c_str!( + "async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n" + ), + pyo3::ffi::c_str!("host_op.py"), + pyo3::ffi::c_str!("host_op"), + )?; + Ok(Some( + module + .getattr("answer")? + .call1((matches!(self.op, OpScript::AwaitFailure),))? + .unbind(), + )) + } + + fn resume_host_call( + &mut self, + py: Python<'_>, + result: PyResult>, + ) -> Result>, InvokeError> { + let answer = result?.extract::(py)?; + self.pending_reply.take().unwrap().send(answer); + Ok(None) + } +} + +impl PythonOwned for SyntheticBinding { + fn close(&mut self, _: Python<'_>) { + self.log.push("host.close"); + } + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } +} + +#[derive(Clone, Copy, Debug)] +pub enum HookScript { + Plain, + RewriteArguments, + ObserveArguments, + FailBegin, + ReplaceResponse, + FailAfterSuccess, + CancelTerminal, + FailTerminal, +} + +pub struct SyntheticHooks { + log: Log, + script: HookScript, +} + +pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks { + SyntheticHooks { + log: log.clone(), + script, + } +} + +impl CallHooks for SyntheticHooks { + fn prepare_arguments( + &mut self, + _: Python<'_>, + arguments: Py, + _: f64, + ) -> PyResult>> { + self.log.push("begin"); + if matches!(self.script, HookScript::FailBegin) { + return Err(PyValueError::new_err("begin failed")); + } + if matches!(self.script, HookScript::RewriteArguments) { + let prepared = Python::attach(|py| -> PyResult> { + let prepared = arguments.bind(py).copy()?; + prepared.set_item("prepared", "hook")?; + prepared.set_item("api_key", "hook-key")?; + Ok(prepared.unbind()) + })?; + return Ok(HookStep::Ready(prepared)); + } + Ok(HookStep::Ready(arguments)) + } + + fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { + if matches!(self.script, HookScript::ObserveArguments) { + let value: String = arguments + .bind(py) + .get_item("prepared")? + .unwrap() + .extract()?; + self.log.push(format!("adopted:{value}")); + } + Ok(()) + } + + fn before_provider_request( + &mut self, + _: Python<'_>, + wire: Box, + _: &RequestContext, + ) -> PyResult>> { + self.log.push("before_provider_request"); + Ok(HookStep::Ready(Box::new(WireRequest { + url: "rewritten".into(), + ..*wire + }))) + } + + fn transform_response( + &mut self, + py: Python<'_>, + response: Py, + _: Timing, + ) -> PyResult>> { + self.log.push("after_success"); + match self.script { + HookScript::ReplaceResponse => Ok(HookStep::Ready( + "replaced".into_pyobject(py)?.into_any().unbind(), + )), + HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")), + _ => Ok(HookStep::Ready(response)), + } + } + + fn on_event( + &mut self, + py: Python<'_>, + event: PythonCallEvent<'_>, + ) -> PyResult> { + let terminal = matches!( + event, + PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. } + ); + if terminal && matches!(self.script, HookScript::CancelTerminal) { + return Err(CancelledError::new_err("callback cancelled")); + } + if terminal && matches!(self.script, HookScript::FailTerminal) { + return Err(PyValueError::new_err("callback failed")); + } + self.log.push(match event { + PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(), + PythonCallEvent::Started { .. } => "started".into(), + PythonCallEvent::Cancelled { .. } => "cancelled".into(), + PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { + format!("response:{}", raw.body) + } + PythonCallEvent::Succeeded { response, .. } => { + format!("succeeded:{}", response.bind(py)) + } + PythonCallEvent::Failed { origin, error, .. } => { + format!("failed:{origin:?}:{}", error.value(py)) + } + }); + Ok(HookStep::Ready(())) + } + + fn on_stream_open(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + self.log.push("opened"); + Ok(()) + } + + fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + self.log.push("delivered"); + Ok(()) + } +} + +impl PythonOwned for SyntheticHooks { + fn close(&mut self, _: Python<'_>) { + self.log.push("adapter.close"); + } + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } +} + +pub type MachineFactory = Box CallMachine + Send + Sync>; + +/// Answers to projection, to the route op and to `before_provider_request` all reach the +/// response, so a driver that misroutes a reply changes what the call returns. +pub fn success_machine() -> MachineFactory { + Box::new(move |projected| { + CallMachine::new(None, move |host| { + Box::pin(async move { + let signed = host.services.call(|reply| ("sign", reply)).await?; + let wire = host + .interceptors + .before_provider_request(wire(), context()) + .await?; + host.interceptors + .after_provider_response(RawResponse { body: "raw".into() }) + .await?; + Ok(format!("{projected}|{signed}|{}", wire.url)) + }) + }) + }) +} + +pub fn failing_machine() -> MachineFactory { + Box::new(move |_| { + CallMachine::new(None, |_| { + Box::pin(async move { Err(Error("provider exploded".into())) }) + }) + }) +} + +/// Drives an asynchronous call's coroutine to its settled value from a sync test. A +/// synchronous call is already settled. +pub fn settled( + py: Python<'_>, + result: PyResult>, + asynchronous: bool, +) -> PyResult> { + if !asynchronous { + return result; + } + let coroutine = result?; + let completed = coroutine + .call_method1(py, "send", (py.None(),)) + .unwrap_err(); + if !completed.is_instance_of::(py) { + return Err(completed); + } + completed.value(py).getattr("value").map(Bound::unbind) +} + +pub fn run_scripted( + py: Python<'_>, + machine: MachineFactory, + op: OpScript, + script: HookScript, + asynchronous: bool, +) -> (PyResult>, Vec) { + let log = Log::default(); + run_composed( + py, + machine, + binding(&log, op), + hooks(&log, script), + std::convert::identity, + call_options(asynchronous), + ) +} + +pub fn run_composed( + py: Python<'_>, + machine: MachineFactory, + binding: SyntheticBinding, + hooks: SyntheticHooks, + compose: impl FnOnce(SyntheticHooks) -> L, + options: CallOptions, +) -> (PyResult>, Vec) { + let asynchronous = options.asynchronous; + let log = binding.log.clone(); + let arguments = PyDict::new(py); + arguments.set_item("model", "m").unwrap(); + let result = run_call( + py, + move |_, _, request| Ok(machine(request)), + binding, + compose(hooks), + arguments.unbind(), + options, + ); + (settled(py, result, asynchronous), log.entries()) +} + +pub struct Streaming; + +impl Protocol for Streaming { + type Response = (); + type Error = Error; + type Request = (); + type HostCall = std::convert::Infallible; + type Chunk = &'static str; + type StreamHead = Vec<(&'static str, &'static str)>; +} + +pub struct StreamingBinding; + +impl PythonBinding for StreamingBinding { + type Protocol = Streaming; + type Failure = Classified; + + fn decode_request( + &mut self, + _: Python<'_>, + _: &Bound<'_, PyDict>, + ) -> Result<(), InvokeError> { + Ok(()) + } + + fn encode_stream_head( + &mut self, + py: Python<'_>, + head: Vec<(&'static str, &'static str)>, + ) -> PyResult> { + let headers = PyDict::new(py); + for (name, value) in head { + headers.set_item(name, value)?; + } + let hidden = PyDict::new(py); + hidden.set_item("additional_headers", headers)?; + Ok(hidden.into_any().unbind()) + } + + fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult> { + Ok(PyString::new(py, chunk).into_any().unbind()) + } + + fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult> { + Ok(py.None()) + } + + fn map_error(&self, _: Python<'_>, error: Error) -> PyResult { + Ok(Classified(error.0)) + } + + fn host_error(error: &PyErr) -> Error { + Error(error.to_string()) + } +} + +impl PythonHostCalls for StreamingBinding { + fn handle_host_call( + &mut self, + _: Python<'_>, + op: std::convert::Infallible, + ) -> Result<(), InvokeError> { + match op {} + } +} + +impl PythonOwned for StreamingBinding { + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } +} + +pub fn streaming_machine() -> impl FnOnce(()) -> HostedMachine + Send + Sync { + move |()| { + hosted_call((), None, |(), _, _, _observations| async { + Ok(CallOutput::Stream { + head: vec![("request-id", "req_1")], + chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])), + }) + }) + } +} + +/// Starts a streaming call and returns the caller's stream object. +pub fn open_stream( + py: Python<'_>, + hooks: SyntheticHooks, + options: CallOptions, +) -> Bound<'_, PyAny> { + let asynchronous = options.asynchronous; + let handed = run_call( + py, + |_, _, request| Ok(streaming_machine()(request)), + StreamingBinding, + hooks, + PyDict::new(py).unbind(), + options, + ) + .unwrap(); + if asynchronous { + let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err(); + stop.value(py).getattr("value").unwrap() + } else { + handed.into_bound(py) + } +} + +/// Drives a `Stream` (async) or `SyncStream` to completion from a sync test. +pub fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec { + if !asynchronous { + return stream + .try_iter() + .unwrap() + .map(|chunk| chunk.unwrap().extract().unwrap()) + .collect(); + } + std::iter::from_fn(|| { + let stop = stream + .call_method0("__anext__") + .unwrap() + .call_method1("send", (py.None(),)) + .unwrap_err(); + if stop.is_instance_of::(py) { + return None; + } + assert!(stop.is_instance_of::(py)); + Some(stop.value(py).getattr("value").unwrap().extract().unwrap()) + }) + .collect() +}