From 3a5b7c12ef119474a596cd58f59807cce5854fb5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 15:35:20 -0700 Subject: [PATCH] refactor(rust): separate machine events from the Python lifecycle's events --- .../crates/callbacks-legacy/src/adapter.rs | 57 +++++++-------- .../tests/deployment_hooks.rs | 11 ++- .../crates/callbacks-legacy/tests/payload.rs | 8 +-- .../crates/callbacks-legacy/tests/terminal.rs | 14 ++-- litellm-rust/crates/callbacks/src/event.rs | 16 +++-- litellm-rust/crates/callbacks/src/host.rs | 4 +- litellm-rust/crates/callbacks/src/run.rs | 5 +- litellm-rust/crates/core/src/machine/mod.rs | 4 +- .../crates/core/src/messages/route.rs | 4 +- litellm-rust/crates/core/src/ocr/handler.rs | 4 +- .../tests/azure_document_intelligence_ocr.rs | 8 +-- litellm-rust/crates/core/tests/ocr.rs | 9 ++- litellm-rust/crates/core/tests/reducto_ocr.rs | 8 +-- .../crates/host-python/src/adapter.rs | 33 ++++++--- litellm-rust/crates/host-python/src/driver.rs | 69 ++++++++++--------- litellm-rust/crates/host-python/src/lib.rs | 2 +- 16 files changed, 142 insertions(+), 114 deletions(-) diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index 9d28db92add..a67da2188fb 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -3,10 +3,10 @@ //! `@client` path makes them. use litellm_callbacks::event::{ - CallEvent, FailureOrigin, RequestContext, Timing, WireRequest, epoch_seconds, + FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds, }; use litellm_host_python::{ - LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, + LifecycleEvent, LifecycleStep, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -346,28 +346,11 @@ impl PythonLifecycle for LegacyLogging { fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult { - match (event, public) { - (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), - (CallEvent::Opened, _) => { - Streaming::Opened.call(py, (self.logger()?.object(py),))?; - self.stream = Some(DeliveredStream { - chunks: PyList::empty(py).unbind(), - first_chunk: None, - }); - Ok(LifecycleStep::Done) - } - (CallEvent::Delivered, Some(PublicValue::Chunk(chunk))) => { - let stream = self.stream.as_mut().ok_or_else(missing_state)?; - if stream.first_chunk.is_none() { - stream.first_chunk = Some(datetime(py, epoch_seconds())?); - } - stream.chunks.bind(py).append(chunk)?; - Ok(LifecycleStep::Done) - } - (CallEvent::ResponseReceived { raw }, _) => { + match event { + LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { let api_key = self .context .as_ref() @@ -382,7 +365,7 @@ impl PythonLifecycle for LegacyLogging { )?; Ok(LifecycleStep::Done) } - (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { + LifecycleEvent::Succeeded { timing, response } => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); match &self.stream { @@ -391,13 +374,17 @@ impl PythonLifecycle for LegacyLogging { } Ok(LifecycleStep::Done) } - (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Failed { + timing, + origin, + error, + } => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); if self.stream.is_some() { return self.stream_failure(py); } - if *origin == FailureOrigin::Call + if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() { @@ -412,10 +399,26 @@ impl PythonLifecycle for LegacyLogging { } self.dispatch_failure(py) } - _ => Err(missing_state()), } } + fn opened(&mut self, py: Python<'_>) -> PyResult<()> { + Streaming::Opened.call(py, (self.logger()?.object(py),))?; + self.stream = Some(DeliveredStream { + chunks: PyList::empty(py).unbind(), + first_chunk: None, + }); + Ok(()) + } + + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { + let stream = self.stream.as_mut().ok_or_else(missing_state)?; + if stream.first_chunk.is_none() { + stream.first_chunk = Some(datetime(py, epoch_seconds())?); + } + stream.chunks.bind(py).append(chunk) + } + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 7bad09c7890..f4602739548 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; +use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -217,13 +217,12 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle .resume(py, Ok(local(&locals, "kwargs").unbind())) .unwrap(); let failure = PyErr::from_value(local(&locals, "failure")); - let failed = CallEvent::Failed { + let failed = LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Call, + error: &failure, }; - let step = logging - .emit(py, &failed, Some(PublicValue::Error(&failure))) - .unwrap(); + let step = logging.emit(py, failed).unwrap(); assert!(awaits_deployment_hook(&step)); let hook_result = if cancelled { Err(CancelledError::new_err("cancelled")) diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 43128bc38ea..67ad4ab8a2a 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,8 +1,8 @@ use std::ffi::CStr; use litellm_auth::SecretValue; -use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleStep, PythonLifecycle, to_py}; +use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; use proptest::prelude::*; use pyo3::prelude::*; use rstest::rstest; @@ -94,13 +94,13 @@ fn before_send_bound( body, }; let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = CallEvent::ResponseReceived { + let raw = MachineEvent::ResponseReceived { raw: RawResponse { body: "raw response".into(), }, }; assert!(matches!( - logging.emit(py, &raw, None).unwrap(), + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), LifecycleStep::Done )); run(py, &locals, c"check()"); diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 3094d7b88d2..5688f70f387 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; +use litellm_callbacks::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -33,8 +33,10 @@ fn succeed( logging .emit( py, - &CallEvent::Succeeded { timing: TIMING }, - Some(PublicValue::Response(&response)), + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, ) .unwrap() } @@ -44,11 +46,11 @@ fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) logging .emit( py, - &CallEvent::Failed { + LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Host, + error: &failure, }, - Some(PublicValue::Error(&failure)), ) .unwrap() } diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index d19e973b812..182dab657d3 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -52,18 +52,20 @@ pub enum FailureOrigin { Host, } +/// What a machine reports while it runs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MachineEvent { + ResponseReceived { raw: RawResponse }, +} + +/// What an in-process host observes: the machine's own events between the driver's +/// start and terminal ones. #[derive(Clone, Debug, PartialEq)] pub enum CallEvent { Started { start_time: f64, }, - ResponseReceived { - raw: RawResponse, - }, - /// The call streams and its stream was handed to the caller. - Opened, - /// One chunk of an open stream reached the caller. - Delivered, + Machine(MachineEvent), Succeeded { timing: Timing, }, diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/callbacks/src/host.rs index eef3e1da8d5..aba35185a18 100644 --- a/litellm-rust/crates/callbacks/src/host.rs +++ b/litellm-rust/crates/callbacks/src/host.rs @@ -1,6 +1,6 @@ use std::future::Future; -use crate::event::{CallEvent, RequestContext, WireRequest}; +use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest}; use crate::route::Route; /// One suspension point of a native call, performed by the host. @@ -10,7 +10,7 @@ pub enum HostOp { wire: Box, context: Box, }, - Emit(CallEvent), + Emit(MachineEvent), /// The response streams: the host hands the caller a stream and answers once the /// caller asks for the first chunk or goes away. Open(R::StreamHead), diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 705c5504c01..6a0c08fba68 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -25,7 +25,10 @@ where .before_send(*wire, &context) .await .map(|wire| HostResult::BeforeSend(Box::new(wire))), - HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted), + HostOp::Emit(event) => host + .emit(&CallEvent::Machine(event)) + .await + .map(|()| HostResult::Emitted), HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), }; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/core/src/machine/mod.rs index 929f0a423c4..d6db488159e 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/core/src/machine/mod.rs @@ -8,7 +8,7 @@ use std::{future::Future, pin::Pin}; pub use auth::{HostTokenProvider, TokenRoute}; use litellm_callbacks::{ - event::{CallEvent, RequestContext, WireRequest}, + event::{MachineEvent, RequestContext, WireRequest}, host::{Demand, HostOp, HostResult}, machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, route::Route, @@ -82,7 +82,7 @@ where } } - pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> { + pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { match self.invoke(HostOp::Emit(event)).await? { HostResult::Emitted => Ok(()), _ => Err(MachineFault::Mismatch.into()), diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index a680b5ce1ec..2fd2f79907a 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -3,7 +3,7 @@ use std::{sync::Mutex, time::Duration}; use bytes::Bytes; use litellm_auth::SecretValue; use litellm_callbacks::{ - event::{CallEvent, RawResponse, RequestContext, WireRequest}, + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, route::Route, }; @@ -172,7 +172,7 @@ async fn execute(host: MessagesHost) -> Result { return relay(&host, response).await; } let text = response.text().await.map_err(network)?; - host.emit(CallEvent::ResponseReceived { + host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: text.clone() }, }) .await?; diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index a6af190eb91..aff1eded2cc 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,6 +1,6 @@ use futures_util::future::BoxFuture; use litellm_auth::SecretValue; -use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; +use litellm_callbacks::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -71,7 +71,7 @@ impl CallHooks for OcrCallHooks { } fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(self.host.emit(CallEvent::ResponseReceived { + Box::pin(self.host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: String::from_utf8_lossy(body).into_owned(), }, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 1fc4d6c2b9e..62544931dfc 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::CallEvent; +use litellm_callbacks::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; @@ -263,7 +263,7 @@ async fn accepted_response_emits_response_received_before_polling() { json!({}), )) .with_observer(move |event| { - let CallEvent::ResponseReceived { raw } = event else { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { return; }; match request_count.lock().unwrap().len() { @@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { mod transformation { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::CallEvent; + use litellm_callbacks::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::transformation::OcrDocument; use serde_json::{Value, json}; @@ -646,7 +646,7 @@ mod transformation { json!({}), )) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed .lock() .unwrap() diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 779d2037bb3..a6001951361 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use litellm_callbacks::{ - event::{CallEvent, WireRequest}, + event::{CallEvent, MachineEvent, WireRequest}, host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; @@ -195,11 +195,9 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { CallEvent::Started { .. } => "started", - CallEvent::ResponseReceived { .. } => "response", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", - CallEvent::Opened => "opened", - CallEvent::Delivered => "delivered", } } @@ -377,6 +375,7 @@ async fn drive_until( intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) } HostOp::Emit(event) => { + let event = CallEvent::Machine(event); ops.push(event_name(&event)); host.emit(&event) .await @@ -420,7 +419,7 @@ async fn invalid_provider_response_emits_response_received_before_normalization_ let observed = responses_received.clone(); let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed.lock().unwrap().push(raw.body.clone()); } }, diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 59891b16e90..8c037889a17 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, WireRequest}; +use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; use rstest::rstest; use serde_json::{Value, json}; @@ -139,7 +139,7 @@ async fn response_received_stays_after_reducto_upload_and_parse() { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } @@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() { } mod transformation { - use litellm_callbacks::event::{CallEvent, WireRequest}; + use litellm_callbacks::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::{ base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, reducto::ocr::transformation::*, @@ -506,7 +506,7 @@ mod transformation { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 795977438fd..f946dbc763a 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, RequestContext, Timing, WireRequest}; +use litellm_callbacks::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; use litellm_callbacks::route::Route; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; @@ -19,11 +19,22 @@ pub enum LifecycleStep { Done, } -/// The host-typed value the driver attaches to a terminal event. -pub enum PublicValue<'a> { - Response(&'a Py), - Error(&'a PyErr), - Chunk(&'a Py), +/// What a lifecycle observes: the driver's start, the machine's own events, and one +/// terminal event carrying the public value the caller receives. +pub enum LifecycleEvent<'a> { + Started { + start_time: f64, + }, + Machine(&'a MachineEvent), + Succeeded { + timing: Timing, + response: &'a Py, + }, + Failed { + timing: Timing, + origin: FailureOrigin, + error: &'a PyErr, + }, } /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in @@ -61,10 +72,16 @@ pub trait PythonLifecycle: Send + Sync { fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult; + /// The call streams and its stream was handed to the caller. The caller is not + /// inside an await here, so this step and `delivered` cannot suspend. + fn opened(&mut self, py: Python<'_>) -> PyResult<()>; + + /// One chunk of an open stream is about to reach the caller. + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()>; + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index c59ba1925dc..25f78d7013e 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; +use litellm_callbacks::event::{FailureOrigin, Timing, epoch_seconds}; use litellm_callbacks::host::{Demand, HostOp, HostResult, HostStep}; use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; use litellm_callbacks::route::Route; @@ -13,7 +13,7 @@ use pyo3::types::PyDict; use tokio::sync::Mutex; use crate::adapter::{ - HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, + HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -159,10 +159,10 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let started = CallEvent::Started { + let started = LifecycleEvent::Started { start_time: self.started_at, }; - match self.adapter.emit(py, &started, None) { + match self.adapter.emit(py, started) { Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } @@ -303,7 +303,7 @@ where } HostOp::Open(_) => return self.opened(py).map(Next::Return), HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), - HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { + HostOp::Emit(event) => match self.adapter.emit(py, LifecycleEvent::Machine(&event)) { Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); @@ -321,12 +321,11 @@ where fn opened(&mut self, py: Python<'_>) -> PyResult { self.stage = Stage::Streaming; - match self.adapter.emit(py, &CallEvent::Opened, None) { - Ok(LifecycleStep::Done) => { + match self.adapter.opened(py) { + Ok(()) => { self.pending = Some(Pending::Consumer); Ok(ExecutionStep::Open) } - Ok(_) => Err(missing_state()), Err(error) => self.interrupt(py, error), } } @@ -340,15 +339,11 @@ where Ok(chunk) => chunk, Err(error) => return self.interrupt(py, error), }; - let observed = - self.adapter - .emit(py, &CallEvent::Delivered, Some(PublicValue::Chunk(&chunk))); - match observed { - Ok(LifecycleStep::Done) => { + match self.adapter.delivered(py, &chunk) { + Ok(()) => { self.pending = Some(Pending::Consumer); Ok(ExecutionStep::Yield(chunk)) } - Ok(_) => Err(missing_state()), Err(error) => self.interrupt(py, error), } } @@ -461,12 +456,11 @@ where } fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { - let event = CallEvent::Succeeded { + let event = LifecycleEvent::Succeeded { timing: self.timing(), + response: &response, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Response(&response)))?; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Succeeded(response); self.on_adapter(py, step, Expect::Terminal) } @@ -481,13 +475,12 @@ where if is_cancellation(py, &error) { return Err(error); } - let event = CallEvent::Failed { + let event = LifecycleEvent::Failed { timing: self.timing(), origin, + error: &error, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Error(&error)))?; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -541,7 +534,7 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{RequestContext, WireRequest}; + use litellm_callbacks::event::{MachineEvent, RequestContext, WireRequest}; use litellm_callbacks::machine::{Interrupted, Step}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -803,23 +796,33 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn emit( &mut self, py: Python<'_>, - event: &CallEvent, - public: Option>, + event: LifecycleEvent<'_>, ) -> PyResult { - self.log.push(match (event, public) { - (CallEvent::Started { .. }, None) => "started".into(), - (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), - (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { - format!("succeeded:{}", value.bind(py)) + self.log.push(match event { + LifecycleEvent::Started { .. } => "started".into(), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + format!("response:{}", raw.body) } - (CallEvent::Failed { origin, .. }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Succeeded { response, .. } => { + format!("succeeded:{}", response.bind(py)) + } + LifecycleEvent::Failed { origin, error, .. } => { format!("failed:{origin:?}:{}", error.value(py)) } - _ => "unexpected".into(), }); Ok(LifecycleStep::Done) } + fn opened(&mut self, _: Python<'_>) -> PyResult<()> { + self.log.push("opened"); + Ok(()) + } + + fn delivered(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + self.log.push("delivered"); + Ok(()) + } + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -899,7 +902,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri wire: Box::new(wire()), context: Box::new(context()), }, - HostOp::Emit(CallEvent::ResponseReceived { + HostOp::Emit(MachineEvent::ResponseReceived { raw: litellm_callbacks::event::RawResponse { body: "raw".into() }, }), ], diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8889b1513db..3738a9b1c3c 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -13,7 +13,7 @@ mod handle; mod marshal; pub use adapter::{ - HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, + HostOpError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, }; pub use argument::lookup; pub use callable::wrap_failure;