scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Revert "feat(python-bridge): run deployment hooks through the native cursor"

Browse source 
This reverts commit b204dec9f9.
This commit is contained in:
Yujong Lee 2026-09-16 19:06:47 -07:00
parent c26b96d1ba
commit b2c17ae5ac
8 changed files with 278 additions and 984 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

262
litellm-rust/crates/core/src/call_lifecycle/callbacks.rs
View file

@ -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<CallbackKind>,
}
pub fn plan_success(facts: &SuccessFacts) -> Vec<Dispatch> {
pub fn plan_success(facts: &SuccessFacts) -> Vec<SuccessDispatch> {
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<Dispatch> {
.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<Dispatch> {
background.into_iter().chain(worker).collect()
}
pub fn plan_failure(phase: HostPhase, asynchronous: bool, internal: bool) -> Option<Dispatch> {
pub fn plan_failure(
phase: HostPhase,
asynchronous: bool,
internal: bool,
) -> Option<CallbackFamily> {
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<CallbackId>,
facts: CursorFacts,
stream: bool,
position: Position,
}
impl DispatchCursor {
pub fn start(dispatch: Dispatch, targets: Vec<CallbackId>, facts: CursorFacts) -> Self {
let family = dispatch.family;
let position = if family.marker().is_some() && facts.already_logged {
pub fn start(
family: CallbackFamily,
targets: Vec<CallbackId>,
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<CallbackId>,
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);
}
}

7
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -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,

40
litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs
View file

@ -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<bool> {
.call_method0("get")?
.extract()
}
pub(super) struct DeploymentHooks;
impl DeploymentHooks {
pub(super) fn before_call(
py: Python<'_>,
kwargs: &Py<PyDict>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
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<PyDict>,
response: &Option<Py<PyAny>>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
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<PyDict>,
error: &Py<PyBaseException>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_post_call_failure_deployment_hook")?
.call1((kwargs, error, call_type))
.map(Bound::unbind)
}
}

257
litellm-rust/crates/python-bridge/src/lifecycle/dispatch.rs
View file

@ -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<CallbackId>,
pub dispatch: Dispatch,
pub family: CallbackFamily,
pub response: Option<Py<PyAny>>,
pub error: Option<Py<PyBaseException>>,
pub start: Py<PyAny>,
pub end: Py<PyAny>,
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<Self> {
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::<bool>()?,
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<PyAny>,
exception: Py<PyBaseException>,
fallback_depth: Py<PyAny>,
},
}
pub(super) struct DeploymentBody {
logger: PythonLogger,
targets: Targets,
cursor: DispatchCursor,
event: DeploymentEvent,
call_type: &'static str,
current: Py<PyAny>,
kwargs: Py<pyo3::types::PyDict>,
pending: Option<CallbackId>,
}
impl DeploymentBody {
pub(super) fn start(
py: Python<'_>,
logger: &PythonLogger,
family: CallbackFamily,
call_type: &'static str,
kwargs: &Py<pyo3::types::PyDict>,
current: Py<PyAny>,
error: Option<&Py<PyBaseException>>,
) -> PyResult<Self> {
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<PyAny>, Py<PyAny>, Py<PyAny>) =
view.extract()?;
DeploymentEvent::Failure {
request,
exception: snapshot
.into_bound(py)
.cast_into::<PyBaseException>()?
.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<Py<PyAny>> {
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<Py<PyAny>>,
) -> 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::<PyException>(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<Py<PyAny>>>,
) -> PyResult<super::handle::ExecutionStep> {
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<Py<PyAny>> {
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<CallbackId> = 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(),

70
litellm-rust/crates/python-bridge/src/lifecycle/mod.rs
View file

@ -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<PyAny>,
) -> PyResult<Py<PyAny>> {
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<PyAny>) -> PyResult<()> {
match phase {
HostPhase::DeploymentPreCall => {
@ -479,19 +458,14 @@ impl PythonCallState {
}
}
fn job(&self, py: Python<'_>, selected: Dispatch) -> PyResult<dispatch::Job> {
fn job(&self, py: Python<'_>, family: CallbackFamily) -> PyResult<dispatch::Job> {
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()),

300
litellm/rust_bridge/leaves.py
View file

@ -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 _:

282
tests/test_litellm/rust_bridge/test_leaves.py
View file

@ -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)

44
tests/test_litellm_rust/ocr/test_callbacks.py
View file

@ -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"]