From 898fbaf985e25942e526e35bd84a4c8b8bdaaf03 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 18 Oct 2026 14:25:30 -0700 Subject: [PATCH] refactor(host): declare each hook point once and fold Python hook stages generically Adding one interception point meant a route-facing trait method, an InterceptRequest variant, channel forwarding, a match arm in every driver, a CallHooks method and a hand-written suspend/resume trio in HookChain. InterceptRequest is now the only list of points. Interceptors has one handle method that answers the points an implementer cares about and passes the rest through, routes raise points through the typed InterceptPoints calls, the channel yields any point to the host, and the native driver answers every point with one arm. HookChain runs its argument, wire-request and response stages through one generic fold. No behavior changes Co-Authored-By: Claude Opus 5.5 --- .../crates/gateway-inference/src/caching.rs | 26 +- litellm-rust/crates/host-http/tests/serve.rs | 19 ++ litellm-rust/crates/host-native/src/driver.rs | 23 +- .../crates/host-native/tests/driver.rs | 36 +++ litellm-rust/crates/host-python/AGENTS.md | 2 +- .../host-python/src/hooks/chain/dispatch.rs | 274 +++++++++++------- .../crates/host-python/tests/support/mod.rs | 3 +- litellm-rust/crates/host/AGENTS.md | 11 +- litellm-rust/crates/host/src/interceptors.rs | 129 ++++++--- .../crates/host/src/machine/context.rs | 39 +-- litellm-rust/crates/host/src/protocol.rs | 40 ++- .../crates/inference-chat/src/handler.rs | 19 ++ .../crates/inference-chat/tests/caching.rs | 20 ++ .../inference-chat/tests/chat_completions.rs | 18 ++ .../inference-chat/tests/support/mod.rs | 18 ++ .../crates/inference-messages/src/handler.rs | 1 + .../inference-messages/tests/caching.rs | 20 ++ .../inference-messages/tests/live/support.rs | 20 ++ .../inference-messages/tests/messages/host.rs | 23 ++ .../inference-messages/tests/messages/main.rs | 21 ++ .../tests/messages/stream.rs | 21 ++ .../inference-messages/tests/support/mod.rs | 18 ++ .../crates/inference-ocr/src/handler.rs | 1 + .../inference-ocr/tests/ocr/lifecycle.rs | 38 +++ .../crates/inference-ocr/tests/ocr/machine.rs | 3 +- .../crates/inference-ocr/tests/ocr/main.rs | 21 ++ .../crates/inference-ocr/tests/support/mod.rs | 18 ++ .../crates/inference-responses/src/handler.rs | 1 + .../inference-responses/tests/caching.rs | 20 ++ .../inference-responses/tests/support/mod.rs | 18 ++ litellm-rust/crates/inference/src/caching.rs | 3 +- litellm-rust/crates/inference/src/context.rs | 2 +- .../crates/inference/tests/caching.rs | 20 ++ 33 files changed, 722 insertions(+), 224 deletions(-) diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 1abea61b164..64bceec8e48 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -73,28 +73,16 @@ fn duration(seconds: f64) -> Result { pub(crate) struct CacheHeaders(std::sync::Arc>); impl litellm_host::interceptors::Interceptors for CacheHeaders { - async fn before_provider_request( + async fn handle( &self, - wire: litellm_host::interceptors::WireRequest, - _: litellm_host::interceptors::RequestContext, - ) -> Result { - Ok(wire) - } - - async fn after_provider_response( - &self, - _: litellm_host::interceptors::RawResponse, + request: litellm_host::protocol::InterceptRequest, ) -> Result<(), litellm_inference::RouteError> { - Ok(()) - } - - async fn result_ready( - &self, - facts: litellm_host::interceptors::ExecutionFacts, - ) -> Result<(), litellm_inference::RouteError> { - if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source { - let _ = self.0.set(key); + if let litellm_host::protocol::InterceptRequest::ResultReady { facts, .. } = &request + && let litellm_host::interceptors::ResultSource::Cache { key } = &facts.source + { + let _ = self.0.set(key.clone()); } + request.pass_through(); Ok(()) } } diff --git a/litellm-rust/crates/host-http/tests/serve.rs b/litellm-rust/crates/host-http/tests/serve.rs index 2ba39ad4cd2..9a07a6814c9 100644 --- a/litellm-rust/crates/host-http/tests/serve.rs +++ b/litellm-rust/crates/host-http/tests/serve.rs @@ -1,4 +1,6 @@ +use litellm_host::interceptors::InterceptPoints; use litellm_host::lifecycle::ExecutionEvent; +use litellm_host::protocol::InterceptRequest; use std::sync::{ Arc, Mutex, atomic::{AtomicBool, AtomicUsize, Ordering}, @@ -151,6 +153,23 @@ struct Hooks { } impl Interceptors for Hooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), TestError> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl Hooks { async fn before_provider_request( &self, wire: WireRequest, diff --git a/litellm-rust/crates/host-native/src/driver.rs b/litellm-rust/crates/host-native/src/driver.rs index 7254721da87..0f83bb8ac59 100644 --- a/litellm-rust/crates/host-native/src/driver.rs +++ b/litellm-rust/crates/host-native/src/driver.rs @@ -3,7 +3,7 @@ use std::ops::ControlFlow; use litellm_host::{ interceptors::Interceptors, machine::{HostFailure, Machine, MachineStep}, - protocol::{HostRequest, InterceptRequest, Protocol, Reply, StreamDelivery}, + protocol::{HostRequest, Protocol, Reply, StreamDelivery}, }; use crate::services::HostCallHandler; @@ -66,27 +66,8 @@ where MachineStep::Suspended(request) => request, }; let answered = match request { - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => self - .interceptors - .result_ready(facts) - .await - .map(|()| reply.send(())), + HostRequest::Intercept(request) => self.interceptors.handle(request).await, HostRequest::HostCall(call) => self.services.handle_host_call(call).await, - HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { - wire, - context, - reply, - }) => self - .interceptors - .before_provider_request(*wire, *context) - .await - .map(|wire| reply.send(wire)), - HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => { - self.interceptors - .after_provider_response(raw) - .await - .map(|()| reply.send(())) - } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.demand = Some(reply); return Ok(Boundary::Open(head)); diff --git a/litellm-rust/crates/host-native/tests/driver.rs b/litellm-rust/crates/host-native/tests/driver.rs index c41c1e5a43b..40b354a692f 100644 --- a/litellm-rust/crates/host-native/tests/driver.rs +++ b/litellm-rust/crates/host-native/tests/driver.rs @@ -1,4 +1,6 @@ +use litellm_host::interceptors::InterceptPoints; use litellm_host::lifecycle::ExecutionEvent; +use litellm_host::protocol::InterceptRequest; use std::{ ops::ControlFlow, sync::{ @@ -109,6 +111,23 @@ struct Hooks { } impl Interceptors for Hooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), TestError> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl Hooks { async fn before_provider_request( &self, wire: WireRequest, @@ -630,6 +649,23 @@ struct ResponseGate { } impl Interceptors for ResponseGate { + async fn handle(&self, request: InterceptRequest) -> Result<(), TestError> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl ResponseGate { async fn before_provider_request( &self, wire: WireRequest, diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 02b8d4f2dcb..ab681a9ab14 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -14,7 +14,7 @@ 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` composes Python runtime hooks in order. 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 +`HookChain` composes Python runtime hooks in order. 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. Every transforming stage runs through the one generic fold in `hooks/chain/dispatch.rs`, so a new stage adds a `Transform` impl rather than its own suspend and resume path 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 diff --git a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs index 4939995cefa..eec4d9034d0 100644 --- a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs +++ b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs @@ -19,12 +19,146 @@ type OwnedEvent = CallEvent, Py, RawResponse>; #[derive(Default)] pub struct HookChain { hooks: Vec>, - arguments: Option<(usize, f64)>, - wire: Option<(usize, RequestContext)>, - response: Option<(usize, Timing)>, + pending: Option, event: Option<(usize, OwnedEvent)>, } +/// Where a suspended transformation resumes: the hook it waits on and the context the rest of +/// the chain still needs. +enum Pending { + Arguments(usize, f64), + Wire(usize, RequestContext), + Response(usize, Timing), +} + +/// A stage that feeds each hook's result to the next hook. +trait Transform { + type Value; + type Context; + + fn call( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + value: Self::Value, + context: &Self::Context, + ) -> PyResult>; + + fn resume( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>; + + fn suspend(index: usize, context: Self::Context) -> Pending; + + fn resumed(pending: Pending) -> Option<(usize, Self::Context)>; +} + +struct Arguments; + +impl Transform for Arguments { + type Value = Py; + type Context = f64; + + fn call( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + value: Py, + started_at: &f64, + ) -> PyResult>> { + hooks.prepare_arguments(py, value, *started_at) + } + + fn resume( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + hooks.resume_arguments(py, result) + } + + fn suspend(index: usize, started_at: f64) -> Pending { + Pending::Arguments(index, started_at) + } + + fn resumed(pending: Pending) -> Option<(usize, f64)> { + match pending { + Pending::Arguments(index, started_at) => Some((index, started_at)), + _ => None, + } + } +} + +struct Wire; + +impl Transform for Wire { + type Value = Box; + type Context = RequestContext; + + fn call( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + value: Box, + context: &RequestContext, + ) -> PyResult>> { + hooks.before_provider_request(py, value, context) + } + + fn resume( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + hooks.resume_wire(py, result) + } + + fn suspend(index: usize, context: RequestContext) -> Pending { + Pending::Wire(index, context) + } + + fn resumed(pending: Pending) -> Option<(usize, RequestContext)> { + match pending { + Pending::Wire(index, context) => Some((index, context)), + _ => None, + } + } +} + +struct Response; + +impl Transform for Response { + type Value = Py; + type Context = Timing; + + fn call( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + value: Py, + timing: &Timing, + ) -> PyResult>> { + hooks.transform_response(py, value, *timing) + } + + fn resume( + hooks: &mut dyn ChainHooks, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + hooks.resume_response(py, result) + } + + fn suspend(index: usize, timing: Timing) -> Pending { + Pending::Response(index, timing) + } + + fn resumed(pending: Pending) -> Option<(usize, Timing)> { + match pending { + Pending::Response(index, timing) => Some((index, timing)), + _ => None, + } + } +} + impl HookChain { pub fn new() -> Self { Self::default() @@ -35,127 +169,49 @@ impl HookChain { self } - fn arguments_from( + fn transform_from( &mut self, py: Python<'_>, index: usize, - arguments: Py, - started_at: f64, - ) -> PyResult>> { + value: T::Value, + context: T::Context, + ) -> PyResult> { let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(arguments)); + return Ok(HookStep::Ready(value)); }; - let step = hooks.prepare_arguments(py, arguments, started_at)?; - self.arguments_step(py, index, step, started_at) + let step = T::call(hooks.as_mut(), py, value, &context)?; + self.transform_step::(py, index, step, context) } - fn arguments_step( + fn transform_step( &mut self, py: Python<'_>, index: usize, - step: ChainStep>, - started_at: f64, - ) -> PyResult>> { + step: ChainStep, + context: T::Context, + ) -> PyResult> { match step { - ChainStep::Ready(value) => self.arguments_from(py, index + 1, value, started_at), + ChainStep::Ready(value) => self.transform_from::(py, index + 1, value, context), ChainStep::Await(awaitable) => { - self.arguments = Some((index, started_at)); - Ok(HookStep::Await(awaitable, Self::resume_arguments)) + self.pending = Some(T::suspend(index, context)); + Ok(HookStep::Await(awaitable, Self::resume_transform::)) } } } - fn resume_arguments( + fn resume_transform( &mut self, py: Python<'_>, result: PyResult>, - ) -> PyResult>> { - let (index, context) = self.arguments.take().ok_or(Error::MissingState)?; + ) -> PyResult> { + let (index, context) = self + .pending + .take() + .and_then(T::resumed) + .ok_or(Error::MissingState)?; let result = resume_unless_cancelled(py, result)?; - let step = self.hooks[index].resume_arguments(py, result)?; - self.arguments_step(py, index, step, context) - } - - fn wire_from( - &mut self, - py: Python<'_>, - index: usize, - wire: Box, - context: &RequestContext, - ) -> PyResult>> { - let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(wire)); - }; - let step = hooks.before_provider_request(py, wire, context)?; - self.wire_step(py, index, step, context) - } - - fn wire_step( - &mut self, - py: Python<'_>, - index: usize, - step: ChainStep>, - context: &RequestContext, - ) -> PyResult>> { - match step { - ChainStep::Ready(value) => self.wire_from(py, index + 1, value, context), - ChainStep::Await(awaitable) => { - self.wire = Some((index, context.clone())); - Ok(HookStep::Await(awaitable, Self::resume_wire)) - } - } - } - - fn resume_wire( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let (index, context) = self.wire.take().ok_or(Error::MissingState)?; - let result = resume_unless_cancelled(py, result)?; - let step = self.hooks[index].resume_wire(py, result)?; - self.wire_step(py, index, step, &context) - } - - fn response_from( - &mut self, - py: Python<'_>, - index: usize, - response: Py, - timing: Timing, - ) -> PyResult>> { - let Some(hooks) = self.hooks.get_mut(index) else { - return Ok(HookStep::Ready(response)); - }; - let step = hooks.transform_response(py, response, timing)?; - self.response_step(py, index, step, timing) - } - - fn response_step( - &mut self, - py: Python<'_>, - index: usize, - step: ChainStep>, - timing: Timing, - ) -> PyResult>> { - match step { - ChainStep::Ready(value) => self.response_from(py, index + 1, value, timing), - ChainStep::Await(awaitable) => { - self.response = Some((index, timing)); - Ok(HookStep::Await(awaitable, Self::resume_response)) - } - } - } - - fn resume_response( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - let (index, context) = self.response.take().ok_or(Error::MissingState)?; - let result = resume_unless_cancelled(py, result)?; - let step = self.hooks[index].resume_response(py, result)?; - self.response_step(py, index, step, context) + let step = T::resume(self.hooks[index].as_mut(), py, result)?; + self.transform_step::(py, index, step, context) } fn event_from( @@ -312,7 +368,7 @@ impl CallHooks for HookChain { arguments: Py, started_at: f64, ) -> PyResult>> { - self.arguments_from(py, 0, arguments, started_at) + self.transform_from::(py, 0, arguments, started_at) } fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { @@ -327,7 +383,7 @@ impl CallHooks for HookChain { wire: Box, context: &RequestContext, ) -> PyResult>> { - self.wire_from(py, 0, wire, context) + self.transform_from::(py, 0, wire, context.clone()) } fn transform_response( @@ -336,7 +392,7 @@ impl CallHooks for HookChain { response: Py, timing: Timing, ) -> PyResult>> { - self.response_from(py, 0, response, timing) + self.transform_from::(py, 0, response, timing) } fn on_event( @@ -372,9 +428,7 @@ impl CallHooks for HookChain { impl PythonOwned for HookChain { fn close(&mut self, py: Python<'_>) { - self.arguments = None; - self.wire = None; - self.response = None; + self.pending = None; self.event = None; for hooks in &mut self.hooks { hooks.close(py); diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index ef09135698e..84778e0fdde 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -1,11 +1,12 @@ #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset +use litellm_host::interceptors::InterceptPoints; use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; use litellm_host::{ call::{CallOutput, HostedMachine, hosted_call}, hooks::CallHooks, - interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}, + interceptors::{RawResponse, RequestContext, WireRequest}, lifecycle::{ExecutionEvent, Timing}, machine::{CallMachine, MachineFault}, protocol::{Protocol, Reply}, diff --git a/litellm-rust/crates/host/AGENTS.md b/litellm-rust/crates/host/AGENTS.md index 0f06f0e57e1..5349ed14e40 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::handle` over `InterceptRequest`, typed calls through `InterceptPoints`, and `hooks::CallHooks` | Request interception and fallible execution callbacks | `PythonCallHooks` also prepares arguments, transforms public responses and receives stream callbacks | | 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 | @@ -18,7 +18,12 @@ Each driver owns terminal dispatch. Hooks can change or fail execution. Passive `hooks.rs` owns the call-stage interface and its runtime-associated types. It contains no Python types or legacy callback policy. A runtime supplies its context and continuation representation through `HookRuntime` -`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` +`protocol.rs` owns `Protocol` and suspension messages, including `InterceptRequest` and `StreamDelivery`. + +`InterceptRequest` is the only list of hook points. Each variant carries its payload and a typed reply, and is either a transform (answered with a value of its input's type) or a notification (answered with `()`). Adding a point means one variant, its `pass_through` and `relay` arms, and its typed call on `InterceptPoints`. Nothing else in this crate or the drivers changes per point + +An `Interceptors` implementer writes one `handle` that answers the points it cares about and ends with `other => other.pass_through()`. Routes raise points through `InterceptPoints` and never build an `InterceptRequest` themselves. `ChannelInterceptors` yields each point to the host as `HostRequest::Intercept`, and a driver answers it by passing the whole request to its own `handle`, with no arm per point + `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` Rust handlers answer suspensions through `litellm-host-native::Driver`, which `litellm-host-http` and `litellm_host_native::in_process` share. `in_process::Host` is an assembly of services, interceptors, stream consumer and optional observation publisher. It is not a trait mirroring every suspension. Use `run_hosted` to preserve the distinction between stream completion and detachment @@ -26,4 +31,4 @@ Keep API policy in gateway-inference and python-bridge, and legacy callback poli Interceptors run inline and can rewrite values or fail execution. Observers consume owned `CallEvent` snapshots from `observation_channel`. Its bounded `ObservationSender` never waits for delivery and counts events dropped when the queue is full or closed. The host owns receiver processing and draining. Pass the same publisher to machine construction and the driver when one receiver should collect execution and lifecycle events. Legacy Python callbacks retain their existing awaited, fallible behavior through the Python adapter -`ExecutionFacts` and `ResultSource` describe execution without pricing or budget policy. `Interceptors::result_ready` delivers these facts through an awaited `InterceptRequest::ResultReady`. Hosts receive them before response transformation or stream delivery. `ExecutionEvent::ResultReady` is the matching lifecycle event and can also be published as a passive snapshot. Accounting must consume the awaited path rather than a lossy observation queue +`ExecutionFacts` and `ResultSource` describe execution without pricing or budget policy. `InterceptPoints::result_ready` delivers these facts through an awaited `InterceptRequest::ResultReady`. Hosts receive them before response transformation or stream delivery. `ExecutionEvent::ResultReady` is the matching lifecycle event and can also be published as a passive snapshot. Accounting must consume the awaited path rather than a lossy observation queue diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index f0fd7434c6d..c66551a8a93 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -1,5 +1,10 @@ use std::future::Future; +use crate::{ + machine::MachineFault, + protocol::{InterceptRequest, Reply, reply}, +}; + use serde_json::Value; /// The provider request as it is about to leave, offered to the host for rewriting. @@ -70,57 +75,79 @@ pub struct ExecutionFacts { pub cost: Cost, } +/// Answers the hook points a call raises. `InterceptRequest` lists every point; an +/// implementer handles the ones it cares about and passes the rest through. pub trait Interceptors: Send + Sync { - fn result_ready(&self, _facts: ExecutionFacts) -> impl Future> + Send { + fn handle(&self, request: InterceptRequest) -> impl Future> + Send { + request.pass_through(); async { Ok(()) } } - fn before_provider_request( + /// Raises one point and waits for its answer. A host channel yields it to the host. + fn intercept( &self, - wire: WireRequest, - context: RequestContext, - ) -> impl Future> + Send; - - fn after_provider_response( - &self, - raw: RawResponse, - ) -> impl Future> + Send; + ask: impl FnOnce(Reply) -> InterceptRequest + Send, + ) -> impl Future> + Send + where + E: From, + { + let (reply, answer) = reply(); + let handled = self.handle(ask(reply)); + async move { + handled.await?; + answer.await.map_err(|_| MachineFault::Abandoned.into()) + } + } } impl + ?Sized> Interceptors for &T { - fn result_ready(&self, facts: ExecutionFacts) -> impl Future> + Send { - (**self).result_ready(facts) + fn handle(&self, request: InterceptRequest) -> impl Future> + Send { + (**self).handle(request) } + fn intercept( + &self, + ask: impl FnOnce(Reply) -> InterceptRequest + Send, + ) -> impl Future> + Send + where + E: From, + { + (**self).intercept(ask) + } +} + +impl Interceptors for () {} + +/// The typed form of each point, for the routes that raise them. +pub trait InterceptPoints: Interceptors +where + E: From, +{ fn before_provider_request( &self, wire: WireRequest, context: RequestContext, ) -> impl Future> + Send { - (**self).before_provider_request(wire, context) + self.intercept(|reply| InterceptRequest::BeforeProviderRequest { + wire: Box::new(wire), + context: Box::new(context), + reply, + }) } fn after_provider_response( &self, raw: RawResponse, ) -> impl Future> + Send { - (**self).after_provider_response(raw) + self.intercept(|reply| InterceptRequest::AfterProviderResponse { raw, reply }) + } + + fn result_ready(&self, facts: ExecutionFacts) -> impl Future> + Send { + self.intercept(|reply| InterceptRequest::ResultReady { facts, reply }) } } -impl Interceptors for () { - async fn before_provider_request( - &self, - wire: WireRequest, - _: RequestContext, - ) -> Result { - Ok(wire) - } - - async fn after_provider_response(&self, _: RawResponse) -> Result<(), E> { - Ok(()) - } -} +impl, T: Interceptors + ?Sized> InterceptPoints for T {} #[cfg(test)] mod tests { @@ -129,7 +156,6 @@ mod tests { use serde_json::json; use super::*; - use crate::protocol::InterceptRequest; use crate::{ machine::{CallMachine, Machine, MachineFault, MachineStep}, protocol::{HostRequest, Protocol}, @@ -178,13 +204,13 @@ mod tests { async fn the_channel_yields_each_hook_as_its_op_and_returns_the_answer() { let mut machine = CallMachine::::new(None, |channel| { Box::pin(async move { - let sent = Interceptors::before_provider_request( + let sent = InterceptPoints::before_provider_request( &channel.interceptors, wire("prepared"), context(), ) .await?; - Interceptors::after_provider_response( + InterceptPoints::after_provider_response( &channel.interceptors, RawResponse { body: "raw".into() }, ) @@ -220,12 +246,49 @@ mod tests { assert_eq!(sent.url, "rewritten"); } + #[rstest::rstest] + #[tokio::test] + async fn handling_on_the_channel_suspends_for_the_host_and_answers_the_original_reply() { + let mut machine = CallMachine::::new(None, |channel| { + Box::pin(async move { + let (reply, answer) = reply(); + channel + .interceptors + .handle(InterceptRequest::BeforeProviderRequest { + wire: Box::new(wire("prepared")), + context: Box::new(context()), + reply, + }) + .await?; + Ok((answer.await.map_err(|_| Fault)?, ())) + }) + }); + + let Ok(MachineStep::Suspended(HostRequest::Intercept( + InterceptRequest::BeforeProviderRequest { wire, reply, .. }, + ))) = machine.resume().await + else { + panic!("handling on the channel yields the request to the host"); + }; + assert_eq!(wire.url, "prepared"); + reply.send(WireRequest { + url: "rewritten".into(), + ..*wire + }); + + let Ok(MachineStep::Complete((sent, ()))) = machine.resume().await else { + panic!("the call completes with the host's answer"); + }; + assert_eq!(sent.url, "rewritten"); + } + #[rstest::rstest] #[tokio::test] async fn no_hooks_pass_the_wire_request_through() { - let sent = Interceptors::::before_provider_request(&(), wire("prepared"), context()) - .await - .unwrap(); + let sent = + InterceptPoints::::before_provider_request(&(), wire("prepared"), context()) + .await + .unwrap(); assert_eq!(sent.url, "prepared"); } } diff --git a/litellm-rust/crates/host/src/machine/context.rs b/litellm-rust/crates/host/src/machine/context.rs index b2c9e95c563..b376deed1b6 100644 --- a/litellm-rust/crates/host/src/machine/context.rs +++ b/litellm-rust/crates/host/src/machine/context.rs @@ -4,10 +4,7 @@ use std::ops::ControlFlow; use litellm_coroutine::Co; use super::coroutine::MachineFault; -use crate::{ - interceptors::{RawResponse, RequestContext, WireRequest}, - protocol::{HostRequest, InterceptRequest, Protocol, Reply, StreamDelivery}, -}; +use crate::protocol::{HostRequest, InterceptRequest, Protocol, Reply, StreamDelivery}; pub struct CallContext { pub services: HostServices

, @@ -85,38 +82,16 @@ impl crate::interceptors::Interceptors for ChannelInterce where P::Error: From, { - async fn before_provider_request( - &self, - wire: WireRequest, - context: RequestContext, - ) -> Result { - self.0 - .request_reply(|reply| { - HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { - wire: Box::new(wire), - context: Box::new(context), - reply, - }) - }) - .await + async fn handle(&self, request: InterceptRequest) -> Result<(), P::Error> { + request.relay(self).await } - async fn result_ready( + async fn intercept( &self, - facts: crate::interceptors::ExecutionFacts, - ) -> Result<(), P::Error> { + ask: impl FnOnce(Reply) -> InterceptRequest + Send, + ) -> Result { self.0 - .request_reply(|reply| { - HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) - }) - .await - } - - async fn after_provider_response(&self, raw: RawResponse) -> Result<(), P::Error> { - self.0 - .request_reply(|reply| { - HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) - }) + .request_reply(|reply| HostRequest::Intercept(ask(reply))) .await } } diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs index 9e5aff8aae8..91d15891cf9 100644 --- a/litellm-rust/crates/host/src/protocol.rs +++ b/litellm-rust/crates/host/src/protocol.rs @@ -2,7 +2,7 @@ use std::ops::ControlFlow; pub use litellm_coroutine::{Abandoned, Answer, Reply, reply}; -use crate::interceptors::{RawResponse, RequestContext, WireRequest}; +use crate::interceptors::{InterceptPoints, RawResponse, RequestContext, WireRequest}; pub trait Protocol: Send + Sync + 'static { type Request: Send + 'static; @@ -35,6 +35,44 @@ pub enum InterceptRequest { }, } +impl InterceptRequest { + /// Answers as if no hook ran: a transform returns its input, a notification returns. + pub fn pass_through(self) { + match self { + Self::ResultReady { reply, .. } | Self::AfterProviderResponse { reply, .. } => { + reply.send(()); + } + Self::BeforeProviderRequest { wire, reply, .. } => reply.send(*wire), + } + } + + /// Raises this request again through `interceptors` and answers its reply with the result. + pub async fn relay(self, interceptors: &I) -> Result<(), E> + where + E: From, + I: InterceptPoints + ?Sized, + { + match self { + Self::ResultReady { facts, reply } => { + reply.send(interceptors.result_ready(facts).await?) + } + Self::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send( + interceptors + .before_provider_request(*wire, *context) + .await?, + ), + Self::AfterProviderResponse { raw, reply } => { + reply.send(interceptors.after_provider_response(raw).await?); + } + } + Ok(()) + } +} + pub enum StreamDelivery { Open(P::StreamHead, Reply>), Chunk(P::Chunk, Reply>), diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 95dd05e90d6..f3fc17f1bc6 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -1,3 +1,4 @@ +use litellm_host::interceptors::InterceptPoints; use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender}; use std::time::Duration; @@ -170,6 +171,7 @@ pub(super) fn outbound_request( #[cfg(test)] mod tests { + use litellm_host::protocol::InterceptRequest; use std::sync::Mutex; use rstest::rstest; @@ -192,6 +194,23 @@ mod tests { } impl Interceptors for RecordingHooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } + } + + impl RecordingHooks { async fn before_provider_request( &self, wire: WireRequest, diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 3fc432f8b46..a08670a9b92 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -1,4 +1,5 @@ mod support; +use litellm_host::protocol::InterceptRequest; use std::{ num::NonZeroUsize, @@ -167,6 +168,25 @@ struct ChangingHooks { } impl Interceptors for ChangingHooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), RouteError> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + InterceptRequest::ResultReady { facts, reply } => { + reply.send(self.result_ready(facts).await?); + } + } + Ok(()) + } +} + +impl ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, diff --git a/litellm-rust/crates/inference-chat/tests/chat_completions.rs b/litellm-rust/crates/inference-chat/tests/chat_completions.rs index 461d4146a63..e8016a05b2e 100644 --- a/litellm-rust/crates/inference-chat/tests/chat_completions.rs +++ b/litellm-rust/crates/inference-chat/tests/chat_completions.rs @@ -1,4 +1,5 @@ use litellm_host::interceptors::RawResponse; +use litellm_host::protocol::InterceptRequest; use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, @@ -334,6 +335,23 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry( use litellm_host::interceptors::{Interceptors, RequestContext, WireRequest}; struct FailingHook; impl Interceptors for FailingHook { + async fn handle(&self, request: InterceptRequest) -> Result<(), Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } + } + + impl FailingHook { async fn before_provider_request( &self, wire: WireRequest, diff --git a/litellm-rust/crates/inference-chat/tests/support/mod.rs b/litellm-rust/crates/inference-chat/tests/support/mod.rs index 74098bf21a2..081d49d3bcf 100644 --- a/litellm-rust/crates/inference-chat/tests/support/mod.rs +++ b/litellm-rust/crates/inference-chat/tests/support/mod.rs @@ -2,6 +2,7 @@ //! secret source. #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset +use litellm_host::protocol::InterceptRequest; use std::{ ops::ControlFlow, @@ -173,6 +174,23 @@ impl RecordingCall

{ impl litellm_host::interceptors::Interceptors for RecordingCall

{ + async fn handle(&self, request: InterceptRequest) -> Result<(), P::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl RecordingCall

{ async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index c4032baa98f..522dd57a9f1 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -1,3 +1,4 @@ +use litellm_host::interceptors::InterceptPoints; use std::time::Duration; use bytes::Bytes; diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 32c962091ad..3725ddaac7c 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -1,4 +1,5 @@ mod support; +use litellm_host::protocol::InterceptRequest; use std::{ sync::{ @@ -134,6 +135,25 @@ struct ChangingHooks { } impl Interceptors for ChangingHooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), RouteError> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + InterceptRequest::ResultReady { facts, reply } => { + reply.send(self.result_ready(facts).await?); + } + } + Ok(()) + } +} + +impl ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 81e295e34ad..5493ee47f59 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -1,3 +1,4 @@ +use litellm_host::protocol::InterceptRequest; use std::{sync::OnceLock, time::Duration}; use futures_util::{FutureExt, TryStreamExt}; @@ -40,6 +41,25 @@ impl LiveCall { } impl Interceptors for LiveCall { + async fn handle(&self, request: InterceptRequest) -> Result<(), Error> { + match request { + InterceptRequest::ResultReady { facts, reply } => { + reply.send(self.result_ready(facts).await?); + } + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + } + Ok(()) + } +} + +impl LiveCall { async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { assert!( self.facts.set(facts).is_ok(), diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index ba20fb9eec4..c4637f75056 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -1,4 +1,5 @@ use litellm_host::lifecycle::ExecutionEvent; +use litellm_host::protocol::InterceptRequest; use std::sync::Mutex; use litellm_host::{ @@ -78,6 +79,28 @@ impl litellm_host::lifecycle::CallObserver for RecordingHost { impl litellm_host::interceptors::Interceptors<::Error> for RecordingHost { + async fn handle( + &self, + request: InterceptRequest, + ) -> Result<(), ::Error> { + match request { + InterceptRequest::ResultReady { facts, reply } => { + reply.send(self.result_ready(facts).await?); + } + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + } + Ok(()) + } +} + +impl RecordingHost { async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); if self.reject_result { diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index fcc432acb76..fdd1540b125 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -1,3 +1,4 @@ +use litellm_host::protocol::InterceptRequest; use std::{ sync::{Arc, Mutex}, time::Duration, @@ -162,6 +163,26 @@ impl litellm_host::lifecycle::CallObserver for LocalMessagesHost { impl litellm_host::interceptors::Interceptors<::Error> for LocalMessagesHost { + async fn handle( + &self, + request: InterceptRequest, + ) -> Result<(), ::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl LocalMessagesHost { async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index cda688107b5..669cb9ebfa5 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -1,3 +1,4 @@ +use litellm_host::protocol::InterceptRequest; use std::{ ops::ControlFlow, sync::{Mutex, mpsc}, @@ -101,6 +102,26 @@ impl litellm_host::lifecycle::CallObserver for RecordingStreamHost { impl litellm_host::interceptors::Interceptors<::Error> for RecordingStreamHost { + async fn handle( + &self, + request: InterceptRequest, + ) -> Result<(), ::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl RecordingStreamHost { async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, diff --git a/litellm-rust/crates/inference-messages/tests/support/mod.rs b/litellm-rust/crates/inference-messages/tests/support/mod.rs index 3cecc865070..fc74d7e79f2 100644 --- a/litellm-rust/crates/inference-messages/tests/support/mod.rs +++ b/litellm-rust/crates/inference-messages/tests/support/mod.rs @@ -2,6 +2,7 @@ //! secret source. #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset +use litellm_host::protocol::InterceptRequest; use std::{ ops::ControlFlow, @@ -174,6 +175,23 @@ impl RecordingCall

{ impl litellm_host::interceptors::Interceptors for RecordingCall

{ + async fn handle(&self, request: InterceptRequest) -> Result<(), P::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl RecordingCall

{ async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index 51b739f47ba..24549456117 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -1,4 +1,5 @@ use futures_util::future::BoxFuture; +use litellm_host::interceptors::InterceptPoints; use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}; use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender}; use litellm_llms::base_llm::ocr::{ diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index d10a2e23c9a..b98dfea6a92 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -1,5 +1,6 @@ use litellm_host::interceptors::RawResponse; use litellm_host::lifecycle::ExecutionEvent; +use litellm_host::protocol::InterceptRequest; use std::sync::{ Arc, Mutex, atomic::{AtomicUsize, Ordering}, @@ -243,6 +244,26 @@ impl litellm_host::lifecycle::CallObserver for CallerTokenHost { impl litellm_host::interceptors::Interceptors<::Error> for CallerTokenHost { + async fn handle( + &self, + request: InterceptRequest, + ) -> Result<(), ::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl CallerTokenHost { async fn before_provider_request( &self, wire: WireRequest, @@ -325,6 +346,23 @@ async fn direct_execution_uses_hooks_without_a_machine() { struct Hooks; impl Interceptors for Hooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } + } + + impl Hooks { async fn before_provider_request( &self, wire: WireRequest, diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs index 4f9266c4a8d..a3002db409f 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs @@ -1,3 +1,4 @@ +use litellm_host::interceptors::InterceptPoints; use std::{ sync::{ Arc, @@ -7,7 +8,7 @@ use std::{ }; use litellm_host::{ - interceptors::{Interceptors, WireRequest}, + interceptors::WireRequest, machine::{HostFailure, Machine, MachineStep}, protocol::{HostRequest, InterceptRequest}, }; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a905f907d2..64fefaa759b 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -1,3 +1,4 @@ +use litellm_host::protocol::InterceptRequest; use litellm_host::{ interceptors::{RequestContext, WireRequest}, lifecycle::CallEvent, @@ -240,6 +241,26 @@ impl litellm_host::lifecycle::CallObserver for LocalOcrHost { impl litellm_host::interceptors::Interceptors<::Error> for LocalOcrHost { + async fn handle( + &self, + request: InterceptRequest, + ) -> Result<(), ::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl LocalOcrHost { async fn before_provider_request( &self, wire: WireRequest, diff --git a/litellm-rust/crates/inference-ocr/tests/support/mod.rs b/litellm-rust/crates/inference-ocr/tests/support/mod.rs index 5b1b16bc483..43332494ae0 100644 --- a/litellm-rust/crates/inference-ocr/tests/support/mod.rs +++ b/litellm-rust/crates/inference-ocr/tests/support/mod.rs @@ -2,6 +2,7 @@ //! secret source. #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset +use litellm_host::protocol::InterceptRequest; use std::{ ops::ControlFlow, @@ -185,6 +186,23 @@ impl RecordingCall

{ impl litellm_host::interceptors::Interceptors for RecordingCall

{ + async fn handle(&self, request: InterceptRequest) -> Result<(), P::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl RecordingCall

{ async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index eadedca7d3f..0bf8aa974e6 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -1,3 +1,4 @@ +use litellm_host::interceptors::InterceptPoints; use litellm_host::lifecycle::ExecutionEvent; use litellm_host::observation::ObservationSender; use std::time::Duration; diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index f45cefdc13b..b9206ffae96 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -1,4 +1,5 @@ mod support; +use litellm_host::protocol::InterceptRequest; use std::{ sync::{ @@ -125,6 +126,25 @@ struct ChangingHooks { } impl Interceptors for ChangingHooks { + async fn handle(&self, request: InterceptRequest) -> Result<(), RouteError> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + InterceptRequest::ResultReady { facts, reply } => { + reply.send(self.result_ready(facts).await?); + } + } + Ok(()) + } +} + +impl ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, diff --git a/litellm-rust/crates/inference-responses/tests/support/mod.rs b/litellm-rust/crates/inference-responses/tests/support/mod.rs index 81f133a0471..88339230150 100644 --- a/litellm-rust/crates/inference-responses/tests/support/mod.rs +++ b/litellm-rust/crates/inference-responses/tests/support/mod.rs @@ -2,6 +2,7 @@ //! secret source. #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset +use litellm_host::protocol::InterceptRequest; use std::{ ops::ControlFlow, @@ -171,6 +172,23 @@ impl RecordingCall

{ impl litellm_host::interceptors::Interceptors for RecordingCall

{ + async fn handle(&self, request: InterceptRequest) -> Result<(), P::Error> { + match request { + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + other => other.pass_through(), + } + Ok(()) + } +} + +impl RecordingCall

{ async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index 6dfda20ead4..d6350f654b5 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -14,7 +14,8 @@ use litellm_cache_response::{ use litellm_host::{ call::{CallOutput, OutputOf}, interceptors::{ - Cost, ExecutionFacts, Interceptors, ProviderIdentity, ResultSource, WireRequest, + Cost, ExecutionFacts, InterceptPoints, Interceptors, ProviderIdentity, ResultSource, + WireRequest, }, lifecycle::{CallEvent, ExecutionEvent}, observation::ObservationSender, diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 6fcbb3051b4..a097baba4d4 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -1,6 +1,6 @@ use litellm_cache_response::CachePolicy; use litellm_host::{ - interceptors::{ExecutionFacts, Interceptors, RawResponse}, + interceptors::{ExecutionFacts, InterceptPoints, Interceptors, RawResponse}, lifecycle::{CallEvent, ExecutionEvent}, observation::ObservationSender, }; diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index 1bd53695c30..41180d4a53f 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -1,3 +1,4 @@ +use litellm_host::protocol::InterceptRequest; use std::{ convert::Infallible, num::NonZeroUsize, @@ -468,6 +469,25 @@ struct CacheHitAccounting { } impl Interceptors for CacheHitAccounting { + async fn handle(&self, request: InterceptRequest) -> Result<(), RouteError> { + match request { + InterceptRequest::ResultReady { facts, reply } => { + reply.send(self.result_ready(facts).await?); + } + InterceptRequest::BeforeProviderRequest { + wire, + context, + reply, + } => reply.send(self.before_provider_request(*wire, *context).await?), + InterceptRequest::AfterProviderResponse { raw, reply } => { + reply.send(self.after_provider_response(raw).await?); + } + } + Ok(()) + } +} + +impl CacheHitAccounting { async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { self.calls.fetch_add(1, Ordering::SeqCst); let ResultSource::Cache { key } = facts.source else {