From c63d46becc016a9f2b570a8d6a6818be9d3c930e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 14 Oct 2026 10:16:21 -0700 Subject: [PATCH] refactor(host): split call hooks into facets with the policy in the return type One trait per payload replaces CallHooks and its on_event bag: ArgumentHooks, ProviderHooks, ResponseHooks and OutcomeHooks over a Surface type bag. Transforms fold through the groups and stop on the first error, outcomes reach every group and only a Cancelled group stops the rest, so ChainRuntime, HookGroupAdapter, DynCallHooks and the runtime terminal match are gone. Hooks composes mounted groups, request stages in mount order, response and terminal stages in reverse host-python keeps PythonSurface, the per-stage cancellation guard behind PythonSurface::group, and report_unless_cancelled as the one place an ordinary outcome error becomes unraisable. The legacy adapter implements each facet in its own file under facets/, which also lists the Python callbacks each stage reaches, and mapping.rs only collects them. SdkPolicy does its work in prepare, so the logging-before-policy ordering rule is no longer load-bearing --- .../crates/callbacks-legacy-python/AGENTS.md | 6 +- .../callbacks-legacy-python/src/adapter.rs | 183 ++----- .../src/facets/arguments.rs | 33 ++ .../callbacks-legacy-python/src/facets/mod.rs | 13 + .../src/facets/outcome.rs | 89 ++++ .../src/facets/provider.rs | 46 ++ .../src/facets/response.rs | 39 ++ .../crates/callbacks-legacy-python/src/lib.rs | 8 +- .../callbacks-legacy-python/src/mapping.rs | 206 +------- litellm-rust/crates/host-python/AGENTS.md | 14 +- .../crates/host-python/src/binding.rs | 2 +- litellm-rust/crates/host-python/src/driver.rs | 81 ++-- litellm-rust/crates/host-python/src/hooks.rs | 45 +- .../crates/host-python/src/hooks/chain.rs | 171 ------- .../crates/host-python/src/hooks/guard.rs | 226 +++++++++ litellm-rust/crates/host-python/src/lib.rs | 4 +- .../crates/host-python/tests/driver.rs | 98 +++- .../crates/host-python/tests/handle.rs | 64 ++- .../crates/host-python/tests/hook_chain.rs | 227 +++++---- .../crates/host-python/tests/support/mod.rs | 134 ++++-- litellm-rust/crates/host/AGENTS.md | 4 +- litellm-rust/crates/host/src/hooks.rs | 168 +++++-- litellm-rust/crates/host/src/hooks/adapter.rs | 91 ---- litellm-rust/crates/host/src/hooks/chain.rs | 449 ++++++++++-------- litellm-rust/crates/host/src/hooks/group.rs | 164 +++++++ .../crates/python-bridge/src/preflight.rs | 54 ++- .../crates/python-bridge/src/routes/mod.rs | 15 +- 27 files changed, 1505 insertions(+), 1129 deletions(-) create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/facets/arguments.rs create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/facets/mod.rs create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/facets/outcome.rs create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/facets/provider.rs create mode 100644 litellm-rust/crates/callbacks-legacy-python/src/facets/response.rs delete mode 100644 litellm-rust/crates/host-python/src/hooks/chain.rs create mode 100644 litellm-rust/crates/host-python/src/hooks/guard.rs delete mode 100644 litellm-rust/crates/host/src/hooks/adapter.rs create mode 100644 litellm-rust/crates/host/src/hooks/group.rs diff --git a/litellm-rust/crates/callbacks-legacy-python/AGENTS.md b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md index 1dd09575c08..d7d82bdca2d 100644 --- a/litellm-rust/crates/callbacks-legacy-python/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy-python/AGENTS.md @@ -1,8 +1,8 @@ - Target invariants, not completion claims -- This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `mapping.rs` owns the executable call bindings and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) +- This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `facets/` holds `LegacyHooks` on each facet of the call, one file per facet, each listing the Python callbacks its stages reach. `mapping.rs` collects those lists and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) - Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here - - SDK request policy (credential inheritance, the budget and retry-count limits) is a separate hook supplied by `python-bridge`. Compose it after this adapter so logging adopts the final keyword view before policy mutates or rejects it - - The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks` using the shared `CallEvent`. They never learn which Python objects consume a call + - SDK request policy (credential inheritance, the budget and retry-count limits) is a separate argument group supplied by `python-bridge` + - The driver in `litellm-host-python`, the routes and core see one `PythonGroup`. They never learn which Python objects consume a call - Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json` - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 85fb65230c0..bab1960c02c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -9,7 +9,7 @@ use litellm_host::{ interceptors::{RawResponse, RequestContext, WireRequest}, lifecycle::{FailureOrigin, Timing, epoch_seconds}, }; -use litellm_host_python::{Awaiter, Error as HostError, from_py, to_py}; +use litellm_host_python::{Error as HostError, from_py, to_py}; use parking_lot::Mutex; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -64,7 +64,7 @@ fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult> { PyDateTime::from_timestamp(py, epoch_seconds, None).map(|value| value.into_any().unbind()) } -fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { +pub(crate) fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { !error.is_instance_of::(py) } @@ -76,14 +76,6 @@ pub(crate) enum FailureStage { 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<'_>, @@ -158,14 +150,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> { + pub(crate) 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(self.call.kwargs().clone_ref(py)) } - fn finalize(&mut self, py: Python<'_>) -> PyResult> { + pub(crate) fn finalize(&mut self, py: Python<'_>) -> PyResult> { finalize( py, &self.response, @@ -290,7 +282,7 @@ impl LegacyLogging { /// 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>> { + pub(crate) fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult>> { let (Some(logger), Some(error)) = (&self.logger, &self.error) else { return Ok(None); }; @@ -314,7 +306,7 @@ impl LegacyLogging { } impl LegacyLogging { - fn adopt_hooked_arguments( + pub(crate) fn adopt_hooked_arguments( &mut self, py: Python<'_>, result: PyResult>, @@ -324,7 +316,7 @@ impl LegacyLogging { Ok(()) } - fn adopt_hooked_response(&mut self, result: PyResult>) -> PyResult<()> { + pub(crate) fn adopt_hooked_response(&mut self, result: PyResult>) -> PyResult<()> { self.response = Some(result?); Ok(()) } @@ -332,7 +324,7 @@ impl LegacyLogging { impl LegacyLogging { /// Sets the call up, returning the deployment pre-call hook to await when it runs. - fn begin_call( + pub(crate) fn begin_call( &mut self, py: Python<'_>, arguments: Py, @@ -405,7 +397,7 @@ impl LegacyLogging { /// Takes the public response, returning the deployment post-call hook to await when /// it runs. - fn respond( + pub(crate) fn respond( &mut self, py: Python<'_>, response: Py, @@ -489,7 +481,7 @@ impl LegacyLogging { } } - fn failed( + pub(crate) fn failed( &mut self, py: Python<'_>, timing: Timing, @@ -536,100 +528,16 @@ 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); +pub struct LegacyHooks(pub(crate) Mutex); impl LegacyHooks { pub fn new(logging: LegacyLogging) -> Self { Self(Mutex::new(logging)) } - fn with(&self, step: impl FnOnce(&mut LegacyLogging, Python<'_>) -> T) -> T { + pub(crate) 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 { @@ -680,9 +588,8 @@ impl PythonOwned for LegacyLogging { mod deployment_hooks_tests { use std::ffi::CStr; - use litellm_host::hooks::CallHooks; + use litellm_host::hooks::{ArgumentHooks, OutcomeHooks, ResponseHooks}; use litellm_host::lifecycle::{FailureOrigin, Timing}; - use litellm_host_python::PythonCallEvent; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -718,7 +625,7 @@ kwargs = {'logger': logger, 'document': document} task: &Scripted, ) -> (LegacyHooks, PyResult>) { let hooks = hooks(legacy_call(py, locals, asynchronous)); - let prepared = finish(hooks.prepare_arguments(task, kwargs(locals), 0.0)); + let prepared = finish(hooks.prepare(task, kwargs(locals), 0.0)); (hooks, prepared) } @@ -744,7 +651,7 @@ kwargs = {'logger': logger, 'document': document} ..legacy_call(py, &locals, asynchronous) }); let task = Scripted::answering([Ok(kwargs(&locals).into_any())]); - finish(hooks.prepare_arguments(&task, kwargs(&locals), 0.0)).unwrap(); + finish(hooks.prepare(&task, kwargs(&locals), 0.0)).unwrap(); assert_eq!(task.count() == 1, asynchronous); locals.set_item("expected", expected).unwrap(); locals.set_item("asynchronous", asynchronous).unwrap(); @@ -859,12 +766,9 @@ logger.hooks = {'pre': lambda kwargs: kwargs} ]); let (hooks, prepared) = begin(py, &locals, true, &task); prepared.unwrap(); - let returned = finish(hooks.transform_response( - &task, - local(&locals, "response").unbind(), - TIMING, - )) - .unwrap(); + let returned = + finish(hooks.transform(&task, local(&locals, "response").unbind(), TIMING)) + .unwrap(); assert_eq!(task.count(), 2); assert!(returned.bind(py).is(local(&locals, "replacement"))); run( @@ -896,7 +800,7 @@ assert finalized is replacement 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)) + finish(hooks.transform(&task, local(&locals, "response").unbind(), TIMING)) .unwrap_err() } else { prepared.unwrap_err() @@ -934,12 +838,7 @@ assert finalized is replacement 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(); + finish(hooks.failed(&task, TIMING, FailureOrigin::Call, &failure)).unwrap(); assert_eq!(task.count(), 3); run( py, @@ -1718,9 +1617,9 @@ def check(): mod terminal_tests { use std::ffi::CStr; - use litellm_host::hooks::CallHooks; - use litellm_host::lifecycle::{FailureOrigin, Timing}; - use litellm_host_python::{PythonCallEvent, PythonOwned}; + use litellm_host::hooks::{Cancelled, OutcomeHooks, ResponseHooks}; + use litellm_host::lifecycle::{CallEvent, FailureOrigin, Timing}; + use litellm_host_python::PythonOwned; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -1750,28 +1649,15 @@ mod terminal_tests { fn succeed(locals: &Bound<'_, PyDict>, hooks: &LegacyHooks) -> usize { let response = local(locals, "response").unbind(); let task = Scripted::default(); - finish(hooks.on_event( - &task, - PythonCallEvent::Succeeded { - timing: TIMING, - response: &response, - }, - )) - .unwrap(); + finish(hooks.succeeded(&task, TIMING, &response)).unwrap(); task.count() } /// 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, - }, - )) + finish(hooks.failed(task, TIMING, FailureOrigin::Host, &failure)) + .map_err(|Cancelled(error)| error) } #[rstest] @@ -1821,17 +1707,16 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy ); let hooks = hooks(logged(py, &locals, true)); let response = local(&locals, "response").unbind(); - let event = PythonCallEvent::Succeeded { - timing: TIMING, - response: &response, - }; let (sender, receiver) = litellm_host::observation::observation_channel( std::num::NonZeroUsize::new(1).unwrap(), ); drop(receiver); - sender.emit(event.snapshot()); + sender.emit(CallEvent::Succeeded { + timing: TIMING, + response: (), + }); let task = Scripted::default(); - finish(hooks.on_event(&task, event)).unwrap(); + finish(hooks.succeeded(&task, TIMING, &response)).unwrap(); assert_eq!(task.count(), 0); assert_eq!(sender.dropped_events(), 1); run(py, &locals, c" @@ -1855,13 +1740,13 @@ assert logger.calls[1][1] is response ..logged(py, &locals, true) }); hooks - .on_stream_open(&local(&locals, "head").unbind()) + .stream_opened(&local(&locals, "head").unbind()) .unwrap(); hooks - .on_stream_chunk(&local(&locals, "first").unbind()) + .stream_chunk(&local(&locals, "first").unbind()) .unwrap(); hooks - .on_stream_chunk(&local(&locals, "last").unbind()) + .stream_chunk(&local(&locals, "last").unbind()) .unwrap(); assert_eq!(succeed(&locals, &hooks), 0); run( diff --git a/litellm-rust/crates/callbacks-legacy-python/src/facets/arguments.rs b/litellm-rust/crates/callbacks-legacy-python/src/facets/arguments.rs new file mode 100644 index 00000000000..b2ceb512354 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/facets/arguments.rs @@ -0,0 +1,33 @@ +use litellm_host::hooks::ArgumentHooks; +use litellm_host_python::{Awaiter, PythonSurface}; +use pyo3::{prelude::*, types::PyDict}; + +use crate::{CallBoundary, LegacyHooks}; + +pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[( + CallBoundary::PrepareArguments, + &["async_pre_call_deployment_hook"], +)]; + +impl ArgumentHooks for LegacyHooks { + /// `function_setup`, then the deployment pre-call hook, whose returned kwargs become + /// the keyword view the rest of the call reads. + async fn prepare<'a>( + &'a self, + task: &'a 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)) + } + + fn prepared(&self, arguments: &Py) -> PyResult<()> { + self.with(|logging, py| logging.adopt_arguments(py, arguments)); + Ok(()) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/facets/mod.rs b/litellm-rust/crates/callbacks-legacy-python/src/facets/mod.rs new file mode 100644 index 00000000000..c4c280c9a43 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/facets/mod.rs @@ -0,0 +1,13 @@ +//! [`LegacyHooks`](crate::LegacyHooks) on each facet of the call, one file per facet, so the +//! stage-to-callback mapping is the file layout. Every facet file also lists the Python +//! callbacks its stages reach, which `mapping.rs` collects. + +mod arguments; +mod outcome; +mod provider; +mod response; + +pub(crate) use arguments::CALLBACKS as ARGUMENT_CALLBACKS; +pub(crate) use outcome::CALLBACKS as OUTCOME_CALLBACKS; +pub(crate) use provider::CALLBACKS as PROVIDER_CALLBACKS; +pub(crate) use response::CALLBACKS as RESPONSE_CALLBACKS; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/facets/outcome.rs b/litellm-rust/crates/callbacks-legacy-python/src/facets/outcome.rs new file mode 100644 index 00000000000..98ecb14f03d --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/facets/outcome.rs @@ -0,0 +1,89 @@ +use litellm_host::{ + hooks::{Cancelled, OutcomeHooks}, + lifecycle::{FailureOrigin, Timing}, +}; +use litellm_host_python::{Awaiter, PythonSurface, report_unless_cancelled}; +use pyo3::{exceptions::PyBaseException, prelude::*}; + +use crate::{CallBoundary, LegacyHooks, adapter::FailureStage, is_cancellation}; + +pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[ + ( + CallBoundary::Succeeded, + &[ + "log_success_event", + "async_log_success_event", + "logging_hook", + "async_logging_hook", + "redact_standard_logging_payload_from_model_call_details", + "log_event", + "async_log_event", + ], + ), + ( + CallBoundary::Failed, + &[ + "async_post_call_failure_deployment_hook", + "log_failure_event", + "async_log_failure_event", + "log_model_group_rate_limit_error", + "log_event", + "async_log_event", + ], + ), +]; + +/// 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 LegacyHooks { + async fn fail( + &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(()), + } + } +} + +impl OutcomeHooks for LegacyHooks { + /// The sync and async success fan-out with the selected response. + async fn succeeded<'a>( + &'a self, + _: &'a dyn Awaiter, + timing: Timing, + response: &Py, + ) -> Result<(), Cancelled> { + report_unless_cancelled(self.with(|logging, py| logging.succeeded(py, timing, response))) + } + + /// The deployment failure hook, then the sync and async failure fan-out with the + /// selected exception. + async fn failed<'a>( + &'a self, + task: &'a dyn Awaiter, + timing: Timing, + origin: FailureOrigin, + error: &Py, + ) -> Result<(), Cancelled> { + report_unless_cancelled(self.fail(task, timing, origin, error).await) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/facets/provider.rs b/litellm-rust/crates/callbacks-legacy-python/src/facets/provider.rs new file mode 100644 index 00000000000..7d155a64df0 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/facets/provider.rs @@ -0,0 +1,46 @@ +use litellm_host::{ + hooks::ProviderHooks, + interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest}, +}; +use litellm_host_python::{Awaiter, PythonSurface}; +use pyo3::prelude::*; + +use crate::{CallBoundary, LegacyHooks}; + +pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[ + ( + CallBoundary::BeforeProviderRequest, + &["log_pre_api_call", "log_input_event"], + ), + (CallBoundary::AfterProviderResponse, &["log_post_api_call"]), +]; + +impl ProviderHooks for LegacyHooks { + /// `Logging.pre_call` over the wire request, which the callbacks may edit in place. + async fn before_provider_request<'a>( + &'a self, + _: &'a dyn Awaiter, + wire: Box, + context: &RequestContext, + ) -> PyResult> { + self.with(|logging, py| logging.pre_call(py, wire, context)) + } + + /// `Logging.post_call` with the raw provider response. + async fn after_provider_response<'a>( + &'a self, + _: &'a dyn Awaiter, + raw: &RawResponse, + ) -> PyResult<()> { + self.with(|logging, py| logging.post_call(py, raw)) + } + + /// The provider, cost and cache facts the spend record is built from. + async fn result_ready<'a>( + &'a self, + _: &'a dyn Awaiter, + facts: &ExecutionFacts, + ) -> PyResult<()> { + self.with(|logging, py| logging.result_ready(py, facts)) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/facets/response.rs b/litellm-rust/crates/callbacks-legacy-python/src/facets/response.rs new file mode 100644 index 00000000000..9cee494333d --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/src/facets/response.rs @@ -0,0 +1,39 @@ +use litellm_host::{hooks::ResponseHooks, lifecycle::Timing}; +use litellm_host_python::{Awaiter, PythonSurface}; +use pyo3::prelude::*; + +use crate::{CallBoundary, LegacyHooks}; + +pub(crate) const CALLBACKS: &[(CallBoundary, &[&str])] = &[ + ( + CallBoundary::TransformResponse, + &["async_post_call_success_deployment_hook"], + ), + (CallBoundary::StreamOpened, &[]), + (CallBoundary::StreamChunk, &[]), +]; + +impl ResponseHooks for LegacyHooks { + /// The deployment success hook, whose returned response is what the caller receives. + async fn transform<'a>( + &'a self, + task: &'a 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)) + } + + fn stream_opened(&self, head: &Py) -> PyResult<()> { + self.with(|logging, py| logging.stream_opened(py, head)) + } + + fn stream_chunk(&self, chunk: &Py) -> PyResult<()> { + self.with(|logging, py| logging.stream_chunk(py, chunk)) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 35ad3b637b5..90371e22ea8 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -1,8 +1,8 @@ //! The legacy `@client` wrapper as the native call sees it: litellm's `Logging` object, the //! sync and async callback registries it fans out to, the deployment hooks and the deferred -//! proxy release. All of it sits behind one -//! [`PythonCallHooks`](litellm_host_python::PythonCallHooks), so the driver, the routes and -//! core never learn which Python object is on the other end. +//! proxy release. All of it sits behind one [`LegacyHooks`] group on the Python surface, one +//! facet per file under `facets/`, so the driver, the routes and core never learn which +//! Python object is on the other end. //! //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`] //! is where those objects live. @@ -11,9 +11,11 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod facets; mod logger; mod mapping; mod python; +pub(crate) use adapter::is_cancellation; pub use adapter::{LegacyHooks, LegacyLogging}; pub use call::PublicCall; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs index 8feac1a34bf..2fab8bffb3f 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs @@ -1,15 +1,9 @@ -use std::future::{Future, ready}; +//! Which Python callbacks each neutral call stage reaches, and which callbacks Python +//! still invokes itself. -use futures_util::future::BoxFuture; -use litellm_host::{ - hooks::CallHooks, - interceptors::{RawResponse, RequestContext, WireRequest}, - lifecycle::{ExecutionEvent, FailureOrigin, Timing}, +use crate::facets::{ + ARGUMENT_CALLBACKS, OUTCOME_CALLBACKS, PROVIDER_CALLBACKS, RESPONSE_CALLBACKS, }; -use litellm_host_python::{Awaiter, PythonCallEvent, PythonRuntime}; -use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; - -use crate::LegacyHooks; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CallBoundary { @@ -36,194 +30,24 @@ pub struct CallbackMapping { pub dispatch: Dispatch, } -struct Binding { - boundary: CallBoundary, - invoke: H, - callbacks: &'static [&'static str], -} - -impl Binding { - fn mappings(&self) -> impl Iterator { - self.callbacks.iter().map(|callback| CallbackMapping { +fn call_mappings( + stages: &'static [(CallBoundary, &'static [&'static str])], +) -> impl Iterator { + stages.iter().flat_map(|(boundary, callbacks)| { + callbacks.iter().map(|callback| CallbackMapping { callback, - dispatch: Dispatch::Call(self.boundary), + dispatch: Dispatch::Call(*boundary), }) - } + }) } -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: |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: LegacyHooks::pre_call, - callbacks: &["log_pre_api_call", "log_input_event"], -}; - -const AFTER: Binding = Binding { - boundary: CallBoundary::AfterProviderResponse, - invoke: LegacyHooks::post_call, - callbacks: &["log_post_api_call"], -}; - -const TRANSFORM: Binding = Binding { - boundary: CallBoundary::TransformResponse, - 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: LegacyHooks::succeeded, - callbacks: &[ - "log_success_event", - "async_log_success_event", - "logging_hook", - "async_logging_hook", - "redact_standard_logging_payload_from_model_call_details", - "log_event", - "async_log_event", - ], -}; - -const FAILURE: Binding = Binding { - boundary: CallBoundary::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", - "async_log_failure_event", - "log_model_group_rate_limit_error", - "log_event", - "async_log_event", - ], -}; - -const OPEN: Binding = Binding { - boundary: CallBoundary::StreamOpened, - invoke: LegacyHooks::stream_opened, - callbacks: &[], -}; - -const CHUNK: Binding = Binding { - boundary: CallBoundary::StreamChunk, - invoke: LegacyHooks::stream_chunk, - callbacks: &[], -}; - pub fn callback_mappings() -> impl Iterator { - PREPARE - .mappings() - .chain(BEFORE.mappings()) - .chain(AFTER.mappings()) - .chain(TRANSFORM.mappings()) - .chain(SUCCESS.mappings()) - .chain(FAILURE.mappings()) - .chain(OPEN.mappings()) - .chain(CHUNK.mappings()) + call_mappings(ARGUMENT_CALLBACKS) + .chain(call_mappings(PROVIDER_CALLBACKS)) + .chain(call_mappings(RESPONSE_CALLBACKS)) + .chain(call_mappings(OUTCOME_CALLBACKS)) .chain(PYTHON_CALLBACKS.iter().copied()) } - -impl CallHooks for LegacyHooks { - fn prepare_arguments<'a>( - &'a self, - task: &'a dyn Awaiter, - arguments: Py, - started_at: f64, - ) -> impl Future>> + Send { - (PREPARE.invoke)(self, task, arguments, started_at) - } - - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { - self.adopt_arguments(arguments); - Ok(()) - } - - fn before_provider_request<'a>( - &'a self, - _: &'a dyn Awaiter, - wire: Box, - context: &RequestContext, - ) -> impl Future>> + Send { - ready((BEFORE.invoke)(self, wire, context)) - } - - fn transform_response<'a>( - &'a self, - task: &'a dyn Awaiter, - response: Py, - timing: Timing, - ) -> impl Future>> + Send { - (TRANSFORM.invoke)(self, task, response, timing) - } - - async fn on_event<'a>( - &'a self, - task: &'a dyn Awaiter, - event: PythonCallEvent<'_>, - ) -> PyResult<()> { - match event { - PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => Ok(()), - PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }) => { - self.result_ready(&facts) - } - PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - (AFTER.invoke)(self, raw) - } - PythonCallEvent::Succeeded { timing, response } => { - (SUCCESS.invoke)(self, timing, response) - } - PythonCallEvent::Failed { - timing, - origin, - error, - } => (FAILURE.invoke)(self, task, timing, origin, error).await, - } - } - - fn on_stream_open(&self, head: &Py) -> PyResult<()> { - (OPEN.invoke)(self, head) - } - - fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { - (CHUNK.invoke)(self, chunk) - } -} - macro_rules! python_callbacks { ($($dispatch:expr => [$($callback:literal),* $(,)?]),* $(,)?) => { const PYTHON_CALLBACKS: &[CallbackMapping] = &[ diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index ae7b243c296..49ee5de966e 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -14,19 +14,19 @@ The driver owns ordering: start, argument preparation, prepared-argument hooks, Boundary tests exercise behavior with a supplied lifecycle binding without importing the LiteLLM Python package. Pin inline awaiting, awaitable final values, exception identity, cancellation, re-entry and release of retained objects, rather than module names or source layout -`HookChain` specializes the shared host chain for Python runtime hooks. The Python adapter owns cancellation guards, close and GC traversal. `PythonRuntime` reports an ordinary terminal notification error as unraisable and lets cancellation stop the chain. Each argument, wire-request and response transformation feeds its result to the next hook. After all argument transformations, the driver calls `arguments_prepared` on every hook in order. Retained callback views must adopt that dictionary before later policy hooks can mutate or reject it. SDK policy is supplied by bridge composition as a hook, never a separate driver phase or parameter. Hooks implement only the stages they need. Default stages preserve the supplied values +`PythonSurface` is the shared `Surface` for Python: `Py` arguments, `Py` responses and chunks, `PyErr` errors, the caller's task as `&dyn Awaiter`. A group implements the four facets (`ArgumentHooks`, `ProviderHooks`, `ResponseHooks`, `OutcomeHooks`) plus `PythonOwned` and is mounted with `PythonSurface::group`, which guards every stage: a cancellation the task resumes a hook with ends that stage there. `Hooks` is the shared chain on this surface, and the bridge composes it; SDK policy is a mounted group, never a driver phase. A group implements only the stages it needs, default stages preserve the supplied values. The driver calls `prepared` on every group after all `prepare` stages, with the final keyword view -Terminal notifications share the selected response or exception. An ordinary notification error is reported as unraisable and does not skip the next hook or replace the selected outcome. Preparation, interception and transformation errors stop the chain. Cancellation stops all further hook dispatch. Suspensions stay inline in the existing driver, and the chain traverses retained event values for GC +Terminal outcomes share the selected response or exception. A Python outcome hook settles its own result through `report_unless_cancelled`: an ordinary exception is written unraisable and does not skip the next group or replace the selected outcome, anything else is `Cancelled` and stops all further dispatch. Preparation, interception and transformation errors stop the chain. Suspensions stay inline in the existing driver ## Existing runtime invariants -- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonBinding`, `PythonHostCalls` and `PythonOwned` traits, and the `PythonRuntime` specialization of `host::hooks::CallHooks` +- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonBinding`, `PythonHostCalls` and `PythonOwned` traits, and the `PythonSurface` specialization of the `host::hooks` facets - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - - `PythonCallHooks` only constrains the shared call-stage interface to `PythonRuntime` and Python ownership. It must not redeclare the stages - - `PythonCallEvent` is a specialization of the shared `CallEvent`, never a separately defined lifecycle. The driver emits `Succeeded` or `Failed` exactly once and never dispatches callbacks after a cancellation. Which Python objects consume those events is the legacy adapter's business + - `PythonGroup` only constrains the shared facets to `PythonSurface` and Python ownership. It must not redeclare the stages + - The driver delivers `succeeded` or `failed` exactly once and never dispatches callbacks after a cancellation. Which Python objects consume those outcomes is the legacy adapter's business - `CallOptions` can publish snapshots independently of callback delivery. Terminal snapshots follow completed hook dispatch, and cancellation never calls a Python callback. For Python-driven calls, leave the machine's observation publisher unset so the driver is the sole publisher of intercepted provider-response snapshots - - `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites - - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is + - `PythonBinding::decode_request` receives the keyword view returned by `prepare` and seen by `prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites + - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare` or `transform`, is raised as is - A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention. Binding/unbinding does not copy payloads diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index 3b3f4c5080c..9707ac0009a 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -11,7 +11,7 @@ pub trait PythonBinding: PythonOwned { /// raises it. type Failure: Into; - /// Decodes the keyword view returned by `prepare_arguments`, including preflight rewrites. + /// Decodes the keyword view returned by `prepare`, including preflight rewrites. fn decode_request( &self, py: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index bc51242dbf2..931eaaf5818 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use std::task::{Context, Poll, Waker}; use parking_lot::{Mutex, MutexGuard}; -use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; +use pyo3::exceptions::{PyBaseException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -11,6 +11,8 @@ use pyo3::types::PyDict; use litellm_coroutine::{Coroutine, CoroutineState}; use litellm_host::{ call::HostedCompletion, + hooks::Cancelled, + interceptors::RawResponse, lifecycle::{CallEvent, ExecutionEvent, FailureOrigin, Timing, epoch_seconds}, machine::{HostFailure, Machine, MachineStep}, observation::ObservationSender, @@ -20,7 +22,7 @@ use litellm_host::{ use crate::PythonHostCalls; use crate::caller::{Answer, PythonCaller, Yielded}; use crate::handle::{Execution, ExecutionBody, ExecutionStep, StepTypes}; -use crate::hooks::{PythonCallEvent, PythonCallHooks}; +use crate::hooks::{PythonGroup, is_cancellation}; use crate::native::{NativeMachine, NativePoll}; use crate::{Error, InvokeError, PythonBinding}; @@ -34,6 +36,8 @@ type StartMachine = Box< + Sync, >; type Body = Coroutine>>; +/// The call's events as the observers see them, borrowing the Python values the call holds. +type PythonCallEvent<'a> = CallEvent<&'a Py, &'a Py, &'a RawResponse>; pub struct CallOptions { pub asynchronous: bool, @@ -51,10 +55,6 @@ impl CallOptions { } } -fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { - !error.is_instance_of::(py) -} - struct Observation { observers: Option, observing: bool, @@ -148,7 +148,7 @@ where impl CallBody where - L: PythonCallHooks + 'static, + L: PythonGroup + 'static, H: PythonBinding + PythonHostCalls + 'static, M: Machine + 'static, M::Complete: Into>>, @@ -188,9 +188,9 @@ where parts.observation.emit(&started); } self.hooks() - .on_event(&self.caller, started) + .started(&self.caller, started_at) .await - .map_err(Setup::Host)?; + .map_err(|Cancelled(error)| Setup::Host(error))?; let arguments = self .parts() .arguments @@ -198,12 +198,10 @@ where .ok_or_else(|| Setup::Host(Error::MissingState.into()))?; let prepared = self .hooks() - .prepare_arguments(&self.caller, arguments, started_at) + .prepare(&self.caller, arguments, started_at) .await .map_err(Setup::Host)?; - self.hooks() - .arguments_prepared(&prepared) - .map_err(Setup::Host)?; + self.hooks().prepared(&prepared).map_err(Setup::Host)?; let decoded = Python::attach(|py| self.binding().decode_request(py, prepared.bind(py))); self.parts().arguments = Some(prepared); let request = decoded.map_err(Setup::Decode)?; @@ -289,42 +287,41 @@ where Ok(()) } HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => { - self.notify(PythonCallEvent::Execution(ExecutionEvent::ResultReady { - facts, - })) - .await?; + self.parts() + .observation + .emit(&CallEvent::Execution(ExecutionEvent::ResultReady { + facts: facts.clone(), + })); + self.hooks().result_ready(&self.caller, &facts).await?; reply.send(()); Ok(()) } HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { - self.notify(PythonCallEvent::Execution( + self.parts().observation.emit(&CallEvent::Execution( ExecutionEvent::ProviderResponseReceived { raw: &raw }, - )) - .await?; + )); + self.hooks() + .after_provider_response(&self.caller, &raw) + .await?; reply.send(()); Ok(()) } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.streaming = true; let head = Python::attach(|py| self.binding().encode_stream_head(py, head))?; - self.hooks().on_stream_open(&head)?; + self.hooks().stream_opened(&head)?; reply.send(self.caller.open_stream(head).await); Ok(()) } HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => { let chunk = Python::attach(|py| self.binding().encode_chunk(py, chunk))?; - self.hooks().on_stream_chunk(&chunk)?; + self.hooks().stream_chunk(&chunk)?; reply.send(self.caller.yield_chunk(chunk).await); Ok(()) } } } - async fn notify(&self, event: PythonCallEvent<'_>) -> PyResult<()> { - self.parts().observation.emit(&event); - self.hooks().on_event(&self.caller, event).await - } - async fn completed( &mut self, completion: HostedCompletion>, @@ -345,11 +342,7 @@ where return self.succeeded(public).await; } let timing = self.parts().observation.timing(); - match self - .hooks() - .transform_response(&self.caller, public, timing) - .await - { + match self.hooks().transform(&self.caller, public, timing).await { Ok(response) => self.succeeded(response).await, Err(error) => self.failure(error, FailureOrigin::Host).await, } @@ -379,11 +372,14 @@ where parts.observation.retain_terminal(&event); parts.selected = Some(Selected::Response(Arc::clone(&response))); } - let notified = self.hooks().on_event(&self.caller, event).await; + let notified = self + .hooks() + .succeeded(&self.caller, timing, &response) + .await; { let mut parts = self.parts(); parts.selected = None; - notified?; + notified.map_err(|Cancelled(error)| error)?; parts.observation.finish(); } Arc::into_inner(response).ok_or_else(|| Error::MissingState.into()) @@ -411,11 +407,14 @@ where parts.observation.retain_terminal(&event); parts.selected = Some(Selected::Exception(Arc::clone(&exception))); } - let notified = self.hooks().on_event(&self.caller, event).await; + let notified = self + .hooks() + .failed(&self.caller, timing, origin, &exception) + .await; { let mut parts = self.parts(); parts.selected = None; - notified?; + notified.map_err(|Cancelled(error)| error)?; parts.observation.finish(); } let exception = Arc::into_inner(exception).ok_or(Error::MissingState)?; @@ -428,7 +427,7 @@ where struct CallDriver where H: PythonBinding, - L: PythonCallHooks, + L: PythonGroup, { body: Mutex, reply: Option>, @@ -438,7 +437,7 @@ where impl CallDriver where H: PythonBinding + 'static, - L: PythonCallHooks + 'static, + L: PythonGroup + 'static, { /// A failure the body raised past its terminal hooks, or a cancellation, is the last /// thing the observers hear about the call. @@ -469,7 +468,7 @@ where impl ExecutionBody for CallDriver where H: PythonBinding + 'static, - L: PythonCallHooks + 'static, + L: PythonGroup + 'static, { fn resume(&mut self, answer: Option) -> PyResult { match (self.reply.take(), answer) { @@ -521,7 +520,7 @@ where impl Drop for CallDriver where H: PythonBinding, - L: PythonCallHooks, + L: PythonGroup, { fn drop(&mut self) { self.body.get_mut().cancel(); @@ -563,7 +562,7 @@ pub fn run_call( options: CallOptions, ) -> PyResult> where - L: PythonCallHooks + 'static, + L: PythonGroup + 'static, H: PythonBinding + PythonHostCalls + 'static, M: Machine + 'static, M::Complete: Into>>, diff --git a/litellm-rust/crates/host-python/src/hooks.rs b/litellm-rust/crates/host-python/src/hooks.rs index 409b37aa63a..91047ae352b 100644 --- a/litellm-rust/crates/host-python/src/hooks.rs +++ b/litellm-rust/crates/host-python/src/hooks.rs @@ -1,32 +1,55 @@ -mod chain; -pub use chain::HookChain; +mod guard; -use crate::PythonOwned; use futures_util::future::BoxFuture; -use litellm_host::hooks::{CallHooks, HookRuntime, RuntimeCallEvent}; -use pyo3::exceptions::PyBaseException; +use litellm_host::hooks::{AnyGroup, Cancelled, Surface}; +use pyo3::exceptions::{PyBaseException, PyException}; use pyo3::prelude::*; use pyo3::types::PyDict; +use crate::PythonOwned; + /// 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; +/// The Python public surface: keyword arguments in, Python objects out, `PyErr` for every +/// failure, and the caller's task to await on. +pub struct PythonSurface; -impl HookRuntime for PythonRuntime { - type Context<'a> = &'a dyn Awaiter; +impl Surface for PythonSurface { + type Task<'a> = &'a dyn Awaiter; type Arguments = Py; type Response = Py; type Chunk = Py; type Error = PyErr; type Failure = Py; + type Group = dyn PythonGroup; } -pub type PythonCallEvent<'a> = RuntimeCallEvent<'a, PythonRuntime>; +/// A group mounted on the Python surface: the four facets plus Python ownership. +pub trait PythonGroup: AnyGroup + PythonOwned {} -pub trait PythonCallHooks: CallHooks + PythonOwned {} +impl + PythonOwned> PythonGroup for H {} -impl + PythonOwned> PythonCallHooks for H {} +pub type Hooks = litellm_host::hooks::Hooks; + +pub(crate) fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { + !error.is_instance_of::(py) +} + +/// How a Python outcome hook settles its own result: an ordinary exception is reported as +/// unraisable and the call goes on, anything else is a cancellation that stops it. +pub fn report_unless_cancelled(result: PyResult<()>) -> Result<(), Cancelled> { + let Err(error) = result else { + return Ok(()); + }; + Python::attach(|py| { + if is_cancellation(py, &error) { + return Err(Cancelled(error)); + } + error.write_unraisable(py, None); + Ok(()) + }) +} diff --git a/litellm-rust/crates/host-python/src/hooks/chain.rs b/litellm-rust/crates/host-python/src/hooks/chain.rs deleted file mode 100644 index 20a6f2afe7d..00000000000 --- a/litellm-rust/crates/host-python/src/hooks/chain.rs +++ /dev/null @@ -1,171 +0,0 @@ -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, ChainRuntime, DynCallHooks, HookGroupAdapter}, - interceptors::{RequestContext, WireRequest}, - lifecycle::Timing, -}; -use parking_lot::Mutex; -use pyo3::{ - exceptions::PyException, - gc::{PyTraverseError, PyVisit}, - prelude::*, - types::PyDict, -}; - -use crate::{Awaiter, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime}; - -pub type HookChain = litellm_host::hooks::HookChain; - -pub trait PythonHookGroup: DynCallHooks + PythonOwned {} -impl + PythonOwned> PythonHookGroup for H {} - -impl ChainRuntime for PythonRuntime { - type Group = dyn PythonHookGroup; - - fn terminal_notification_failed(error: PyErr) -> PyResult<()> { - if is_cancellation(&error) { - return Err(error); - } - Python::attach(|py| error.write_unraisable(py, None)); - Ok(()) - } -} - -impl HookGroupAdapter for PythonRuntime { - fn adapt_group(hooks: H) -> Box { - Box::new(PythonHookAdapter(hooks)) - } -} - -struct PythonHookAdapter(H); - -impl PythonOwned for PythonHookAdapter { - fn close(&mut self, py: Python<'_>) { - self.0.close(py); - } - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.0.traverse(visit) - } -} - -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, - } - }) - } -} - -impl CallHooks for PythonHookAdapter { - async fn prepare_arguments<'a>( - &'a self, - task: &'a dyn Awaiter, - arguments: Py, - started_at: f64, - ) -> PyResult> { - let stage = Stage::new(task); - stage - .guard(self.0.prepare_arguments(&stage, arguments, started_at)) - .await - } - - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { - self.0.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); - stage - .guard(self.0.before_provider_request(&stage, wire, context)) - .await - } - - async fn transform_response<'a>( - &'a self, - task: &'a dyn Awaiter, - response: Py, - timing: Timing, - ) -> PyResult> { - let stage = Stage::new(task); - stage - .guard(self.0.transform_response(&stage, response, timing)) - .await - } - - async fn on_event<'a>( - &'a self, - task: &'a dyn Awaiter, - event: PythonCallEvent<'_>, - ) -> PyResult<()> { - let stage = Stage::new(task); - stage.guard(self.0.on_event(&stage, event)).await - } - - fn on_stream_open(&self, head: &Py) -> PyResult<()> { - self.0.on_stream_open(head) - } - - fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { - self.0.on_stream_chunk(chunk) - } -} - -impl PythonOwned for HookChain { - fn close(&mut self, py: Python<'_>) { - for hooks in self.groups_mut() { - hooks.close(py); - } - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.groups().try_for_each(|hooks| hooks.traverse(visit)) - } -} diff --git a/litellm-rust/crates/host-python/src/hooks/guard.rs b/litellm-rust/crates/host-python/src/hooks/guard.rs new file mode 100644 index 00000000000..e29619a0b18 --- /dev/null +++ b/litellm-rust/crates/host-python/src/hooks/guard.rs @@ -0,0 +1,226 @@ +use std::future::{Future, pending, poll_fn}; +use std::pin::pin; +use std::task::Poll; + +use futures_util::future::BoxFuture; +use litellm_host::{ + hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks}, + interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest}, + lifecycle::{FailureOrigin, Timing}, +}; +use parking_lot::Mutex; +use pyo3::{ + exceptions::PyBaseException, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; + +use crate::hooks::is_cancellation; +use crate::{Awaiter, Hooks, PythonGroup, PythonOwned, PythonSurface}; + +impl PythonSurface { + /// Mounts a group on the Python surface. A cancellation the caller's task resumes a + /// hook with ends that hook's stage there, whatever the hook would have done with it. + pub fn group(hooks: impl PythonGroup + 'static) -> Box { + Box::new(Guarded(hooks)) + } +} + +struct Guarded(H); + +impl PythonOwned for Guarded { + fn close(&mut self, py: Python<'_>) { + self.0.close(py); + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.0.traverse(visit) + } +} + +/// The caller's task as one stage hands it to its hook. +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>, + cancelled: impl FnOnce(PyErr) -> E, + ) -> Result { + let mut hook = pin!(hook); + let mut cancelled = Some(cancelled); + poll_fn(|context| match hook.as_mut().poll(context) { + Poll::Ready(result) => Poll::Ready(result), + Poll::Pending => match self.cancelled.lock().take() { + Some(error) => Poll::Ready(Err(cancelled.take().expect("settled once")(error))), + None => Poll::Pending, + }, + }) + .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 Python::attach(|py| is_cancellation(py, &error)) => { + *self.cancelled.lock() = Some(error); + pending().await + } + answer => answer, + } + }) + } +} + +impl ArgumentHooks for Guarded { + async fn prepare<'a>( + &'a self, + task: &'a dyn Awaiter, + arguments: Py, + started_at: f64, + ) -> PyResult> { + let stage = Stage::new(task); + stage + .guard(self.0.prepare(&stage, arguments, started_at), |error| error) + .await + } + + fn prepared(&self, arguments: &Py) -> PyResult<()> { + self.0.prepared(arguments) + } +} + +impl ProviderHooks for Guarded { + async fn before_provider_request<'a>( + &'a self, + task: &'a dyn Awaiter, + wire: Box, + context: &RequestContext, + ) -> PyResult> { + let stage = Stage::new(task); + stage + .guard( + self.0.before_provider_request(&stage, wire, context), + |error| error, + ) + .await + } + + async fn after_provider_response<'a>( + &'a self, + task: &'a dyn Awaiter, + raw: &RawResponse, + ) -> PyResult<()> { + let stage = Stage::new(task); + stage + .guard(self.0.after_provider_response(&stage, raw), |error| error) + .await + } + + async fn result_ready<'a>( + &'a self, + task: &'a dyn Awaiter, + facts: &ExecutionFacts, + ) -> PyResult<()> { + let stage = Stage::new(task); + stage + .guard(self.0.result_ready(&stage, facts), |error| error) + .await + } +} + +impl ResponseHooks for Guarded { + async fn transform<'a>( + &'a self, + task: &'a dyn Awaiter, + response: Py, + timing: Timing, + ) -> PyResult> { + let stage = Stage::new(task); + stage + .guard(self.0.transform(&stage, response, timing), |error| error) + .await + } + + fn stream_opened(&self, head: &Py) -> PyResult<()> { + self.0.stream_opened(head) + } + + fn stream_chunk(&self, chunk: &Py) -> PyResult<()> { + self.0.stream_chunk(chunk) + } +} + +impl OutcomeHooks for Guarded { + async fn started<'a>( + &'a self, + task: &'a dyn Awaiter, + start_time: f64, + ) -> Result<(), Cancelled> { + let stage = Stage::new(task); + stage + .guard(self.0.started(&stage, start_time), Cancelled) + .await + } + + async fn succeeded<'a>( + &'a self, + task: &'a dyn Awaiter, + timing: Timing, + response: &Py, + ) -> Result<(), Cancelled> { + let stage = Stage::new(task); + stage + .guard(self.0.succeeded(&stage, timing, response), Cancelled) + .await + } + + async fn failed<'a>( + &'a self, + task: &'a dyn Awaiter, + timing: Timing, + origin: FailureOrigin, + error: &Py, + ) -> Result<(), Cancelled> { + let stage = Stage::new(task); + stage + .guard(self.0.failed(&stage, timing, origin, error), Cancelled) + .await + } + + async fn cancelled<'a>( + &'a self, + task: &'a dyn Awaiter, + timing: Timing, + ) -> Result<(), Cancelled> { + let stage = Stage::new(task); + stage + .guard(self.0.cancelled(&stage, timing), Cancelled) + .await + } +} + +impl PythonOwned for Hooks { + fn close(&mut self, py: Python<'_>) { + for group in self.groups_mut() { + group.close(py); + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.groups().try_for_each(|group| group.traverse(visit)) + } +} diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index babc870655f..53b82203629 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,6 +1,6 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and //! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine) -//! against a Python binding, host services and active call hooks. Everything here is Python-specific by +//! against a Python binding, host services and the mounted hook groups. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. mod argument; @@ -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::{Awaiter, HookChain, PythonCallEvent, PythonCallHooks, PythonRuntime}; +pub use hooks::{Awaiter, Hooks, PythonGroup, PythonSurface, report_unless_cancelled}; 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 3365eabade8..ed487716dac 100644 --- a/litellm-rust/crates/host-python/tests/driver.rs +++ b/litellm-rust/crates/host-python/tests/driver.rs @@ -5,12 +5,12 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use litellm_host::{ - hooks::CallHooks, + hooks::{ArgumentHooks, OutcomeHooks, ProviderHooks, ResponseHooks}, lifecycle::{CallEvent, ExecutionEvent, FailureOrigin}, machine::CallMachine, observation::observation_channel, }; -use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call}; +use litellm_host_python::{CallOptions, Hooks, PythonOwned, PythonSurface, run_call}; use pyo3::{ exceptions::{ PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError, @@ -79,10 +79,19 @@ fn composed_hooks_share_prepared_arguments_and_one_execution( ) { Python::attach(|py| { let log = Log::default(); - let chain = HookChain::new() - .with(hooks(&log, HookScript::ObserveArguments)) - .with(hooks(&log, HookScript::RewriteArguments)) - .with(hooks(&log, HookScript::ReplaceResponse)); + let chain = Hooks::new() + .with(PythonSurface::group(hooks( + &log, + HookScript::ObserveArguments, + ))) + .with(PythonSurface::group(hooks( + &log, + HookScript::RewriteArguments, + ))) + .with(PythonSurface::group(hooks( + &log, + HookScript::ReplaceResponse, + ))); let result = handed( py, run_call( @@ -183,7 +192,6 @@ fn observation_delivery_preserves_callback_order_and_response( #[case::preparation(HookScript::FailBegin, false, 1)] #[case::transformation(HookScript::FailAfterSuccess, false, 1)] #[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)] -#[case::terminal_failure(HookScript::FailTerminal, false, 0)] fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks( #[from(interpreter)] _python: Interpreter, #[case] script: HookScript, @@ -245,6 +253,51 @@ fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks( }); } +#[rstest] +fn an_ordinary_terminal_hook_error_is_reported_and_the_call_still_succeeds( + #[from(interpreter)] _python: Interpreter, + #[values(false, true)] asynchronous: bool, +) { + Python::attach(|py| { + let unraisable = PyDict::new(py); + py.run( + c"import sys\nseen = []\nsys.unraisablehook = lambda event: seen.append(event.exc_value)", + Some(&unraisable), + Some(&unraisable), + ) + .unwrap(); + let (options, mut events) = observed(4); + let log = Log::default(); + let (result, log) = run_composed( + py, + success_machine(), + binding(&log, OpScript::Answer), + hooks(&log, HookScript::FailTerminal), + std::convert::identity, + CallOptions { + asynchronous, + ..options + }, + ); + assert_eq!( + result.unwrap().extract::(py).unwrap(), + "project:1|sign|rewritten" + ); + assert!(matches!(events().last(), Some(CallEvent::Succeeded { .. }))); + assert!(!log.iter().any(|entry| entry.starts_with("failed:"))); + let seen: Vec = unraisable + .get_item("seen") + .unwrap() + .unwrap() + .extract::>>() + .unwrap() + .iter() + .map(|error| error.to_string()) + .collect(); + assert_eq!(seen, ["callback failed"]); + }); +} + #[rstest] #[case::preparation(HookScript::FailBegin, OpScript::Answer)] #[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)] @@ -534,8 +587,8 @@ enum ArgumentPolicy { Reject(Py), } -impl CallHooks for ArgumentPolicy { - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { +impl ArgumentHooks for ArgumentPolicy { + fn 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())), @@ -543,6 +596,10 @@ impl CallHooks for ArgumentPolicy { } } +impl ProviderHooks for ArgumentPolicy {} +impl ResponseHooks for ArgumentPolicy {} +impl OutcomeHooks for ArgumentPolicy {} + impl PythonOwned for ArgumentPolicy { fn close(&mut self, _: Python<'_>) {} @@ -581,9 +638,12 @@ fn resource_setup_uses_prepared_arguments_and_failures_are_terminal( Ok(success_machine()(request)) }, binding(&log, OpScript::Answer), - HookChain::new() - .with(hooks(&log, HookScript::RewriteArguments)) - .with(ArgumentPolicy::Inherit), + Hooks::new() + .with(PythonSurface::group(hooks( + &log, + HookScript::RewriteArguments, + ))) + .with(PythonSurface::group(ArgumentPolicy::Inherit)), arguments.clone().unbind(), call_options(asynchronous), ), @@ -675,9 +735,11 @@ fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_start binding(&log, OpScript::Answer), hooks(&log, HookScript::Plain), |hooks| { - HookChain::new() - .with(hooks) - .with(ArgumentPolicy::Reject(raised.clone_ref(py))) + Hooks::new() + .with(PythonSurface::group(hooks)) + .with(PythonSurface::group(ArgumentPolicy::Reject( + raised.clone_ref(py), + ))) }, call_options(asynchronous), ); @@ -707,7 +769,11 @@ fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote( success_machine(), binding(&log, OpScript::Answer), hooks(&log, HookScript::Plain), - |hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit), + |hooks| { + Hooks::new() + .with(PythonSurface::group(hooks)) + .with(PythonSurface::group(ArgumentPolicy::Inherit)) + }, call_options(asynchronous), ); assert_eq!( diff --git a/litellm-rust/crates/host-python/tests/handle.rs b/litellm-rust/crates/host-python/tests/handle.rs index e8bf19e9476..26876196a5f 100644 --- a/litellm-rust/crates/host-python/tests/handle.rs +++ b/litellm-rust/crates/host-python/tests/handle.rs @@ -1,9 +1,12 @@ mod support; -use litellm_host::{hooks::CallHooks, lifecycle::Timing}; +use litellm_host::{ + hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks}, + lifecycle::{FailureOrigin, Timing}, +}; use litellm_host_python::{ - Awaiter, Execution, ExecutionBody, ExecutionStep, PythonCallEvent, PythonOwned, PythonRuntime, - Suspension, run_call, + Awaiter, Execution, ExecutionBody, ExecutionStep, PythonOwned, PythonSurface, Suspension, + report_unless_cancelled, run_call, }; use parking_lot::Mutex; use pyo3::{ @@ -120,8 +123,22 @@ struct PausingHooks { holder: Holder, } -impl CallHooks for PausingHooks { - async fn transform_response<'a>( +impl PausingHooks { + /// Pauses the call on the test's awaitable, once, at the stage the holder names. + async fn pause_at(&self, task: &dyn Awaiter, here: bool) -> Result<(), Cancelled> { + let pause = self.pause.lock().take_if(|_| here); + match pause { + Some(pause) => report_unless_cancelled(task.awaited(pause).await.map(drop)), + None => Ok(()), + } + } +} + +impl ArgumentHooks for PausingHooks {} +impl ProviderHooks for PausingHooks {} + +impl ResponseHooks for PausingHooks { + async fn transform<'a>( &'a self, _: &'a dyn Awaiter, response: Py, @@ -136,22 +153,33 @@ impl CallHooks for PausingHooks { })), } } +} - async fn on_event<'a>( +impl OutcomeHooks for PausingHooks { + async fn started<'a>(&'a self, task: &'a dyn Awaiter, _: f64) -> Result<(), Cancelled> { + self.pause_at(task, matches!(self.holder, Holder::Hooks)) + .await + } + + async fn succeeded<'a>( &'a self, task: &'a dyn Awaiter, - event: PythonCallEvent<'_>, - ) -> PyResult<()> { - let pauses_here = match self.holder { - Holder::Hooks => true, - Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }), - Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }), - }; - let pause = self.pause.lock().take_if(|_| pauses_here); - if let Some(pause) = pause { - task.awaited(pause).await?; - } - Ok(()) + _: Timing, + _: &Py, + ) -> Result<(), Cancelled> { + self.pause_at(task, matches!(self.holder, Holder::Response)) + .await + } + + async fn failed<'a>( + &'a self, + task: &'a dyn Awaiter, + _: Timing, + _: FailureOrigin, + _: &Py, + ) -> Result<(), Cancelled> { + self.pause_at(task, matches!(self.holder, Holder::Exception)) + .await } } diff --git a/litellm-rust/crates/host-python/tests/hook_chain.rs b/litellm-rust/crates/host-python/tests/hook_chain.rs index b7ad24fd784..04909856841 100644 --- a/litellm-rust/crates/host-python/tests/hook_chain.rs +++ b/litellm-rust/crates/host-python/tests/hook_chain.rs @@ -4,14 +4,15 @@ use std::future::{Future, ready}; use futures_util::{FutureExt, future::BoxFuture}; use litellm_host::{ - hooks::CallHooks, + hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks}, interceptors::{RequestContext, WireRequest}, lifecycle::{FailureOrigin, Timing}, }; use litellm_host_python::{ - Awaiter, HookChain, PythonCallEvent, PythonOwned, PythonRuntime, run_call, + Awaiter, Hooks, PythonGroup, PythonOwned, PythonSurface, report_unless_cancelled, run_call, }; use pyo3::{ + exceptions::PyBaseException, gc::{PyTraverseError, PyVisit}, prelude::*, types::{PyDict, PyTuple}, @@ -59,8 +60,28 @@ impl PythonOwned for ScriptHooks { } } -impl CallHooks for ScriptHooks { - async fn prepare_arguments<'a>( +impl ScriptHooks { + async fn event( + &self, + task: &dyn Awaiter, + name: &str, + value: impl FnOnce(Python<'_>) -> Py, + ) -> Result<(), Cancelled> { + let result = self + .invoke(task, "event", |py| { + Ok( + PyTuple::new(py, [name.into_pyobject(py)?.into_any().unbind(), value(py)])? + .into_any() + .unbind(), + ) + }) + .await; + report_unless_cancelled(result.map(drop)) + } +} + +impl ArgumentHooks for ScriptHooks { + async fn prepare<'a>( &'a self, task: &'a dyn Awaiter, arguments: Py, @@ -72,11 +93,13 @@ impl CallHooks for ScriptHooks { Python::attach(|py| Ok(value.into_bound(py).cast_into::()?.unbind())) } - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + fn prepared(&self, arguments: &Py) -> PyResult<()> { Python::attach(|py| self.object.call_method1(py, "adopt", (arguments,)))?; Ok(()) } +} +impl ProviderHooks for ScriptHooks { async fn before_provider_request<'a>( &'a self, task: &'a dyn Awaiter, @@ -94,8 +117,10 @@ impl CallHooks for ScriptHooks { ..*wire })) } +} - async fn transform_response<'a>( +impl ResponseHooks for ScriptHooks { + async fn transform<'a>( &'a self, task: &'a dyn Awaiter, response: Py, @@ -104,42 +129,44 @@ impl CallHooks for ScriptHooks { self.invoke(task, "transform", |_| Ok(response)).await } - async fn on_event<'a>( - &'a self, - task: &'a dyn Awaiter, - event: PythonCallEvent<'_>, - ) -> 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(&self, _head: &Py) -> PyResult<()> { + fn stream_opened(&self, _head: &Py) -> PyResult<()> { Python::attach(|py| self.object.call_method1(py, "stream", (py.None(),)))?; Ok(()) } - fn on_stream_chunk(&self, chunk: &Py) -> PyResult<()> { + fn stream_chunk(&self, chunk: &Py) -> PyResult<()> { Python::attach(|py| self.object.call_method1(py, "stream", (chunk,)))?; Ok(()) } } +impl OutcomeHooks for ScriptHooks { + async fn started<'a>(&'a self, task: &'a dyn Awaiter, _: f64) -> Result<(), Cancelled> { + self.event(task, "started", |py| py.None()).await + } + + async fn succeeded<'a>( + &'a self, + task: &'a dyn Awaiter, + _: Timing, + response: &Py, + ) -> Result<(), Cancelled> { + self.event(task, "success", |py| response.clone_ref(py)) + .await + } + + async fn failed<'a>( + &'a self, + task: &'a dyn Awaiter, + _: Timing, + _: FailureOrigin, + error: &Py, + ) -> Result<(), Cancelled> { + self.event(task, "failure", |py| error.clone_ref(py).into_any()) + .await + } +} + /// Runs every awaitable a hook hands it to completion on a fresh event loop. struct RunToCompletion; @@ -154,7 +181,7 @@ impl Awaiter for RunToCompletion { } } -fn finish(hook: impl Future>) -> PyResult { +fn finish(hook: impl Future) -> T { hook.now_or_never() .expect("every awaitable runs to completion inline") } @@ -209,12 +236,33 @@ fourth = Hooks('d') }) } -fn chain(py: Python<'_>, scripts: &Py, asynchronous: bool) -> HookChain { - let hook = |name| ScriptHooks { +fn script_hooks( + py: Python<'_>, + scripts: &Py, + name: &str, + asynchronous: bool, +) -> ScriptHooks { + ScriptHooks { object: scripts.bind(py).get_item(name).unwrap().unwrap().unbind(), asynchronous, - }; - HookChain::new().with(hook("first")).with(hook("second")) + } +} + +fn chain(py: Python<'_>, scripts: &Py, asynchronous: bool) -> Hooks { + ["first", "second"] + .into_iter() + .fold(Hooks::new(), |hooks, name| { + hooks.with(PythonSurface::group(script_hooks( + py, + scripts, + name, + asynchronous, + ))) + }) +} + +fn mounted(hooks: Hooks, group: impl PythonGroup + 'static) -> Hooks { + hooks.with(PythonSurface::group(group)) } const TIMING: Timing = Timing { @@ -224,12 +272,16 @@ const TIMING: Timing = Timing { struct PreparedPolicy; -impl CallHooks for PreparedPolicy { - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { +impl ArgumentHooks for PreparedPolicy { + fn prepared(&self, arguments: &Py) -> PyResult<()> { Python::attach(|py| arguments.bind(py).set_item("policy", "configured")) } } +impl ProviderHooks for PreparedPolicy {} +impl ResponseHooks for PreparedPolicy {} +impl OutcomeHooks for PreparedPolicy {} + impl PythonOwned for PreparedPolicy { fn close(&mut self, _: Python<'_>) {} @@ -241,16 +293,16 @@ impl PythonOwned for PreparedPolicy { #[rstest] #[case::sync(false)] #[case::suspended(true)] -fn transformations_feed_each_other_and_notifications_share_final_values( +fn request_stages_feed_forward_and_response_stages_unwind_over_the_final_values( scripts: Py, #[case] asynchronous: bool, ) { Python::attach(|py| { - let hooks = chain(py, &scripts, asynchronous).with(PreparedPolicy); + let hooks = mounted(chain(py, &scripts, asynchronous), PreparedPolicy); let original = PyDict::new(py).unbind(); let arguments = - finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap(); - hooks.arguments_prepared(&arguments).unwrap(); + finish(hooks.prepare(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap(); + hooks.prepared(&arguments).unwrap(); let context = RequestContext { model: "model".into(), custom_llm_provider: "provider".into(), @@ -267,18 +319,10 @@ fn transformations_feed_each_other_and_notifications_share_final_values( assert_eq!(wire.url, "urlab"); let response = PyDict::new(py).unbind().into_any(); 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(); + finish(hooks.transform(&RunToCompletion, response.clone_ref(py), TIMING)).unwrap(); + finish(hooks.succeeded(&RunToCompletion, TIMING, &final_response)).unwrap(); + hooks.stream_opened(&py.None()).unwrap(); + hooks.stream_chunk(&response).unwrap(); let locals = scripts.bind(py); locals.set_item("arguments", arguments).unwrap(); locals.set_item("original", original).unwrap(); @@ -291,10 +335,10 @@ assert original == {} assert first.adopted is arguments and second.adopted is arguments assert first.adopted['policy'] == second.adopted['policy'] == 'configured' assert first.model == second.model == 'model' -assert final_response == ((response, 'a'), 'b') +assert final_response == ((response, 'b'), 'a') assert [(name, kind) for name, kind, value in log] == [ - ('a', 'prepare'), ('b', 'prepare'), ('a', 'success'), ('b', 'success'), - ('a', 'stream'), ('b', 'stream'), ('a', 'stream'), ('b', 'stream'), + ('a', 'prepare'), ('b', 'prepare'), ('b', 'success'), ('a', 'success'), + ('b', 'stream'), ('a', 'stream'), ('b', 'stream'), ('a', 'stream'), ] assert log[2][2] is final_response and log[3][2] is final_response assert log[6][2] is response and log[7][2] is response @@ -318,15 +362,10 @@ fn terminal_failure_does_not_skip_later_hooks( #[values("first", "second")] failing_hook: &str, ) { Python::attach(|py| { - let hooks = chain(py, &scripts, asynchronous).with(ScriptHooks { - object: scripts - .bind(py) - .get_item("third") - .unwrap() - .unwrap() - .unbind(), - asynchronous, - }); + let hooks = mounted( + chain(py, &scripts, asynchronous), + script_hooks(py, &scripts, "third", asynchronous), + ); let locals = scripts.bind(py); locals .get_item(failing_hook) @@ -339,25 +378,17 @@ fn terminal_failure_does_not_skip_later_hooks( .unwrap(); let error = pyo3::exceptions::PyValueError::new_err("provider failed").into_value(py); let response = py.None(); - let event = if failed { - PythonCallEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Call, - error: &error, - } + if failed { + finish(hooks.failed(&RunToCompletion, TIMING, FailureOrigin::Call, &error)).unwrap(); } else { - PythonCallEvent::Succeeded { - timing: TIMING, - response: &response, - } - }; - finish(hooks.on_event(&RunToCompletion, event)).unwrap(); + finish(hooks.succeeded(&RunToCompletion, TIMING, &response)).unwrap(); + } locals .set_item("selected", if failed { error.into_any() } else { response }) .unwrap(); py.run( c" -assert [name for name, kind, value in log] == ['a', 'b', 'c'] +assert [name for name, kind, value in log] == ['c', 'b', 'a'] assert all(kind == log[0][1] for name, kind, value in log) assert all(value is selected for name, kind, value in log) ", @@ -381,8 +412,7 @@ fn transformation_failure_stops_the_chain(scripts: Py, #[case] asynchron Some(locals), ) .unwrap(); - let result = - finish(hooks.prepare_arguments(&RunToCompletion, PyDict::new(py).unbind(), 0.0)); + let result = finish(hooks.prepare(&RunToCompletion, PyDict::new(py).unbind(), 0.0)); assert!( result.unwrap_err().value(py).is(locals .get_item("first") @@ -408,26 +438,17 @@ fn cancellation_stops_notification_dispatch(scripts: Py, #[case] asynchr let hooks = chain(py, &scripts, asynchronous); let locals = scripts.bind(py); py.run( - c"first.error = asyncio.CancelledError()", + c"second.error = asyncio.CancelledError()", Some(locals), Some(locals), ) .unwrap(); let response = py.None(); - let result = finish(hooks.on_event( - &RunToCompletion, - PythonCallEvent::Succeeded { - timing: TIMING, - response: &response, - }, - )); - assert!( - result - .unwrap_err() - .is_instance_of::(py) - ); + let Cancelled(error) = + finish(hooks.succeeded(&RunToCompletion, TIMING, &response)).unwrap_err(); + assert!(error.is_instance_of::(py)); py.run( - c"assert len(log) == 1 and log[0][0] == 'a'", + c"assert len(log) == 1 and log[0][0] == 'b'", Some(locals), Some(locals), ) @@ -504,7 +525,7 @@ asyncio.run(exercise()) #[case::pair(2, "ab")] #[case::three(3, "abc")] #[case::four(4, "abcd")] -fn builder_runs_hooks_in_append_order( +fn argument_groups_run_in_mount_order( scripts: Py, #[case] count: usize, #[case] expected: &str, @@ -514,15 +535,11 @@ fn builder_runs_hooks_in_append_order( 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, - }) + .fold(Hooks::new(), |chain, name| { + mounted(chain, script_hooks(py, &scripts, name, asynchronous)) }); let original = PyDict::new(py).unbind(); - let result = - finish(hooks.prepare_arguments(&RunToCompletion, original.clone_ref(py), 0.0)).unwrap(); + let result = finish(hooks.prepare(&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 0dd7b4ce9b1..f68fa6545a6 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -4,18 +4,18 @@ use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; use litellm_host::{ call::{CallOutput, HostedMachine, hosted_call}, - hooks::CallHooks, - interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}, - lifecycle::{ExecutionEvent, Timing}, + hooks::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks}, + interceptors::{ExecutionFacts, Interceptors, RawResponse, RequestContext, WireRequest}, + lifecycle::{FailureOrigin, Timing}, machine::{CallMachine, MachineFault}, protocol::{Protocol, Reply}, }; use litellm_host_python::{ - Awaiter, CallOptions, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks, - PythonHostCalls, PythonOwned, PythonRuntime, run_call, + Awaiter, CallOptions, InvokeError, PythonBinding, PythonGroup, PythonHostCalls, PythonOwned, + PythonSurface, report_unless_cancelled, run_call, }; use pyo3::{ - exceptions::{PyTypeError, PyValueError, asyncio::CancelledError}, + exceptions::{PyBaseException, PyTypeError, PyValueError, asyncio::CancelledError}, gc::{PyTraverseError, PyVisit}, prelude::*, types::{PyDict, PyString}, @@ -313,8 +313,27 @@ pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks { } } -impl CallHooks for SyntheticHooks { - async fn prepare_arguments<'a>( +impl SyntheticHooks { + /// What a terminal hook does with the outcome it was handed: cancels the call, raises + /// an ordinary error that is only reported, or records it. + fn outcome(&self, entry: String) -> Result<(), Cancelled> { + match self.script { + HookScript::CancelTerminal => { + Err(Cancelled(CancelledError::new_err("callback cancelled"))) + } + HookScript::FailTerminal => { + report_unless_cancelled(Err(PyValueError::new_err("callback failed"))) + } + _ => { + self.log.push(entry); + Ok(()) + } + } + } +} + +impl ArgumentHooks for SyntheticHooks { + async fn prepare<'a>( &'a self, _: &'a dyn Awaiter, arguments: Py, @@ -336,7 +355,7 @@ impl CallHooks for SyntheticHooks { Ok(arguments) } - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { + fn prepared(&self, arguments: &Py) -> PyResult<()> { if matches!(self.script, HookScript::ObserveArguments) { let value: String = Python::attach(|py| arguments.bind(py).get_item("prepared")?.unwrap().extract())?; @@ -344,7 +363,9 @@ impl CallHooks for SyntheticHooks { } Ok(()) } +} +impl ProviderHooks for SyntheticHooks { async fn before_provider_request<'a>( &'a self, _: &'a dyn Awaiter, @@ -358,7 +379,23 @@ impl CallHooks for SyntheticHooks { })) } - async fn transform_response<'a>( + async fn after_provider_response<'a>( + &'a self, + _: &'a dyn Awaiter, + raw: &RawResponse, + ) -> PyResult<()> { + self.log.push(format!("response:{}", raw.body)); + Ok(()) + } + + async fn result_ready<'a>(&'a self, _: &'a dyn Awaiter, _: &ExecutionFacts) -> PyResult<()> { + self.log.push("cache_hit"); + Ok(()) + } +} + +impl ResponseHooks for SyntheticHooks { + async fn transform<'a>( &'a self, _: &'a dyn Awaiter, response: Py, @@ -374,49 +411,56 @@ impl CallHooks for SyntheticHooks { } } - async fn on_event<'a>( - &'a self, - _: &'a dyn Awaiter, - event: PythonCallEvent<'_>, - ) -> PyResult<()> { - let terminal = matches!( - event, - PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. } - ); - if terminal && matches!(self.script, HookScript::CancelTerminal) { - return Err(CancelledError::new_err("callback cancelled")); - } - if terminal && matches!(self.script, HookScript::FailTerminal) { - return Err(PyValueError::new_err("callback failed")); - } - self.log.push(Python::attach(|py| match event { - PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(), - PythonCallEvent::Started { .. } => "started".into(), - PythonCallEvent::Cancelled { .. } => "cancelled".into(), - PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - format!("response:{}", raw.body) - } - PythonCallEvent::Succeeded { response, .. } => { - format!("succeeded:{}", response.bind(py)) - } - PythonCallEvent::Failed { origin, error, .. } => { - format!("failed:{origin:?}:{}", error.bind(py)) - } - })); - Ok(()) - } - - fn on_stream_open(&self, _: &Py) -> PyResult<()> { + fn stream_opened(&self, _: &Py) -> PyResult<()> { self.log.push("opened"); Ok(()) } - fn on_stream_chunk(&self, _: &Py) -> PyResult<()> { + fn stream_chunk(&self, _: &Py) -> PyResult<()> { self.log.push("delivered"); Ok(()) } } +impl OutcomeHooks for SyntheticHooks { + async fn started<'a>(&'a self, _: &'a dyn Awaiter, _: f64) -> Result<(), Cancelled> { + self.log.push("started"); + Ok(()) + } + + async fn succeeded<'a>( + &'a self, + _: &'a dyn Awaiter, + _: Timing, + response: &Py, + ) -> Result<(), Cancelled> { + self.outcome(Python::attach(|py| { + format!("succeeded:{}", response.bind(py)) + })) + } + + async fn failed<'a>( + &'a self, + _: &'a dyn Awaiter, + _: Timing, + origin: FailureOrigin, + error: &Py, + ) -> Result<(), Cancelled> { + self.outcome(Python::attach(|py| { + format!("failed:{origin:?}:{}", error.bind(py)) + })) + } + + async fn cancelled<'a>( + &'a self, + _: &'a dyn Awaiter, + _: Timing, + ) -> Result<(), Cancelled> { + self.log.push("cancelled"); + Ok(()) + } +} + impl PythonOwned for SyntheticHooks { fn close(&mut self, _: Python<'_>) { self.log.push("adapter.close"); @@ -495,7 +539,7 @@ pub fn run_scripted( ) } -pub fn run_composed( +pub fn run_composed( py: Python<'_>, machine: MachineFactory, binding: SyntheticBinding, diff --git a/litellm-rust/crates/host/AGENTS.md b/litellm-rust/crates/host/AGENTS.md index d3387544d9e..59bd06a3c12 100644 --- a/litellm-rust/crates/host/AGENTS.md +++ b/litellm-rust/crates/host/AGENTS.md @@ -4,7 +4,7 @@ | --- | --- | --- | --- | | Input and output conversion | `Protocol::{Request, Response, StreamHead, Chunk, Error}` | Typed input. `ResponseEncoder` and `StreamEncoder` produce HTTP values | `PythonBinding` decodes prepared arguments, encodes public values and maps native errors | | Host services | `Protocol::HostCall`, `HostServices::call` | `host-native::services::HostCallHandler` answers typed calls. `()` handles protocols without host calls | `PythonHostCalls` invokes retained Python objects. It may share an owner with the binding | -| Active hooks | `Interceptors::{before_provider_request, after_provider_response}` and `hooks::CallHooks` | Request interception and fallible execution callbacks | `PythonCallHooks` also prepares arguments, transforms public responses and receives stream callbacks | +| Active hooks | `Interceptors::{before_provider_request, after_provider_response}` and the `hooks` facets (`ArgumentHooks`, `ProviderHooks`, `ResponseHooks`, `OutcomeHooks`) over a `Surface` | Request interception and fallible execution callbacks | `PythonGroup` implements the four facets on `PythonSurface`, mounted through `Hooks` | | Passive observation | `observation::ObservationSender` | Queued execution and lifecycle snapshots, retained by the response body | Public Python callbacks remain active hooks with their existing failure policy | | Runtime driving | `Machine`, `HostRequest::{HostCall, Intercept, Stream}` | Body polling controls demand | Native polling and inline caller-task Python awaits control progress | @@ -16,7 +16,7 @@ Each driver owns terminal dispatch. Hooks can change or fail execution. Passive `interceptors.rs` owns `Interceptors` and its request/response payload types. `lifecycle.rs` owns `CallObserver`, `CallEvent`, `ExecutionEvent`, timing, failure origin, and the observation wrappers. Event payloads are generic so a runtime can retain its own response, exception and raw-response references without introducing a language dependency. `snapshot()` projects them into the owned observation contract without retaining runtime objects. Pass interceptors and observers separately at direct route and HTTP entrypoints. Routes publish execution events independently of interception -`hooks.rs` owns the async call-stage interface and its runtime-associated types. `hooks::HookChain` composes stages in order, using `DynCallHooks` for boxed storage and `CallHooks` for callback authors. A terminal notification reaches every group, and `ChainRuntime::terminal_notification_failed` decides whether a group's error is reported or stops the chain. A runtime supplies its context through `HookRuntime` and its stored hook type through `ChainRuntime`. Python ownership and cancellation guards stay in the Python adapter +`hooks.rs` owns the call's facets, one trait per payload, each with its own failure policy in its return type: `ArgumentHooks` and `ProviderHooks` fold a value through the groups in mount order and stop on the first error, `ResponseHooks` unwinds in reverse mount order, `OutcomeHooks` reaches every group in reverse and only a `Cancelled` group stops the rest. What a group does with an ordinary outcome error is the group's business. `Surface` names a host's public values (`Task`, `Arguments`, `Response`, `Chunk`, `Error`, `Failure`) and how a group is stored (`Group`). `hooks::Hooks` composes mounted groups and implements every facet; `AnyGroup` is the object-safe mirror a group is boxed through, never the trait a group author implements. Python ownership and cancellation guards stay in the Python adapter `protocol.rs` owns `Protocol` and suspension messages, including `InterceptRequest` and `StreamDelivery`. `call.rs` owns route outputs and their adaptation into a hosted machine. Rust service handling belongs in `host-native::services`. Coroutine channel handles stay in `machine/context.rs` diff --git a/litellm-rust/crates/host/src/hooks.rs b/litellm-rust/crates/host/src/hooks.rs index bfb1ef1cf83..e4318bff8eb 100644 --- a/litellm-rust/crates/host/src/hooks.rs +++ b/litellm-rust/crates/host/src/hooks.rs @@ -1,73 +1,141 @@ -mod adapter; -mod chain; +//! What a call lets its host hook into, one trait per payload. A host picks the facets +//! its public surface has and composes any number of groups through [`Hooks`]. -pub use adapter::DynCallHooks; -pub use chain::{ChainRuntime, HookChain, HookGroupAdapter}; +mod chain; +mod group; use std::future::{Future, ready}; +pub use chain::Hooks; +pub use group::AnyGroup; + use crate::{ - interceptors::{RawResponse, RequestContext, WireRequest}, - lifecycle::{CallEvent, Timing}, + interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest}, + lifecycle::{FailureOrigin, Timing}, }; -pub trait HookRuntime { - type Context<'a>: Send; +/// A host's public values as its hooks see them. +pub trait Surface: Sized + 'static { + /// The caller's task a hook awaits on, or `()` for a host that never suspends. + type Task<'a>: Copy + Send; type Arguments: Send; type Response: Send + Sync; type Chunk; type Error: Send + Sync; + /// What a failed outcome carries: the host's retained form of its error. type Failure: Sync; + /// How [`Hooks`] stores a mounted group. + type Group: ?Sized + AnyGroup + Send + Sync; } -pub type RuntimeCallEvent<'a, R> = - CallEvent<&'a ::Response, &'a ::Failure, &'a RawResponse>; +/// A hook stopped the call: nothing after it runs and the error is raised as it is. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Cancelled(pub E); -pub trait CallHooks: Sync { - fn prepare_arguments<'a>( +/// The caller's arguments before the route decodes them. Groups run in mount order and +/// each receives the previous group's output; the first error stops the chain. +pub trait ArgumentHooks: Send + Sync { + fn prepare<'a>( &'a self, - _runtime: R::Context<'a>, - arguments: R::Arguments, + _task: S::Task<'a>, + arguments: S::Arguments, _started_at: f64, - ) -> impl Future> + Send { + ) -> impl Future> + Send { ready(Ok(arguments)) } - fn arguments_prepared(&self, _arguments: &R::Arguments) -> Result<(), R::Error> { - Ok(()) - } - - fn before_provider_request<'a>( - &'a self, - _runtime: R::Context<'a>, - wire: Box, - _context: &RequestContext, - ) -> impl Future, R::Error>> + Send { - ready(Ok(wire)) - } - - fn transform_response<'a>( - &'a self, - _runtime: R::Context<'a>, - response: R::Response, - _timing: Timing, - ) -> impl Future> + Send { - ready(Ok(response)) - } - - fn on_event<'a>( - &'a self, - _runtime: R::Context<'a>, - _event: RuntimeCallEvent<'_, R>, - ) -> impl Future> + Send { - ready(Ok(())) - } - - fn on_stream_open(&self, _head: &R::Response) -> Result<(), R::Error> { - Ok(()) - } - - fn on_stream_chunk(&self, _chunk: &R::Chunk) -> Result<(), R::Error> { + fn prepared(&self, _arguments: &S::Arguments) -> Result<(), S::Error> { Ok(()) } } + +/// The provider exchange: the wire request about to leave, the raw response that came +/// back, and the facts of the execution. Groups run in mount order; the first error +/// interrupts the call. +pub trait ProviderHooks: Send + Sync { + fn before_provider_request<'a>( + &'a self, + _task: S::Task<'a>, + wire: Box, + _context: &RequestContext, + ) -> impl Future, S::Error>> + Send { + ready(Ok(wire)) + } + + fn after_provider_response<'a>( + &'a self, + _task: S::Task<'a>, + _raw: &RawResponse, + ) -> impl Future> + Send { + ready(Ok(())) + } + + fn result_ready<'a>( + &'a self, + _task: S::Task<'a>, + _facts: &ExecutionFacts, + ) -> impl Future> + Send { + ready(Ok(())) + } +} + +/// The public response and stream. Transforms unwind in reverse mount order, so the first +/// group mounted sees the response last; the first error stops the chain. +pub trait ResponseHooks: Send + Sync { + fn transform<'a>( + &'a self, + _task: S::Task<'a>, + response: S::Response, + _timing: Timing, + ) -> impl Future> + Send { + ready(Ok(response)) + } + + fn stream_opened(&self, _head: &S::Response) -> Result<(), S::Error> { + Ok(()) + } + + fn stream_chunk(&self, _chunk: &S::Chunk) -> Result<(), S::Error> { + Ok(()) + } +} + +/// How the call ended. Every group hears the outcome, terminal ones in reverse mount +/// order, and only a [`Cancelled`] group stops the rest. What a group does with its own +/// errors is its business. +pub trait OutcomeHooks: Send + Sync { + fn started<'a>( + &'a self, + _task: S::Task<'a>, + _start_time: f64, + ) -> impl Future>> + Send { + ready(Ok(())) + } + + fn succeeded<'a>( + &'a self, + _task: S::Task<'a>, + _timing: Timing, + _response: &S::Response, + ) -> impl Future>> + Send { + ready(Ok(())) + } + + fn failed<'a>( + &'a self, + _task: S::Task<'a>, + _timing: Timing, + _origin: FailureOrigin, + _error: &S::Failure, + ) -> impl Future>> + Send { + ready(Ok(())) + } + + fn cancelled<'a>( + &'a self, + _task: S::Task<'a>, + _timing: Timing, + ) -> impl Future>> + Send { + ready(Ok(())) + } +} diff --git a/litellm-rust/crates/host/src/hooks/adapter.rs b/litellm-rust/crates/host/src/hooks/adapter.rs deleted file mode 100644 index 3c0b0e295f0..00000000000 --- a/litellm-rust/crates/host/src/hooks/adapter.rs +++ /dev/null @@ -1,91 +0,0 @@ -use futures_util::future::BoxFuture; - -use super::{CallHooks, HookRuntime, RuntimeCallEvent}; -use crate::{ - interceptors::{RequestContext, WireRequest}, - lifecycle::Timing, -}; - -pub trait DynCallHooks: Sync { - fn prepare_arguments<'a>( - &'a self, - runtime: R::Context<'a>, - arguments: R::Arguments, - started_at: f64, - ) -> BoxFuture<'a, Result>; - fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error>; - fn before_provider_request<'a: 'b, 'b>( - &'a self, - runtime: R::Context<'a>, - wire: Box, - context: &'b RequestContext, - ) -> BoxFuture<'b, Result, R::Error>>; - fn transform_response<'a>( - &'a self, - runtime: R::Context<'a>, - response: R::Response, - timing: Timing, - ) -> BoxFuture<'a, Result>; - fn on_event<'a: 'b, 'b>( - &'a self, - runtime: R::Context<'a>, - event: RuntimeCallEvent<'b, R>, - ) -> BoxFuture<'b, Result<(), R::Error>>; - fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error>; - fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error>; -} - -impl> DynCallHooks for H { - fn prepare_arguments<'a>( - &'a self, - runtime: R::Context<'a>, - arguments: R::Arguments, - started_at: f64, - ) -> BoxFuture<'a, Result> { - Box::pin(CallHooks::prepare_arguments( - self, runtime, arguments, started_at, - )) - } - - fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error> { - CallHooks::arguments_prepared(self, arguments) - } - - fn before_provider_request<'a: 'b, 'b>( - &'a self, - runtime: R::Context<'a>, - wire: Box, - context: &'b RequestContext, - ) -> BoxFuture<'b, Result, R::Error>> { - Box::pin(CallHooks::before_provider_request( - self, runtime, wire, context, - )) - } - - fn transform_response<'a>( - &'a self, - runtime: R::Context<'a>, - response: R::Response, - timing: Timing, - ) -> BoxFuture<'a, Result> { - Box::pin(CallHooks::transform_response( - self, runtime, response, timing, - )) - } - - fn on_event<'a: 'b, 'b>( - &'a self, - runtime: R::Context<'a>, - event: RuntimeCallEvent<'b, R>, - ) -> BoxFuture<'b, Result<(), R::Error>> { - Box::pin(CallHooks::on_event(self, runtime, event)) - } - - fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error> { - CallHooks::on_stream_open(self, head) - } - - fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error> { - CallHooks::on_stream_chunk(self, chunk) - } -} diff --git a/litellm-rust/crates/host/src/hooks/chain.rs b/litellm-rust/crates/host/src/hooks/chain.rs index 01cff5c4d8c..dc99378027b 100644 --- a/litellm-rust/crates/host/src/hooks/chain.rs +++ b/litellm-rust/crates/host/src/hooks/chain.rs @@ -1,133 +1,175 @@ -use super::{CallHooks, DynCallHooks, HookRuntime, RuntimeCallEvent}; +use super::{ + AnyGroup, ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks, Surface, +}; use crate::{ - interceptors::{RequestContext, WireRequest}, - lifecycle::{CallEvent, Timing}, + interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest}, + lifecycle::{FailureOrigin, Timing}, }; -pub trait ChainRuntime: HookRuntime + Sized + 'static { - type Group: DynCallHooks + ?Sized; - - /// `Ok` lets the next group receive the terminal event, `Err` stops the chain. - fn terminal_notification_failed(error: Self::Error) -> Result<(), Self::Error>; +/// The groups mounted on a call, run as one group of each facet. Request-side stages run +/// in mount order, response-side and terminal stages unwind in reverse, like layers. +pub struct Hooks { + groups: Vec>, } -pub trait HookGroupAdapter: ChainRuntime { - fn adapt_group(hooks: H) -> Box; -} - -pub struct HookChain { - hooks: Vec>, -} - -impl HookChain { - pub fn groups(&self) -> impl Iterator { - self.hooks.iter().map(Box::as_ref) - } - - pub fn groups_mut(&mut self) -> impl Iterator { - self.hooks.iter_mut().map(Box::as_mut) - } - +impl Hooks { pub fn new() -> Self { Self::default() } - pub fn with(mut self, hooks: H) -> Self - where - R: HookGroupAdapter, - { - self.hooks.push(R::adapt_group(hooks)); + pub fn with(mut self, group: Box) -> Self { + self.groups.push(group); self } + + pub fn groups(&self) -> impl Iterator { + self.groups.iter().map(Box::as_ref) + } + + pub fn groups_mut(&mut self) -> impl Iterator { + self.groups.iter_mut().map(Box::as_mut) + } } -impl CallHooks for HookChain -where - for<'a> R::Context<'a>: Copy, -{ - async fn prepare_arguments<'a>( +impl Default for Hooks { + fn default() -> Self { + Self { groups: Vec::new() } + } +} + +impl ArgumentHooks for Hooks { + async fn prepare<'a>( &'a self, - runtime: R::Context<'a>, - arguments: R::Arguments, + task: S::Task<'a>, + arguments: S::Arguments, started_at: f64, - ) -> Result { + ) -> Result { let mut arguments = arguments; - for hooks in &self.hooks { - arguments = hooks - .prepare_arguments(runtime, arguments, started_at) - .await?; + for group in &self.groups { + arguments = group.prepare(task, arguments, started_at).await?; } Ok(arguments) } - fn arguments_prepared(&self, arguments: &R::Arguments) -> Result<(), R::Error> { - self.hooks + fn prepared(&self, arguments: &S::Arguments) -> Result<(), S::Error> { + self.groups .iter() - .try_for_each(|hooks| hooks.arguments_prepared(arguments)) + .try_for_each(|group| group.prepared(arguments)) } +} +impl ProviderHooks for Hooks { async fn before_provider_request<'a>( &'a self, - runtime: R::Context<'a>, + task: S::Task<'a>, wire: Box, context: &RequestContext, - ) -> Result, R::Error> { + ) -> Result, S::Error> { let mut wire = wire; - for hooks in &self.hooks { - wire = hooks - .before_provider_request(runtime, wire, context) - .await?; + for group in &self.groups { + wire = group.before_provider_request(task, wire, context).await?; } Ok(wire) } - async fn transform_response<'a>( + async fn after_provider_response<'a>( &'a self, - runtime: R::Context<'a>, - response: R::Response, - timing: Timing, - ) -> Result { - let mut response = response; - for hooks in &self.hooks { - response = hooks.transform_response(runtime, response, timing).await?; - } - Ok(response) - } - - async fn on_event<'a>( - &'a self, - runtime: R::Context<'a>, - event: RuntimeCallEvent<'_, R>, - ) -> Result<(), R::Error> { - let terminal = matches!( - event, - CallEvent::Succeeded { .. } | CallEvent::Failed { .. } - ); - for hooks in &self.hooks { - match hooks.on_event(runtime, event.clone()).await { - Err(error) if terminal => R::terminal_notification_failed(error)?, - result => result?, - } + task: S::Task<'a>, + raw: &RawResponse, + ) -> Result<(), S::Error> { + for group in &self.groups { + group.after_provider_response(task, raw).await?; } Ok(()) } - fn on_stream_open(&self, head: &R::Response) -> Result<(), R::Error> { - self.hooks - .iter() - .try_for_each(|hooks| hooks.on_stream_open(head)) - } - - fn on_stream_chunk(&self, chunk: &R::Chunk) -> Result<(), R::Error> { - self.hooks - .iter() - .try_for_each(|hooks| hooks.on_stream_chunk(chunk)) + async fn result_ready<'a>( + &'a self, + task: S::Task<'a>, + facts: &ExecutionFacts, + ) -> Result<(), S::Error> { + for group in &self.groups { + group.result_ready(task, facts).await?; + } + Ok(()) } } -impl Default for HookChain { - fn default() -> Self { - Self { hooks: Vec::new() } +impl ResponseHooks for Hooks { + async fn transform<'a>( + &'a self, + task: S::Task<'a>, + response: S::Response, + timing: Timing, + ) -> Result { + let mut response = response; + for group in self.groups.iter().rev() { + response = group.transform(task, response, timing).await?; + } + Ok(response) + } + + fn stream_opened(&self, head: &S::Response) -> Result<(), S::Error> { + self.groups + .iter() + .rev() + .try_for_each(|group| group.stream_opened(head)) + } + + fn stream_chunk(&self, chunk: &S::Chunk) -> Result<(), S::Error> { + self.groups + .iter() + .rev() + .try_for_each(|group| group.stream_chunk(chunk)) + } +} + +impl OutcomeHooks for Hooks { + async fn started<'a>( + &'a self, + task: S::Task<'a>, + start_time: f64, + ) -> Result<(), Cancelled> { + for group in &self.groups { + group.started(task, start_time).await?; + } + Ok(()) + } + + async fn succeeded<'a>( + &'a self, + task: S::Task<'a>, + timing: Timing, + response: &S::Response, + ) -> Result<(), Cancelled> { + for group in self.groups.iter().rev() { + group.succeeded(task, timing, response).await?; + } + Ok(()) + } + + async fn failed<'a>( + &'a self, + task: S::Task<'a>, + timing: Timing, + origin: FailureOrigin, + error: &S::Failure, + ) -> Result<(), Cancelled> { + for group in self.groups.iter().rev() { + group.failed(task, timing, origin, error).await?; + } + Ok(()) + } + + async fn cancelled<'a>( + &'a self, + task: S::Task<'a>, + timing: Timing, + ) -> Result<(), Cancelled> { + for group in self.groups.iter().rev() { + group.cancelled(task, timing).await?; + } + Ok(()) } } @@ -138,36 +180,21 @@ mod tests { use rstest::{fixture, rstest}; use super::*; - use crate::lifecycle::FailureOrigin; - struct TestRuntime; + struct TestSurface; - impl HookRuntime for TestRuntime { - type Context<'a> = &'a Mutex>; + impl Surface for TestSurface { + type Task<'a> = &'a Mutex>; type Arguments = String; type Response = String; type Chunk = u32; type Error = &'static str; type Failure = &'static str; + type Group = dyn AnyGroup; } - impl ChainRuntime for TestRuntime { - type Group = dyn DynCallHooks; - - fn terminal_notification_failed(error: &'static str) -> Result<(), &'static str> { - match error { - "cancelled" => Err(error), - _ => Ok(()), - } - } - } - - impl + 'static> HookGroupAdapter for TestRuntime { - fn adapt_group(hooks: H) -> Box { - Box::new(hooks) - } - } - + /// A group that tags every value it transforms with its name, logs every stage it + /// runs, and fails one stage when told to. struct Group { name: &'static str, suspensions: u8, @@ -175,6 +202,10 @@ mod tests { } impl Group { + fn mounted(self) -> Box> { + Box::new(self) + } + async fn value( &self, trace: &Mutex>, @@ -192,8 +223,8 @@ mod tests { } } - impl CallHooks for Group { - async fn prepare_arguments<'a>( + impl ArgumentHooks for Group { + async fn prepare<'a>( &'a self, trace: &'a Mutex>, arguments: String, @@ -201,7 +232,9 @@ mod tests { ) -> Result { self.value(trace, "prepare", arguments).await } + } + impl ProviderHooks for Group { async fn before_provider_request<'a>( &'a self, trace: &'a Mutex>, @@ -211,8 +244,10 @@ mod tests { let url = self.value(trace, "wire", wire.url.clone()).await?; Ok(Box::new(WireRequest { url, ..*wire })) } + } - async fn transform_response<'a>( + impl ResponseHooks for Group { + async fn transform<'a>( &'a self, trace: &'a Mutex>, response: String, @@ -220,18 +255,43 @@ mod tests { ) -> Result { self.value(trace, "response", response).await } + } - async fn on_event<'a>( + impl OutcomeHooks for Group { + async fn started<'a>( &'a self, trace: &'a Mutex>, - event: RuntimeCallEvent<'_, TestRuntime>, - ) -> Result<(), &'static str> { - let label = match event { - CallEvent::Succeeded { response, .. } => response.as_str(), - CallEvent::Failed { error, .. } => *error, - _ => "event", - }; - self.value(trace, label, String::new()).await.map(|_| ()) + _: f64, + ) -> Result<(), Cancelled<&'static str>> { + self.value(trace, "started", String::new()) + .await + .map(drop) + .map_err(Cancelled) + } + + async fn succeeded<'a>( + &'a self, + trace: &'a Mutex>, + _: Timing, + response: &String, + ) -> Result<(), Cancelled<&'static str>> { + self.value(trace, response, String::new()) + .await + .map(drop) + .map_err(Cancelled) + } + + async fn failed<'a>( + &'a self, + trace: &'a Mutex>, + _: Timing, + _: FailureOrigin, + error: &&'static str, + ) -> Result<(), Cancelled<&'static str>> { + self.value(trace, error, String::new()) + .await + .map(drop) + .map_err(Cancelled) } } @@ -259,147 +319,128 @@ mod tests { } } + fn chain( + suspensions: u8, + rejecting: Option<(&'static str, &'static str)>, + ) -> Hooks { + ["a", "b", "c"] + .into_iter() + .fold(Hooks::new(), |hooks, name| { + hooks.with( + Group { + name, + suspensions, + reject: rejecting + .filter(|(group, _)| *group == name) + .map(|(_, error)| error), + } + .mounted(), + ) + }) + } + #[rstest] #[tokio::test] - async fn groups_compose_without_python( + async fn request_stages_run_in_mount_order_and_response_stages_unwind( trace: Mutex>, timing: Timing, context: RequestContext, #[values(0, 1, 2)] suspensions: u8, ) { - let hooks = HookChain::::new() - .with(Group { - name: "a", - suspensions, - reject: None, - }) - .with(Group { - name: "b", - suspensions, - reject: None, - }); - let arguments = CallHooks::prepare_arguments(&hooks, &trace, "input".into(), 0.0) + let hooks = chain(suspensions, None); + OutcomeHooks::started(&hooks, &trace, 0.0).await.unwrap(); + let arguments = ArgumentHooks::prepare(&hooks, &trace, "input".into(), 0.0) .await .unwrap(); - assert_eq!(arguments, "inputab"); + assert_eq!(arguments, "inputabc"); let wire = Box::new(WireRequest { url: "url".into(), headers: Vec::new(), body: serde_json::Value::Null, }); - let rewritten = CallHooks::before_provider_request(&hooks, &trace, wire, &context) + let rewritten = ProviderHooks::before_provider_request(&hooks, &trace, wire, &context) .await .unwrap(); - assert_eq!(rewritten.url, "urlab"); - let response = CallHooks::transform_response(&hooks, &trace, "output".into(), timing) + assert_eq!(rewritten.url, "urlabc"); + let response = ResponseHooks::transform(&hooks, &trace, "output".into(), timing) + .await + .unwrap(); + assert_eq!(response, "outputcba"); + OutcomeHooks::succeeded(&hooks, &trace, timing, &response) .await .unwrap(); - assert_eq!(response, "outputab"); - CallHooks::on_event( - &hooks, - &trace, - CallEvent::Succeeded { - timing, - response: &response, - }, - ) - .await - .unwrap(); assert_eq!( *trace.lock().unwrap(), [ + "started:a", + "started:b", + "started:c", "prepare:a", "prepare:b", + "prepare:c", "wire:a", "wire:b", - "response:a", + "wire:c", + "response:c", "response:b", - "outputab:a", - "outputab:b", + "response:a", + "outputcba:c", + "outputcba:b", + "outputcba:a", ] ); } - fn rejecting_chain(suspensions: u8, error: &'static str) -> HookChain { - ["a", "b", "c"] - .into_iter() - .fold(HookChain::new(), |hooks, name| { - hooks.with(Group { - name, - suspensions, - reject: (name == "b").then_some(error), - }) - }) - } - #[rstest] #[tokio::test] - async fn transformation_failure_stops_later_groups( + async fn a_transformation_error_stops_later_groups( trace: Mutex>, #[values(0, 2)] suspensions: u8, ) { - let hooks = rejecting_chain(suspensions, "rejected"); - let result = CallHooks::prepare_arguments(&hooks, &trace, String::new(), 0.0).await; + let hooks = chain(suspensions, Some(("b", "rejected"))); + let result = ArgumentHooks::prepare(&hooks, &trace, String::new(), 0.0).await; assert_eq!(result, Err("rejected")); assert_eq!(*trace.lock().unwrap(), ["prepare:a", "prepare:b"]); } #[rstest] - #[case::success_reaches_every_group(true, "rejected", Ok(()), &["selected:a", "selected:b", "selected:c"])] - #[case::failure_reaches_every_group(false, "rejected", Ok(()), &["failed:a", "failed:b", "failed:c"])] - #[case::cancellation_stops_success(true, "cancelled", Err("cancelled"), &["selected:a", "selected:b"])] - #[case::cancellation_stops_failure(false, "cancelled", Err("cancelled"), &["failed:a", "failed:b"])] + #[case::success(true)] + #[case::failure(false)] #[tokio::test] - async fn terminal_notification_errors_follow_the_runtime( + async fn a_cancelled_outcome_stops_the_groups_mounted_before_it( trace: Mutex>, timing: Timing, #[values(0, 2)] suspensions: u8, #[case] succeeded: bool, - #[case] error: &'static str, - #[case] expected: Result<(), &'static str>, - #[case] reached: &[&str], ) { - let hooks = rejecting_chain(suspensions, error); + let hooks = chain(suspensions, Some(("b", "cancelled"))); let response = "selected".to_owned(); - let event = if succeeded { - CallEvent::Succeeded { - timing, - response: &response, - } + let result = if succeeded { + OutcomeHooks::succeeded(&hooks, &trace, timing, &response).await } else { - CallEvent::Failed { - timing, - origin: FailureOrigin::Host, - error: &"failed", - } + OutcomeHooks::failed(&hooks, &trace, timing, FailureOrigin::Host, &"failed").await + }; + assert_eq!(result, Err(Cancelled("cancelled"))); + let reached = if succeeded { + ["selected:c", "selected:b"] + } else { + ["failed:c", "failed:b"] }; - let result = CallHooks::on_event(&hooks, &trace, event).await; - assert_eq!(result, expected); assert_eq!(*trace.lock().unwrap(), reached); } #[rstest] #[tokio::test] - async fn non_terminal_notification_error_stops_later_groups(trace: Mutex>) { - let hooks = rejecting_chain(0, "rejected"); - let result = - CallHooks::on_event(&hooks, &trace, CallEvent::Started { start_time: 0.0 }).await; - assert_eq!(result, Err("rejected")); - assert_eq!(*trace.lock().unwrap(), ["event:a", "event:b"]); - } - - #[rstest] - #[tokio::test] - async fn empty_chain_preserves_values(trace: Mutex>, timing: Timing) { - let hooks = HookChain::::new(); + async fn an_empty_chain_preserves_values(trace: Mutex>, timing: Timing) { + let hooks = Hooks::::new(); assert_eq!( - CallHooks::prepare_arguments(&hooks, &trace, "input".into(), 0.0) + ArgumentHooks::prepare(&hooks, &trace, "input".into(), 0.0) .await .unwrap(), "input" ); assert_eq!( - CallHooks::transform_response(&hooks, &trace, "output".into(), timing) + ResponseHooks::transform(&hooks, &trace, "output".into(), timing) .await .unwrap(), "output" diff --git a/litellm-rust/crates/host/src/hooks/group.rs b/litellm-rust/crates/host/src/hooks/group.rs new file mode 100644 index 00000000000..58b1ed15438 --- /dev/null +++ b/litellm-rust/crates/host/src/hooks/group.rs @@ -0,0 +1,164 @@ +use futures_util::future::BoxFuture; + +use super::{ArgumentHooks, Cancelled, OutcomeHooks, ProviderHooks, ResponseHooks, Surface}; +use crate::{ + interceptors::{ExecutionFacts, RawResponse, RequestContext, WireRequest}, + lifecycle::{FailureOrigin, Timing}, +}; + +/// Every facet behind one object-safe trait, so [`super::Hooks`] can store groups of +/// different types. Implemented for anything that implements the four facets. +pub trait AnyGroup: Send + Sync { + fn prepare<'a>( + &'a self, + task: S::Task<'a>, + arguments: S::Arguments, + started_at: f64, + ) -> BoxFuture<'a, Result>; + fn prepared(&self, arguments: &S::Arguments) -> Result<(), S::Error>; + fn before_provider_request<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + wire: Box, + context: &'b RequestContext, + ) -> BoxFuture<'b, Result, S::Error>>; + fn after_provider_response<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + raw: &'b RawResponse, + ) -> BoxFuture<'b, Result<(), S::Error>>; + fn result_ready<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + facts: &'b ExecutionFacts, + ) -> BoxFuture<'b, Result<(), S::Error>>; + fn transform<'a>( + &'a self, + task: S::Task<'a>, + response: S::Response, + timing: Timing, + ) -> BoxFuture<'a, Result>; + fn stream_opened(&self, head: &S::Response) -> Result<(), S::Error>; + fn stream_chunk(&self, chunk: &S::Chunk) -> Result<(), S::Error>; + fn started<'a>( + &'a self, + task: S::Task<'a>, + start_time: f64, + ) -> BoxFuture<'a, Result<(), Cancelled>>; + fn succeeded<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + timing: Timing, + response: &'b S::Response, + ) -> BoxFuture<'b, Result<(), Cancelled>>; + fn failed<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + timing: Timing, + origin: FailureOrigin, + error: &'b S::Failure, + ) -> BoxFuture<'b, Result<(), Cancelled>>; + fn cancelled<'a>( + &'a self, + task: S::Task<'a>, + timing: Timing, + ) -> BoxFuture<'a, Result<(), Cancelled>>; +} + +impl AnyGroup for H +where + S: Surface, + H: ArgumentHooks + ProviderHooks + ResponseHooks + OutcomeHooks, +{ + fn prepare<'a>( + &'a self, + task: S::Task<'a>, + arguments: S::Arguments, + started_at: f64, + ) -> BoxFuture<'a, Result> { + Box::pin(ArgumentHooks::prepare(self, task, arguments, started_at)) + } + + fn prepared(&self, arguments: &S::Arguments) -> Result<(), S::Error> { + ArgumentHooks::prepared(self, arguments) + } + + fn before_provider_request<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + wire: Box, + context: &'b RequestContext, + ) -> BoxFuture<'b, Result, S::Error>> { + Box::pin(ProviderHooks::before_provider_request( + self, task, wire, context, + )) + } + + fn after_provider_response<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + raw: &'b RawResponse, + ) -> BoxFuture<'b, Result<(), S::Error>> { + Box::pin(ProviderHooks::after_provider_response(self, task, raw)) + } + + fn result_ready<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + facts: &'b ExecutionFacts, + ) -> BoxFuture<'b, Result<(), S::Error>> { + Box::pin(ProviderHooks::result_ready(self, task, facts)) + } + + fn transform<'a>( + &'a self, + task: S::Task<'a>, + response: S::Response, + timing: Timing, + ) -> BoxFuture<'a, Result> { + Box::pin(ResponseHooks::transform(self, task, response, timing)) + } + + fn stream_opened(&self, head: &S::Response) -> Result<(), S::Error> { + ResponseHooks::stream_opened(self, head) + } + + fn stream_chunk(&self, chunk: &S::Chunk) -> Result<(), S::Error> { + ResponseHooks::stream_chunk(self, chunk) + } + + fn started<'a>( + &'a self, + task: S::Task<'a>, + start_time: f64, + ) -> BoxFuture<'a, Result<(), Cancelled>> { + Box::pin(OutcomeHooks::started(self, task, start_time)) + } + + fn succeeded<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + timing: Timing, + response: &'b S::Response, + ) -> BoxFuture<'b, Result<(), Cancelled>> { + Box::pin(OutcomeHooks::succeeded(self, task, timing, response)) + } + + fn failed<'a: 'b, 'b>( + &'a self, + task: S::Task<'a>, + timing: Timing, + origin: FailureOrigin, + error: &'b S::Failure, + ) -> BoxFuture<'b, Result<(), Cancelled>> { + Box::pin(OutcomeHooks::failed(self, task, timing, origin, error)) + } + + fn cancelled<'a>( + &'a self, + task: S::Task<'a>, + timing: Timing, + ) -> BoxFuture<'a, Result<(), Cancelled>> { + Box::pin(OutcomeHooks::cancelled(self, task, timing)) + } +} diff --git a/litellm-rust/crates/python-bridge/src/preflight.rs b/litellm-rust/crates/python-bridge/src/preflight.rs index ced78582906..27c9fe80232 100644 --- a/litellm-rust/crates/python-bridge/src/preflight.rs +++ b/litellm-rust/crates/python-bridge/src/preflight.rs @@ -1,5 +1,5 @@ -use litellm_host::hooks::CallHooks; -use litellm_host_python::{PythonOwned, PythonRuntime}; +use litellm_host::hooks::{ArgumentHooks, OutcomeHooks, ProviderHooks, ResponseHooks}; +use litellm_host_python::{Awaiter, PythonOwned, PythonSurface}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -33,22 +33,33 @@ impl PythonPreflight { #[cfg(test)] pub(crate) const PYTHON_CONTRACT: &str = include_str!("../preflight_contract.json"); +/// The SDK's request policy as one argument group: inherited credentials are written into +/// the keyword view and a call over its limits is rejected before any provider work. pub(crate) struct SdkPolicy; -impl CallHooks for SdkPolicy { - fn arguments_prepared(&self, arguments: &Py) -> PyResult<()> { +impl ArgumentHooks for SdkPolicy { + async fn prepare<'a>( + &'a self, + _: &'a dyn Awaiter, + arguments: Py, + _: f64, + ) -> PyResult> { Python::attach(|py| { inherit_credentials(py, arguments.bind(py), || { Ok(PythonPreflight::CredentialList .call(py, ())? .cast_into::()?) })?; - PythonPreflight::CheckLimits.call(py, (arguments,))?; - Ok(()) + PythonPreflight::CheckLimits.call(py, (&arguments,))?; + Ok(arguments) }) } } +impl ProviderHooks for SdkPolicy {} +impl ResponseHooks for SdkPolicy {} +impl OutcomeHooks for SdkPolicy {} + impl PythonOwned for SdkPolicy { fn close(&mut self, _: Python<'_>) {} @@ -108,7 +119,11 @@ fn inherit_credentials<'py>( #[cfg(test)] mod tests { use std::collections::BTreeSet; + use std::pin::pin; use std::sync::Mutex; + use std::task::{Context, Poll, Waker}; + + use futures_util::future::BoxFuture; use super::*; use strum::VariantArray; @@ -427,9 +442,7 @@ preflight.credential_list = lambda: [Credential(), Credential()] arguments = {'litellm_credential_name': 'missing'} ", ); - SdkPolicy - .arguments_prepared(&argument_dict(&locals).unbind()) - .unwrap(); + prepared(&argument_dict(&locals).unbind()).unwrap(); py.run( c" assert arguments == {'litellm_credential_name': 'missing'}, arguments @@ -465,9 +478,7 @@ preflight.check_limits = check_limits arguments = {'litellm_credential_name': 'ocr-test'} ", ); - let error = SdkPolicy - .arguments_prepared(&argument_dict(&locals).unbind()) - .unwrap_err(); + let error = prepared(&argument_dict(&locals).unbind()).unwrap_err(); assert!( error .value(py) @@ -485,6 +496,25 @@ assert arguments['api_key'] == 'inherited' }); } + /// Runs the policy over `arguments`. It never awaits, so the task is never asked. + fn prepared(arguments: &Py) -> PyResult> { + struct NeverAwaited; + impl Awaiter for NeverAwaited { + fn awaited(&self, _: Py) -> BoxFuture<'_, PyResult>> { + unreachable!("the SDK policy never awaits") + } + } + let arguments = Python::attach(|py| arguments.clone_ref(py)); + let mut prepare = pin!(SdkPolicy.prepare(&NeverAwaited, arguments, 0.0)); + match prepare + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + { + Poll::Ready(result) => result, + Poll::Pending => unreachable!("the SDK policy never suspends"), + } + } + fn argument_dict<'py>(locals: &Bound<'py, PyDict>) -> Bound<'py, PyDict> { locals .get_item("arguments") diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index dafca72ef0a..2dd45dbc938 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -11,10 +11,11 @@ pub(crate) mod traces; 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::{ - CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args, + CallOptions, Hooks, PythonBinding, PythonGroup, PythonHostCalls, PythonSurface, + effective_py_args, }; +use litellm_host_python::{Error as HostError, present}; use parking_lot::Mutex; use pyo3::{ gc::{PyTraverseError, PyVisit}, @@ -119,7 +120,7 @@ fn call_hooks( operation: LoggingOperation, call: &NativeCall<'_>, asynchronous: bool, -) -> PyResult<(Py, impl PythonCallHooks + use<>)> { +) -> PyResult<(Py, impl PythonGroup + use<>)> { let call = PublicCall::capture(&call.base, &call.args, &call.kwargs)?; let arguments = call.arguments(py); Ok(( @@ -144,7 +145,7 @@ fn run_public_call( + Sync + 'static, host: H, - hooks: impl PythonCallHooks + 'static, + hooks: impl PythonGroup + 'static, asynchronous: bool, ) -> PyResult> where @@ -158,9 +159,9 @@ where start(py, arguments, request).map(crate::logger::LoggedMachine::new) }, host, - HookChain::new() - .with(hooks) - .with(crate::preflight::SdkPolicy), + Hooks::new() + .with(PythonSurface::group(hooks)) + .with(PythonSurface::group(crate::preflight::SdkPolicy)), arguments, CallOptions::new(asynchronous, execution_steps), )