From b2c17ae5accff95d242afd9386945ad34ea454bb Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:06:47 -0700 Subject: [PATCH] Revert "feat(python-bridge): run deployment hooks through the native cursor" This reverts commit b204dec9f9cfd5e68b9d56e0bd4a7ae198afdd6f. --- .../core/src/call_lifecycle/callbacks.rs | 262 +++++---------- .../crates/core/src/call_lifecycle/mod.rs | 7 +- .../python-bridge/src/lifecycle/bindings.rs | 40 +++ .../python-bridge/src/lifecycle/dispatch.rs | 257 +-------------- .../crates/python-bridge/src/lifecycle/mod.rs | 70 ++-- litellm/rust_bridge/leaves.py | 300 +++++++----------- tests/test_litellm/rust_bridge/test_leaves.py | 282 ---------------- tests/test_litellm_rust/ocr/test_callbacks.py | 44 +-- 8 files changed, 278 insertions(+), 984 deletions(-) delete mode 100644 tests/test_litellm/rust_bridge/test_leaves.py diff --git a/litellm-rust/crates/core/src/call_lifecycle/callbacks.rs b/litellm-rust/crates/core/src/call_lifecycle/callbacks.rs index c1b523cef6d..82d33898bd0 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/callbacks.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/callbacks.rs @@ -90,6 +90,18 @@ pub enum CallbackFamily { } impl CallbackFamily { + pub const fn delivery(self) -> Delivery { + match self { + Self::RequestPreCall | Self::RequestPostCall | Self::SyncFailure => Delivery::Inline, + Self::DeploymentPreCall + | Self::DeploymentPostCall + | Self::DeploymentFailure + | Self::AsyncFailure => Delivery::Await, + Self::SyncSuccess => Delivery::Worker, + Self::AsyncSuccess => Delivery::Background, + } + } + pub const fn dispatch_method(self) -> CallbackMethod { match self { Self::RequestPreCall => CallbackMethod::LogPreApiCall, @@ -166,22 +178,12 @@ pub enum ReleaseGate { } #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Dispatch { +pub struct SuccessDispatch { pub family: CallbackFamily, pub delivery: Delivery, pub gate: ReleaseGate, } -impl Dispatch { - pub const fn immediate(family: CallbackFamily, delivery: Delivery) -> Self { - Self { - family, - delivery, - gate: ReleaseGate::Immediate, - } - } -} - #[derive(Clone, Debug, PartialEq, Eq)] pub struct SuccessFacts { pub asynchronous: bool, @@ -191,15 +193,15 @@ pub struct SuccessFacts { pub sync_target_kinds: Vec, } -pub fn plan_success(facts: &SuccessFacts) -> Vec { +pub fn plan_success(facts: &SuccessFacts) -> Vec { if !facts.asynchronous { - return vec![Dispatch { + return vec![SuccessDispatch { family: CallbackFamily::SyncSuccess, delivery: Delivery::Worker, gate: ReleaseGate::Immediate, }]; } - let background = (!facts.internal && !facts.fallbacks).then_some(Dispatch { + let background = (!facts.internal && !facts.fallbacks).then_some(SuccessDispatch { family: CallbackFamily::AsyncSuccess, delivery: Delivery::Background, gate: if facts.deferred { @@ -212,7 +214,7 @@ pub fn plan_success(facts: &SuccessFacts) -> Vec { .sync_target_kinds .iter() .any(|kind| kind.runs_sync_handler_for_async_call()) - .then_some(Dispatch { + .then_some(SuccessDispatch { family: CallbackFamily::SyncSuccess, delivery: Delivery::Worker, gate: ReleaseGate::Immediate, @@ -220,45 +222,21 @@ pub fn plan_success(facts: &SuccessFacts) -> Vec { background.into_iter().chain(worker).collect() } -pub fn plan_failure(phase: HostPhase, asynchronous: bool, internal: bool) -> Option { +pub fn plan_failure( + phase: HostPhase, + asynchronous: bool, + internal: bool, +) -> Option { if asynchronous && internal { return None; } match phase { - HostPhase::Failure => Some(Dispatch::immediate( - CallbackFamily::SyncFailure, - Delivery::Inline, - )), - HostPhase::AsyncFailure => Some(Dispatch::immediate( - CallbackFamily::AsyncFailure, - Delivery::Await, - )), + HostPhase::Failure => Some(CallbackFamily::SyncFailure), + HostPhase::AsyncFailure => Some(CallbackFamily::AsyncFailure), _ => None, } } -pub fn plan_request(family: CallbackFamily) -> Dispatch { - let delivery = match family { - CallbackFamily::RequestPreCall | CallbackFamily::RequestPostCall => Delivery::Inline, - _ => Delivery::Await, - }; - Dispatch::immediate(family, delivery) -} - -pub fn object_target_eligible( - sync_request: bool, - method: CallbackMethod, - kind: CallbackKind, -) -> bool { - match (method, kind) { - ( - CallbackMethod::LogSuccessEvent | CallbackMethod::LogFailureEvent, - CallbackKind::CustomLogger | CallbackKind::Callable { .. }, - ) => sync_request, - _ => true, - } -} - pub trait DispatchFacts { fn eligible(&mut self, target: CallbackId, method: CallbackMethod) -> bool; } @@ -286,24 +264,21 @@ enum Position { Complete { aborted: bool }, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct CursorFacts { - pub already_logged: bool, - pub stream: bool, - pub sync_request: bool, -} - pub struct DispatchCursor { - dispatch: Dispatch, + family: CallbackFamily, targets: Vec, - facts: CursorFacts, + stream: bool, position: Position, } impl DispatchCursor { - pub fn start(dispatch: Dispatch, targets: Vec, facts: CursorFacts) -> Self { - let family = dispatch.family; - let position = if family.marker().is_some() && facts.already_logged { + pub fn start( + family: CallbackFamily, + targets: Vec, + already_logged: bool, + stream: bool, + ) -> Self { + let position = if family.marker().is_some() && already_logged { Position::Complete { aborted: false } } else if family.prepares_logging() { Position::Prepare @@ -311,23 +286,15 @@ impl DispatchCursor { Position::Dispatch(0) }; Self { - dispatch, + family, targets, - facts, + stream, position, } } - pub fn object_target_eligible(&self, method: CallbackMethod, kind: CallbackKind) -> bool { - object_target_eligible(self.facts.sync_request, method, kind) - } - pub const fn family(&self) -> CallbackFamily { - self.dispatch.family - } - - pub const fn delivery(&self) -> Delivery { - self.dispatch.delivery + self.family } pub fn targets(&self) -> &[CallbackId] { @@ -336,7 +303,7 @@ impl DispatchCursor { pub fn accept(&mut self, outcome: InvocationOutcome) { if outcome == InvocationOutcome::Failed - && self.dispatch.family.error_policy() == TargetErrorPolicy::Propagate + && self.family.error_policy() == TargetErrorPolicy::Propagate { self.position = Position::Complete { aborted: true }; } @@ -350,7 +317,7 @@ impl DispatchCursor { return DispatchStep::PrepareLogging; } Position::Hook(index) => { - let Some(method) = self.dispatch.family.hook_method() else { + let Some(method) = self.family.hook_method() else { self.position = Position::Mark; continue; }; @@ -365,10 +332,8 @@ impl DispatchCursor { } Position::Mark => { self.position = Position::Dispatch(0); - match self.dispatch.family.marker() { - Some(marker) if !self.facts.stream => { - return DispatchStep::MarkLogged(marker); - } + match self.family.marker() { + Some(marker) if !self.stream => return DispatchStep::MarkLogged(marker), _ => continue, } } @@ -378,7 +343,7 @@ impl DispatchCursor { continue; }; self.position = Position::Dispatch(index + 1); - let method = self.dispatch.family.dispatch_method(); + let method = self.family.dispatch_method(); if facts.eligible(target, method) { return DispatchStep::Invoke(self.invocation(target, method)); } @@ -389,7 +354,7 @@ impl DispatchCursor { } fn after_prepare(&self) -> Position { - if self.dispatch.family.hook_method().is_some() { + if self.family.hook_method().is_some() { Position::Hook(0) } else { Position::Mark @@ -400,7 +365,7 @@ impl DispatchCursor { CallbackInvocation { target, method, - delivery: self.dispatch.delivery, + delivery: self.family.delivery(), } } } @@ -442,34 +407,6 @@ mod tests { values.iter().copied().map(CallbackId).collect() } - fn delivery_for(family: CallbackFamily) -> Delivery { - match family { - CallbackFamily::SyncSuccess => Delivery::Worker, - CallbackFamily::AsyncSuccess => Delivery::Background, - CallbackFamily::SyncFailure - | CallbackFamily::RequestPreCall - | CallbackFamily::RequestPostCall => Delivery::Inline, - _ => Delivery::Await, - } - } - - fn start( - family: CallbackFamily, - targets: Vec, - already_logged: bool, - stream: bool, - ) -> DispatchCursor { - DispatchCursor::start( - Dispatch::immediate(family, delivery_for(family)), - targets, - CursorFacts { - already_logged, - stream, - sync_request: true, - }, - ) - } - #[test] fn terminal_families_order_dynamic_before_global_and_keep_first_duplicate() { let combined = CallbackFamily::SyncSuccess.targets(&ids(&[3, 1, 4]), Some(&ids(&[1, 2]))); @@ -490,7 +427,8 @@ mod tests { #[test] fn success_runs_every_hook_before_any_dispatch_and_marks_between_passes() { - let mut cursor = start(CallbackFamily::SyncSuccess, ids(&[1, 2]), false, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::SyncSuccess, ids(&[1, 2]), false, false); let steps = drain(&mut cursor, &mut AllEligible); let invocation = |target, method| { DispatchStep::Invoke(CallbackInvocation { @@ -515,7 +453,8 @@ mod tests { #[test] fn async_success_uses_async_leaf_methods_and_background_delivery() { - let mut cursor = start(CallbackFamily::AsyncSuccess, ids(&[7]), false, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::AsyncSuccess, ids(&[7]), false, false); let steps = drain(&mut cursor, &mut AllEligible); let methods: Vec<_> = steps .iter() @@ -550,7 +489,7 @@ mod tests { CallbackMethod::AsyncLogFailureEvent, ), ] { - let mut cursor = start(family, ids(&[1, 2]), false, false); + let mut cursor = DispatchCursor::start(family, ids(&[1, 2]), false, false); let steps = drain(&mut cursor, &mut AllEligible); assert_eq!(steps[0], DispatchStep::PrepareLogging); assert!(matches!(steps[1], DispatchStep::MarkLogged(_))); @@ -571,7 +510,8 @@ mod tests { ] ); } - let mut cursor = start(CallbackFamily::RequestPreCall, ids(&[1]), false, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::RequestPreCall, ids(&[1]), false, false); let steps = drain(&mut cursor, &mut AllEligible); assert_eq!( steps, @@ -588,12 +528,14 @@ mod tests { #[test] fn already_logged_marker_skips_the_whole_terminal_family_but_not_request_families() { - let mut cursor = start(CallbackFamily::AsyncSuccess, ids(&[1]), true, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::AsyncSuccess, ids(&[1]), true, false); assert_eq!( cursor.next(&mut AllEligible), DispatchStep::Complete { aborted: false } ); - let mut cursor = start(CallbackFamily::RequestPostCall, ids(&[1]), true, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::RequestPostCall, ids(&[1]), true, false); assert!(matches!( cursor.next(&mut AllEligible), DispatchStep::Invoke(_) @@ -602,7 +544,7 @@ mod tests { #[test] fn streaming_skips_the_marker_write_but_still_dispatches() { - let mut cursor = start(CallbackFamily::SyncSuccess, ids(&[1]), false, true); + let mut cursor = DispatchCursor::start(CallbackFamily::SyncSuccess, ids(&[1]), false, true); let steps = drain(&mut cursor, &mut AllEligible); assert!( !steps @@ -620,7 +562,8 @@ mod tests { #[test] fn ineligible_targets_are_skipped_per_method_without_affecting_others() { - let mut cursor = start(CallbackFamily::SyncSuccess, ids(&[1, 2]), false, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::SyncSuccess, ids(&[1, 2]), false, false); let mut facts = Gate(|target, method| { !(target == CallbackId(1) && method == CallbackMethod::LoggingHook) && !(target == CallbackId(2) && method == CallbackMethod::LogSuccessEvent) @@ -643,7 +586,8 @@ mod tests { #[test] fn contained_failures_continue_and_propagating_failures_abort() { - let mut cursor = start(CallbackFamily::SyncFailure, ids(&[1, 2]), false, false); + let mut cursor = + DispatchCursor::start(CallbackFamily::SyncFailure, ids(&[1, 2]), false, false); assert_eq!(cursor.next(&mut AllEligible), DispatchStep::PrepareLogging); assert!(matches!( cursor.next(&mut AllEligible), @@ -662,7 +606,7 @@ mod tests { }) )); - let mut cursor = start( + let mut cursor = DispatchCursor::start( CallbackFamily::DeploymentPreCall, ids(&[1, 2]), false, @@ -690,7 +634,7 @@ mod tests { }); assert_eq!( plan, - [Dispatch { + [SuccessDispatch { family: CallbackFamily::SyncSuccess, delivery: Delivery::Worker, gate: ReleaseGate::Immediate, @@ -712,7 +656,7 @@ mod tests { }; assert_eq!( plan_success(&base), - [Dispatch { + [SuccessDispatch { family: CallbackFamily::AsyncSuccess, delivery: Delivery::Background, gate: ReleaseGate::Immediate, @@ -729,12 +673,12 @@ mod tests { assert_eq!( plan_success(&with_external), [ - Dispatch { + SuccessDispatch { family: CallbackFamily::AsyncSuccess, delivery: Delivery::Background, gate: ReleaseGate::Deferred, }, - Dispatch { + SuccessDispatch { family: CallbackFamily::SyncSuccess, delivery: Delivery::Worker, gate: ReleaseGate::Immediate, @@ -748,7 +692,7 @@ mod tests { }; assert_eq!( plan_success(&internal_or_fallback), - [Dispatch { + [SuccessDispatch { family: CallbackFamily::SyncSuccess, delivery: Delivery::Worker, gate: ReleaseGate::Immediate, @@ -775,83 +719,29 @@ mod tests { fn failure_families_follow_the_phase_and_skip_internal_async_calls() { assert_eq!( plan_failure(HostPhase::Failure, false, true), - Some(Dispatch::immediate( - CallbackFamily::SyncFailure, - Delivery::Inline - )) + Some(CallbackFamily::SyncFailure) ); assert_eq!( plan_failure(HostPhase::AsyncFailure, true, false), - Some(Dispatch::immediate( - CallbackFamily::AsyncFailure, - Delivery::Await - )) + Some(CallbackFamily::AsyncFailure) ); assert_eq!(plan_failure(HostPhase::Failure, true, true), None); assert_eq!(plan_failure(HostPhase::Success, false, false), None); } #[test] - fn delivery_is_selected_by_the_plan_and_carried_on_every_invocation() { - let mut cursor = DispatchCursor::start( - Dispatch::immediate(CallbackFamily::SyncFailure, Delivery::Worker), - ids(&[1]), - CursorFacts { - already_logged: false, - stream: false, - sync_request: true, - }, - ); - let deliveries: Vec<_> = drain(&mut cursor, &mut AllEligible) - .into_iter() - .filter_map(|step| match step { - DispatchStep::Invoke(invocation) => Some(invocation.delivery), - _ => None, - }) - .collect(); - assert_eq!(deliveries, [Delivery::Worker]); + fn delivery_is_a_property_of_the_family_not_of_the_callable() { + assert_eq!(CallbackFamily::RequestPreCall.delivery(), Delivery::Inline); assert_eq!( - plan_request(CallbackFamily::RequestPreCall).delivery, - Delivery::Inline - ); - assert_eq!( - plan_request(CallbackFamily::DeploymentPreCall).delivery, + CallbackFamily::DeploymentPreCall.delivery(), Delivery::Await ); - } - - #[test] - fn sync_leaf_methods_skip_object_targets_on_async_requests() { - let asynchronous = DispatchCursor::start( - Dispatch::immediate(CallbackFamily::SyncSuccess, Delivery::Worker), - ids(&[1]), - CursorFacts { - already_logged: false, - stream: false, - sync_request: false, - }, - ); - for kind in [ - CallbackKind::CustomLogger, - CallbackKind::Callable { internal: false }, - ] { - assert!(!asynchronous.object_target_eligible(CallbackMethod::LogSuccessEvent, kind)); - assert!(!asynchronous.object_target_eligible(CallbackMethod::LogFailureEvent, kind)); - assert!(asynchronous.object_target_eligible(CallbackMethod::LoggingHook, kind)); - assert!( - asynchronous.object_target_eligible(CallbackMethod::AsyncLogSuccessEvent, kind) - ); - } - assert!(asynchronous.object_target_eligible( - CallbackMethod::LogSuccessEvent, - CallbackKind::Named { known: false } - )); - let synchronous = start(CallbackFamily::SyncSuccess, ids(&[1]), false, false); - assert!( - synchronous.object_target_eligible( - CallbackMethod::LogSuccessEvent, - CallbackKind::CustomLogger - ) + assert_eq!(CallbackFamily::SyncSuccess.delivery(), Delivery::Worker); + assert_eq!( + CallbackFamily::AsyncSuccess.delivery(), + Delivery::Background ); + assert_eq!(CallbackFamily::SyncFailure.delivery(), Delivery::Inline); + assert_eq!(CallbackFamily::AsyncFailure.delivery(), Delivery::Await); } } diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 8788e0926b0..f5625ac6537 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -7,10 +7,9 @@ pub mod registration; pub mod types; pub use callbacks::{ - CallbackFamily, CallbackId, CallbackInvocation, CallbackKind, CallbackMethod, CursorFacts, - Delivery, Dispatch, DispatchCursor, DispatchFacts, DispatchStep, InvocationOutcome, - LoggedMarker, ReleaseGate, SuccessFacts, TargetErrorPolicy, object_target_eligible, - plan_failure, plan_request, plan_success, + CallbackFamily, CallbackId, CallbackInvocation, CallbackKind, CallbackMethod, Delivery, + DispatchCursor, DispatchFacts, DispatchStep, InvocationOutcome, LoggedMarker, ReleaseGate, + SuccessDispatch, SuccessFacts, TargetErrorPolicy, plan_failure, plan_success, }; pub use types::{ CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest, diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs index 1a1ecbce116..f78d4f87c62 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -1,3 +1,4 @@ +use pyo3::exceptions::PyBaseException; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -55,3 +56,42 @@ pub(super) fn is_internal_call(py: Python<'_>) -> PyResult { .call_method0("get")? .extract() } + +pub(super) struct DeploymentHooks; + +impl DeploymentHooks { + pub(super) fn before_call( + py: Python<'_>, + kwargs: &Py, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_pre_call_deployment_hook")? + .call1((kwargs, call_type)) + .map(Bound::unbind) + } + + pub(super) fn after_success( + py: Python<'_>, + kwargs: &Py, + response: &Option>, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_post_call_success_deployment_hook")? + .call1((kwargs, response, call_type)) + .map(Bound::unbind) + } + + pub(super) fn after_failure( + py: Python<'_>, + kwargs: &Py, + error: &Py, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_post_call_failure_deployment_hook")? + .call1((kwargs, error, call_type)) + .map(Bound::unbind) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs b/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs index ca204d1d01b..4215b8e6c2d 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs @@ -1,7 +1,6 @@ use litellm_core::call_lifecycle::{ - CallbackFamily, CallbackId, CallbackInvocation, CallbackKind, CallbackMethod, CursorFacts, - Delivery, Dispatch, DispatchCursor, DispatchFacts, DispatchStep, InvocationOutcome, - LoggedMarker, object_target_eligible, plan_request, + CallbackFamily, CallbackId, CallbackInvocation, CallbackKind, CallbackMethod, Delivery, + DispatchCursor, DispatchFacts, DispatchStep, InvocationOutcome, LoggedMarker, }; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; @@ -126,13 +125,12 @@ pub(super) struct Job { pub logger: PythonLogger, pub targets: Targets, pub ids: Vec, - pub dispatch: Dispatch, + pub family: CallbackFamily, pub response: Option>, pub error: Option>, pub start: Py, pub end: Py, pub stream: bool, - pub sync_request: bool, } impl Job { @@ -146,7 +144,7 @@ impl Job { } fn family_name(&self) -> &'static str { - match self.dispatch.family { + match self.family { CallbackFamily::SyncSuccess => "sync_success", CallbackFamily::AsyncSuccess => "async_success", CallbackFamily::SyncFailure => "sync_failure", @@ -156,7 +154,7 @@ impl Job { } fn outcome(&self) -> Outcome { - match self.dispatch.family { + match self.family { CallbackFamily::SyncSuccess | CallbackFamily::AsyncSuccess => Outcome::Success, _ => Outcome::Failure, } @@ -173,9 +171,6 @@ impl DispatchFacts for Eligibility<'_, '_> { fn eligible(&mut self, target: CallbackId, method: CallbackMethod) -> bool { let object = self.job.targets.object(self.py, target); let kind = self.job.targets.kind(target); - if !object_target_eligible(self.job.sync_request, method, kind) { - return false; - } let result = match method { CallbackMethod::LoggingHook | CallbackMethod::AsyncLoggingHook => { if kind != CallbackKind::CustomLogger { @@ -232,22 +227,14 @@ pub(super) struct Runner { impl Runner { pub(super) fn start(py: Python<'_>, job: Job) -> PyResult { let leaves = leaves(py)?; - let already_logged = match job.dispatch.family.marker() { + let already = match job.family.marker() { Some(marker) => leaves .getattr("already_logged")? .call1((job.logger.object(py), marker.key()))? .extract::()?, None => false, }; - let cursor = DispatchCursor::start( - job.dispatch, - job.ids.clone(), - CursorFacts { - already_logged, - stream: job.stream, - sync_request: job.sync_request, - }, - ); + let cursor = DispatchCursor::start(job.family, job.ids.clone(), already, job.stream); let result = job.response.as_ref().map(|value| value.clone_ref(py)); Ok(Self { job, @@ -353,6 +340,7 @@ impl Runner { let logger = self.job.logger.object(py); let target = self.job.targets.object(py, invocation.target); let kind = self.job.targets.kind(invocation.target); + let awaits = matches!(invocation.delivery, Delivery::Await | Delivery::Background); let value = match (invocation.method, kind) { (CallbackMethod::LoggingHook, CallbackKind::CustomLogger) => { let replaced = @@ -414,19 +402,10 @@ impl Runner { ))?, _ => return Ok(None), }; - match invocation.delivery { - Delivery::Inline | Delivery::Worker => Ok(None), - Delivery::Await | Delivery::Background if value.is_none() => Ok(None), - Delivery::Await | Delivery::Background if value.hasattr("__await__")? => { - Ok(Some(value.unbind())) - } - Delivery::Await | Delivery::Background => Err(PyRuntimeError::new_err(format!( - "{:?} leaf for {:?} returned a non-awaitable {}", - invocation.method, - self.job.dispatch.family, - value.get_type().name()? - ))), + if awaits && !value.is_none() { + return Ok(Some(value.unbind())); } + Ok(None) } fn accept( @@ -565,15 +544,7 @@ pub(super) struct RequestJob<'a> { pub(super) fn dispatch_request(py: Python<'_>, job: RequestJob<'_>) -> PyResult<()> { let leaves = leaves(py)?; let (targets, ids) = family_targets(py, job.logger, job.family)?; - let mut cursor = DispatchCursor::start( - plan_request(job.family), - ids, - CursorFacts { - already_logged: false, - stream: false, - sync_request: true, - }, - ); + let mut cursor = DispatchCursor::start(job.family, ids, false, false); let logger = job.logger.object(py); let event = match job.family { CallbackFamily::RequestPreCall => "pre_api_call", @@ -677,196 +648,6 @@ pub(super) fn family_targets( Ok((targets, ordered)) } -pub(super) enum DeploymentEvent { - PreCall, - PostCall, - Failure { - request: Py, - exception: Py, - fallback_depth: Py, - }, -} - -pub(super) struct DeploymentBody { - logger: PythonLogger, - targets: Targets, - cursor: DispatchCursor, - event: DeploymentEvent, - call_type: &'static str, - current: Py, - kwargs: Py, - pending: Option, -} - -impl DeploymentBody { - pub(super) fn start( - py: Python<'_>, - logger: &PythonLogger, - family: CallbackFamily, - call_type: &'static str, - kwargs: &Py, - current: Py, - error: Option<&Py>, - ) -> PyResult { - let (targets, ids) = family_targets(py, logger, family)?; - let event = match family { - CallbackFamily::DeploymentPreCall => DeploymentEvent::PreCall, - CallbackFamily::DeploymentPostCall => DeploymentEvent::PostCall, - CallbackFamily::DeploymentFailure => { - let exception = error.ok_or_else(super::missing_state)?; - let view = leaves(py)? - .getattr("failure_deployment_hook_view")? - .call1((kwargs, exception))?; - let (request, snapshot, fallback_depth): (Py, Py, Py) = - view.extract()?; - DeploymentEvent::Failure { - request, - exception: snapshot - .into_bound(py) - .cast_into::()? - .unbind(), - fallback_depth, - } - } - _ => return Err(super::missing_state()), - }; - Ok(Self { - logger: logger.clone_ref(py), - targets, - cursor: DispatchCursor::start( - plan_request(family), - ids, - CursorFacts { - already_logged: false, - stream: false, - sync_request: true, - }, - ), - event, - call_type, - current, - kwargs: kwargs.clone_ref(py), - pending: None, - }) - } - - fn invoke(&self, py: Python<'_>, target: CallbackId) -> PyResult> { - let leaves = leaves(py)?; - let object = self.targets.object(py, target); - let awaitable = match &self.event { - DeploymentEvent::PreCall => leaves.getattr("pre_call_deployment_hook")?.call1(( - object, - &self.current, - self.call_type, - ))?, - DeploymentEvent::PostCall => leaves - .getattr("post_call_success_deployment_hook")? - .call1((object, &self.kwargs, &self.current, self.call_type))?, - DeploymentEvent::Failure { - request, - exception, - fallback_depth, - } => leaves - .getattr("post_call_failure_deployment_hook")? - .call1((object, request, exception, self.call_type, fallback_depth))?, - }; - Ok(awaitable.unbind()) - } - - fn accept( - &mut self, - py: Python<'_>, - target: CallbackId, - result: PyResult>, - ) -> PyResult<()> { - match result { - Ok(value) => { - if !value.is_none(py) && !matches!(self.event, DeploymentEvent::Failure { .. }) { - self.current = value; - } - self.cursor.accept(InvocationOutcome::Completed); - Ok(()) - } - Err(error) - if matches!(self.event, DeploymentEvent::Failure { .. }) - && error.is_instance_of::(py) => - { - let object = self.targets.object(py, target); - leaves(py)? - .getattr("report_deployment_failure_hook_error")? - .call1((object, error.value(py)))?; - self.cursor.accept(InvocationOutcome::Failed); - Ok(()) - } - Err(error) => Err(error), - } - } -} - -impl super::handle::ExecutionBody for DeploymentBody { - fn resume( - &mut self, - result: Option>>, - ) -> PyResult { - Python::attach(|py| { - if let Some(result) = result { - let target = self.pending.take().ok_or_else(super::missing_state)?; - self.accept(py, target, result)?; - } - let mut facts = DeploymentEligibility(&self.targets); - loop { - match self.cursor.next(&mut facts) { - DispatchStep::Invoke(invocation) => { - let awaitable = self.invoke(py, invocation.target)?; - self.pending = Some(invocation.target); - return Ok(super::handle::ExecutionStep::Await(awaitable)); - } - DispatchStep::Complete { .. } => { - return Ok(super::handle::ExecutionStep::Return( - self.current.clone_ref(py), - )); - } - DispatchStep::PrepareLogging | DispatchStep::MarkLogged(_) => {} - } - } - }) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.logger.traverse(visit)?; - self.targets.traverse(visit)?; - visit.call(&self.current)?; - visit.call(&self.kwargs)?; - if let DeploymentEvent::Failure { - request, - exception, - fallback_depth, - } = &self.event - { - visit.call(request)?; - visit.call(exception)?; - visit.call(fallback_depth)?; - } - Ok(()) - } -} - -struct DeploymentEligibility<'a>(&'a Targets); - -impl DispatchFacts for DeploymentEligibility<'_> { - fn eligible(&mut self, target: CallbackId, _: CallbackMethod) -> bool { - self.0.kind(target) == CallbackKind::CustomLogger - } -} - -pub(super) fn deployment_coroutine(py: Python<'_>, body: DeploymentBody) -> PyResult> { - let execution = Py::new(py, super::handle::Execution::new(body))?; - py.import("litellm.rust_bridge.lifecycle")? - .getattr("drive")? - .call1((execution,)) - .map(Bound::unbind) -} - #[cfg(test)] mod tests { use super::*; @@ -901,23 +682,13 @@ target = CustomLogger() let list = PyList::new(py, [&target]).unwrap().into_any(); let (targets, ids) = Targets::read(py, &[list]).unwrap(); let ids: Vec = ids.into_iter().flatten().collect(); - let dispatch = Dispatch::immediate(family, Delivery::Worker); Runner { - cursor: DispatchCursor::start( - dispatch, - ids.clone(), - CursorFacts { - already_logged: false, - stream: false, - sync_request: true, - }, - ), + cursor: DispatchCursor::start(family, ids.clone(), false, false), job: Job { logger: logger.extract().unwrap(), targets, ids, - dispatch, - sync_request: true, + family, response: Some(target.clone().unbind()), error: None, start: py.None(), diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index bf6c537e240..ca247138083 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -13,7 +13,7 @@ use litellm_core::call_lifecycle::host::{ HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, }; use litellm_core::call_lifecycle::{ - CallbackFamily, Delivery, Dispatch, ReleaseGate, SuccessFacts, plan_failure, plan_success, + CallbackFamily, Delivery, ReleaseGate, SuccessFacts, plan_failure, plan_success, }; mod arguments; @@ -26,6 +26,7 @@ mod setup; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; pub(crate) use arguments::{BoundArguments, Signature}; +use bindings::DeploymentHooks; pub(crate) use bindings::PythonLogger; use handle::{Execution, ExecutionBody, ExecutionStep}; @@ -323,34 +324,30 @@ impl PythonCallState { match phase { HostPhase::Setup => self.setup(py)?, HostPhase::DeploymentPreCall => { - let copied = self.kwargs.bind(py).copy()?.into_any().unbind(); - return Ok(HostStep::Suspend(self.deployment( + return Ok(HostStep::Suspend(DeploymentHooks::before_call( py, - CallbackFamily::DeploymentPreCall, - copied, + &self.kwargs, + self.call_type, )?)); } HostPhase::Prepare => self.prepare(py)?, HostPhase::DeploymentPostCall => { - let response = self - .response - .as_ref() - .map(|value| value.clone_ref(py)) - .unwrap_or_else(|| py.None()); - return Ok(HostStep::Suspend(self.deployment( + return Ok(HostStep::Suspend(DeploymentHooks::after_success( py, - CallbackFamily::DeploymentPostCall, - response, + &self.kwargs, + &self.response, + self.call_type, )?)); } HostPhase::Finalize => self.finalize(py)?, HostPhase::Success => self.dispatch_success(py)?, HostPhase::DeploymentFailure => { - if self.error.is_some() { - return Ok(HostStep::Suspend(self.deployment( + if let Some(error) = &self.error { + return Ok(HostStep::Suspend(DeploymentHooks::after_failure( py, - CallbackFamily::DeploymentFailure, - py.None(), + &self.kwargs, + error, + self.call_type, )?)); } } @@ -369,24 +366,6 @@ impl PythonCallState { Ok(HostStep::Ready(py.None())) } - fn deployment( - &self, - py: Python<'_>, - family: CallbackFamily, - current: Py, - ) -> PyResult> { - let body = dispatch::DeploymentBody::start( - py, - self.logger()?, - family, - self.call_type, - &self.kwargs, - current, - self.error.as_ref(), - )?; - dispatch::deployment_coroutine(py, body) - } - fn accept(&mut self, py: Python<'_>, phase: HostPhase, value: Py) -> PyResult<()> { match phase { HostPhase::DeploymentPreCall => { @@ -479,19 +458,14 @@ impl PythonCallState { } } - fn job(&self, py: Python<'_>, selected: Dispatch) -> PyResult { + fn job(&self, py: Python<'_>, family: CallbackFamily) -> PyResult { let logger = self.logger()?; - let (targets, ids) = dispatch::family_targets(py, logger, selected.family)?; - let sync_request = dispatch::leaves(py)? - .getattr("is_sync_request")? - .call1((logger.object(py),))? - .extract()?; + let (targets, ids) = dispatch::family_targets(py, logger, family)?; Ok(dispatch::Job { logger: logger.clone_ref(py), targets, ids, - dispatch: selected, - sync_request, + family, response: self.response.as_ref().map(|value| value.clone_ref(py)), error: self.error.as_ref().map(|value| value.clone_ref(py)), start: self.start.clone_ref(py), @@ -523,7 +497,7 @@ impl PythonCallState { sync_target_kinds: sync_targets.kinds(&sync_ids), }; for selected in plan_success(&facts) { - let runner = dispatch::Runner::start(py, self.job(py, selected)?)?; + let runner = dispatch::Runner::start(py, self.job(py, selected.family)?)?; match (selected.delivery, selected.gate) { (Delivery::Worker, _) => { let job = Py::new(py, dispatch::WorkerJob::new(runner))?; @@ -564,14 +538,14 @@ impl PythonCallState { } else { HostPhase::Failure }; - let Some(selected) = plan_failure(phase, self.asynchronous, self.internal) else { + let Some(family) = plan_failure(phase, self.asynchronous, self.internal) else { return Ok(None); }; if self.supplied { - return compat::dispatch_failure(py, self, selected.family); + return compat::dispatch_failure(py, self, family); } - let mut runner = dispatch::Runner::start(py, self.job(py, selected)?)?; - match selected.delivery { + let mut runner = dispatch::Runner::start(py, self.job(py, family)?)?; + match family.delivery() { Delivery::Inline => match runner.resume(py, None)? { dispatch::Step::Done => Ok(None), dispatch::Step::Await(_) => Err(missing_state()), diff --git a/litellm/rust_bridge/leaves.py b/litellm/rust_bridge/leaves.py index eef5a280430..6b9020aba78 100644 --- a/litellm/rust_bridge/leaves.py +++ b/litellm/rust_bridge/leaves.py @@ -31,15 +31,15 @@ from litellm.integrations.custom_logger import CustomLogger if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import StandardCallbackDynamicParams TerminalFamily: TypeAlias = Literal["sync_success", "async_success", "sync_failure", "async_failure"] -Family: TypeAlias = Literal["request", "deployment", TerminalFamily] +Family: TypeAlias = Literal["request", TerminalFamily] Details: TypeAlias = dict[ str, object ] # mutable-ok: model_call_details is the shared mutable envelope callbacks write to Timestamp: TypeAlias = datetime.datetime LegacyCall: TypeAlias = Callable[..., object] +LegacyAsyncCall: TypeAlias = Callable[..., Awaitable[None]] class LoggerView(Protocol): @@ -51,7 +51,7 @@ class LoggerView(Protocol): completion_start_time: Timestamp | None model_call_details: Details log_raw_request_response: bool - standard_callback_dynamic_params: StandardCallbackDynamicParams + standard_callback_dynamic_params: object standard_built_in_tools_params: object def record_api_call_start_time(self) -> None: ... @@ -176,10 +176,12 @@ def _integration(callback: CustomLogger) -> IntegrationView: return cast(IntegrationView, callback) # cast-ok: legacy CustomLogger methods are untyped -def _singleton(name: str) -> object: +def _legacy_module() -> Mapping[str, object]: from litellm.litellm_core_utils import litellm_logging - return getattr(litellm_logging, name, None) # pyright: ignore[reportAny] # legacy module globals are rebound at runtime + return cast( # cast-ok: module globals hold the legacy integration singletons + Mapping[str, object], vars(litellm_logging) + ) def _print_verbose() -> LegacyCall: @@ -188,6 +190,14 @@ def _print_verbose() -> LegacyCall: return cast(LegacyCall, litellm_logging.print_verbose) # cast-ok: legacy debug printer is untyped +def _method(target: object, name: str) -> LegacyCall: + return _call_of(_attribute(target, name)) + + +def _async_method(target: object, name: str) -> LegacyAsyncCall: + return cast(LegacyAsyncCall, _attribute(target, name)) # cast-ok: legacy integration singletons are untyped + + def _redact_string(value: str) -> str: from litellm.litellm_core_utils import litellm_logging @@ -293,27 +303,25 @@ def log_post_api_call(logger: Logging, callback: CustomLogger) -> None: def dispatch_named_request(logger: Logging, name: str, event: Literal["pre_api_call", "post_api_call"]) -> None: - from litellm.integrations.supabase import Supabase - view: Final = _logger(logger) - details: Final = view.model_call_details - match name: - case "supabase" if event == "pre_api_call" and isinstance(client := _singleton("supabaseClient"), Supabase): - client.input_log_event( - model=view.model, - messages=view.messages, - end_user=details.get("user", "default"), - litellm_call_id=details["litellm_call_id"], - print_verbose=_print_verbose(), - ) - case "sentry" if callable(add_breadcrumb := _singleton("add_breadcrumb")): - add_breadcrumb(category="litellm.llm_call", message=f"Model Call Details {event}: {details}", level="info") - case _: - return + module: Final = _legacy_module() + if name == "supabase" and event == "pre_api_call" and (client := module.get("supabaseClient")) is not None: + details: Final = view.model_call_details + _method(client, "input_log_event")( + model=view.model, + messages=view.messages, + end_user=details.get("user", "default"), + litellm_call_id=details["litellm_call_id"], + print_verbose=_print_verbose(), + ) + if name == "sentry" and (add_breadcrumb := module.get("add_breadcrumb")) is not None: + cast(LegacyCall, add_breadcrumb)( # cast-ok: legacy sentry hook + category="litellm.llm_call", message=f"Model Call Details {event}: {view.model_call_details}", level="info" + ) def dispatch_callable_request(logger: Logging, callback: LegacyCall) -> None: - custom: Final = _singleton("customLogger") + custom: Final = _legacy_module().get("customLogger") if not isinstance(custom, CustomLogger): return view: Final = _logger(logger) @@ -326,66 +334,6 @@ def dispatch_callable_request(logger: Logging, callback: LegacyCall) -> None: ) -def _typed_call_type(call_type: str) -> object: - from litellm.types.utils import CallTypes - - try: - return CallTypes(call_type) - except ValueError: - return None - - -def _awaitable(value: object) -> Awaitable[object]: - return cast(Awaitable[object], value) # cast-ok: legacy async hooks return coroutines - - -def pre_call_deployment_hook(callback: CustomLogger, kwargs: Details, call_type: str) -> Awaitable[object]: - hook: Final = _call_of(_attribute(callback, "async_pre_call_deployment_hook")) - return _awaitable(hook(kwargs, _typed_call_type(call_type))) - - -def post_call_success_deployment_hook( - callback: CustomLogger, kwargs: Details, response: object, call_type: str -) -> Awaitable[object]: - hook: Final = _call_of(_attribute(callback, "async_post_call_success_deployment_hook")) - return _awaitable(hook(kwargs, response, _typed_call_type(call_type))) - - -def failure_deployment_hook_view( - kwargs: Details, exception: BaseException -) -> tuple[Mapping[str, object], BaseException, int | None]: - from litellm import utils - - raw_depth: Final = kwargs.get("fallback_depth") - depth: Final = raw_depth if isinstance(raw_depth, int) else None - safe_request: Final = MappingProxyType({key: value for key, value in kwargs.items() if key != "attempted_targets"}) - snapshot: Final = _call_of(utils._snapshot_exception_for_hook)(exception) # pyright: ignore[reportPrivateUsage] # legacy snapshot helper - return safe_request, cast(BaseException, snapshot), depth # cast-ok: snapshot is a same-class copy of the exception - - -def post_call_failure_deployment_hook( - callback: CustomLogger, - request: Mapping[str, object], - exception: BaseException, - call_type: str, - fallback_depth: int | None, -) -> Awaitable[object]: - from litellm import utils - - hook: Final = _call_of(_attribute(callback, "async_post_call_failure_deployment_hook")) - accepts_depth: Final = _call_of(utils._accepts_fallback_depth_kwarg_for_class)(type(callback)) # pyright: ignore[reportPrivateUsage] # legacy signature probe - typed: Final = _typed_call_type(call_type) - if accepts_depth: - return _awaitable(hook(request, exception, typed, fallback_depth=fallback_depth)) - return _awaitable(hook(request, exception, typed)) - - -def report_deployment_failure_hook_error(callback: object, error: BaseException) -> None: - from litellm._logging import verbose_logger - - verbose_logger.debug("async_post_call_failure_deployment_hook error in %s: %s", type(callback).__name__, error) - - def report_target_failure(logger: Logging, callback: object, family: Family, error: BaseException) -> None: from litellm._logging import verbose_logger @@ -395,10 +343,10 @@ def report_target_failure(logger: Logging, callback: object, family: Family, err callback, "".join(traceback.format_exception(error)), ) - capture: Final = _singleton("capture_exception") - if callable(capture) and family in ("request", "sync_success", "sync_failure"): - capture(error) - if family not in ("request", "deployment", "sync_failure"): + capture: Final = _legacy_module().get("capture_exception") + if capture is not None and family in ("request", "sync_success", "sync_failure"): + cast(LegacyCall, capture)(error) # cast-ok: legacy sentry hook + if family not in ("request", "sync_failure"): _logger(logger)._handle_callback_failure(callback=callback) # pyright: ignore[reportPrivateUsage] # legacy prometheus counter @@ -525,14 +473,6 @@ async def async_logging_hook(logger: Logging, callback: CustomLogger, result: ob return replaced -def is_sync_request(logger: Logging) -> bool: - from litellm.litellm_core_utils.litellm_logging import Logging as LoggingClass - - params: Final = _details_of(_logger(logger).model_call_details.get("litellm_params") or _EMPTY) - decide: Final = _call_of(LoggingClass._is_sync_litellm_request) # pyright: ignore[reportPrivateUsage] # legacy predicate - return decide(params) is True - - def mark_logged(logger: Logging, marker: str) -> None: _logger(logger).model_call_details[marker] = True @@ -584,7 +524,7 @@ def async_log_failure_event( def _custom_logger_singleton() -> IntegrationView: from litellm.litellm_core_utils import litellm_logging - existing: Final = _singleton("customLogger") + existing: Final = _legacy_module().get("customLogger") if isinstance(existing, CustomLogger): return _integration(existing) created: Final = CustomLogger() @@ -624,71 +564,63 @@ def dispatch_callable( ) +_SUCCESS_SINGLETONS: Final[Mapping[str, str]] = MappingProxyType( + { + "promptlayer": "promptLayerLogger", + "supabase": "supabaseClient", + "wandb": "weightsBiasesLogger", + "logfire": "logfireLogger", + "lunary": "lunaryLogger", + "helicone": "heliconeLogger", + "greenscale": "greenscaleLogger", + "athina": "athinaLogger", + "traceloop": "traceloopLogger", + "s3": "s3Logger", + "openmeter": "openMeterLogger", + } +) + + def dispatch_named_success( logger: Logging, name: str, response: object, start_time: Timestamp, end_time: Timestamp ) -> Awaitable[None] | None: - from litellm.integrations.athina import AthinaLogger - from litellm.integrations.greenscale import GreenscaleLogger - from litellm.integrations.helicone import HeliconeLogger - from litellm.integrations.logfire_logger import LogfireLevel, LogfireLogger - from litellm.integrations.lunary import LunaryLogger - from litellm.integrations.openmeter import OpenMeterLogger - from litellm.integrations.prompt_layer import PromptLayerLogger - from litellm.integrations.s3 import S3Logger - from litellm.integrations.supabase import Supabase - from litellm.integrations.traceloop import TraceloopLogger - from litellm.integrations.weights_biases import WeightsBiasesLogger - view: Final = _logger(logger) details: Final = view.model_call_details print_verbose: Final = _print_verbose() + integration: Final = _legacy_module().get(_SUCCESS_SINGLETONS.get(name, "")) without_response: Final = { # mutable-ok: legacy integrations receive a private mutable copy key: value for key, value in details.items() if key != "original_response" } match name: - case "promptlayer" if isinstance(promptlayer := _singleton("promptLayerLogger"), PromptLayerLogger): - promptlayer.log_event( + case "promptlayer" | "wandb" | "athina" if integration is not None: + _method(integration, "log_event")( kwargs=details, response_obj=response, start_time=start_time, end_time=end_time, print_verbose=print_verbose, ) - case "wandb" if isinstance(wandb := _singleton("weightsBiasesLogger"), WeightsBiasesLogger): - wandb.log_event( - kwargs=details, - response_obj=response, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - case "athina" if isinstance(athina := _singleton("athinaLogger"), AthinaLogger): - athina.log_event( - kwargs=details, - response_obj=response, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - case "logfire" if isinstance(logfire := _singleton("logfireLogger"), LogfireLogger): - logfire.log_event( + case "logfire" if integration is not None: + from litellm.integrations.logfire_logger import LogfireLevel + + _method(integration, "log_event")( kwargs=without_response, response_obj=response, start_time=start_time, end_time=end_time, print_verbose=print_verbose, - level=LogfireLevel.INFO, + level=LogfireLevel.INFO.value, ) - case "greenscale" if isinstance(greenscale := _singleton("greenscaleLogger"), GreenscaleLogger): - greenscale.log_event( + case "greenscale" if integration is not None: + _method(integration, "log_event")( kwargs=without_response, response_obj=response, start_time=start_time, end_time=end_time, print_verbose=print_verbose, ) - case "supabase" if isinstance(supabase := _singleton("supabaseClient"), Supabase): - supabase.log_event( + case "supabase" if integration is not None: + _method(integration, "log_event")( model=view.model, messages=view.messages, end_user=details.get("user", "default"), @@ -698,8 +630,8 @@ def dispatch_named_success( litellm_call_id=details["litellm_call_id"], print_verbose=print_verbose, ) - case "lunary" if isinstance(lunary := _singleton("lunaryLogger"), LunaryLogger): - lunary.log_event( + case "lunary" if integration is not None: + _method(integration, "log_event")( kwargs=details, type="llm", event="end", @@ -712,8 +644,8 @@ def dispatch_named_success( run_id=view.litellm_call_id, print_verbose=print_verbose, ) - case "helicone" if isinstance(helicone := _singleton("heliconeLogger"), HeliconeLogger): - helicone.log_success( + case "helicone" if integration is not None: + _method(integration, "log_success")( model=view.model, messages=view.messages, response_obj=response, @@ -726,8 +658,8 @@ def dispatch_named_success( _langfuse( logger, response=response, start_time=start_time, end_time=end_time, level=None, status_message=None ) - case "traceloop" if isinstance(traceloop := _singleton("traceloopLogger"), TraceloopLogger): - traceloop.log_event( + case "traceloop" if integration is not None: + _method(integration, "log_event")( kwargs=details, response_obj=response, start_time=start_time, @@ -735,16 +667,16 @@ def dispatch_named_success( user_id=details.get("user", None), print_verbose=print_verbose, ) - case "s3" if isinstance(s3 := _singleton("s3Logger"), S3Logger): - s3.log_event( + case "s3" if integration is not None: + _method(integration, "log_event")( kwargs=details, response_obj=response, start_time=start_time, end_time=end_time, print_verbose=print_verbose, ) - case "openmeter" if isinstance(openmeter := _singleton("openMeterLogger"), OpenMeterLogger): - return openmeter.async_log_success_event( + case "openmeter" if integration is not None: + return _async_method(integration, "async_log_success_event")( kwargs=details, response_obj=response, start_time=start_time, end_time=end_time ) case "dynamodb": @@ -760,10 +692,10 @@ def _dynamodb( from litellm.integrations.dynamodb import DyanmoDBLogger from litellm.litellm_core_utils import litellm_logging - existing: Final = _singleton("dynamoLogger") + existing: Final = _legacy_module().get("dynamoLogger") dynamo: Final = existing if isinstance(existing, DyanmoDBLogger) else DyanmoDBLogger() litellm_logging.dynamoLogger = dynamo # pyright: ignore[reportAttributeAccessIssue] # legacy module global - return dynamo._async_log_event( # pyright: ignore[reportPrivateUsage] # legacy async entry point + return _async_method(dynamo, "_async_log_event")( kwargs=details, response_obj=response, start_time=start_time, end_time=end_time, print_verbose=print_verbose ) @@ -777,34 +709,39 @@ def _langfuse( level: str | None, status_message: str | None, ) -> None: - from litellm.integrations.langfuse.langfuse import LangFuseLogger - from litellm.integrations.langfuse.langfuse_handler import LangFuseHandler - from litellm.litellm_core_utils import litellm_logging - from litellm.types.utils import ModelResponse + from litellm.integrations.langfuse import langfuse_handler view: Final = _logger(logger) + module: Final = _legacy_module() kwargs: Final = { # mutable-ok: langfuse receives a private mutable copy key: value for key, value in view.model_call_details.items() if key != "original_response" } - global_logger: Final = _singleton("langFuseLogger") - handler: Final = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=global_logger if isinstance(global_logger, LangFuseLogger) else None, - standard_callback_dynamic_params=view.standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=litellm_logging.in_memory_dynamic_logger_cache, + select: Final = cast( # cast-ok: legacy factory + LegacyCall, langfuse_handler.LangFuseHandler.get_langfuse_logger_for_request ) - user: Final = kwargs.get("user") - result: Final = handler.log_event_on_langfuse( + handler: Final = select( + globalLangfuseLogger=module.get("langFuseLogger"), + standard_callback_dynamic_params=view.standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=module["in_memory_dynamic_logger_cache"], + ) + if handler is None: + return + extra: Final[Mapping[str, object]] = ( + MappingProxyType({"level": level, "status_message": status_message}) if level is not None else _EMPTY + ) + result: Final = _method(handler, "log_event_on_langfuse")( kwargs=kwargs, - response_obj=cast(ModelResponse, response), # cast-ok: OCR responses sit outside the legacy union + response_obj=response, start_time=start_time, end_time=end_time, - user_id=user if isinstance(user, str) else None, - level="DEFAULT" if level is None else level, - status_message=status_message, + user_id=kwargs.get("user", None), + **extra, ) - trace_id: Final = result.get("trace_id") - if isinstance(trace_id, str): - litellm_logging.in_memory_trace_id_cache.set_cache( + trace_id: Final = ( + cast(Details, result).get("trace_id") if isinstance(result, dict) else None # cast-ok: legacy response dict + ) # cast-ok: legacy response dict + if trace_id is not None: + _method(module["in_memory_trace_id_cache"], "set_cache")( litellm_call_id=view.litellm_call_id, service_name="langfuse", trace_id=trace_id ) @@ -817,17 +754,16 @@ def dispatch_named_failure( start_time: Timestamp, end_time: Timestamp, ) -> None: - from litellm.integrations.logfire_logger import LogfireLevel, LogfireLogger - from litellm.integrations.lunary import LunaryLogger - from litellm.integrations.supabase import Supabase - from litellm.integrations.traceloop import TraceloopLogger - view: Final = _logger(logger) + module: Final = _legacy_module() details: Final = view.model_call_details print_verbose: Final = _print_verbose() + without_response: Final = MappingProxyType( + {key: value for key, value in details.items() if key != "original_response"} + ) match name: - case "lunary" if isinstance(lunary := _singleton("lunaryLogger"), LunaryLogger): - lunary.log_event( + case "lunary" if (lunary := module.get("lunaryLogger")) is not None: + _method(lunary, "log_event")( kwargs=details, type="llm", event="error", @@ -840,10 +776,10 @@ def dispatch_named_failure( end_time=end_time, print_verbose=print_verbose, ) - case "sentry" if callable(capture := _singleton("capture_exception")): - capture(exception) - case "supabase" if isinstance(supabase := _singleton("supabaseClient"), Supabase): - supabase.log_event( + case "sentry" if (capture := module.get("capture_exception")) is not None: + cast(LegacyCall, capture)(exception) # cast-ok: legacy sentry hook + case "supabase" if (supabase := module.get("supabaseClient")) is not None: + _method(supabase, "log_event")( model=view.model, messages=view.messages, end_user=details.get("user", "default"), @@ -862,8 +798,8 @@ def dispatch_named_failure( level="ERROR", status_message=str(exception), ) - case "traceloop" if isinstance(traceloop := _singleton("traceloopLogger"), TraceloopLogger): - traceloop.log_event( + case "traceloop" if (traceloop := module.get("traceloopLogger")) is not None: + _method(traceloop, "log_event")( start_time=start_time, end_time=end_time, response_obj=None, @@ -873,16 +809,18 @@ def dispatch_named_failure( level="ERROR", kwargs=details, ) - case "logfire" if isinstance(logfire := _singleton("logfireLogger"), LogfireLogger): - logfire.log_event( + case "logfire" if (logfire := module.get("logfireLogger")) is not None: + from litellm.integrations.logfire_logger import LogfireLevel + + _method(logfire, "log_event")( kwargs={ # mutable-ok: logfire receives a private mutable copy - key: value for key, value in details.items() if key != "original_response" - } - | {"exception": exception}, # mutable-ok: merged into the private copy above + **without_response, + "exception": exception, + }, response_obj=None, start_time=start_time, end_time=end_time, - level=LogfireLevel.ERROR, + level=LogfireLevel.ERROR.value, print_verbose=print_verbose, ) case _: diff --git a/tests/test_litellm/rust_bridge/test_leaves.py b/tests/test_litellm/rust_bridge/test_leaves.py deleted file mode 100644 index a8f594da3f4..00000000000 --- a/tests/test_litellm/rust_bridge/test_leaves.py +++ /dev/null @@ -1,282 +0,0 @@ -import asyncio -import datetime -import inspect -from collections.abc import Iterator -from typing import Final -from unittest.mock import Mock - -import pytest - -from litellm.integrations.athina import AthinaLogger -from litellm.integrations.dynamodb import DyanmoDBLogger -from litellm.integrations.greenscale import GreenscaleLogger -from litellm.integrations.helicone import HeliconeLogger -from litellm.integrations.langfuse.langfuse import LangFuseLogger -from litellm.integrations.logfire_logger import LogfireLevel, LogfireLogger -from litellm.integrations.lunary import LunaryLogger -from litellm.integrations.openmeter import OpenMeterLogger -from litellm.integrations.prompt_layer import PromptLayerLogger -from litellm.integrations.s3 import S3Logger -from litellm.integrations.supabase import Supabase -from litellm.integrations.traceloop import TraceloopLogger -from litellm.integrations.weights_biases import WeightsBiasesLogger -from litellm.litellm_core_utils import litellm_logging -from litellm.rust_bridge import leaves - -START: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) -END: Final = START + datetime.timedelta(seconds=1) -DETAILS: Final[dict[str, object]] = { - "litellm_call_id": "call-id", - "input": "document", - "user": "user-1", - "original_response": "raw", - "litellm_params": {}, -} - - -def _logger() -> Mock: - return Mock( - model="mistral/mistral-ocr-latest", - messages="document", - litellm_call_id="call-id", - model_call_details=dict(DETAILS), - standard_callback_dynamic_params={}, - ) - - -def _stub(cls: type, method: str) -> tuple[object, Mock]: - instance: Final[object] = object.__new__(cls) - recorder: Final = Mock(name=f"{cls.__name__}.{method}") - setattr(instance, method, recorder) - return instance, recorder - - -def _assert_binds(cls: type, method: str, recorder: Mock, instance: object) -> None: - recorder.assert_called_once() - target: Final[object] = getattr(cls, method) - assert callable(target) - signature: Final = inspect.signature(target) - bound: Final = signature.bind(instance, *recorder.call_args.args, **recorder.call_args.kwargs) - required: Final = frozenset( - name for name, parameter in signature.parameters.items() if parameter.default is inspect.Parameter.empty - ) - assert required <= set(bound.arguments) - - -SUCCESS_CASES: Final[tuple[tuple[str, str, type, str], ...]] = ( - ("promptlayer", "promptLayerLogger", PromptLayerLogger, "log_event"), - ("wandb", "weightsBiasesLogger", WeightsBiasesLogger, "log_event"), - ("athina", "athinaLogger", AthinaLogger, "log_event"), - ("logfire", "logfireLogger", LogfireLogger, "log_event"), - ("greenscale", "greenscaleLogger", GreenscaleLogger, "log_event"), - ("supabase", "supabaseClient", Supabase, "log_event"), - ("lunary", "lunaryLogger", LunaryLogger, "log_event"), - ("helicone", "heliconeLogger", HeliconeLogger, "log_success"), - ("traceloop", "traceloopLogger", TraceloopLogger, "log_event"), - ("s3", "s3Logger", S3Logger, "log_event"), -) - -FAILURE_CASES: Final[tuple[tuple[str, str, type, str], ...]] = ( - ("lunary", "lunaryLogger", LunaryLogger, "log_event"), - ("supabase", "supabaseClient", Supabase, "log_event"), - ("traceloop", "traceloopLogger", TraceloopLogger, "log_event"), - ("logfire", "logfireLogger", LogfireLogger, "log_event"), -) - - -def _install(monkeypatch: pytest.MonkeyPatch, global_name: str, value: object) -> None: - monkeypatch.setattr(litellm_logging, global_name, value, raising=False) - - -@pytest.mark.parametrize(("name", "global_name", "cls", "method"), SUCCESS_CASES) -def test_named_success_calls_integration_with_its_real_signature( - monkeypatch: pytest.MonkeyPatch, name: str, global_name: str, cls: type, method: str -) -> None: - instance, recorder = _stub(cls, method) - _install(monkeypatch, global_name, instance) - response: Final = object() - - assert leaves.dispatch_named_success(_logger(), name, response, START, END) is None - - _assert_binds(cls, method, recorder, instance) - kwargs: Final = recorder.call_args.kwargs - assert kwargs["start_time"] is START and kwargs["end_time"] is END - assert kwargs["response_obj"] is response - assert callable(kwargs["print_verbose"]) - - -@pytest.mark.parametrize(("name", "global_name", "cls", "method"), FAILURE_CASES) -def test_named_failure_calls_integration_with_its_real_signature( - monkeypatch: pytest.MonkeyPatch, name: str, global_name: str, cls: type, method: str -) -> None: - instance, recorder = _stub(cls, method) - _install(monkeypatch, global_name, instance) - error: Final = RuntimeError("boom") - - leaves.dispatch_named_failure(_logger(), name, error, "formatted", START, END) - - _assert_binds(cls, method, recorder, instance) - kwargs: Final = recorder.call_args.kwargs - assert kwargs["start_time"] is START and kwargs["end_time"] is END - assert kwargs.get("response_obj") is None - - -def test_logfire_receives_enum_levels_and_exception_on_failure(monkeypatch: pytest.MonkeyPatch) -> None: - instance, recorder = _stub(LogfireLogger, "log_event") - _install(monkeypatch, "logfireLogger", instance) - error: Final = RuntimeError("boom") - - leaves.dispatch_named_success(_logger(), "logfire", object(), START, END) - assert recorder.call_args.kwargs["level"] is LogfireLevel.INFO - assert "original_response" not in recorder.call_args.kwargs["kwargs"] - - leaves.dispatch_named_failure(_logger(), "logfire", error, "formatted", START, END) - assert recorder.call_args.kwargs["level"] is LogfireLevel.ERROR - assert recorder.call_args.kwargs["kwargs"]["exception"] is error - assert "original_response" not in recorder.call_args.kwargs["kwargs"] - - -def test_lunary_marks_success_and_error_events(monkeypatch: pytest.MonkeyPatch) -> None: - instance, recorder = _stub(LunaryLogger, "log_event") - _install(monkeypatch, "lunaryLogger", instance) - - leaves.dispatch_named_success(_logger(), "lunary", object(), START, END) - assert recorder.call_args.kwargs["event"] == "end" - assert recorder.call_args.kwargs["run_id"] == "call-id" - - leaves.dispatch_named_failure(_logger(), "lunary", RuntimeError("boom"), "formatted", START, END) - assert recorder.call_args.kwargs["event"] == "error" - assert recorder.call_args.kwargs["error"] == "formatted" - - -def test_supabase_request_hook_only_fires_pre_call(monkeypatch: pytest.MonkeyPatch) -> None: - instance, recorder = _stub(Supabase, "input_log_event") - _install(monkeypatch, "supabaseClient", instance) - - leaves.dispatch_named_request(_logger(), "supabase", "post_api_call") - recorder.assert_not_called() - - leaves.dispatch_named_request(_logger(), "supabase", "pre_api_call") - _assert_binds(Supabase, "input_log_event", recorder, instance) - assert recorder.call_args.kwargs["end_user"] == "user-1" - - -def test_sentry_hooks_are_called_when_configured(monkeypatch: pytest.MonkeyPatch) -> None: - breadcrumb: Final = Mock() - capture: Final = Mock() - _install(monkeypatch, "add_breadcrumb", breadcrumb) - _install(monkeypatch, "capture_exception", capture) - error: Final = RuntimeError("boom") - - leaves.dispatch_named_request(_logger(), "sentry", "pre_api_call") - assert breadcrumb.call_args.kwargs["category"] == "litellm.llm_call" - - leaves.dispatch_named_failure(_logger(), "sentry", error, "formatted", START, END) - capture.assert_called_once_with(error) - - -@pytest.mark.parametrize(("name", "global_name"), [(case[0], case[1]) for case in SUCCESS_CASES]) -def test_unconfigured_or_foreign_singleton_is_skipped( - monkeypatch: pytest.MonkeyPatch, name: str, global_name: str -) -> None: - _install(monkeypatch, global_name, None) - assert leaves.dispatch_named_success(_logger(), name, object(), START, END) is None - - foreign: Final = Mock() - _install(monkeypatch, global_name, foreign) - assert leaves.dispatch_named_success(_logger(), name, object(), START, END) is None - assert not foreign.method_calls - - -def _run(awaitable: object) -> None: - async def consume() -> None: - assert inspect.isawaitable(awaitable) - _ = await awaitable - - asyncio.run(consume()) - - -def test_openmeter_returns_the_async_coroutine(monkeypatch: pytest.MonkeyPatch) -> None: - instance: Final = object.__new__(OpenMeterLogger) - recorder: Final = Mock() - - async def async_log_success_event(**kwargs: object) -> None: - recorder(**kwargs) - - monkeypatch.setattr(instance, "async_log_success_event", async_log_success_event) - _install(monkeypatch, "openMeterLogger", instance) - response: Final = object() - - _run(leaves.dispatch_named_success(_logger(), "openmeter", response, START, END)) - - _assert_binds(OpenMeterLogger, "async_log_success_event", recorder, instance) - assert recorder.call_args.kwargs["response_obj"] is response - - -def test_dynamodb_reuses_the_existing_singleton(monkeypatch: pytest.MonkeyPatch) -> None: - instance: Final = object.__new__(DyanmoDBLogger) - recorder: Final = Mock() - - async def _async_log_event(**kwargs: object) -> None: - recorder(**kwargs) - - monkeypatch.setattr(instance, "_async_log_event", _async_log_event) - _install(monkeypatch, "dynamoLogger", instance) - - _run(leaves.dispatch_named_success(_logger(), "dynamodb", object(), START, END)) - - _assert_binds(DyanmoDBLogger, "_async_log_event", recorder, instance) - assert litellm_logging.dynamoLogger is instance - - -@pytest.fixture -def langfuse_handler(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[LangFuseLogger, Mock, Mock]]: - handler: Final = object.__new__(LangFuseLogger) - recorder: Final = Mock(return_value={"trace_id": "trace-1"}) - monkeypatch.setattr(handler, "log_event_on_langfuse", recorder) - _install(monkeypatch, "langFuseLogger", handler) - cache: Final = Mock() - _install(monkeypatch, "in_memory_trace_id_cache", cache) - yield handler, recorder, cache - - -def test_langfuse_success_logs_and_caches_trace_id(langfuse_handler: tuple[LangFuseLogger, Mock, Mock]) -> None: - handler, recorder, cache = langfuse_handler - response: Final = object() - - leaves.dispatch_named_success(_logger(), "langfuse", response, START, END) - - _assert_binds(LangFuseLogger, "log_event_on_langfuse", recorder, handler) - kwargs: Final = recorder.call_args.kwargs - assert kwargs["response_obj"] is response - assert kwargs["user_id"] == "user-1" - assert kwargs["level"] == "DEFAULT" and kwargs["status_message"] is None - assert "original_response" not in kwargs["kwargs"] - cache.set_cache.assert_called_once_with(litellm_call_id="call-id", service_name="langfuse", trace_id="trace-1") - - -def test_langfuse_failure_logs_error_level(langfuse_handler: tuple[LangFuseLogger, Mock, Mock]) -> None: - handler, recorder, _ = langfuse_handler - - leaves.dispatch_named_failure(_logger(), "langfuse", RuntimeError("boom"), "formatted", START, END) - - _assert_binds(LangFuseLogger, "log_event_on_langfuse", recorder, handler) - kwargs: Final = recorder.call_args.kwargs - assert kwargs["response_obj"] is None - assert kwargs["level"] == "ERROR" and kwargs["status_message"] == "boom" - - -def test_langfuse_skips_cache_without_trace_id(langfuse_handler: tuple[LangFuseLogger, Mock, Mock]) -> None: - _, recorder, cache = langfuse_handler - recorder.return_value = {} - - leaves.dispatch_named_success(_logger(), "langfuse", object(), START, END) - - cache.set_cache.assert_not_called() - - -def test_signature_binding_rejects_drift() -> None: - instance, recorder = _stub(Supabase, "log_event") - recorder(model="m", end_userr="typo", response_obj=None, start_time=START, end_time=END) - with pytest.raises(TypeError): - _assert_binds(Supabase, "log_event", recorder, instance) diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 904504150a7..9487550c80e 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -506,52 +506,20 @@ def legacy_orchestration_disabled(monkeypatch: pytest.MonkeyPatch) -> list[str]: for name in FORBIDDEN_ORCHESTRATION: monkeypatch.setattr(Logging, name, forbid(name)) - def forbid_utils(name: str): - async def hook(*args, **kwargs): - reached.append(name) - raise AssertionError(f"legacy orchestration reached: {name}") + def forbidden_setup(*args, **kwargs): + reached.append("function_setup") + raise AssertionError("legacy orchestration reached: function_setup") - def setup(*args, **kwargs): - reached.append(name) - raise AssertionError(f"legacy orchestration reached: {name}") - - return setup if name == "function_setup" else hook - - for name in ( - "function_setup", - "async_pre_call_deployment_hook", - "async_post_call_success_deployment_hook", - "async_post_call_failure_deployment_hook", - ): - monkeypatch.setattr(utils, name, forbid_utils(name)) + monkeypatch.setattr(utils, "function_setup", forbidden_setup) return reached -class DeploymentRecorder(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.hooks: list[str] = [] - - async def async_pre_call_deployment_hook(self, kwargs, call_type): - self.hooks.append(f"pre:{call_type.value}") - return {**kwargs, "metadata": {**(kwargs.get("metadata") or {}), "deployment": "seen"}} - - async def async_post_call_success_deployment_hook(self, request_data, response, call_type): - self.hooks.append(f"post:{request_data['metadata']['deployment']}") - return response - - async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, fallback_depth=None): - self.hooks.append(f"failure:{type(exception).__name__}:{fallback_depth}") - - @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_ocr_success_runs_integrations_without_legacy_orchestration( ocr_server: RecordingServer, legacy_orchestration_disabled: list[str], asynchronous: bool ) -> None: recorder: Final = RecordingLogger() - deployment: Final = DeploymentRecorder() - litellm.callbacks.append(deployment) arguments: Final = {"callbacks": [recorder]} response: Final = ( await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) @@ -560,7 +528,6 @@ async def test_native_ocr_success_runs_integrations_without_legacy_orchestration success_event: Final = "async_log_success_event" if asynchronous else "log_success_event" events: Final = await recorder.wait_for_async(success_event) assert legacy_orchestration_disabled == [] - assert deployment.hooks == (["pre:aocr", "post:seen"] if asynchronous else []) assert recorder.names.count("log_pre_api_call") == 1 assert events[0].kwargs["standard_logging_object"]["status"] == "success" assert events[0].kwargs["response_cost"] is not None @@ -574,13 +541,10 @@ async def test_native_ocr_failure_runs_integrations_without_legacy_orchestration ) -> None: ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) recorder: Final = RecordingLogger() - deployment: Final = DeploymentRecorder() - litellm.callbacks.append(deployment) arguments: Final = {"callbacks": [recorder]} with pytest.raises(litellm.InternalServerError) as caught: await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) assert legacy_orchestration_disabled == [] - assert deployment.hooks == (["pre:aocr", "failure:InternalServerError:None"] if asynchronous else []) failures: Final = tuple(event for event in recorder.events if event.name.endswith("log_failure_event")) assert [event.name for event in failures] == ( ["log_failure_event", "async_log_failure_event"] if asynchronous else ["log_failure_event"]