diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index dedd41b2efc..09896a96952 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3631,9 +3631,11 @@ dependencies = [ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ + "futures-util", "litellm-auth", "litellm-host", "litellm-host-python", + "parking_lot", "proptest", "pyo3", "rstest", diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index ed5e0fb9691..0bde2250700 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -9,6 +9,9 @@ repository.workspace = true litellm-host.workspace = true litellm-host-python.workspace = true +futures-util.workspace = true +parking_lot.workspace = true + pyo3.workspace = true strum.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index d891b5a1f0c..85fb65230c0 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -9,7 +9,8 @@ use litellm_host::{ interceptors::{RawResponse, RequestContext, WireRequest}, lifecycle::{FailureOrigin, Timing, epoch_seconds}, }; -use litellm_host_python::{Error as HostError, HookStep, from_py, to_py}; +use litellm_host_python::{Awaiter, Error as HostError, from_py, to_py}; +use parking_lot::Mutex; use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, @@ -67,6 +68,22 @@ fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { !error.is_instance_of::(py) } +/// What a failure runs before the failure handlers: the deployment failure hook, whose +/// outcome never stops them, or the handlers themselves, with the async family's +/// awaitable when there is one. +pub(crate) enum FailureStage { + DeploymentHook(Py), + Handlers(Option>), +} + +/// An awaited failure handler ends the call however it ends, unless it was cancelled. +fn async_failure_result(result: PyResult>) -> PyResult<()> { + match result { + Err(error) if Python::attach(|py| is_cancellation(py, &error)) => Err(error), + _ => Ok(()), + } +} + impl LegacyLogging { pub fn new( py: Python<'_>, @@ -141,14 +158,14 @@ impl LegacyLogging { /// The keyword view the rest of the call reads: a copy, so the deployment hook's own /// dict is left as the hook returned it, carrying the logger as `@client` injects it. - fn prepare(&mut self, py: Python<'_>) -> PyResult>> { + fn prepare(&mut self, py: Python<'_>) -> PyResult> { let prepared = self.call.kwargs().bind(py).copy()?; prepared.set_item("litellm_logging_obj", self.logger()?.object(py))?; self.call.set_kwargs(prepared.unbind()); - Ok(HookStep::Ready(self.call.kwargs().clone_ref(py))) + Ok(self.call.kwargs().clone_ref(py)) } - fn finalize(&mut self, py: Python<'_>) -> PyResult>> { + fn finalize(&mut self, py: Python<'_>) -> PyResult> { finalize( py, &self.response, @@ -159,7 +176,7 @@ impl LegacyLogging { )?; self.response .as_ref() - .map(|response| HookStep::Ready(response.clone_ref(py))) + .map(|response| response.clone_ref(py)) .ok_or_else(|| HostError::MissingState.into()) } @@ -238,14 +255,14 @@ impl LegacyLogging { /// A failure after the stream reached the caller bills the delivered chunks as /// partial usage. The sync path has no loop to schedule that on, so it falls back to /// the plain failure handler. - fn stream_failure(&mut self, py: Python<'_>) -> PyResult> { + fn stream_failure(&mut self, py: Python<'_>) -> PyResult>> { let (Some(logger), Some(error), Some(stream), Some(billing)) = ( &self.logger, &self.error, &self.stream, self.stream_billing(), ) else { - return Ok(HookStep::Ready(())); + return Ok(None); }; if !self.asynchronous { return self.dispatch_failure(py); @@ -264,24 +281,21 @@ impl LegacyLogging { ), ); match scheduled { - Ok(awaitable) => Ok(HookStep::Await( - awaitable.unbind(), - Self::resume_async_failure, - )), + Ok(awaitable) => Ok(Some(awaitable.unbind())), Err(failure) if is_cancellation(py, &failure) => Err(failure), - Err(_) => Ok(HookStep::Ready(())), + Err(_) => Ok(None), } } /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. - fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult> { + fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult>> { let (Some(logger), Some(error)) = (&self.logger, &self.error) else { - return Ok(HookStep::Ready(())); + return Ok(None); }; if self.asynchronous && self.internal { - return Ok(HookStep::Ready(())); + return Ok(None); } if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false) && is_cancellation(py, &failure) @@ -289,64 +303,41 @@ impl LegacyLogging { return Err(failure); } if !self.asynchronous { - return Ok(HookStep::Ready(())); + return Ok(None); } match logger.failure(py, error, &self.start, &self.end, true) { - Ok(Some(awaitable)) => Ok(HookStep::Await(awaitable, Self::resume_async_failure)), - Ok(None) => Ok(HookStep::Ready(())), + Ok(awaitable) => Ok(awaitable), Err(failure) if is_cancellation(py, &failure) => Err(failure), - Err(_) => Ok(HookStep::Ready(())), + Err(_) => Ok(None), } } } impl LegacyLogging { - fn resume_begin( + fn adopt_hooked_arguments( &mut self, py: Python<'_>, result: PyResult>, - ) -> PyResult>> { + ) -> PyResult<()> { self.call .set_kwargs(result?.into_bound(py).cast_into::()?.unbind()); - self.prepare(py) + Ok(()) } - fn resume_after_success( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { + fn adopt_hooked_response(&mut self, result: PyResult>) -> PyResult<()> { self.response = Some(result?); - self.finalize(py) - } - - fn resume_deployment_failure( - &mut self, - py: Python<'_>, - _: PyResult>, - ) -> PyResult> { - self.dispatch_failure(py) - } - - fn resume_async_failure( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult> { - match result { - Err(error) if is_cancellation(py, &error) => Err(error), - _ => Ok(HookStep::Ready(())), - } + Ok(()) } } impl LegacyLogging { - pub(crate) fn prepare_call( + /// Sets the call up, returning the deployment pre-call hook to await when it runs. + fn begin_call( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult>> { + ) -> PyResult>> { self.call.set_kwargs(arguments); self.start = datetime(py, started_at)?; self.internal = is_internal_call(py)?; @@ -360,13 +351,10 @@ impl LegacyLogging { )?; self.logger = Some(result.logger()?); self.call.set_kwargs(result.kwargs()?); - if self.runs_deployment_hooks() { - return Ok(HookStep::Await( - DeploymentHooks::before_call(py, self.call.kwargs(), self.call_type())?, - Self::resume_begin, - )); + if !self.runs_deployment_hooks() { + return Ok(None); } - self.prepare(py) + DeploymentHooks::before_call(py, self.call.kwargs(), self.call_type()).map(Some) } pub(crate) fn pre_call( @@ -374,7 +362,7 @@ impl LegacyLogging { py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult>> { + ) -> PyResult> { let logger = self.logger()?; logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; let body = to_py(py, &wire.body)? @@ -408,40 +396,35 @@ impl LegacyLogging { .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; - Ok(HookStep::Ready(Box::new(WireRequest { + Ok(Box::new(WireRequest { body: from_py(&body)?, headers, ..*wire - }))) + })) } - pub(crate) fn transform_public_response( + /// Takes the public response, returning the deployment post-call hook to await when + /// it runs. + fn respond( &mut self, py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult>> { + ) -> PyResult>> { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.runs_deployment_hooks() { - return Ok(HookStep::Await( - DeploymentHooks::after_success( - py, - self.call.kwargs(), - &self.response, - self.call_type(), - )?, - Self::resume_after_success, - )); + if !self.runs_deployment_hooks() { + return Ok(None); } - self.finalize(py) + DeploymentHooks::after_success(py, self.call.kwargs(), &self.response, self.call_type()) + .map(Some) } pub(crate) fn result_ready( &mut self, py: Python<'_>, facts: &litellm_host::interceptors::ExecutionFacts, - ) -> PyResult> { + ) -> PyResult<()> { use litellm_host::interceptors::{Cost, ResultSource}; let logger = self.logger()?.object(py); @@ -474,14 +457,10 @@ impl LegacyLogging { }; details.set_item("cache_hit", self.cache_key.is_some())?; details.set_item("cache_key", self.cache_key.as_deref())?; - Ok(HookStep::Ready(())) + Ok(()) } - pub(crate) fn post_call( - &mut self, - py: Python<'_>, - raw: &RawResponse, - ) -> PyResult> { + pub(crate) fn post_call(&mut self, py: Python<'_>, raw: &RawResponse) -> PyResult<()> { let api_key = self .request .as_ref() @@ -493,8 +472,7 @@ impl LegacyLogging { api_key, self.request.as_ref().map(|request| &request.body), self.request.as_ref().map(|request| &request.headers), - )?; - Ok(HookStep::Ready(())) + ) } pub(crate) fn succeeded( @@ -502,36 +480,33 @@ impl LegacyLogging { py: Python<'_>, timing: Timing, response: &Py, - ) -> PyResult> { + ) -> PyResult<()> { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { - Some(stream) => self.stream_success(py, stream)?, - None => self.dispatch_success(py)?, + Some(stream) => self.stream_success(py, stream), + None => self.dispatch_success(py), } - Ok(HookStep::Ready(())) } - pub(crate) fn failed( + fn failed( &mut self, py: Python<'_>, timing: Timing, origin: FailureOrigin, - error: &PyErr, - ) -> PyResult> { + error: &Py, + ) -> PyResult { self.end = Some(datetime(py, timing.end_time)?); - self.error = Some(error.clone_ref(py).into_value(py)); + self.error = Some(error.clone_ref(py)); if self.stream.is_some() { - return self.stream_failure(py); + return self.stream_failure(py).map(FailureStage::Handlers); } if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() { let error = self.error.as_ref().ok_or(HostError::MissingState)?; - return Ok(HookStep::Await( - DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?, - Self::resume_deployment_failure, - )); + return DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type()) + .map(FailureStage::DeploymentHook); } - self.dispatch_failure(py) + self.dispatch_failure(py).map(FailureStage::Handlers) } pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py) -> PyResult<()> { @@ -559,6 +534,117 @@ impl LegacyLogging { } } +/// [`LegacyLogging`] as the call's hooks. The state is locked only around Python work and +/// never across an await, so the call's handle can traverse it while a hook is suspended. +pub struct LegacyHooks(Mutex); + +impl LegacyHooks { + pub fn new(logging: LegacyLogging) -> Self { + Self(Mutex::new(logging)) + } + + fn with(&self, step: impl FnOnce(&mut LegacyLogging, Python<'_>) -> T) -> T { + Python::attach(|py| step(&mut self.0.lock(), py)) + } + + pub(crate) async fn prepare_call( + &self, + task: &dyn Awaiter, + arguments: Py, + started_at: f64, + ) -> PyResult> { + let hook = self.with(|logging, py| logging.begin_call(py, arguments, started_at))?; + if let Some(hook) = hook { + let hooked = task.awaited(hook).await; + self.with(|logging, py| logging.adopt_hooked_arguments(py, hooked))?; + } + self.with(|logging, py| logging.prepare(py)) + } + + pub(crate) fn adopt_arguments(&self, arguments: &Py) { + self.with(|logging, py| logging.adopt_arguments(py, arguments)); + } + + pub(crate) fn pre_call( + &self, + wire: Box, + context: &RequestContext, + ) -> PyResult> { + self.with(|logging, py| logging.pre_call(py, wire, context)) + } + + pub(crate) async fn transform_public_response( + &self, + task: &dyn Awaiter, + response: Py, + timing: Timing, + ) -> PyResult> { + let hook = self.with(|logging, py| logging.respond(py, response, timing))?; + if let Some(hook) = hook { + let hooked = task.awaited(hook).await; + self.with(|logging, _| logging.adopt_hooked_response(hooked))?; + } + self.with(|logging, py| logging.finalize(py)) + } + + pub(crate) fn result_ready( + &self, + facts: &litellm_host::interceptors::ExecutionFacts, + ) -> PyResult<()> { + self.with(|logging, py| logging.result_ready(py, facts)) + } + + pub(crate) fn post_call(&self, raw: &RawResponse) -> PyResult<()> { + self.with(|logging, py| logging.post_call(py, raw)) + } + + pub(crate) fn succeeded(&self, timing: Timing, response: &Py) -> PyResult<()> { + self.with(|logging, py| logging.succeeded(py, timing, response)) + } + + pub(crate) async fn failed( + &self, + task: &dyn Awaiter, + timing: Timing, + origin: FailureOrigin, + error: &Py, + ) -> PyResult<()> { + let stage = self.with(|logging, py| logging.failed(py, timing, origin, error))?; + let handler = match stage { + FailureStage::DeploymentHook(hook) => { + drop(task.awaited(hook).await); + self.with(|logging, py| logging.dispatch_failure(py))? + } + FailureStage::Handlers(handler) => handler, + }; + match handler { + Some(handler) => async_failure_result(task.awaited(handler).await), + None => Ok(()), + } + } + + pub(crate) fn stream_opened(&self, head: &Py) -> PyResult<()> { + self.with(|logging, py| logging.stream_opened(py, head)) + } + + pub(crate) fn stream_chunk(&self, chunk: &Py) -> PyResult<()> { + self.with(|logging, py| logging.stream_chunk(py, chunk)) + } +} + +impl PythonOwned for LegacyHooks { + fn close(&mut self, py: Python<'_>) { + self.0.get_mut().close(py); + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + match self.0.try_lock() { + Some(logging) => logging.traverse(visit), + None => Ok(()), + } + } +} + impl PythonOwned for LegacyLogging { fn close(&mut self, py: Python<'_>) { if let Some(logger) = self.logger.take() @@ -596,26 +682,16 @@ mod deployment_hooks_tests { use litellm_host::hooks::CallHooks; use litellm_host::lifecycle::{FailureOrigin, Timing}; - use litellm_host_python::{HookStep, PythonCallEvent}; + use litellm_host_python::PythonCallEvent; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; use rstest::rstest; - use super::LegacyLogging; - use crate::test_support::{legacy_call, local, namespace, run}; - - fn resume( - logging: &mut LegacyLogging, - step: HookStep, - py: Python<'_>, - value: PyResult>, - ) -> PyResult> { - let HookStep::Await(_, continuation) = step else { - panic!("expected suspension") - }; - continuation(logging, py, value) - } + use super::{LegacyHooks, LegacyLogging}; + use crate::test_support::{ + Scripted, failure, finish, hooks, legacy_call, local, namespace, run, + }; const CALL: &CStr = c" document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} @@ -627,32 +703,23 @@ kwargs = {'logger': logger, 'document': document} end_time: 1.0, }; - fn begin<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - asynchronous: bool, - ) -> (LegacyLogging, HookStep>) { - let mut logging = legacy_call(py, locals, asynchronous); - let kwargs = local(locals, "kwargs") + fn kwargs(locals: &Bound<'_, PyDict>) -> Py { + local(locals, "kwargs") .cast_into::() .unwrap() - .unbind(); - let step = logging.prepare_arguments(py, kwargs, 0.0).unwrap(); - (logging, step) + .unbind() } - fn arguments<'py>( - py: Python<'py>, - step: HookStep>, - ) -> Bound<'py, PyDict> { - let HookStep::Ready(arguments) = step else { - panic!("expected the prepared arguments"); - }; - arguments.into_bound(py) - } - - fn awaits_deployment_hook(step: &HookStep) -> bool { - matches!(step, HookStep::Await(_, _)) + /// Prepares the call's `kwargs`, answering the deployment pre-call hook from `task`. + fn begin( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + asynchronous: bool, + task: &Scripted, + ) -> (LegacyHooks, PyResult>) { + let hooks = hooks(legacy_call(py, locals, asynchronous)); + let prepared = finish(hooks.prepare_arguments(task, kwargs(locals), 0.0)); + (hooks, prepared) } #[rstest] @@ -672,16 +739,13 @@ kwargs = {'logger': logger, 'document': document} Python::initialize(); Python::attach(|py| { let locals = namespace(py, CALL); - let mut logging = LegacyLogging { + let hooks = hooks(LegacyLogging { operation, ..legacy_call(py, &locals, asynchronous) - }; - let kwargs = local(&locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let step = logging.prepare_arguments(py, kwargs, 0.0).unwrap(); - assert_eq!(awaits_deployment_hook(&step), asynchronous); + }); + let task = Scripted::answering([Ok(kwargs(&locals).into_any())]); + finish(hooks.prepare_arguments(&task, kwargs(&locals), 0.0)).unwrap(); + assert_eq!(task.count() == 1, asynchronous); locals.set_item("expected", expected).unwrap(); locals.set_item("asynchronous", asynchronous).unwrap(); run( @@ -703,8 +767,9 @@ if asynchronous: Python::initialize(); Python::attach(|py| { let locals = namespace(py, CALL); - let (_, step) = begin(py, &locals, asynchronous); - assert_eq!(awaits_deployment_hook(&step), asynchronous); + let task = Scripted::answering([Ok(kwargs(&locals).into_any())]); + begin(py, &locals, asynchronous, &task).1.unwrap(); + assert_eq!(task.count() == 1, asynchronous); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -727,16 +792,10 @@ kwargs = {'logger': logger, 'document': document} replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} ", ); - let (mut logging, step) = begin(py, &locals, true); - assert!(awaits_deployment_hook(&step)); - let step = resume( - &mut logging, - step, - py, - Ok(local(&locals, "replaced_kwargs").unbind()), - ) - .unwrap(); - locals.set_item("prepared", arguments(py, step)).unwrap(); + let task = Scripted::answering([Ok(local(&locals, "replaced_kwargs").unbind())]); + let (_, prepared) = begin(py, &locals, true, &task); + assert_eq!(task.count(), 1); + locals.set_item("prepared", prepared.unwrap()).unwrap(); run( py, &locals, @@ -767,14 +826,8 @@ logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) kwargs = {'logger': logger, 'vendor_extension': opaque} ", ); - let (mut logging, step) = begin(py, &locals, asynchronous); - let step = match step { - HookStep::Await(hook_result, resume) => { - resume(&mut logging, py, Ok(hook_result)).unwrap() - } - step => step, - }; - locals.set_item("prepared", arguments(py, step)).unwrap(); + let (_, prepared) = begin(py, &locals, asynchronous, &Scripted::default()); + locals.set_item("prepared", prepared.unwrap()).unwrap(); locals.set_item("asynchronous", asynchronous).unwrap(); run( py, @@ -800,28 +853,19 @@ replacement = object() logger.hooks = {'pre': lambda kwargs: kwargs} ", ); - let (mut logging, step) = begin(py, &locals, true); - resume( - &mut logging, - step, - py, + let task = Scripted::answering([ Ok(local(&locals, "kwargs").unbind()), - ) - .unwrap(); - let step = logging - .transform_response(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - assert!(awaits_deployment_hook(&step)); - let step = resume( - &mut logging, - step, - py, Ok(local(&locals, "replacement").unbind()), - ) + ]); + let (hooks, prepared) = begin(py, &locals, true, &task); + prepared.unwrap(); + let returned = finish(hooks.transform_response( + &task, + local(&locals, "response").unbind(), + TIMING, + )) .unwrap(); - let HookStep::Ready(returned) = step else { - panic!("expected the finalized response"); - }; + assert_eq!(task.count(), 2); assert!(returned.bind(py).is(local(&locals, "replacement"))); run( py, @@ -841,27 +885,21 @@ assert finalized is replacement Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); - let (mut logging, step) = begin(py, &locals, true); let cancellation = CancelledError::new_err("cancelled"); let cancelled = cancellation.value(py).clone(); - let error = if post_call { - resume( - &mut logging, - step, - py, - Ok(local(&locals, "kwargs").unbind()), - ) - .unwrap(); - let step = logging - .transform_response(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - resume(&mut logging, step, py, Err(cancellation)) - .err() - .unwrap() + let answers = if post_call { + vec![Ok(local(&locals, "kwargs").unbind()), Err(cancellation)] } else { - resume(&mut logging, step, py, Err(cancellation)) - .err() - .unwrap() + vec![Err(cancellation)] + }; + let task = Scripted::answering(answers); + let (hooks, prepared) = begin(py, &locals, true, &task); + let error = if post_call { + prepared.unwrap(); + finish(hooks.transform_response(&task, local(&locals, "response").unbind(), TIMING)) + .unwrap_err() + } else { + prepared.unwrap_err() }; assert!(error.value(py).is(&cancelled)); let names: Vec = local(&locals, "logger") @@ -883,31 +921,26 @@ assert finalized is replacement py, c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", ); - let (mut logging, step) = begin(py, &locals, true); - resume( - &mut logging, - step, - py, - Ok(local(&locals, "kwargs").unbind()), - ) - .unwrap(); - let failure = PyErr::from_value(local(&locals, "failure")); - let failed = PythonCallEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Call, - error: &failure, - }; - let step = logging.on_event(py, failed).unwrap(); - assert!(awaits_deployment_hook(&step)); let hook_result = if cancelled { Err(CancelledError::new_err("cancelled")) } else { Ok(py.None()) }; - assert!(matches!( - resume(&mut logging, step, py, hook_result).unwrap(), - HookStep::Await(_, _) - )); + let task = Scripted::answering([ + Ok(local(&locals, "kwargs").unbind()), + hook_result, + Ok(py.None()), + ]); + let (hooks, prepared) = begin(py, &locals, true, &task); + prepared.unwrap(); + let failure = failure(&locals); + let failed = PythonCallEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Call, + error: &failure, + }; + finish(hooks.on_event(&task, failed)).unwrap(); + assert_eq!(task.count(), 3); run( py, &locals, @@ -925,10 +958,8 @@ mod payload_tests { use std::ffi::CStr; use litellm_auth::SecretValue; - use litellm_host::hooks::CallHooks; use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; - use litellm_host::lifecycle::ExecutionEvent; - use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py}; + use litellm_host_python::{PythonOwned, to_py}; use proptest::prelude::*; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -1022,8 +1053,7 @@ logger._provider_reported_cost = None .unwrap() .set_item("response_cost", existing_cost) .unwrap(); - let step = logging.result_ready(py, &facts).unwrap(); - assert!(matches!(step, HookStep::Ready(()))); + logging.result_ready(py, &facts).unwrap(); locals.set_item("reported", reported).unwrap(); locals.set_item("cache_hit", cache_hit).unwrap(); run( @@ -1093,44 +1123,26 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update headers: vec![("x-route".into(), "route".into())], body, }; - let (_, step) = send_and_receive(py, &mut logging, wire, &context); + let wire = send_and_receive(py, &mut logging, wire, &context); run(py, &locals, c"check()"); - let HookStep::Ready(wire) = step else { - panic!("before_provider_request did not hand back the wire request"); - }; *wire }) } /// `before_provider_request` over `wire`, then the provider's raw response the way the driver /// delivers it, so `pre_call` and `post_call` have both seen the retained payload. - fn send_and_receive<'a>( + fn send_and_receive( py: Python<'_>, - logging: &'a mut LegacyLogging, + logging: &mut LegacyLogging, wire: WireRequest, context: &RequestContext, - ) -> ( - &'a mut LegacyLogging, - HookStep>, - ) { - let step = logging - .before_provider_request(py, Box::new(wire), context) - .unwrap(); + ) -> Box { + let wire = logging.pre_call(py, Box::new(wire), context).unwrap(); let raw = RawResponse { body: "raw response".into(), }; - assert!(matches!( - logging - .on_event( - py, - PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { - raw: &raw - }) - ) - .unwrap(), - HookStep::Ready(()) - )); - (logging, step) + logging.post_call(py, &raw).unwrap(); + wire } fn route_context() -> RequestContext { @@ -1203,13 +1215,12 @@ prepared = {'pages': replacement} .cast_into::() .unwrap() .unbind(); - logging.arguments_prepared(py, &prepared).unwrap(); + logging.adopt_arguments(py, &prepared); let wire = WireRequest { body: json!({"pages": [1]}), ..route_wire() }; - let (_, step) = send_and_receive(py, &mut logging, wire, &route_context()); - assert!(matches!(step, HookStep::Ready(_))); + send_and_receive(py, &mut logging, wire, &route_context()); locals.set_item("description", description).unwrap(); run( py, @@ -1709,16 +1720,18 @@ mod terminal_tests { use litellm_host::hooks::CallHooks; use litellm_host::lifecycle::{FailureOrigin, Timing}; - use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned}; + use litellm_host_python::{PythonCallEvent, PythonOwned}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; use rstest::rstest; - use super::LegacyLogging; + use super::{LegacyHooks, LegacyLogging}; use crate::PythonLogger; - use crate::test_support::{legacy_call, local, namespace, run}; + use crate::test_support::{ + Scripted, failure, finish, hooks, legacy_call, local, namespace, run, + }; const TIMING: Timing = Timing { start_time: 0.0, @@ -1732,39 +1745,33 @@ mod terminal_tests { } } - fn succeed( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - logging: &mut LegacyLogging, - ) -> HookStep { + /// Reports the call's `response` as its success, returning how many awaitables the + /// success fan-out handed back. + fn succeed(locals: &Bound<'_, PyDict>, hooks: &LegacyHooks) -> usize { let response = local(locals, "response").unbind(); - logging - .on_event( - py, - PythonCallEvent::Succeeded { - timing: TIMING, - response: &response, - }, - ) - .unwrap() + let task = Scripted::default(); + finish(hooks.on_event( + &task, + PythonCallEvent::Succeeded { + timing: TIMING, + response: &response, + }, + )) + .unwrap(); + task.count() } - fn fail( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - logging: &mut LegacyLogging, - ) -> HookStep { - let failure = PyErr::from_value(local(locals, "failure")); - logging - .on_event( - py, - PythonCallEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Host, - error: &failure, - }, - ) - .unwrap() + /// Reports the call's `failure`, answering any awaited handler from `task`. + fn fail(locals: &Bound<'_, PyDict>, hooks: &LegacyHooks, task: &Scripted) -> PyResult<()> { + let failure = failure(locals); + finish(hooks.on_event( + task, + PythonCallEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Host, + error: &failure, + }, + )) } #[rstest] @@ -1785,11 +1792,8 @@ mod terminal_tests { Python::attach(|py| { let locals = namespace(py, c"response = object()"); run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - assert!(matches!( - succeed(py, &locals, &mut logging), - HookStep::Ready(()) - )); + let hooks = hooks(logged(py, &locals, asynchronous)); + assert_eq!(succeed(&locals, &hooks), 0); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -1815,7 +1819,7 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy py, c"response = object()\nlogger._defer_async_logging = True", ); - let mut logging = logged(py, &locals, true); + let hooks = hooks(logged(py, &locals, true)); let response = local(&locals, "response").unbind(); let event = PythonCallEvent::Succeeded { timing: TIMING, @@ -1826,10 +1830,9 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy ); drop(receiver); sender.emit(event.snapshot()); - assert!(matches!( - logging.on_event(py, event).unwrap(), - HookStep::Ready(()) - )); + let task = Scripted::default(); + finish(hooks.on_event(&task, event)).unwrap(); + assert_eq!(task.count(), 0); assert_eq!(sender.dropped_events(), 1); run(py, &locals, c" assert logger.names() == ['sync_success_for_async_call'], logger.calls @@ -1847,23 +1850,20 @@ assert logger.calls[1][1] is response Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"first = b'first'\nlast = b'last'\nresponse = None\nhead = {'additional_headers': {'request-id': 'req_native'}}"); - let mut logging = LegacyLogging { + let hooks = hooks(LegacyLogging { operation: crate::LoggingOperation::Messages, ..logged(py, &locals, true) - }; - logging - .on_stream_open(py, &local(&locals, "head").unbind()) + }); + hooks + .on_stream_open(&local(&locals, "head").unbind()) .unwrap(); - logging - .on_stream_chunk(py, &local(&locals, "first").unbind()) + hooks + .on_stream_chunk(&local(&locals, "first").unbind()) .unwrap(); - logging - .on_stream_chunk(py, &local(&locals, "last").unbind()) + hooks + .on_stream_chunk(&local(&locals, "last").unbind()) .unwrap(); - assert!(matches!( - succeed(py, &locals, &mut logging), - HookStep::Ready(()) - )); + assert_eq!(succeed(&locals, &hooks), 0); run( py, &locals, @@ -1889,14 +1889,13 @@ assert chunks[1] is last Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = LegacyLogging { + let hooks = hooks(LegacyLogging { internal: true, ..logged(py, &locals, asynchronous) - }; - assert!(matches!( - fail(py, &locals, &mut logging), - HookStep::Ready(()) - )); + }); + let task = Scripted::default(); + fail(&locals, &hooks, &task).unwrap(); + assert_eq!(task.count(), 0); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -1911,11 +1910,11 @@ assert chunks[1] is last Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"response = object()"); - let mut logging = LegacyLogging { + let hooks = hooks(LegacyLogging { internal: true, ..logged(py, &locals, true) - }; - succeed(py, &locals, &mut logging); + }); + succeed(&locals, &hooks); run( py, &locals, @@ -1941,13 +1940,12 @@ class FailingLogger(StubLogger): logger = FailingLogger() ", ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - succeed(py, &locals, &mut logging), - HookStep::Ready(()) - )); + let hooks = hooks(logged(py, &locals, true)); + assert_eq!(succeed(&locals, &hooks), 0); assert!( - logging + hooks + .0 + .lock() .response .as_ref() .unwrap() @@ -1970,10 +1968,11 @@ logger = FailingLogger() Python::attach(|py| { let locals = namespace(py, c"failure = ValueError('provider')"); run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - let step = fail(py, &locals, &mut logging); + let hooks = hooks(logged(py, &locals, asynchronous)); + let task = Scripted::default(); + fail(&locals, &hooks, &task).unwrap(); let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!(matches!(step, HookStep::Await(_, _)), awaits_async_handler); + assert_eq!(task.count() == 1, awaits_async_handler); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -2005,13 +2004,14 @@ class FailingLogger(StubLogger): logger = FailingLogger() ", ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - fail(py, &locals, &mut logging), - HookStep::Await(_, _) - )); + let hooks = hooks(logged(py, &locals, true)); + let task = Scripted::default(); + fail(&locals, &hooks, &task).unwrap(); + assert_eq!(task.count(), 1); assert!( - logging + hooks + .0 + .lock() .error .as_ref() .unwrap() @@ -2037,18 +2037,18 @@ logger = FailingLogger() Python::initialize(); Python::attach(|py| { let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = logged(py, &locals, true); - let HookStep::Await(_, resume) = fail(py, &locals, &mut logging) else { - panic!("expected async failure handler") - }; + let hooks = hooks(logged(py, &locals, true)); let result = match error { None => Ok(py.None()), Some(false) => Err(PyRuntimeError::new_err("handler failed")), Some(true) => Err(CancelledError::new_err("cancelled")), }; let expected = result.as_ref().err().map(|error| error.value(py).clone()); - match resume(&mut logging, py, result) { - Ok(step) => assert!(done && matches!(step, HookStep::Ready(()))), + let task = Scripted::answering([result]); + let ended = fail(&locals, &hooks, &task); + assert_eq!(task.count(), 1); + match ended { + Ok(()) => assert!(done), Err(propagated) => { assert!(!done); assert!(propagated.value(py).is(expected.unwrap())); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 38c6b1aedbd..35ad3b637b5 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -14,7 +14,7 @@ mod deferred; mod logger; mod mapping; mod python; -pub use adapter::LegacyLogging; +pub use adapter::{LegacyHooks, LegacyLogging}; pub use call::PublicCall; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs index a4aee4eb4da..8feac1a34bf 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs @@ -1,12 +1,15 @@ +use std::future::{Future, ready}; + +use futures_util::future::BoxFuture; use litellm_host::{ hooks::CallHooks, interceptors::{RawResponse, RequestContext, WireRequest}, lifecycle::{ExecutionEvent, FailureOrigin, Timing}, }; -use litellm_host_python::{HookStep, PythonCallEvent, PythonRuntime}; -use pyo3::{prelude::*, types::PyDict}; +use litellm_host_python::{Awaiter, PythonCallEvent, PythonRuntime}; +use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; -use crate::LegacyLogging; +use crate::LegacyHooks; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CallBoundary { @@ -48,44 +51,62 @@ impl Binding { } } -type Step = PyResult>; -type Prepare = fn(&mut LegacyLogging, Python<'_>, Py, f64) -> Step>; -type Before = - fn(&mut LegacyLogging, Python<'_>, Box, &RequestContext) -> Step>; -type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>; -type Transform = fn(&mut LegacyLogging, Python<'_>, Py, Timing) -> Step>; -type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py) -> Step<()>; -type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>; -type Open = fn(&mut LegacyLogging, Python<'_>, &Py) -> PyResult<()>; -type Chunk = fn(&mut LegacyLogging, Python<'_>, &Py) -> PyResult<()>; +type Prepare = for<'a> fn( + &'a LegacyHooks, + &'a dyn Awaiter, + Py, + f64, +) -> BoxFuture<'a, PyResult>>; +type Before = fn(&LegacyHooks, Box, &RequestContext) -> PyResult>; +type After = fn(&LegacyHooks, &RawResponse) -> PyResult<()>; +type Transform = for<'a> fn( + &'a LegacyHooks, + &'a dyn Awaiter, + Py, + Timing, +) -> BoxFuture<'a, PyResult>>; +type Success = fn(&LegacyHooks, Timing, &Py) -> PyResult<()>; +type Failure = for<'a> fn( + &'a LegacyHooks, + &'a dyn Awaiter, + Timing, + FailureOrigin, + &'a Py, +) -> BoxFuture<'a, PyResult<()>>; +type Open = fn(&LegacyHooks, &Py) -> PyResult<()>; +type Chunk = fn(&LegacyHooks, &Py) -> PyResult<()>; const PREPARE: Binding = Binding { boundary: CallBoundary::PrepareArguments, - invoke: LegacyLogging::prepare_call, + invoke: |hooks, task, arguments, started_at| { + Box::pin(hooks.prepare_call(task, arguments, started_at)) + }, callbacks: &["async_pre_call_deployment_hook"], }; const BEFORE: Binding = Binding { boundary: CallBoundary::BeforeProviderRequest, - invoke: LegacyLogging::pre_call, + invoke: LegacyHooks::pre_call, callbacks: &["log_pre_api_call", "log_input_event"], }; const AFTER: Binding = Binding { boundary: CallBoundary::AfterProviderResponse, - invoke: LegacyLogging::post_call, + invoke: LegacyHooks::post_call, callbacks: &["log_post_api_call"], }; const TRANSFORM: Binding = Binding { boundary: CallBoundary::TransformResponse, - invoke: LegacyLogging::transform_public_response, + invoke: |hooks, task, response, timing| { + Box::pin(hooks.transform_public_response(task, response, timing)) + }, callbacks: &["async_post_call_success_deployment_hook"], }; const SUCCESS: Binding = Binding { boundary: CallBoundary::Succeeded, - invoke: LegacyLogging::succeeded, + invoke: LegacyHooks::succeeded, callbacks: &[ "log_success_event", "async_log_success_event", @@ -99,7 +120,9 @@ const SUCCESS: Binding = Binding { const FAILURE: Binding = Binding { boundary: CallBoundary::Failed, - invoke: LegacyLogging::failed, + invoke: |hooks, task, timing, origin, error| { + Box::pin(hooks.failed(task, timing, origin, error)) + }, callbacks: &[ "async_post_call_failure_deployment_hook", "log_failure_event", @@ -112,13 +135,13 @@ const FAILURE: Binding = Binding { const OPEN: Binding = Binding { boundary: CallBoundary::StreamOpened, - invoke: LegacyLogging::stream_opened, + invoke: LegacyHooks::stream_opened, callbacks: &[], }; const CHUNK: Binding = Binding { boundary: CallBoundary::StreamChunk, - invoke: LegacyLogging::stream_chunk, + invoke: LegacyHooks::stream_chunk, callbacks: &[], }; @@ -135,67 +158,69 @@ pub fn callback_mappings() -> impl Iterator { .chain(PYTHON_CALLBACKS.iter().copied()) } -impl CallHooks for LegacyLogging { - fn prepare_arguments( - &mut self, - py: Python<'_>, +impl CallHooks for LegacyHooks { + fn prepare_arguments<'a>( + &'a self, + task: &'a dyn Awaiter, arguments: Py, started_at: f64, - ) -> Step> { - (PREPARE.invoke)(self, py, arguments, started_at) + ) -> impl Future>> + Send { + (PREPARE.invoke)(self, task, arguments, started_at) } - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - self.adopt_arguments(py, arguments); + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + self.adopt_arguments(arguments); Ok(()) } - fn before_provider_request( - &mut self, - py: Python<'_>, + fn before_provider_request<'a>( + &'a self, + _: &'a dyn Awaiter, wire: Box, context: &RequestContext, - ) -> Step> { - (BEFORE.invoke)(self, py, wire, context) + ) -> impl Future>> + Send { + ready((BEFORE.invoke)(self, wire, context)) } - fn transform_response( - &mut self, - py: Python<'_>, + fn transform_response<'a>( + &'a self, + task: &'a dyn Awaiter, response: Py, timing: Timing, - ) -> Step> { - (TRANSFORM.invoke)(self, py, response, timing) + ) -> impl Future>> + Send { + (TRANSFORM.invoke)(self, task, response, timing) } - fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> Step<()> { + async fn on_event<'a>( + &'a self, + task: &'a dyn Awaiter, + event: PythonCallEvent<'_>, + ) -> PyResult<()> { match event { - PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => { - Ok(HookStep::Ready(())) - } + PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => Ok(()), PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }) => { - self.result_ready(py, &facts) + self.result_ready(&facts) } PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - (AFTER.invoke)(self, py, raw) + (AFTER.invoke)(self, raw) } PythonCallEvent::Succeeded { timing, response } => { - (SUCCESS.invoke)(self, py, timing, response) + (SUCCESS.invoke)(self, timing, response) } PythonCallEvent::Failed { timing, origin, error, - } => (FAILURE.invoke)(self, py, timing, origin, error), + } => (FAILURE.invoke)(self, task, timing, origin, error).await, } } - fn on_stream_open(&mut self, py: Python<'_>, head: &Py) -> PyResult<()> { - (OPEN.invoke)(self, py, head) + fn on_stream_open(&self, head: &Py) -> PyResult<()> { + (OPEN.invoke)(self, head) } - fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { - (CHUNK.invoke)(self, py, chunk) + fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { + (CHUNK.invoke)(self, chunk) } } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 093ce60e510..88cb926411d 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -1,9 +1,17 @@ +use std::collections::VecDeque; use std::ffi::CStr; +use std::future::{Future, ready}; +use std::pin::pin; +use std::task::{Context, Poll, Waker}; + +use futures_util::future::BoxFuture; +use litellm_host_python::Awaiter; +use parking_lot::Mutex; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; -use crate::{LegacyLogging, PublicCall}; +use crate::{LegacyHooks, LegacyLogging, PublicCall}; /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. /// `tests/unit/rust_bridge/host/test_callbacks_legacy_python.py` pins this file to the Python @@ -193,3 +201,51 @@ pub(crate) fn legacy_call( let call = PublicCall::capture(&dict("bound"), &PyTuple::empty(py), &dict("kwargs")).unwrap(); LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous) } + +/// Answers each awaitable a hook hands it from `answers` in order, then with the awaitable +/// itself, and counts what it was handed. +#[derive(Default)] +pub(crate) struct Scripted { + answers: Mutex>>>, + awaited: Mutex, +} + +impl Scripted { + pub(crate) fn answering(answers: impl IntoIterator>>) -> Self { + Self { + answers: Mutex::new(answers.into_iter().collect()), + awaited: Mutex::new(0), + } + } + + pub(crate) fn count(&self) -> usize { + *self.awaited.lock() + } +} + +impl Awaiter for Scripted { + fn awaited(&self, awaitable: Py) -> BoxFuture<'_, PyResult>> { + *self.awaited.lock() += 1; + let answer = self.answers.lock().pop_front().unwrap_or(Ok(awaitable)); + Box::pin(ready(answer)) + } +} + +/// Runs a hook whose every await `Scripted` answers inline. +pub(crate) fn finish(hook: impl Future) -> T { + match pin!(hook).poll(&mut Context::from_waker(Waker::noop())) { + Poll::Ready(value) => value, + Poll::Pending => panic!("a scripted hook never suspends"), + } +} + +pub(crate) fn hooks(logging: LegacyLogging) -> LegacyHooks { + LegacyHooks::new(logging) +} + +pub(crate) fn failure(locals: &Bound<'_, PyDict>) -> Py { + local(locals, "failure") + .cast_into::() + .unwrap() + .unbind() +} diff --git a/litellm-rust/crates/host-python/src/caller.rs b/litellm-rust/crates/host-python/src/caller.rs index 5e80b294b88..bbbb2d94ea9 100644 --- a/litellm-rust/crates/host-python/src/caller.rs +++ b/litellm-rust/crates/host-python/src/caller.rs @@ -1,10 +1,11 @@ use std::ops::ControlFlow; +use futures_util::future::BoxFuture; use litellm_coroutine::{Abandoned, Co, Reply}; use pyo3::prelude::*; -use crate::Error; use crate::handle::Suspension; +use crate::{Awaiter, Error}; pub(crate) type Answer = PyResult>; @@ -56,3 +57,9 @@ impl PythonCaller { } } } + +impl Awaiter for PythonCaller { + fn awaited(&self, awaitable: Py) -> BoxFuture<'_, Answer> { + Box::pin(self.await_python(awaitable)) + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index f8e35aa09ba..7a2eac87f13 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -20,7 +20,7 @@ use litellm_host::{ use crate::PythonHostCalls; use crate::caller::{Answer, PythonCaller, Yielded}; use crate::handle::{Execution, ExecutionBody, ExecutionStep, StepTypes}; -use crate::hooks::{HookStep, PythonCallEvent, PythonCallHooks}; +use crate::hooks::{PythonCallEvent, PythonCallHooks}; use crate::native::{NativeMachine, NativePoll}; use crate::{Error, InvokeError, PythonBinding}; @@ -91,18 +91,27 @@ impl Observation { } } -/// 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 lock their own state around Python work, so the handle can traverse them while a +/// hook is suspended. +struct Shared { hooks: L, + parts: Mutex>, +} + +struct Parts { + binding: H, arguments: Option>, interrupted: Option>, - selected: Option>, + selected: Option, observation: Observation, } -type SharedParts = Arc>>; +/// The call's outcome while its terminal hooks run. The body borrows it for the event and +/// holds no Python reference of its own, so the one reference is the traversed one. +enum Selected { + Response(Arc>), + Exception(Arc>), +} enum Setup { Host(PyErr), @@ -121,7 +130,7 @@ fn classified(py: Python<'_>, binding: &H, error: ErrorOf) classifier_error } -/// The call as one async body. Every hook step and host call is awaited in place on the +/// The call as one async body. Every hook and host call is awaited in place on the /// caller's task; the parts lock is never held across an await. struct CallBody where @@ -129,7 +138,7 @@ where M: Machine, { caller: PythonCaller, - parts: SharedParts, + shared: Arc>, native: NativeMachine, start: Option, M>>, streaming: bool, @@ -142,8 +151,12 @@ where M: Machine + 'static, M::Complete: Into>>, { - fn parts(&self) -> MutexGuard<'_, Parts> { - self.parts.lock() + fn parts(&self) -> MutexGuard<'_, Parts> { + self.shared.parts.lock() + } + + fn hooks(&self) -> &L { + &self.shared.hooks } async fn run(mut self) -> PyResult> { @@ -159,45 +172,39 @@ where async fn begin(&mut self) -> Result<(), Setup>> { let started_at = epoch_seconds(); - let started = { + let started = PythonCallEvent::Started { + start_time: started_at, + }; + { 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)?; + parts.observation.emit(&started); + } + self.hooks() + .on_event(&self.caller, started) + .await + .map_err(Setup::Host)?; let arguments = self .parts() .arguments .take() .ok_or_else(|| Setup::Host(Error::MissingState.into()))?; - let prepare = { + let prepared = self + .hooks() + .prepare_arguments(&self.caller, arguments, started_at) + .await + .map_err(Setup::Host)?; + self.hooks() + .arguments_prepared(&prepared) + .map_err(Setup::Host)?; + let decoded = { 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 decoded = Python::attach(|py| parts.binding.decode_request(py, prepared.bind(py))); + parts.arguments = Some(prepared); + decoded }; + let request = decoded.map_err(Setup::Decode)?; let start = self .start .take() @@ -214,22 +221,6 @@ where 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 { @@ -286,11 +277,10 @@ where 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?; + let wire = self + .hooks() + .before_provider_request(&self.caller, wire, &context) + .await?; reply.send(*wire); Ok(()) } @@ -314,26 +304,18 @@ where 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) - })? + Python::attach(|py| parts.binding.encode_stream_head(py, head))? }; + self.hooks().on_stream_open(&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) - })? + Python::attach(|py| parts.binding.encode_chunk(py, chunk))? }; + self.hooks().on_stream_chunk(&chunk)?; reply.send(self.caller.yield_chunk(chunk).await); Ok(()) } @@ -359,12 +341,8 @@ where } 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 + self.parts().observation.emit(&event); + self.hooks().on_event(&self.caller, event).await } async fn completed( @@ -390,16 +368,12 @@ where 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 { + let timing = self.parts().observation.timing(); + match self + .hooks() + .transform_response(&self.caller, public, timing) + .await + { Ok(response) => self.succeeded(response).await, Err(error) => self.failure(error, FailureOrigin::Host).await, } @@ -417,58 +391,60 @@ where 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 + let response = Arc::new(response); + let timing = self.parts().observation.timing(); + let event = PythonCallEvent::Succeeded { + timing, + response: &response, }; - self.settle(step).await?; - let mut parts = self.parts(); - parts.observation.finish(); - parts - .selected - .take() - .ok_or_else(|| Error::MissingState.into()) + { + let mut parts = self.parts(); + parts.observation.retain_terminal(&event); + parts.selected = Some(Selected::Response(Arc::clone(&response))); + } + let notified = self.hooks().on_event(&self.caller, event).await; + { + let mut parts = self.parts(); + parts.selected = None; + notified?; + parts.observation.finish(); + } + Arc::into_inner(response).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 timing = { 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) - })? + parts.observation.timing() }; - self.settle(step).await?; - let mut parts = self.parts(); - parts.observation.finish(); - let selected = parts.selected.take().ok_or(Error::MissingState)?; + let exception = Python::attach(|py| { + if is_cancellation(py, &error) { + return Err(error); + } + Ok(Arc::new(error.into_value(py))) + })?; + let event = PythonCallEvent::Failed { + timing, + origin, + error: &exception, + }; + { + let mut parts = self.parts(); + parts.observation.retain_terminal(&event); + parts.selected = Some(Selected::Exception(Arc::clone(&exception))); + } + let notified = self.hooks().on_event(&self.caller, event).await; + { + let mut parts = self.parts(); + parts.selected = None; + notified?; + parts.observation.finish(); + } + let exception = Arc::into_inner(exception).ok_or(Error::MissingState)?; Err(Python::attach(|py| { - PyErr::from_value(selected.into_bound(py)) + PyErr::from_value(exception.into_bound(py).into_any()) })) } } @@ -480,7 +456,7 @@ where { body: Mutex, reply: Option>, - parts: SharedParts, + shared: Arc>, } impl CallDriver @@ -491,19 +467,20 @@ where /// 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(); + let mut parts = self.shared.parts.lock(); if !parts.observation.observing { return; } let timing = parts.observation.timing(); Python::attach(|py| { + let exception = error.value(py).clone().unbind(); let event = if is_cancellation(py, error) { PythonCallEvent::Cancelled { timing } } else { PythonCallEvent::Failed { timing, origin: FailureOrigin::Host, - error, + error: &exception, } }; parts.observation.emit(&event); @@ -550,14 +527,18 @@ where } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - let Some(parts) = self.parts.try_lock() else { + self.shared.hooks.traverse(visit)?; + let Some(parts) = self.shared.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) + match &parts.selected { + Some(Selected::Response(response)) => visit.call(&**response), + Some(Selected::Exception(exception)) => visit.call(&**exception), + None => Ok(()), + } } } @@ -569,17 +550,22 @@ where fn drop(&mut self) { self.body.get_mut().cancel(); self.reply = None; - let mut parts = self.parts.lock(); + let Some(Shared { hooks, parts }) = Arc::get_mut(&mut self.shared) else { + return; + }; + let Parts { + binding, + observation, + .. + } = parts.get_mut(); Python::attach(|py| { - if parts.observation.observing { - let timing = parts.observation.timing(); - parts - .observation - .emit(&PythonCallEvent::Cancelled { timing }); - parts.observation.observing = false; + if observation.observing { + let timing = observation.timing(); + observation.emit(&PythonCallEvent::Cancelled { timing }); + observation.observing = false; } - parts.hooks.close(py); - parts.binding.close(py); + hooks.close(py); + binding.close(py); }); } } @@ -605,27 +591,29 @@ where M: Machine + 'static, M::Complete: Into>>, { - let parts = Arc::new(Mutex::new(Parts { - binding, + let shared = Arc::new(Shared { hooks, - arguments: Some(arguments), - interrupted: None, - selected: None, - observation: Observation { - observers: options.observers, - observing: false, - started_at: 0.0, - ended_at: None, - terminal: None, - }, - })); + parts: Mutex::new(Parts { + binding, + arguments: Some(arguments), + interrupted: 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 shared = Arc::clone(&shared); let asynchronous = options.asynchronous; move |co| { CallBody { caller: PythonCaller::new(co), - parts, + shared, native: NativeMachine::new(asynchronous), start: Some(Box::new(start)), streaming: false, @@ -636,7 +624,7 @@ where let driver = CallDriver { body: Mutex::new(body), reply: None, - parts, + shared, }; Py::new(py, Execution::new(driver, options.steps)).map(Py::into_any) } diff --git a/litellm-rust/crates/host-python/src/hooks.rs b/litellm-rust/crates/host-python/src/hooks.rs index bab628bb7c6..409b37aa63a 100644 --- a/litellm-rust/crates/host-python/src/hooks.rs +++ b/litellm-rust/crates/host-python/src/hooks.rs @@ -2,32 +2,27 @@ mod chain; pub use chain::HookChain; use crate::PythonOwned; +use futures_util::future::BoxFuture; use litellm_host::hooks::{CallHooks, HookRuntime, RuntimeCallEvent}; +use pyo3::exceptions::PyBaseException; use pyo3::prelude::*; use pyo3::types::PyDict; -/// What a hook step produced: either the value the driver asked for, or a Python -/// awaitable the driver hands back to the caller's task before asking again. -pub type HookResume = fn(&mut L, Python<'_>, PyResult>) -> PyResult>; - -pub enum HookStep { - Await(Py, HookResume), - Ready(T), +/// Where a hook or host call awaits a Python awaitable: the caller's task, which the call +/// suspends to and resumes from with the awaitable's result. +pub trait Awaiter: Sync { + fn awaited(&self, awaitable: Py) -> BoxFuture<'_, PyResult>>; } pub struct PythonRuntime; impl HookRuntime for PythonRuntime { - type Context<'a> = Python<'a>; + type Context<'a> = &'a dyn Awaiter; type Arguments = Py; type Response = Py; type Chunk = Py; type Error = PyErr; - type Step = HookStep; - - fn ready(value: T) -> Self::Step { - HookStep::Ready(value) - } + type Failure = Py; } pub type PythonCallEvent<'a> = RuntimeCallEvent<'a, PythonRuntime>; diff --git a/litellm-rust/crates/host-python/src/hooks/chain.rs b/litellm-rust/crates/host-python/src/hooks/chain.rs new file mode 100644 index 00000000000..a29be6b40f6 --- /dev/null +++ b/litellm-rust/crates/host-python/src/hooks/chain.rs @@ -0,0 +1,279 @@ +use std::future::{Future, pending, poll_fn}; +use std::pin::pin; +use std::task::Poll; + +use futures_util::future::BoxFuture; +use litellm_host::{ + hooks::CallHooks, + interceptors::{RequestContext, WireRequest}, + lifecycle::{CallEvent, Timing}, +}; +use parking_lot::Mutex; +use pyo3::{ + exceptions::PyException, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; + +use crate::{Awaiter, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime}; + +trait ChainHooks: PythonOwned { + fn prepare_arguments<'a>( + &'a self, + task: &'a dyn Awaiter, + arguments: Py, + started_at: f64, + ) -> BoxFuture<'a, PyResult>>; + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()>; + fn before_provider_request<'a>( + &'a self, + task: &'a dyn Awaiter, + wire: Box, + context: &'a RequestContext, + ) -> BoxFuture<'a, PyResult>>; + fn transform_response<'a>( + &'a self, + task: &'a dyn Awaiter, + response: Py, + timing: Timing, + ) -> BoxFuture<'a, PyResult>>; + fn on_event<'a>( + &'a self, + task: &'a dyn Awaiter, + event: PythonCallEvent<'a>, + ) -> BoxFuture<'a, PyResult<()>>; + fn on_stream_open(&self, head: &Py) -> PyResult<()>; + fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()>; +} + +impl ChainHooks for H { + fn prepare_arguments<'a>( + &'a self, + task: &'a dyn Awaiter, + arguments: Py, + started_at: f64, + ) -> BoxFuture<'a, PyResult>> { + Box::pin(CallHooks::prepare_arguments( + self, task, arguments, started_at, + )) + } + + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + CallHooks::arguments_prepared(self, arguments) + } + + fn before_provider_request<'a>( + &'a self, + task: &'a dyn Awaiter, + wire: Box, + context: &'a RequestContext, + ) -> BoxFuture<'a, PyResult>> { + Box::pin(CallHooks::before_provider_request( + self, task, wire, context, + )) + } + + fn transform_response<'a>( + &'a self, + task: &'a dyn Awaiter, + response: Py, + timing: Timing, + ) -> BoxFuture<'a, PyResult>> { + Box::pin(CallHooks::transform_response(self, task, response, timing)) + } + + fn on_event<'a>( + &'a self, + task: &'a dyn Awaiter, + event: PythonCallEvent<'a>, + ) -> BoxFuture<'a, PyResult<()>> { + Box::pin(CallHooks::on_event(self, task, event)) + } + + fn on_stream_open(&self, head: &Py) -> PyResult<()> { + CallHooks::on_stream_open(self, head) + } + + fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { + CallHooks::on_stream_chunk(self, chunk) + } +} + +fn is_cancellation(error: &PyErr) -> bool { + Python::attach(|py| !error.is_instance_of::(py)) +} + +/// The caller's task as one chain stage hands it to its hooks. A cancellation the task +/// resumes with ends the stage there: the hook that awaited it never runs again. +struct Stage<'a> { + task: &'a dyn Awaiter, + cancelled: Mutex>, +} + +impl<'a> Stage<'a> { + fn new(task: &'a dyn Awaiter) -> Self { + Self { + task, + cancelled: Mutex::new(None), + } + } + + async fn guard(&self, hook: impl Future>) -> PyResult { + let mut hook = pin!(hook); + poll_fn(|context| match hook.as_mut().poll(context) { + Poll::Ready(result) => Poll::Ready(result), + Poll::Pending => self + .cancelled + .lock() + .take() + .map_or(Poll::Pending, |error| Poll::Ready(Err(error))), + }) + .await + } +} + +impl Awaiter for Stage<'_> { + fn awaited(&self, awaitable: Py) -> BoxFuture<'_, PyResult>> { + Box::pin(async move { + match self.task.awaited(awaitable).await { + Err(error) if is_cancellation(&error) => { + *self.cancelled.lock() = Some(error); + pending().await + } + answer => answer, + } + }) + } +} + +fn is_terminal(event: &CallEvent) -> bool { + matches!( + event, + CallEvent::Succeeded { .. } | CallEvent::Failed { .. } + ) +} + +/// A terminal notification reaches every hook, so an ordinary error from one is reported +/// and the next hook still runs. +fn notification_result(terminal: bool, result: PyResult<()>) -> PyResult<()> { + match result { + Err(error) if terminal && !is_cancellation(&error) => { + Python::attach(|py| error.write_unraisable(py, None)); + Ok(()) + } + result => result, + } +} + +#[derive(Default)] +pub struct HookChain { + hooks: Vec>, +} + +impl HookChain { + pub fn new() -> Self { + Self::default() + } + + pub fn with(mut self, hooks: impl PythonCallHooks + 'static) -> Self { + self.hooks.push(Box::new(hooks)); + self + } +} + +impl CallHooks for HookChain { + async fn prepare_arguments<'a>( + &'a self, + task: &'a dyn Awaiter, + arguments: Py, + started_at: f64, + ) -> PyResult> { + let stage = Stage::new(task); + let mut arguments = arguments; + for hooks in &self.hooks { + arguments = stage + .guard(hooks.prepare_arguments(&stage, arguments, started_at)) + .await?; + } + Ok(arguments) + } + + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.arguments_prepared(arguments)) + } + + async fn before_provider_request<'a>( + &'a self, + task: &'a dyn Awaiter, + wire: Box, + context: &RequestContext, + ) -> PyResult> { + let stage = Stage::new(task); + let mut wire = wire; + for hooks in &self.hooks { + wire = stage + .guard(hooks.before_provider_request(&stage, wire, context)) + .await?; + } + Ok(wire) + } + + async fn transform_response<'a>( + &'a self, + task: &'a dyn Awaiter, + response: Py, + timing: Timing, + ) -> PyResult> { + let stage = Stage::new(task); + let mut response = response; + for hooks in &self.hooks { + response = stage + .guard(hooks.transform_response(&stage, response, timing)) + .await?; + } + Ok(response) + } + + async fn on_event<'a>( + &'a self, + task: &'a dyn Awaiter, + event: PythonCallEvent<'_>, + ) -> PyResult<()> { + let stage = Stage::new(task); + let terminal = is_terminal(&event); + for hooks in &self.hooks { + let result = stage.guard(hooks.on_event(&stage, event.clone())).await; + notification_result(terminal, result)?; + } + Ok(()) + } + + fn on_stream_open(&self, head: &Py) -> PyResult<()> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.on_stream_open(head)) + } + + fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.on_stream_chunk(chunk)) + } +} + +impl PythonOwned for HookChain { + fn close(&mut self, py: Python<'_>) { + for hooks in &mut self.hooks { + hooks.close(py); + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.hooks + .iter() + .try_for_each(|hooks| hooks.traverse(visit)) + } +} diff --git a/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs b/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs deleted file mode 100644 index c7a5dad3858..00000000000 --- a/litellm-rust/crates/host-python/src/hooks/chain/adapter.rs +++ /dev/null @@ -1,202 +0,0 @@ -use litellm_host::interceptors::{RequestContext, WireRequest}; -use litellm_host::lifecycle::Timing; -use pyo3::{ - gc::{PyTraverseError, PyVisit}, - prelude::*, - types::PyDict, -}; - -use crate::{Error, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned}; - -pub(super) enum ChainStep { - Ready(T), - Await(Py), -} - -enum Continuation { - Arguments(HookResume>), - Wire(HookResume>), - Response(HookResume>), - Event(HookResume), -} - -pub(super) struct HookAdapter { - hooks: H, - continuation: Option>, -} - -impl HookAdapter { - pub(super) fn new(hooks: H) -> Self { - Self { - hooks, - continuation: None, - } - } - - fn step( - &mut self, - step: HookStep, - continuation: impl FnOnce(HookResume) -> Continuation, - ) -> ChainStep { - match step { - HookStep::Ready(value) => ChainStep::Ready(value), - HookStep::Await(awaitable, resume) => { - self.continuation = Some(continuation(resume)); - ChainStep::Await(awaitable) - } - } - } -} - -pub(super) trait ChainHooks: PythonOwned { - fn prepare_arguments( - &mut self, - py: Python<'_>, - arguments: Py, - started_at: f64, - ) -> PyResult>>; - fn resume_arguments( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>>; - fn before_provider_request( - &mut self, - py: Python<'_>, - wire: Box, - context: &RequestContext, - ) -> PyResult>>; - fn resume_wire( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>>; - fn transform_response( - &mut self, - py: Python<'_>, - response: Py, - timing: Timing, - ) -> PyResult>>; - fn resume_response( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>>; - fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> PyResult>; - fn resume_event( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>; - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()>; - fn on_stream_open(&mut self, py: Python<'_>, head: &Py) -> PyResult<()>; - fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()>; -} - -impl ChainHooks for HookAdapter { - fn prepare_arguments( - &mut self, - py: Python<'_>, - arguments: Py, - started_at: f64, - ) -> PyResult>> { - let step = self.hooks.prepare_arguments(py, arguments, started_at)?; - Ok(self.step(step, Continuation::Arguments)) - } - - fn resume_arguments( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let Some(Continuation::Arguments(resume)) = self.continuation.take() else { - return Err(Error::MissingState.into()); - }; - let step = resume(&mut self.hooks, py, result)?; - Ok(self.step(step, Continuation::Arguments)) - } - - fn before_provider_request( - &mut self, - py: Python<'_>, - wire: Box, - context: &RequestContext, - ) -> PyResult>> { - let step = self.hooks.before_provider_request(py, wire, context)?; - Ok(self.step(step, Continuation::Wire)) - } - - fn resume_wire( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let Some(Continuation::Wire(resume)) = self.continuation.take() else { - return Err(Error::MissingState.into()); - }; - let step = resume(&mut self.hooks, py, result)?; - Ok(self.step(step, Continuation::Wire)) - } - - fn transform_response( - &mut self, - py: Python<'_>, - response: Py, - timing: Timing, - ) -> PyResult>> { - let step = self.hooks.transform_response(py, response, timing)?; - Ok(self.step(step, Continuation::Response)) - } - - fn resume_response( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let Some(Continuation::Response(resume)) = self.continuation.take() else { - return Err(Error::MissingState.into()); - }; - let step = resume(&mut self.hooks, py, result)?; - Ok(self.step(step, Continuation::Response)) - } - - fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> PyResult> { - let step = self.hooks.on_event(py, event)?; - Ok(self.step(step, Continuation::Event)) - } - - fn resume_event( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult> { - let Some(Continuation::Event(resume)) = self.continuation.take() else { - return Err(Error::MissingState.into()); - }; - let step = resume(&mut self.hooks, py, result)?; - Ok(self.step(step, Continuation::Event)) - } - - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - self.hooks.arguments_prepared(py, arguments) - } - - fn on_stream_open(&mut self, py: Python<'_>, head: &Py) -> PyResult<()> { - self.hooks.on_stream_open(py, head) - } - - fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { - self.hooks.on_stream_chunk(py, chunk) - } -} - -impl PythonOwned for HookAdapter { - fn close(&mut self, py: Python<'_>) { - self.continuation = None; - self.hooks.close(py); - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.hooks.traverse(visit) - } -} diff --git a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs deleted file mode 100644 index 4939995cefa..00000000000 --- a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs +++ /dev/null @@ -1,394 +0,0 @@ -use litellm_host::{ - hooks::CallHooks, - interceptors::{RawResponse, RequestContext, WireRequest}, - lifecycle::{CallEvent, ExecutionEvent, Timing}, -}; -use pyo3::{ - exceptions::{PyBaseException, PyException}, - gc::{PyTraverseError, PyVisit}, - prelude::*, - types::PyDict, -}; - -use crate::{Error, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime}; - -use super::adapter::{ChainHooks, ChainStep, HookAdapter}; - -type OwnedEvent = CallEvent, Py, RawResponse>; - -#[derive(Default)] -pub struct HookChain { - hooks: Vec>, - arguments: Option<(usize, f64)>, - wire: Option<(usize, RequestContext)>, - response: Option<(usize, Timing)>, - event: Option<(usize, OwnedEvent)>, -} - -impl HookChain { - pub fn new() -> Self { - Self::default() - } - - pub fn with(mut self, hooks: impl PythonCallHooks + 'static) -> Self { - self.hooks.push(Box::new(HookAdapter::new(hooks))); - self - } - - fn arguments_from( - &mut self, - py: Python<'_>, - index: usize, - arguments: Py, - started_at: f64, - ) -> PyResult>> { - let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(arguments)); - }; - let step = hooks.prepare_arguments(py, arguments, started_at)?; - self.arguments_step(py, index, step, started_at) - } - - fn arguments_step( - &mut self, - py: Python<'_>, - index: usize, - step: ChainStep>, - started_at: f64, - ) -> PyResult>> { - match step { - ChainStep::Ready(value) => self.arguments_from(py, index + 1, value, started_at), - ChainStep::Await(awaitable) => { - self.arguments = Some((index, started_at)); - Ok(HookStep::Await(awaitable, Self::resume_arguments)) - } - } - } - - fn resume_arguments( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let (index, context) = self.arguments.take().ok_or(Error::MissingState)?; - let result = resume_unless_cancelled(py, result)?; - let step = self.hooks[index].resume_arguments(py, result)?; - self.arguments_step(py, index, step, context) - } - - fn wire_from( - &mut self, - py: Python<'_>, - index: usize, - wire: Box, - context: &RequestContext, - ) -> PyResult>> { - let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(wire)); - }; - let step = hooks.before_provider_request(py, wire, context)?; - self.wire_step(py, index, step, context) - } - - fn wire_step( - &mut self, - py: Python<'_>, - index: usize, - step: ChainStep>, - context: &RequestContext, - ) -> PyResult>> { - match step { - ChainStep::Ready(value) => self.wire_from(py, index + 1, value, context), - ChainStep::Await(awaitable) => { - self.wire = Some((index, context.clone())); - Ok(HookStep::Await(awaitable, Self::resume_wire)) - } - } - } - - fn resume_wire( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let (index, context) = self.wire.take().ok_or(Error::MissingState)?; - let result = resume_unless_cancelled(py, result)?; - let step = self.hooks[index].resume_wire(py, result)?; - self.wire_step(py, index, step, &context) - } - - fn response_from( - &mut self, - py: Python<'_>, - index: usize, - response: Py, - timing: Timing, - ) -> PyResult>> { - let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(response)); - }; - let step = hooks.transform_response(py, response, timing)?; - self.response_step(py, index, step, timing) - } - - fn response_step( - &mut self, - py: Python<'_>, - index: usize, - step: ChainStep>, - timing: Timing, - ) -> PyResult>> { - match step { - ChainStep::Ready(value) => self.response_from(py, index + 1, value, timing), - ChainStep::Await(awaitable) => { - self.response = Some((index, timing)); - Ok(HookStep::Await(awaitable, Self::resume_response)) - } - } - } - - fn resume_response( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let (index, context) = self.response.take().ok_or(Error::MissingState)?; - let result = resume_unless_cancelled(py, result)?; - let step = self.hooks[index].resume_response(py, result)?; - self.response_step(py, index, step, context) - } - - fn event_from( - &mut self, - py: Python<'_>, - index: usize, - event: OwnedEvent, - ) -> PyResult> { - let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(())); - }; - let result = dispatch(py, hooks.as_mut(), &event); - let step = notification_result(py, is_terminal(&event), result)?; - self.event_step(py, index, step, event) - } - - fn event_step( - &mut self, - py: Python<'_>, - index: usize, - step: ChainStep<()>, - event: OwnedEvent, - ) -> PyResult> { - match step { - ChainStep::Ready(()) => self.event_from(py, index + 1, event), - ChainStep::Await(awaitable) => { - self.event = Some((index, event)); - Ok(HookStep::Await(awaitable, Self::resume_event)) - } - } - } - - fn resume_event( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult> { - let (index, event) = self.event.take().ok_or(Error::MissingState)?; - let result = resume_unless_cancelled(py, result)?; - let result = self.hooks[index].resume_event(py, result); - let step = notification_result(py, is_terminal(&event), result)?; - self.event_step(py, index, step, event) - } -} - -fn resume_unless_cancelled( - py: Python<'_>, - result: PyResult>, -) -> PyResult>> { - match result { - Err(error) if !error.is_instance_of::(py) => Err(error), - result => Ok(result), - } -} - -fn is_terminal(event: &CallEvent) -> bool { - matches!( - event, - CallEvent::Succeeded { .. } | CallEvent::Failed { .. } - ) -} - -fn notification_result( - py: Python<'_>, - terminal: bool, - result: PyResult>, -) -> PyResult> { - match result { - Err(error) if terminal && error.is_instance_of::(py) => { - error.write_unraisable(py, None); - Ok(ChainStep::Ready(())) - } - result => result, - } -} - -fn retain_event(py: Python<'_>, event: PythonCallEvent<'_>) -> OwnedEvent { - match event { - CallEvent::Execution(ExecutionEvent::ResultReady { facts }) => { - CallEvent::Execution(ExecutionEvent::ResultReady { facts }) - } - CallEvent::Started { start_time } => CallEvent::Started { start_time }, - CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }) - } - CallEvent::Succeeded { timing, response } => CallEvent::Succeeded { - timing, - response: response.clone_ref(py), - }, - CallEvent::Failed { - timing, - origin, - error, - } => CallEvent::Failed { - timing, - origin, - error: error.clone_ref(py).into_value(py), - }, - CallEvent::Cancelled { timing } => CallEvent::Cancelled { timing }, - } -} - -fn dispatch( - py: Python<'_>, - hooks: &mut dyn ChainHooks, - event: &OwnedEvent, -) -> PyResult> { - match event { - CallEvent::Execution(ExecutionEvent::ResultReady { facts }) => hooks.on_event( - py, - CallEvent::Execution(ExecutionEvent::ResultReady { - facts: facts.clone(), - }), - ), - CallEvent::Started { start_time } => hooks.on_event( - py, - CallEvent::Started { - start_time: *start_time, - }, - ), - CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => hooks.on_event( - py, - CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }), - ), - CallEvent::Succeeded { timing, response } => hooks.on_event( - py, - CallEvent::Succeeded { - timing: *timing, - response, - }, - ), - CallEvent::Failed { - timing, - origin, - error, - } => hooks.on_event( - py, - CallEvent::Failed { - timing: *timing, - origin: *origin, - error: &PyErr::from_value(error.bind(py).clone().into_any()), - }, - ), - CallEvent::Cancelled { timing } => { - hooks.on_event(py, CallEvent::Cancelled { timing: *timing }) - } - } -} - -impl CallHooks for HookChain { - fn prepare_arguments( - &mut self, - py: Python<'_>, - arguments: Py, - started_at: f64, - ) -> PyResult>> { - self.arguments_from(py, 0, arguments, started_at) - } - - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - self.hooks - .iter_mut() - .try_for_each(|hooks| hooks.arguments_prepared(py, arguments)) - } - - fn before_provider_request( - &mut self, - py: Python<'_>, - wire: Box, - context: &RequestContext, - ) -> PyResult>> { - self.wire_from(py, 0, wire, context) - } - - fn transform_response( - &mut self, - py: Python<'_>, - response: Py, - timing: Timing, - ) -> PyResult>> { - self.response_from(py, 0, response, timing) - } - - fn on_event( - &mut self, - py: Python<'_>, - event: PythonCallEvent<'_>, - ) -> PyResult> { - for (index, hooks) in self.hooks.iter_mut().enumerate() { - let result = hooks.on_event(py, event.clone()); - match notification_result(py, is_terminal(&event), result)? { - ChainStep::Ready(()) => {} - ChainStep::Await(awaitable) => { - self.event = Some((index, retain_event(py, event))); - return Ok(HookStep::Await(awaitable, Self::resume_event)); - } - } - } - Ok(HookStep::Ready(())) - } - - fn on_stream_open(&mut self, py: Python<'_>, head: &Py) -> PyResult<()> { - self.hooks - .iter_mut() - .try_for_each(|hooks| hooks.on_stream_open(py, head)) - } - - fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { - self.hooks - .iter_mut() - .try_for_each(|hooks| hooks.on_stream_chunk(py, chunk)) - } -} - -impl PythonOwned for HookChain { - fn close(&mut self, py: Python<'_>) { - self.arguments = None; - self.wire = None; - self.response = None; - self.event = None; - for hooks in &mut self.hooks { - hooks.close(py); - } - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - for hooks in &self.hooks { - hooks.traverse(visit)?; - } - match &self.event { - Some((_, CallEvent::Succeeded { response, .. })) => visit.call(response), - Some((_, CallEvent::Failed { error, .. })) => visit.call(error), - _ => Ok(()), - } - } -} diff --git a/litellm-rust/crates/host-python/src/hooks/chain/mod.rs b/litellm-rust/crates/host-python/src/hooks/chain/mod.rs deleted file mode 100644 index c6267c385ff..00000000000 --- a/litellm-rust/crates/host-python/src/hooks/chain/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -mod adapter; -mod dispatch; - -pub use dispatch::HookChain; diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8e37ab02669..babc870655f 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -29,7 +29,7 @@ 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, StepTypes, Suspension}; -pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime}; +pub use hooks::{Awaiter, HookChain, 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 157cfc92148..3365eabade8 100644 --- a/litellm-rust/crates/host-python/tests/driver.rs +++ b/litellm-rust/crates/host-python/tests/driver.rs @@ -535,11 +535,11 @@ enum ArgumentPolicy { } impl CallHooks for ArgumentPolicy { - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - match self { + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + Python::attach(|py| 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())), - } + }) } } diff --git a/litellm-rust/crates/host-python/tests/handle.rs b/litellm-rust/crates/host-python/tests/handle.rs index fe900f2c1ec..e8bf19e9476 100644 --- a/litellm-rust/crates/host-python/tests/handle.rs +++ b/litellm-rust/crates/host-python/tests/handle.rs @@ -2,9 +2,10 @@ mod support; use litellm_host::{hooks::CallHooks, lifecycle::Timing}; use litellm_host_python::{ - Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime, + Awaiter, Execution, ExecutionBody, ExecutionStep, PythonCallEvent, PythonOwned, PythonRuntime, Suspension, run_call, }; +use parking_lot::Mutex; use pyo3::{ exceptions::PyBaseException, gc::{PyTraverseError, PyVisit}, @@ -114,55 +115,59 @@ enum Holder { } struct PausingHooks { - retained: Option>, - pause: Option>, + retained: Mutex>>, + pause: Mutex>>, holder: Holder, } impl CallHooks for PausingHooks { - fn transform_response( - &mut self, - py: Python<'_>, + async fn transform_response<'a>( + &'a self, + _: &'a dyn Awaiter, response: Py, _: Timing, - ) -> PyResult>> { + ) -> PyResult> { + let retained = || self.retained.lock().take().unwrap(); 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), - )), + Holder::Hooks => Ok(response), + Holder::Response => Ok(retained()), + Holder::Exception => Err(Python::attach(|py| { + PyErr::from_value(retained().into_bound(py)) + })), } } - fn on_event( - &mut self, - _: Python<'_>, + async fn on_event<'a>( + &'a self, + task: &'a dyn Awaiter, event: PythonCallEvent<'_>, - ) -> PyResult> { + ) -> 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(()), - }) + let pause = self.pause.lock().take_if(|_| pauses_here); + if let Some(pause) = pause { + task.awaited(pause).await?; + } + Ok(()) } } impl PythonOwned for PausingHooks { fn close(&mut self, _: Python<'_>) { - self.retained = None; - self.pause = None; + *self.retained.get_mut() = None; + *self.pause.get_mut() = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.retained)?; - visit.call(&self.pause) + for slot in [&self.retained, &self.pause] { + if let Some(value) = slot.try_lock() { + visit.call(&*value)?; + } + } + Ok(()) } } @@ -180,8 +185,8 @@ fn paused_call( move |_, _, request| Ok(machine(request)), binding(&Log::default(), OpScript::Answer), PausingHooks { - retained: Some(retained), - pause: Some(pause), + retained: Mutex::new(Some(retained)), + pause: Mutex::new(Some(pause)), holder, }, PyDict::new(py).unbind(), diff --git a/litellm-rust/crates/host-python/tests/hook_chain.rs b/litellm-rust/crates/host-python/tests/hook_chain.rs index 4b520ca2665..b7ad24fd784 100644 --- a/litellm-rust/crates/host-python/tests/hook_chain.rs +++ b/litellm-rust/crates/host-python/tests/hook_chain.rs @@ -1,74 +1,57 @@ +mod support; + +use std::future::{Future, ready}; + +use futures_util::{FutureExt, future::BoxFuture}; use litellm_host::{ hooks::CallHooks, interceptors::{RequestContext, WireRequest}, lifecycle::{FailureOrigin, Timing}, }; -use litellm_host_python::{HookChain, HookStep, PythonCallEvent, PythonOwned, PythonRuntime}; +use litellm_host_python::{ + Awaiter, HookChain, PythonCallEvent, PythonOwned, PythonRuntime, run_call, +}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, types::{PyDict, PyTuple}, }; use rstest::{fixture, rstest}; +use support::{ + Interpreter, Log, OpScript, binding, call_options, handed, interpreter, success_machine, +}; struct ScriptHooks { object: Py, asynchronous: bool, - wire: Option>, } impl ScriptHooks { - fn invoke(&self, py: Python<'_>, name: &str, value: Py) -> PyResult> { + async fn invoke( + &self, + task: &dyn Awaiter, + name: &str, + value: impl FnOnce(Python<'_>) -> PyResult>, + ) -> PyResult> { + let result = Python::attach(|py| { + let value = value(py)?; + if self.asynchronous { + self.object.call_method1(py, "invoke", (name, value)) + } else { + self.object.call_method1(py, name, (value,)) + } + })?; if self.asynchronous { - self.object.call_method1(py, "invoke", (name, value)) + task.awaited(result).await } else { - self.object.call_method1(py, name, (value,)) + Ok(result) } } - - fn arguments( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - Ok(HookStep::Ready( - result?.into_bound(py).cast_into::()?.unbind(), - )) - } - - fn response( - &mut self, - _: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - result.map(HookStep::Ready) - } - - fn notification( - &mut self, - _: Python<'_>, - result: PyResult>, - ) -> PyResult> { - result.map(|_| HookStep::Ready(())) - } - - fn request( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let wire = self.wire.take().unwrap(); - Ok(HookStep::Ready(Box::new(WireRequest { - url: result?.extract(py)?, - ..*wire - }))) - } } impl PythonOwned for ScriptHooks { fn close(&mut self, py: Python<'_>) { self.object = py.None(); - self.wire = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -77,98 +60,105 @@ impl PythonOwned for ScriptHooks { } impl CallHooks for ScriptHooks { - fn prepare_arguments( - &mut self, - py: Python<'_>, + async fn prepare_arguments<'a>( + &'a self, + task: &'a dyn Awaiter, arguments: Py, _: f64, - ) -> PyResult>> { - let value = self.invoke(py, "prepare", arguments.into_any())?; - if self.asynchronous { - Ok(HookStep::Await(value, Self::arguments)) - } else { - self.arguments(py, Ok(value)) - } + ) -> PyResult> { + let value = self + .invoke(task, "prepare", |_| Ok(arguments.into_any())) + .await?; + Python::attach(|py| Ok(value.into_bound(py).cast_into::()?.unbind())) } - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - self.object.call_method1(py, "adopt", (arguments,))?; + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + Python::attach(|py| self.object.call_method1(py, "adopt", (arguments,)))?; Ok(()) } - fn before_provider_request( - &mut self, - py: Python<'_>, + async fn before_provider_request<'a>( + &'a self, + task: &'a dyn Awaiter, wire: Box, context: &RequestContext, - ) -> PyResult>> { - self.object.bind(py).setattr("model", &context.model)?; - let value = self.invoke( - py, - "before", - wire.url.clone().into_pyobject(py)?.into_any().unbind(), - )?; - self.wire = Some(wire); - if self.asynchronous { - Ok(HookStep::Await(value, Self::request)) - } else { - self.request(py, Ok(value)) - } + ) -> PyResult> { + Python::attach(|py| self.object.bind(py).setattr("model", &context.model))?; + let url = self + .invoke(task, "before", |py| { + Ok(wire.url.clone().into_pyobject(py)?.into_any().unbind()) + }) + .await?; + Ok(Box::new(WireRequest { + url: Python::attach(|py| url.extract(py))?, + ..*wire + })) } - fn transform_response( - &mut self, - py: Python<'_>, + async fn transform_response<'a>( + &'a self, + task: &'a dyn Awaiter, response: Py, _: Timing, - ) -> PyResult>> { - let value = self.invoke(py, "transform", response)?; - if self.asynchronous { - Ok(HookStep::Await(value, Self::response)) - } else { - self.response(py, Ok(value)) - } + ) -> PyResult> { + self.invoke(task, "transform", |_| Ok(response)).await } - fn on_event( - &mut self, - py: Python<'_>, + async fn on_event<'a>( + &'a self, + task: &'a dyn Awaiter, event: PythonCallEvent<'_>, - ) -> PyResult> { - let (name, value) = match event { - PythonCallEvent::Succeeded { response, .. } => ("success", response.clone_ref(py)), - PythonCallEvent::Failed { error, .. } => { - ("failure", error.clone_ref(py).into_value(py).into_any()) - } - PythonCallEvent::Started { .. } => ("started", py.None()), - PythonCallEvent::Cancelled { .. } => ("cancelled", py.None()), - PythonCallEvent::Execution(_) => ("provider", py.None()), - }; - let value = self.invoke( - py, - "event", - PyTuple::new(py, [name.into_pyobject(py)?.into_any().unbind(), value])? - .into_any() - .unbind(), - )?; - if self.asynchronous { - Ok(HookStep::Await(value, Self::notification)) - } else { - self.notification(py, Ok(value)) - } + ) -> PyResult<()> { + self.invoke(task, "event", |py| { + let (name, value) = match event { + PythonCallEvent::Succeeded { response, .. } => ("success", response.clone_ref(py)), + PythonCallEvent::Failed { error, .. } => { + ("failure", error.clone_ref(py).into_any()) + } + PythonCallEvent::Started { .. } => ("started", py.None()), + PythonCallEvent::Cancelled { .. } => ("cancelled", py.None()), + PythonCallEvent::Execution(_) => ("provider", py.None()), + }; + Ok( + PyTuple::new(py, [name.into_pyobject(py)?.into_any().unbind(), value])? + .into_any() + .unbind(), + ) + }) + .await + .map(drop) } - fn on_stream_open(&mut self, py: Python<'_>, _head: &Py) -> PyResult<()> { - self.object.call_method1(py, "stream", (py.None(),))?; + fn on_stream_open(&self, _head: &Py) -> PyResult<()> { + Python::attach(|py| self.object.call_method1(py, "stream", (py.None(),)))?; Ok(()) } - fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { - self.object.call_method1(py, "stream", (chunk,))?; + fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { + Python::attach(|py| self.object.call_method1(py, "stream", (chunk,)))?; Ok(()) } } +/// Runs every awaitable a hook hands it to completion on a fresh event loop. +struct RunToCompletion; + +impl Awaiter for RunToCompletion { + fn awaited(&self, awaitable: Py) -> BoxFuture<'_, PyResult>> { + let result = Python::attach(|py| { + py.import("asyncio")? + .call_method1("run", (awaitable,)) + .map(Bound::unbind) + }); + Box::pin(ready(result)) + } +} + +fn finish(hook: impl Future>) -> PyResult { + hook.now_or_never() + .expect("every awaitable runs to completion inline") +} + #[fixture] fn scripts() -> Py { Python::initialize(); @@ -223,25 +213,10 @@ fn chain(py: Python<'_>, scripts: &Py, asynchronous: bool) -> HookChain let hook = |name| ScriptHooks { object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(), asynchronous, - wire: None, }; HookChain::new().with(hook("first")).with(hook("second")) } -fn finish(py: Python<'_>, hooks: &mut H, step: HookStep) -> PyResult { - match step { - HookStep::Ready(value) => Ok(value), - HookStep::Await(awaitable, resume) => { - let value = py - .import("asyncio")? - .call_method1("run", (awaitable,)) - .map(Bound::unbind); - let next = resume(hooks, py, value)?; - finish(py, hooks, next) - } - } -} - const TIMING: Timing = Timing { start_time: 0.0, end_time: 1.0, @@ -250,8 +225,8 @@ const TIMING: Timing = Timing { struct PreparedPolicy; impl CallHooks for PreparedPolicy { - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - arguments.bind(py).set_item("policy", "configured") + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + Python::attach(|py| arguments.bind(py).set_item("policy", "configured")) } } @@ -271,13 +246,11 @@ fn transformations_feed_each_other_and_notifications_share_final_values( #[case] asynchronous: bool, ) { Python::attach(|py| { - let mut hooks = chain(py, &scripts, asynchronous).with(PreparedPolicy); + let hooks = chain(py, &scripts, asynchronous).with(PreparedPolicy); let original = PyDict::new(py).unbind(); - let step = hooks - .prepare_arguments(py, original.clone_ref(py), 0.0) - .unwrap(); - let arguments = finish(py, &mut hooks, step).unwrap(); - hooks.arguments_prepared(py, &arguments).unwrap(); + let arguments = + finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap(); + hooks.arguments_prepared(&arguments).unwrap(); let context = RequestContext { model: "model".into(), custom_llm_provider: "provider".into(), @@ -290,25 +263,22 @@ fn transformations_feed_each_other_and_notifications_share_final_values( headers: vec![], body: serde_json::Value::Null, }); - let step = hooks.before_provider_request(py, wire, &context).unwrap(); - assert_eq!(finish(py, &mut hooks, step).unwrap().url, "urlab"); + let wire = finish(hooks.before_provider_request(&RunToCompletion, wire, &context)).unwrap(); + assert_eq!(wire.url, "urlab"); let response = PyDict::new(py).unbind().into_any(); - let step = hooks - .transform_response(py, response.clone_ref(py), TIMING) - .unwrap(); - let final_response = finish(py, &mut hooks, step).unwrap(); - let step = hooks - .on_event( - py, - PythonCallEvent::Succeeded { - timing: TIMING, - response: &final_response, - }, - ) - .unwrap(); - finish(py, &mut hooks, step).unwrap(); - hooks.on_stream_open(py, &py.None()).unwrap(); - hooks.on_stream_chunk(py, &response).unwrap(); + let final_response = + finish(hooks.transform_response(&RunToCompletion, response.clone_ref(py), TIMING)) + .unwrap(); + finish(hooks.on_event( + &RunToCompletion, + PythonCallEvent::Succeeded { + timing: TIMING, + response: &final_response, + }, + )) + .unwrap(); + hooks.on_stream_open(&py.None()).unwrap(); + hooks.on_stream_chunk(&response).unwrap(); let locals = scripts.bind(py); locals.set_item("arguments", arguments).unwrap(); locals.set_item("original", original).unwrap(); @@ -348,7 +318,7 @@ fn terminal_failure_does_not_skip_later_hooks( #[values("first", "second")] failing_hook: &str, ) { Python::attach(|py| { - let mut hooks = chain(py, &scripts, asynchronous).with(ScriptHooks { + let hooks = chain(py, &scripts, asynchronous).with(ScriptHooks { object: scripts .bind(py) .get_item("third") @@ -356,7 +326,6 @@ fn terminal_failure_does_not_skip_later_hooks( .unwrap() .unbind(), asynchronous, - wire: None, }); let locals = scripts.bind(py); locals @@ -368,7 +337,7 @@ fn terminal_failure_does_not_skip_later_hooks( pyo3::exceptions::PyRuntimeError::new_err("callback failed").into_value(py), ) .unwrap(); - let error = pyo3::exceptions::PyValueError::new_err("provider failed"); + let error = pyo3::exceptions::PyValueError::new_err("provider failed").into_value(py); let response = py.None(); let event = if failed { PythonCallEvent::Failed { @@ -382,17 +351,9 @@ fn terminal_failure_does_not_skip_later_hooks( response: &response, } }; - let step = hooks.on_event(py, event).unwrap(); - finish(py, &mut hooks, step).unwrap(); + finish(hooks.on_event(&RunToCompletion, event)).unwrap(); locals - .set_item( - "selected", - if failed { - error.into_value(py).into_any() - } else { - response - }, - ) + .set_item("selected", if failed { error.into_any() } else { response }) .unwrap(); py.run( c" @@ -412,7 +373,7 @@ assert all(value is selected for name, kind, value in log) #[case::async_(true)] fn transformation_failure_stops_the_chain(scripts: Py, #[case] asynchronous: bool) { Python::attach(|py| { - let mut hooks = chain(py, &scripts, asynchronous); + let hooks = chain(py, &scripts, asynchronous); let locals = scripts.bind(py); py.run( c"first.error = ValueError('prepare failed')", @@ -420,9 +381,8 @@ fn transformation_failure_stops_the_chain(scripts: Py, #[case] asynchron Some(locals), ) .unwrap(); - let result = hooks - .prepare_arguments(py, PyDict::new(py).unbind(), 0.0) - .and_then(|step| finish(py, &mut hooks, step)); + let result = + finish(hooks.prepare_arguments(&RunToCompletion, PyDict::new(py).unbind(), 0.0)); assert!( result.unwrap_err().value(py).is(locals .get_item("first") @@ -445,7 +405,7 @@ fn transformation_failure_stops_the_chain(scripts: Py, #[case] asynchron #[case::async_(true)] fn cancellation_stops_notification_dispatch(scripts: Py, #[case] asynchronous: bool) { Python::attach(|py| { - let mut hooks = chain(py, &scripts, asynchronous); + let hooks = chain(py, &scripts, asynchronous); let locals = scripts.bind(py); py.run( c"first.error = asyncio.CancelledError()", @@ -454,15 +414,13 @@ fn cancellation_stops_notification_dispatch(scripts: Py, #[case] asynchr ) .unwrap(); let response = py.None(); - let result = hooks - .on_event( - py, - PythonCallEvent::Succeeded { - timing: TIMING, - response: &response, - }, - ) - .and_then(|step| finish(py, &mut hooks, step)); + let result = finish(hooks.on_event( + &RunToCompletion, + PythonCallEvent::Succeeded { + timing: TIMING, + response: &response, + }, + )); assert!( result .unwrap_err() @@ -477,56 +435,11 @@ fn cancellation_stops_notification_dispatch(scripts: Py, #[case] asynchr }); } -struct Preparing { - hooks: HookChain, - arguments: Option>, - resume: Option>>, -} - -impl litellm_host_python::ExecutionBody for Preparing { - fn resume( - &mut self, - result: Option>>, - ) -> PyResult { - Python::attach(|py| { - let step = match self.resume.take() { - Some(resume) => resume(&mut self.hooks, py, result.unwrap())?, - None => self - .hooks - .prepare_arguments(py, self.arguments.take().unwrap(), 0.0)?, - }; - match step { - HookStep::Ready(arguments) => { - self.hooks.arguments_prepared(py, &arguments)?; - Ok(litellm_host_python::ExecutionStep::Return( - arguments.into_any(), - )) - } - HookStep::Await(awaitable, resume) => { - self.resume = Some(resume); - Ok(litellm_host_python::ExecutionStep::Suspended( - litellm_host_python::Suspension::Await(awaitable), - )) - } - } - }) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.arguments)?; - self.hooks.traverse(visit) - } -} - -fn lifecycle(py: Python<'_>) -> PyResult> { - let source = - std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/execution.py")) - .unwrap(); - PyModule::from_code(py, &source, c"execution.py", c"hook_chain_test_lifecycle") -} - #[rstest] -fn suspended_hooks_keep_the_callers_task_and_context(scripts: Py) { +fn suspended_hooks_keep_the_callers_task_and_context( + #[from(interpreter)] _python: Interpreter, + scripts: Py, +) { Python::attach(|py| { let locals = scripts.bind(py); py.run( @@ -537,8 +450,9 @@ state = contextvars.ContextVar('state') async def invoke(name, value): assert asyncio.current_task() is caller assert threading.get_ident() == thread - if name == 'prepare': - state.set(state.get() + 'x') + if name != 'prepare': + return value + state.set(state.get() + 'x') await asyncio.sleep(0) assert asyncio.current_task() is caller return {**value, 'context': state.get()} @@ -549,24 +463,20 @@ second = Hooks('b', invoke) Some(locals), ) .unwrap(); - let hooks = chain(py, &scripts, true); - let execution = litellm_host_python::Execution::new( - Preparing { - hooks, - arguments: Some(PyDict::new(py).unbind()), - resume: None, - }, - lifecycle, - ); - let module = lifecycle(py).unwrap(); - let call = module - .getattr("drive") - .unwrap() - .call1(( - Py::new(py, execution).unwrap(), - module.getattr("AsyncStream").unwrap(), - )) - .unwrap(); + let machine = success_machine(); + let call = handed( + py, + run_call( + py, + move |_, _, request| Ok(machine(request)), + binding(&Log::default(), OpScript::Answer), + chain(py, &scripts, true), + PyDict::new(py).unbind(), + call_options(true), + ), + true, + ) + .unwrap(); locals.set_item("call", call).unwrap(); py.run( c" @@ -575,10 +485,10 @@ async def exercise(): caller = asyncio.current_task() thread = threading.get_ident() state.set('caller') - result = await call - assert result['context'] == 'callerxx' + await call + assert first.adopted['context'] == 'callerxx' assert state.get() == 'callerxx' - assert first.adopted is result and second.adopted is result + assert second.adopted is first.adopted asyncio.run(exercise()) ", Some(locals), @@ -588,106 +498,6 @@ asyncio.run(exercise()) }); } -#[pyclass(weakref)] -struct HookOwner { - hooks: HookChain, -} - -#[pymethods] -impl HookOwner { - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - self.hooks.traverse(&visit) - } - - fn __clear__(&mut self, py: Python<'_>) { - self.hooks.close(py); - } -} - -#[rstest] -#[case::first_response(false, false)] -#[case::first_exception(true, false)] -#[case::second_response(false, true)] -#[case::second_exception(true, true)] -fn suspended_notification_cycles_are_collectable( - scripts: Py, - #[case] failed: bool, - #[case] first_ready: bool, -) { - Python::attach(|py| { - let hook = |name, asynchronous| ScriptHooks { - object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(), - asynchronous, - wire: None, - }; - let owner = Py::new( - py, - HookOwner { - hooks: HookChain::new() - .with(hook("first", !first_ready)) - .with(hook("second", true)), - }, - ) - .unwrap(); - let locals = scripts.bind(py); - locals.set_item("owner", &owner).unwrap(); - py.run( - c" -import gc -import weakref -class Payload(Exception): - pass -payload = Payload() -payload.owner = owner -owner_ref = weakref.ref(owner) -payload_ref = weakref.ref(payload) -", - Some(locals), - Some(locals), - ) - .unwrap(); - let payload = locals.get_item("payload").unwrap().unwrap().unbind(); - let step = if failed { - owner.borrow_mut(py).hooks.on_event( - py, - PythonCallEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Call, - error: &PyErr::from_value(payload.bind(py).clone()), - }, - ) - } else { - owner.borrow_mut(py).hooks.on_event( - py, - PythonCallEvent::Succeeded { - timing: TIMING, - response: &payload, - }, - ) - } - .unwrap(); - let HookStep::Await(awaitable, _) = step else { - panic!("notification must suspend") - }; - awaitable.call_method0(py, "close").unwrap(); - drop(awaitable); - drop(payload); - drop(owner); - py.run( - c" -log.clear() -del owner, payload -gc.collect() -assert owner_ref() is None -assert payload_ref() is None -", - Some(locals), - Some(locals), - ) - .unwrap(); - }); -} - #[rstest] #[case::empty(0, "")] #[case::single(1, "a")] @@ -701,21 +511,18 @@ fn builder_runs_hooks_in_append_order( #[values(false, true)] asynchronous: bool, ) { Python::attach(|py| { - let mut hooks = ["first", "second", "third", "fourth"] + let hooks = ["first", "second", "third", "fourth"] .into_iter() .take(count) .fold(HookChain::new(), |chain, name| { chain.with(ScriptHooks { object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(), asynchronous, - wire: None, }) }); let original = PyDict::new(py).unbind(); - let step = hooks - .prepare_arguments(py, original.clone_ref(py), 0.0) - .unwrap(); - let result = finish(py, &mut hooks, step).unwrap(); + let result = + finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap(); if count == 0 { assert!(result.bind(py).is(original.bind(py))); } else { diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index ef09135698e..4767cdec19e 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -11,7 +11,7 @@ use litellm_host::{ protocol::{Protocol, Reply}, }; use litellm_host_python::{ - CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks, + Awaiter, CallOptions, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks, PythonHostCalls, PythonOwned, PythonRuntime, run_call, }; use pyo3::{ @@ -335,12 +335,12 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks { } impl CallHooks for SyntheticHooks { - fn prepare_arguments( - &mut self, - _: Python<'_>, + async fn prepare_arguments<'a>( + &'a self, + _: &'a dyn Awaiter, arguments: Py, _: f64, - ) -> PyResult>> { + ) -> PyResult> { self.log.push("begin"); if matches!(self.script, HookScript::FailBegin) { return Err(PyValueError::new_err("begin failed")); @@ -352,57 +352,54 @@ impl CallHooks for SyntheticHooks { prepared.set_item("api_key", "hook-key")?; Ok(prepared.unbind()) })?; - return Ok(HookStep::Ready(prepared)); + return Ok(prepared); } - Ok(HookStep::Ready(arguments)) + Ok(arguments) } - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { if matches!(self.script, HookScript::ObserveArguments) { - let value: String = arguments - .bind(py) - .get_item("prepared")? - .unwrap() - .extract()?; + let value: String = + Python::attach(|py| arguments.bind(py).get_item("prepared")?.unwrap().extract())?; self.log.push(format!("adopted:{value}")); } Ok(()) } - fn before_provider_request( - &mut self, - _: Python<'_>, + async fn before_provider_request<'a>( + &'a self, + _: &'a dyn Awaiter, wire: Box, _: &RequestContext, - ) -> PyResult>> { + ) -> PyResult> { self.log.push("before_provider_request"); - Ok(HookStep::Ready(Box::new(WireRequest { + Ok(Box::new(WireRequest { url: "rewritten".into(), ..*wire - }))) + })) } - fn transform_response( - &mut self, - py: Python<'_>, + async fn transform_response<'a>( + &'a self, + _: &'a dyn Awaiter, response: Py, _: Timing, - ) -> PyResult>> { + ) -> PyResult> { self.log.push("after_success"); match self.script { - HookScript::ReplaceResponse => Ok(HookStep::Ready( - "replaced".into_pyobject(py)?.into_any().unbind(), - )), + HookScript::ReplaceResponse => { + Python::attach(|py| Ok("replaced".into_pyobject(py)?.into_any().unbind())) + } HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")), - _ => Ok(HookStep::Ready(response)), + _ => Ok(response), } } - fn on_event( - &mut self, - py: Python<'_>, + async fn on_event<'a>( + &'a self, + _: &'a dyn Awaiter, event: PythonCallEvent<'_>, - ) -> PyResult> { + ) -> PyResult<()> { let terminal = matches!( event, PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. } @@ -413,7 +410,7 @@ impl CallHooks for SyntheticHooks { if terminal && matches!(self.script, HookScript::FailTerminal) { return Err(PyValueError::new_err("callback failed")); } - self.log.push(match event { + self.log.push(Python::attach(|py| match event { PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(), PythonCallEvent::Started { .. } => "started".into(), PythonCallEvent::Cancelled { .. } => "cancelled".into(), @@ -424,18 +421,18 @@ impl CallHooks for SyntheticHooks { format!("succeeded:{}", response.bind(py)) } PythonCallEvent::Failed { origin, error, .. } => { - format!("failed:{origin:?}:{}", error.value(py)) + format!("failed:{origin:?}:{}", error.bind(py)) } - }); - Ok(HookStep::Ready(())) + })); + Ok(()) } - fn on_stream_open(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + fn on_stream_open(&self, _: &Py) -> PyResult<()> { self.log.push("opened"); Ok(()) } - fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + fn on_stream_chunk(&self, _: &Py) -> PyResult<()> { self.log.push("delivered"); Ok(()) } diff --git a/litellm-rust/crates/host/src/hooks.rs b/litellm-rust/crates/host/src/hooks.rs index da16bef258c..7f78c059ed2 100644 --- a/litellm-rust/crates/host/src/hooks.rs +++ b/litellm-rust/crates/host/src/hooks.rs @@ -1,79 +1,67 @@ +use std::future::{Future, ready}; + use crate::{ interceptors::{RawResponse, RequestContext, WireRequest}, lifecycle::{CallEvent, Timing}, }; pub trait HookRuntime { - type Context<'a>; - type Arguments; - type Response; + type Context<'a>: Send; + type Arguments: Send; + type Response: Send + Sync; type Chunk; - type Error; - type Step; - - fn ready(value: T) -> Self::Step; + type Error: Send + Sync; + type Failure: Sync; } pub type RuntimeCallEvent<'a, R> = - CallEvent<&'a ::Response, &'a ::Error, &'a RawResponse>; + CallEvent<&'a ::Response, &'a ::Failure, &'a RawResponse>; -pub trait CallHooks: Sized { - fn prepare_arguments( - &mut self, - _runtime: R::Context<'_>, +pub trait CallHooks: Sync { + fn prepare_arguments<'a>( + &'a self, + _runtime: R::Context<'a>, arguments: R::Arguments, _started_at: f64, - ) -> Result, R::Error> { - Ok(R::ready(arguments)) + ) -> impl Future> + Send { + ready(Ok(arguments)) } - fn arguments_prepared( - &mut self, - _runtime: R::Context<'_>, - _arguments: &R::Arguments, - ) -> Result<(), R::Error> { + fn arguments_prepared(&self, _arguments: &R::Arguments) -> Result<(), R::Error> { Ok(()) } - fn before_provider_request( - &mut self, - _runtime: R::Context<'_>, + fn before_provider_request<'a>( + &'a self, + _runtime: R::Context<'a>, wire: Box, _context: &RequestContext, - ) -> Result>, R::Error> { - Ok(R::ready(wire)) + ) -> impl Future, R::Error>> + Send { + ready(Ok(wire)) } - fn transform_response( - &mut self, - _runtime: R::Context<'_>, + fn transform_response<'a>( + &'a self, + _runtime: R::Context<'a>, response: R::Response, _timing: Timing, - ) -> Result, R::Error> { - Ok(R::ready(response)) + ) -> impl Future> + Send { + ready(Ok(response)) } - fn on_event( - &mut self, - _runtime: R::Context<'_>, + fn on_event<'a>( + &'a self, + _runtime: R::Context<'a>, _event: RuntimeCallEvent<'_, R>, - ) -> Result, R::Error> { - Ok(R::ready(())) + ) -> impl Future> + Send { + ready(Ok(())) } - fn on_stream_open( - &mut self, - _runtime: R::Context<'_>, - _head: &R::Response, - ) -> Result<(), R::Error> { + fn on_stream_open(&self, _head: &R::Response) -> Result<(), R::Error> { Ok(()) } - fn on_stream_chunk( - &mut self, - _runtime: R::Context<'_>, - _chunk: &R::Chunk, - ) -> Result<(), R::Error> { + fn on_stream_chunk(&self, _chunk: &R::Chunk) -> Result<(), R::Error> { Ok(()) } } diff --git a/litellm-rust/crates/python-bridge/src/preflight.rs b/litellm-rust/crates/python-bridge/src/preflight.rs index bc692498880..ced78582906 100644 --- a/litellm-rust/crates/python-bridge/src/preflight.rs +++ b/litellm-rust/crates/python-bridge/src/preflight.rs @@ -36,14 +36,16 @@ pub(crate) const PYTHON_CONTRACT: &str = include_str!("../preflight_contract.jso pub(crate) struct SdkPolicy; impl CallHooks for SdkPolicy { - fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { - inherit_credentials(py, arguments.bind(py), || { - Ok(PythonPreflight::CredentialList - .call(py, ())? - .cast_into::()?) - })?; - PythonPreflight::CheckLimits.call(py, (arguments,))?; - Ok(()) + fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + Python::attach(|py| { + inherit_credentials(py, arguments.bind(py), || { + Ok(PythonPreflight::CredentialList + .call(py, ())? + .cast_into::()?) + })?; + PythonPreflight::CheckLimits.call(py, (arguments,))?; + Ok(()) + }) } } @@ -426,7 +428,7 @@ arguments = {'litellm_credential_name': 'missing'} ", ); SdkPolicy - .arguments_prepared(py, &argument_dict(&locals).unbind()) + .arguments_prepared(&argument_dict(&locals).unbind()) .unwrap(); py.run( c" @@ -464,7 +466,7 @@ arguments = {'litellm_credential_name': 'ocr-test'} ", ); let error = SdkPolicy - .arguments_prepared(py, &argument_dict(&locals).unbind()) + .arguments_prepared(&argument_dict(&locals).unbind()) .unwrap_err(); assert!( error diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 08c44efc58e..b129f603436 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -9,7 +9,7 @@ pub(crate) mod responses; pub(crate) mod token_counter; pub(crate) mod traces; -use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCall}; +use litellm_callbacks_legacy_python::{LegacyHooks, LegacyLogging, LoggingOperation, PublicCall}; use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol}; use litellm_host_python::{Error as HostError, present}; use litellm_host_python::{ @@ -125,7 +125,7 @@ fn call_hooks( let arguments = call.arguments(py); Ok(( arguments, - LegacyLogging::new(py, operation, call, asynchronous), + LegacyHooks::new(LegacyLogging::new(py, operation, call, asynchronous)), )) }