From f8190bbe80ccefd6b24585dbe885be4e19ca6bd2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 08:30:13 -0700 Subject: [PATCH] refactor(python-bridge): classify callbacks natively instead of function_setup Native OCR setup no longer calls utils.function_setup. The bridge reads registration facts (coroutine-ness, CustomLogger, known names, existing list membership) and core's plan_registration decides every public registry mutation, which the bridge writes back through litellm.rust_bridge.setup. The Logging object is built by a narrow Python factory with the same constructor arguments. Caller-supplied Logging instances keep identity and skip registration. A differential test asserts registry side effects equal function_setup for nine registration shapes; mutating the planner fails four of them. --- .../crates/core/src/call_lifecycle/mod.rs | 1 + .../core/src/call_lifecycle/registration.rs | 374 ++++++++++++++++++ .../crates/python-bridge/src/diagnostics.rs | 24 +- .../python-bridge/src/lifecycle/bindings.rs | 93 +---- .../crates/python-bridge/src/lifecycle/mod.rs | 25 +- .../python-bridge/src/lifecycle/setup.rs | 334 ++++++++++++++++ .../python-bridge/src/routes/ocr/callbacks.rs | 15 +- .../python-bridge/src/routes/ocr/host.rs | 33 +- .../python-bridge/src/routes/ocr/mod.rs | 17 +- .../python-bridge/src/routes/ocr/project.rs | 19 +- litellm/rust_bridge/_native.pyi | 8 + litellm/rust_bridge/lifecycle.py | 30 -- litellm/rust_bridge/setup.py | 262 ++++++++++++ tests/test_litellm/rust_bridge/test_setup.py | 169 ++++++++ tests/test_litellm_rust/ocr/test_lifecycle.py | 20 +- 15 files changed, 1256 insertions(+), 168 deletions(-) create mode 100644 litellm-rust/crates/core/src/call_lifecycle/registration.rs create mode 100644 litellm-rust/crates/python-bridge/src/lifecycle/setup.rs create mode 100644 litellm/rust_bridge/setup.py create mode 100644 tests/test_litellm/rust_bridge/test_setup.py diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 992b29ad63e..f5625ac6537 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -3,6 +3,7 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH}; pub mod callbacks; pub mod host; +pub mod registration; pub mod types; pub use callbacks::{ diff --git a/litellm-rust/crates/core/src/call_lifecycle/registration.rs b/litellm-rust/crates/core/src/call_lifecycle/registration.rs new file mode 100644 index 00000000000..eac0526c9fa --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/registration.rs @@ -0,0 +1,374 @@ +use std::collections::HashSet; + +use super::callbacks::CallbackId; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum Registry { + Input, + AsyncInput, + Success, + AsyncSuccess, + Failure, + AsyncFailure, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Registration { + Named { known: bool, async_only: bool }, + Object { asynchronous: bool }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Entry { + pub id: CallbackId, + pub registration: Registration, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Candidate { + pub resolved: Option, + pub duplicate_type: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq, Default)] +pub struct RegistrationFacts { + pub candidates: Vec, + pub input: Vec, + pub success: Vec, + pub failure: Vec, + pub async_success: Vec, + pub async_failure: Vec, + pub bootstrap_pending: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NamedEvent { + Success, + Failure, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RegistryMutation { + Append(Registry, CallbackId), + Remove(Registry, CallbackId), + ExpandNamed(NamedEvent, CallbackId), + Bootstrap, +} + +struct Planner { + input: Vec, + success: Vec, + failure: Vec, + async_success: HashSet, + async_failure: HashSet, + mutations: Vec, +} + +impl Planner { + fn contains(&self, registry: Registry, id: CallbackId) -> bool { + match registry { + Registry::Input => self.input.iter().any(|entry| entry.id == id), + Registry::Success => self.success.iter().any(|entry| entry.id == id), + Registry::Failure => self.failure.iter().any(|entry| entry.id == id), + Registry::AsyncSuccess => self.async_success.contains(&id), + Registry::AsyncFailure => self.async_failure.contains(&id), + Registry::AsyncInput => false, + } + } + + fn append(&mut self, registry: Registry, entry: Entry) { + if self.contains(registry, entry.id) { + return; + } + match registry { + Registry::Input => self.input.push(entry), + Registry::Success => self.success.push(entry), + Registry::Failure => self.failure.push(entry), + Registry::AsyncSuccess => { + self.async_success.insert(entry.id); + } + Registry::AsyncFailure => { + self.async_failure.insert(entry.id); + } + Registry::AsyncInput => {} + } + self.mutations + .push(RegistryMutation::Append(registry, entry.id)); + } + + fn record(&mut self, mutation: RegistryMutation) { + self.mutations.push(mutation); + } +} + +fn is_asynchronous(registration: Registration) -> bool { + matches!(registration, Registration::Object { asynchronous: true }) +} + +pub fn plan_registration(facts: &RegistrationFacts) -> Vec { + let mut planner = Planner { + input: facts.input.clone(), + success: facts.success.clone(), + failure: facts.failure.clone(), + async_success: facts.async_success.iter().copied().collect(), + async_failure: facts.async_failure.iter().copied().collect(), + mutations: Vec::new(), + }; + + for candidate in &facts.candidates { + let Some(entry) = candidate.resolved else { + continue; + }; + if candidate.duplicate_type { + continue; + } + planner.append(Registry::Input, entry); + if !is_asynchronous(entry.registration) { + planner.append(Registry::Success, entry); + planner.append(Registry::Failure, entry); + } + planner.append(Registry::AsyncSuccess, entry); + planner.append(Registry::AsyncFailure, entry); + } + + if facts.bootstrap_pending + && !(planner.input.is_empty() && planner.success.is_empty() && planner.failure.is_empty()) + { + planner.record(RegistryMutation::Bootstrap); + } + + let input = planner.input.clone(); + for entry in input + .iter() + .filter(|entry| is_asynchronous(entry.registration)) + { + planner.record(RegistryMutation::Append(Registry::AsyncInput, entry.id)); + planner.record(RegistryMutation::Remove(Registry::Input, entry.id)); + } + + let success = planner.success.clone(); + for entry in &success { + match entry.registration { + Registration::Object { asynchronous: true } + | Registration::Named { + async_only: true, .. + } => { + planner.append(Registry::AsyncSuccess, *entry); + planner.record(RegistryMutation::Remove(Registry::Success, entry.id)); + } + Registration::Named { known: true, .. } => { + planner.record(RegistryMutation::ExpandNamed(NamedEvent::Success, entry.id)); + } + _ => {} + } + } + + let failure = planner.failure.clone(); + for entry in &failure { + match entry.registration { + Registration::Object { asynchronous: true } => { + planner.append(Registry::AsyncFailure, *entry); + planner.record(RegistryMutation::Remove(Registry::Failure, entry.id)); + } + Registration::Named { known: true, .. } => { + planner.record(RegistryMutation::ExpandNamed(NamedEvent::Failure, entry.id)); + } + _ => {} + } + } + + planner.mutations +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DynamicSuccessSlot { + Sync, + Async, +} + +pub fn classify_dynamic_success(entry: Entry, named_async: bool) -> DynamicSuccessSlot { + match entry.registration { + Registration::Object { asynchronous: true } => DynamicSuccessSlot::Async, + Registration::Named { .. } if named_async => DynamicSuccessSlot::Async, + _ => DynamicSuccessSlot::Sync, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn object(id: u64, asynchronous: bool) -> Entry { + Entry { + id: CallbackId(id), + registration: Registration::Object { asynchronous }, + } + } + + fn named(id: u64, known: bool, async_only: bool) -> Entry { + Entry { + id: CallbackId(id), + registration: Registration::Named { known, async_only }, + } + } + + fn candidate(entry: Entry) -> Candidate { + Candidate { + resolved: Some(entry), + duplicate_type: false, + } + } + + #[test] + fn sync_callback_in_callbacks_registers_in_every_list_once() { + let facts = RegistrationFacts { + candidates: vec![candidate(object(1, false)), candidate(object(1, false))], + ..RegistrationFacts::default() + }; + assert_eq!( + plan_registration(&facts), + [ + RegistryMutation::Append(Registry::Input, CallbackId(1)), + RegistryMutation::Append(Registry::Success, CallbackId(1)), + RegistryMutation::Append(Registry::Failure, CallbackId(1)), + RegistryMutation::Append(Registry::AsyncSuccess, CallbackId(1)), + RegistryMutation::Append(Registry::AsyncFailure, CallbackId(1)), + ] + ); + } + + #[test] + fn async_callable_in_callbacks_skips_sync_lists_and_moves_out_of_input() { + let facts = RegistrationFacts { + candidates: vec![candidate(object(2, true))], + ..RegistrationFacts::default() + }; + assert_eq!( + plan_registration(&facts), + [ + RegistryMutation::Append(Registry::Input, CallbackId(2)), + RegistryMutation::Append(Registry::AsyncSuccess, CallbackId(2)), + RegistryMutation::Append(Registry::AsyncFailure, CallbackId(2)), + RegistryMutation::Append(Registry::AsyncInput, CallbackId(2)), + RegistryMutation::Remove(Registry::Input, CallbackId(2)), + ] + ); + } + + #[test] + fn unresolved_and_duplicate_type_named_candidates_are_skipped() { + let facts = RegistrationFacts { + candidates: vec![ + Candidate { + resolved: None, + duplicate_type: false, + }, + Candidate { + resolved: Some(object(3, false)), + duplicate_type: true, + }, + ], + ..RegistrationFacts::default() + }; + assert!(plan_registration(&facts).is_empty()); + } + + #[test] + fn already_registered_callbacks_are_not_appended_again() { + let facts = RegistrationFacts { + candidates: vec![candidate(object(1, false))], + input: vec![object(1, false)], + success: vec![object(1, false)], + failure: vec![object(1, false)], + async_success: vec![CallbackId(1)], + async_failure: vec![CallbackId(1)], + ..RegistrationFacts::default() + }; + assert!(plan_registration(&facts).is_empty()); + } + + #[test] + fn bootstrap_runs_once_when_any_public_list_is_populated() { + let empty = RegistrationFacts { + bootstrap_pending: true, + ..RegistrationFacts::default() + }; + assert!(plan_registration(&empty).is_empty()); + let populated = RegistrationFacts { + bootstrap_pending: true, + candidates: vec![candidate(object(1, false))], + ..RegistrationFacts::default() + }; + assert!(plan_registration(&populated).contains(&RegistryMutation::Bootstrap)); + let already = RegistrationFacts { + bootstrap_pending: false, + success: vec![object(1, false)], + ..RegistrationFacts::default() + }; + assert!(!plan_registration(&already).contains(&RegistryMutation::Bootstrap)); + } + + #[test] + fn success_safety_net_moves_async_and_async_only_names_and_expands_known_names() { + let facts = RegistrationFacts { + success: vec![ + object(1, true), + named(2, false, true), + named(3, true, false), + named(4, false, false), + object(5, false), + ], + ..RegistrationFacts::default() + }; + assert_eq!( + plan_registration(&facts), + [ + RegistryMutation::Append(Registry::AsyncSuccess, CallbackId(1)), + RegistryMutation::Remove(Registry::Success, CallbackId(1)), + RegistryMutation::Append(Registry::AsyncSuccess, CallbackId(2)), + RegistryMutation::Remove(Registry::Success, CallbackId(2)), + RegistryMutation::ExpandNamed(NamedEvent::Success, CallbackId(3)), + ] + ); + } + + #[test] + fn failure_safety_net_ignores_async_only_names() { + let facts = RegistrationFacts { + failure: vec![ + object(1, true), + named(2, false, true), + named(3, true, false), + ], + async_failure: vec![CallbackId(1)], + ..RegistrationFacts::default() + }; + assert_eq!( + plan_registration(&facts), + [ + RegistryMutation::Remove(Registry::Failure, CallbackId(1)), + RegistryMutation::ExpandNamed(NamedEvent::Failure, CallbackId(3)), + ] + ); + } + + #[test] + fn dynamic_success_split_follows_async_callables_and_selected_names() { + assert_eq!( + classify_dynamic_success(object(1, true), false), + DynamicSuccessSlot::Async + ); + assert_eq!( + classify_dynamic_success(object(1, false), false), + DynamicSuccessSlot::Sync + ); + assert_eq!( + classify_dynamic_success(named(2, true, false), true), + DynamicSuccessSlot::Async + ); + assert_eq!( + classify_dynamic_success(named(2, true, true), false), + DynamicSuccessSlot::Sync + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index cc153a89b8f..0c2b4c9e731 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,6 +1,6 @@ use litellm_python_interop::release_count; use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyTuple}; #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { @@ -9,6 +9,27 @@ fn gil_stats(py: Python<'_>) -> PyResult> { Ok(stats.into_any().unbind()) } +#[pyfunction] +fn _debug_setup( + py: Python<'_>, + call_type: String, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + start: Py, + asynchronous: bool, +) -> PyResult<(Py, Py)> { + let leaked: &'static str = Box::leak(call_type.into_boxed_str()); + let result = crate::lifecycle::debug_setup( + py, + leaked, + &args.unbind(), + &kwargs.unbind(), + &start, + asynchronous, + )?; + Ok(result) +} + #[cfg(feature = "panic-test")] #[pyfunction] fn _panic_for_test() { @@ -17,6 +38,7 @@ fn _panic_for_test() { pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + module.add_function(wrap_pyfunction!(_debug_setup, module)?)?; #[cfg(feature = "panic-test")] module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; Ok(()) diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs index 23ac0283646..6f95444fc8c 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -1,7 +1,7 @@ use pyo3::exceptions::PyBaseException; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; +use pyo3::types::PyDict; #[derive(FromPyObject)] pub(crate) struct PythonLogger(Py); @@ -126,32 +126,6 @@ impl PythonLogger { } } -pub(super) struct SetupResult<'py>(Bound<'py, PyAny>); - -impl SetupResult<'_> { - pub(super) fn logger(&self) -> PyResult { - self.0.getattr("logger")?.extract() - } - - pub(super) fn kwargs(&self) -> PyResult> { - Ok(self.0.getattr("kwargs")?.extract()?) - } -} - -pub(super) fn setup<'py>( - py: Python<'py>, - call_type: &str, - args: &Py, - kwargs: &Py, - start: &Py, - asynchronous: bool, -) -> PyResult> { - py.import("litellm.rust_bridge.lifecycle")? - .getattr("setup")? - .call1((call_type, args, kwargs, start, asynchronous)) - .map(SetupResult) -} - pub(super) fn finalize( py: Python<'_>, response: &Option>, @@ -215,71 +189,6 @@ impl DeploymentHooks { #[cfg(test)] mod tests { use super::*; - use pyo3::exceptions::PyTypeError; - - #[test] - fn setup_fields_are_checked_in_order_without_eager_logger_method_reads() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - pyo3::ffi::c_str!( - r#" -reads = [] -class Logger: - def __getattribute__(self, name): - reads.append(name) - raise AssertionError('logger methods must remain lazy') -logger = Logger() -class Setup: - @property - def logger(self): - reads.append('logger') - return logger - @property - def kwargs(self): - reads.append('kwargs') - return [] -result = Setup() -"# - ), - Some(&locals), - Some(&locals), - ) - .unwrap(); - let result = SetupResult(locals.get_item("result").unwrap().unwrap()); - let logger = result.logger().unwrap(); - assert!( - logger - .object(py) - .is(locals.get_item("logger").unwrap().unwrap()) - ); - assert_eq!( - locals - .get_item("reads") - .unwrap() - .unwrap() - .extract::>() - .unwrap(), - ["logger"] - ); - assert!( - result - .kwargs() - .unwrap_err() - .is_instance_of::(py) - ); - assert_eq!( - locals - .get_item("reads") - .unwrap() - .unwrap() - .extract::>() - .unwrap(), - ["logger", "kwargs"] - ); - }); - } #[test] fn logger_resolves_each_callback_at_invocation_and_preserves_arguments() { diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index 3a9e305c9f6..5aede7baa13 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -19,6 +19,7 @@ mod arguments; mod bindings; mod handle; mod preparation; +mod setup; pub(crate) use arguments::{BoundArguments, Signature}; use bindings::DeploymentHooks; @@ -93,6 +94,18 @@ pub(crate) fn run_call( } } +pub(crate) fn debug_setup( + py: Python<'_>, + call_type: &'static str, + args: &Py, + kwargs: &Py, + start: &Py, + asynchronous: bool, +) -> PyResult<(Py, Py)> { + let result = setup::setup(py, call_type, args, kwargs, start, asynchronous)?; + Ok((result.logger.object(py).clone().unbind(), result.kwargs)) +} + pub(crate) fn missing_state() -> PyErr { pyo3::exceptions::PyRuntimeError::new_err("missing native call state") } @@ -162,7 +175,11 @@ impl PythonLifecycle { HostFailure::Cancelled(native) }; let state = self.route.state_mut(); - state.retain_first_error(py, error, cancelled && phase != Some(HostPhase::DeploymentFailure)); + state.retain_first_error( + py, + error, + cancelled && phase != Some(HostPhase::DeploymentFailure), + ); let _ = state.finish(py); failure } @@ -385,7 +402,7 @@ impl PythonCallState { pub fn setup(&mut self, py: Python<'_>) -> PyResult<()> { self.start = now(py)?; self.internal = bindings::is_internal_call(py)?; - let result = bindings::setup( + let result = setup::setup( py, self.call_type, &self.args, @@ -393,8 +410,8 @@ impl PythonCallState { &self.start, self.asynchronous, )?; - self.logger = Some(result.logger()?); - self.kwargs = result.kwargs()?; + self.logger = Some(result.logger); + self.kwargs = result.kwargs; Ok(()) } diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/setup.rs b/litellm-rust/crates/python-bridge/src/lifecycle/setup.rs new file mode 100644 index 00000000000..db90a001c52 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/setup.rs @@ -0,0 +1,334 @@ +use litellm_core::call_lifecycle::CallbackId; +use litellm_core::call_lifecycle::registration::{ + Candidate, DynamicSuccessSlot, Entry, NamedEvent, Registration, RegistrationFacts, Registry, + RegistryMutation, classify_dynamic_success, plan_registration, +}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList, PyString, PyTuple}; + +use super::bindings::PythonLogger; + +const SETUP_MODULE: &str = "litellm.rust_bridge.setup"; + +struct Targets<'py> { + objects: Vec>, +} + +impl<'py> Targets<'py> { + fn new() -> Self { + Self { + objects: Vec::new(), + } + } + + fn intern(&mut self, object: &Bound<'py, PyAny>) -> CallbackId { + if let Some(index) = self + .objects + .iter() + .position(|existing| existing.is(object) || existing.eq(object).unwrap_or(false)) + { + return CallbackId(index as u64); + } + self.objects.push(object.clone()); + CallbackId((self.objects.len() - 1) as u64) + } + + fn get(&self, id: CallbackId) -> PyResult<&Bound<'py, PyAny>> { + self.objects + .get(id.0 as usize) + .ok_or_else(super::missing_state) + } +} + +fn registration<'py>( + setup: &Bound<'py, PyModule>, + object: &Bound<'py, PyAny>, +) -> PyResult { + if let Ok(name) = object.cast::() { + let name = name.to_str()?; + let known = setup.getattr("is_known_name")?.call1((name,))?.extract()?; + return Ok(Registration::Named { + known, + async_only: matches!(name, "dynamodb" | "openmeter"), + }); + } + let asynchronous = setup + .getattr("is_async_callable")? + .call1((object,))? + .extract()?; + Ok(Registration::Object { asynchronous }) +} + +fn entries<'py>( + setup: &Bound<'py, PyModule>, + targets: &mut Targets<'py>, + list: &Bound<'py, PyAny>, +) -> PyResult> { + list.try_iter()? + .map(|object| { + let object = object?; + Ok(Entry { + id: targets.intern(&object), + registration: registration(setup, &object)?, + }) + }) + .collect() +} + +fn registry_name(registry: Registry) -> &'static str { + match registry { + Registry::Input => "input", + Registry::AsyncInput => "async_input", + Registry::Success => "success", + Registry::AsyncSuccess => "async_success", + Registry::Failure => "failure", + Registry::AsyncFailure => "async_failure", + } +} + +fn read_registry<'py>(setup: &Bound<'py, PyModule>, name: &str) -> PyResult> { + setup.getattr("registry")?.call1((name,)) +} + +fn read_candidates<'py>( + setup: &Bound<'py, PyModule>, + targets: &mut Targets<'py>, + dynamic: Option>, +) -> PyResult> { + let mut candidates = Vec::new(); + let global = read_registry(setup, "callbacks")?; + let sources = std::iter::once(global).chain(dynamic); + for source in sources { + for object in source.try_iter()? { + let object = object?; + let candidate = if object.is_instance_of::() { + let resolved = setup + .getattr("resolve_named_integration")? + .call1((&object,))?; + if resolved.is_none() { + Candidate { + resolved: None, + duplicate_type: false, + } + } else { + let duplicate_type = setup + .getattr("async_success_registry_has_type")? + .call1((&resolved,))? + .extract()?; + Candidate { + resolved: Some(Entry { + id: targets.intern(&resolved), + registration: registration(setup, &resolved)?, + }), + duplicate_type, + } + } + } else { + Candidate { + resolved: Some(Entry { + id: targets.intern(&object), + registration: registration(setup, &object)?, + }), + duplicate_type: false, + } + }; + candidates.push(candidate); + } + } + Ok(candidates) +} + +fn apply<'py>( + setup: &Bound<'py, PyModule>, + targets: &Targets<'py>, + mutations: &[RegistryMutation], + function_id: Option<&Bound<'py, PyAny>>, +) -> PyResult<()> { + for mutation in mutations { + match mutation { + RegistryMutation::Append(registry, id) => { + setup + .getattr("append_registry")? + .call1((registry_name(*registry), targets.get(*id)?))?; + } + RegistryMutation::Remove(registry, id) => { + setup + .getattr("remove_registry")? + .call1((registry_name(*registry), targets.get(*id)?))?; + } + RegistryMutation::ExpandNamed(event, id) => { + let event = match event { + NamedEvent::Success => "success", + NamedEvent::Failure => "failure", + }; + setup + .getattr("expand_named")? + .call1((targets.get(*id)?, event))?; + } + RegistryMutation::Bootstrap => { + setup.getattr("bootstrap")?.call1((function_id,))?; + } + } + } + Ok(()) +} + +struct DynamicLists<'py> { + success: Option>, + async_success: Option>, + failure: Option>, +} + +fn split_dynamic<'py>( + py: Python<'py>, + setup: &Bound<'py, PyModule>, + targets: &mut Targets<'py>, + kwargs: &Bound<'py, PyDict>, +) -> PyResult> { + let success = match kwargs.get_item("success_callback")? { + Some(value) if value.is_instance_of::() => { + let list = value.cast_into::()?; + let sync = PyList::empty(py); + let asynchronous = PyList::empty(py); + for object in list.iter() { + let entry = Entry { + id: targets.intern(&object), + registration: registration(setup, &object)?, + }; + let named_async = object + .cast::() + .ok() + .and_then(|name| { + name.to_str() + .ok() + .map(|name| matches!(name, "dynamodb" | "s3")) + }) + .unwrap_or(false); + match classify_dynamic_success(entry, named_async) { + DynamicSuccessSlot::Sync => sync.append(&object)?, + DynamicSuccessSlot::Async => asynchronous.append(&object)?, + } + } + kwargs.del_item("success_callback")?; + Some((sync, (!asynchronous.is_empty()).then_some(asynchronous))) + } + _ => None, + }; + let failure = match kwargs.get_item("failure_callback")? { + Some(value) if value.is_instance_of::() => { + kwargs.del_item("failure_callback")?; + Some(value.cast_into::()?) + } + _ => None, + }; + let (success, async_success) = match success { + Some((sync, asynchronous)) => (Some(sync), asynchronous), + None => (None, None), + }; + Ok(DynamicLists { + success, + async_success, + failure, + }) +} + +pub(super) struct Setup { + pub logger: PythonLogger, + pub kwargs: Py, +} + +pub(super) fn setup( + py: Python<'_>, + call_type: &str, + args: &Py, + kwargs: &Py, + start: &Py, + asynchronous: bool, +) -> PyResult { + let setup = py.import(SETUP_MODULE)?; + let kwargs = kwargs.bind(py).copy()?; + if !kwargs.contains("litellm_call_id")? { + let call_id = py.import("uuid")?.call_method0("uuid4")?.str()?; + kwargs.set_item("litellm_call_id", call_id)?; + } + if let Some(supplied) = kwargs.get_item("litellm_logging_obj")? { + let logging_class = py + .import("litellm.litellm_core_utils.litellm_logging")? + .getattr("Logging")?; + if supplied.is_instance(&logging_class)? { + return Ok(Setup { + logger: supplied.extract()?, + kwargs: kwargs.unbind(), + }); + } + } + + setup.getattr("prepare_environment")?.call0()?; + let guardrails = setup.getattr("applied_guardrails")?.call1((&kwargs,))?; + let function_id = kwargs.get_item("id")?; + + let mut targets = Targets::new(); + let dynamic = match kwargs.get_item("callbacks")? { + Some(value) => { + kwargs.del_item("callbacks")?; + (!value.is_none()).then_some(value) + } + None => None, + }; + let candidates = read_candidates(&setup, &mut targets, dynamic)?; + let facts = RegistrationFacts { + candidates, + input: entries(&setup, &mut targets, &read_registry(&setup, "input")?)?, + success: entries(&setup, &mut targets, &read_registry(&setup, "success")?)?, + failure: entries(&setup, &mut targets, &read_registry(&setup, "failure")?)?, + async_success: entries( + &setup, + &mut targets, + &read_registry(&setup, "async_success")?, + )? + .into_iter() + .map(|entry| entry.id) + .collect(), + async_failure: entries( + &setup, + &mut targets, + &read_registry(&setup, "async_failure")?, + )? + .into_iter() + .map(|entry| entry.id) + .collect(), + bootstrap_pending: setup.getattr("bootstrap_pending")?.call0()?.extract()?, + }; + apply( + &setup, + &targets, + &plan_registration(&facts), + function_id.as_ref(), + )?; + + let dynamic = split_dynamic(py, &setup, &mut targets, &kwargs)?; + setup.getattr("breadcrumb")?.call1((&kwargs,))?; + if let Some(logger_fn) = kwargs.get_item("logger_fn")? { + setup.getattr("logger_fn")?.call1((logger_fn,))?; + } + let model = match args.bind(py).get_item(0) { + Ok(model) => Some(model), + Err(_) => kwargs.get_item("model")?, + }; + let build = setup.getattr("build_logging")?; + let build_kwargs = PyDict::new(py); + build_kwargs.set_item("call_type", call_type)?; + build_kwargs.set_item("model", model)?; + build_kwargs.set_item("kwargs", &kwargs)?; + build_kwargs.set_item("start_time", start)?; + build_kwargs.set_item("asynchronous", asynchronous)?; + build_kwargs.set_item("dynamic_success", dynamic.success)?; + build_kwargs.set_item("dynamic_async_success", dynamic.async_success)?; + build_kwargs.set_item("dynamic_failure", dynamic.failure)?; + build_kwargs.set_item("guardrails", guardrails)?; + let logger = build.call((), Some(&build_kwargs))?; + Ok(Setup { + logger: logger.extract()?, + kwargs: kwargs.unbind(), + }) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index de70f7ce6f0..deb0a4e7f4c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -38,7 +38,13 @@ pub(super) fn pre_call( ) -> PyResult<()> { py.import("litellm.rust_bridge.ocr")? .getattr("pre_call")? - .call1((logger.object(py), request.api_key.as_deref(), &payload.body, &payload.headers, &request.url))?; + .call1(( + logger.object(py), + request.api_key.as_deref(), + &payload.body, + &payload.headers, + &request.url, + ))?; Ok(()) } @@ -50,7 +56,12 @@ pub(super) fn post_call( ) -> PyResult<()> { py.import("litellm.rust_bridge.ocr")? .getattr("post_call")? - .call1((logger.object(py), to_py(py, original_response)?, &payload.body, &payload.headers))?; + .call1(( + logger.object(py), + to_py(py, original_response)?, + &payload.body, + &payload.headers, + ))?; Ok(()) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 5ccb65780df..72971bce2d4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -40,17 +40,28 @@ pub(super) struct PythonPayload { impl PythonPayload { fn from_request(py: Python<'_>, request: &OcrDuringCallRequest) -> PyResult { - let body = to_py(py, &request.body)?.into_bound(py).cast_into::()?; + let body = to_py(py, &request.body)? + .into_bound(py) + .cast_into::()?; let headers = PyDict::new(py); for (name, value) in &request.headers { headers.set_item(name, value)?; } - Ok(Self { body: body.unbind(), headers: headers.unbind() }) + Ok(Self { + body: body.unbind(), + headers: headers.unbind(), + }) } - fn write_back(&self, py: Python<'_>, mut request: OcrDuringCallRequest) -> PyResult { + fn write_back( + &self, + py: Python<'_>, + mut request: OcrDuringCallRequest, + ) -> PyResult { request.body = from_py(self.body.bind(py))?; - request.headers = self.headers.bind(py) + request.headers = self + .headers + .bind(py) .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; @@ -81,7 +92,9 @@ impl PythonOcrHost { } fn project(&mut self, py: Python<'_>) -> PyResult { - let arguments = self.signature.bind(self.state.args.bind(py), self.state.kwargs.bind(py))?; + let arguments = self + .signature + .bind(self.state.args.bind(py), self.state.kwargs.bind(py))?; let Projection { native, retained } = project(py, &arguments)?; let host_token_provider = retained.azure_ad_token_provider.is_some(); self.retained = Some(retained); @@ -126,7 +139,11 @@ impl PythonOcrHost { py: Python<'_>, request: OcrPostCallRequest, ) -> PyResult { - let payload = self.retained()?.payload.as_ref().ok_or_else(missing_state)?; + let payload = self + .retained()? + .payload + .as_ref() + .ok_or_else(missing_state)?; callbacks::post_call( py, self.state.logger()?, @@ -202,9 +219,7 @@ impl PythonRoute for PythonOcrHost { OcrHostOperation::AcquireAzureAdToken => { OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?)) } - OcrHostOperation::PreCall(request) => { - OcrHostResult::PreCall(Ok(request)) - } + OcrHostOperation::PreCall(request) => OcrHostResult::PreCall(Ok(request)), OcrHostOperation::DuringCall(request) => { OcrHostResult::DuringCall(Ok(self.during_call(py, request)?)) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index cbd32792843..d1fd977c5b6 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -62,13 +62,16 @@ fn call( ..OcrAdmission::all() }, ))?; - let host = PythonOcrHost::new(PythonCallState::new( - py, - args.unbind(), - kwargs.copy()?.unbind(), - asynchronous, - signature.name, - )?, signature); + let host = PythonOcrHost::new( + PythonCallState::new( + py, + args.unbind(), + kwargs.copy()?.unbind(), + asynchronous, + signature.name, + )?, + signature, + ); run_call(py, call, host) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 57cddac48cb..5157f7683ed 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -20,14 +20,18 @@ use crate::marshal::{BoundRouteInputs, Projection}; /// `optional_params`. const BOUND_FIELDS: &[&str] = &["model", "document", "timeout", "input_sources"]; -fn project_document(document: &Bound<'_, PyAny>) -> PyResult> { +fn project_document( + document: &Bound<'_, PyAny>, +) -> PyResult> { let kind: String = document.get_item("type")?.extract()?; if kind != "file" { let value: serde_json::Value = from_py(document)?; - return Ok(OcrDocument::try_from(value).map(|document| FileDocumentInput { - input: document.into(), - reader: None, - })); + return Ok( + OcrDocument::try_from(value).map(|document| FileDocumentInput { + input: document.into(), + reader: None, + }), + ); } document.extract().map(Ok) } @@ -154,10 +158,7 @@ mod tests { let document = py .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None) .unwrap(); - let error = project_document(&document) - .unwrap() - .err() - .unwrap(); + let error = project_document(&document).unwrap().err().unwrap(); assert!(error.to_string().contains("document")); }); } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7d7e4216c40..bb114336208 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,3 +1,4 @@ +import datetime from asyncio import Future from collections.abc import Coroutine, Mapping @@ -111,3 +112,10 @@ class TokenCounter: def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... def gil_stats() -> dict[str, int]: ... +def _debug_setup( + call_type: str, + args: tuple[object, ...], + kwargs: dict[str, object], + start: datetime.datetime, + asynchronous: bool, +) -> tuple[object, dict[str, object]]: ... diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index 35429ea7ccc..432ece7e093 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,7 +1,6 @@ from __future__ import annotations import datetime -import uuid from collections.abc import Awaitable, Mapping from dataclasses import dataclass from typing import ( @@ -64,35 +63,6 @@ class MetadataUpdater(Protocol): ) -> None: ... -@dataclass(frozen=True, slots=True) -class CallSetup: - logger: Logging - kwargs: dict[str, object] - - -def setup( - call_type: str, - args: tuple[object, ...], - kwargs: Mapping[str, object], - start_time: datetime.datetime, - asynchronous: bool, -) -> CallSetup: - from litellm import utils - from litellm.litellm_core_utils.litellm_logging import Logging - - arguments: Final = { # mutable-ok: function_setup consumes an owned kwargs dict - "litellm_call_id": str(uuid.uuid4()), - **kwargs, - } - supplied: Final = arguments.get("litellm_logging_obj") - if isinstance(supplied, Logging): - return CallSetup(supplied, arguments) - logger, prepared = utils.function_setup( - call_type, utils.Rules(), start_time, *args, is_async_call=asynchronous, **arguments - ) - return CallSetup(logger, prepared) - - def check_limits(kwargs: Mapping[str, object]) -> None: import litellm from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit diff --git a/litellm/rust_bridge/setup.py b/litellm/rust_bridge/setup.py new file mode 100644 index 00000000000..820d39c7b3e --- /dev/null +++ b/litellm/rust_bridge/setup.py @@ -0,0 +1,262 @@ +from __future__ import annotations + +import datetime +from collections.abc import Callable, Mapping, MutableSequence, Sequence +from typing import ( # noqa: TID251 # narrows legacy untyped registries at the boundary + TYPE_CHECKING, + Final, + Literal, + Protocol, + cast, +) + +from litellm.integrations.custom_logger import CustomLogger + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + +CallbackTarget = str | Callable[..., object] | CustomLogger +RegistryName = Literal["input", "async_input", "success", "async_success", "failure", "async_failure", "callbacks"] + +_REGISTRY_ATTRIBUTES: Final[Mapping[RegistryName, str]] = { + "input": "input_callback", + "async_input": "_async_input_callback", + "success": "success_callback", + "async_success": "_async_success_callback", + "failure": "failure_callback", + "async_failure": "_async_failure_callback", + "callbacks": "callbacks", +} + + +class _CallbackManager(Protocol): + def add_litellm_success_callback(self, callback: CallbackTarget) -> None: ... + + def add_litellm_failure_callback(self, callback: CallbackTarget) -> None: ... + + def add_litellm_async_success_callback(self, callback: CallbackTarget) -> None: ... + + def add_litellm_async_failure_callback(self, callback: CallbackTarget) -> None: ... + + +class _LoggingFactory(Protocol): + def __call__( + self, + *, + model: str | None, + messages: object, + stream: bool, + litellm_call_id: str, + litellm_trace_id: str | None, + function_id: str, + call_type: str, + start_time: datetime.datetime, + dynamic_success_callbacks: list[CallbackTarget] | None, + dynamic_failure_callbacks: list[CallbackTarget] | None, + dynamic_async_success_callbacks: list[CallbackTarget] | None, + dynamic_async_failure_callbacks: list[CallbackTarget] | None, + kwargs: dict[str, object], + applied_guardrails: list[str], + supports_correlation_logging: bool, + ) -> Logging: ... + + +class _EnvironmentUpdater(Protocol): + def __call__( + self, + *, + model: str | None, + user: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + stream_options: object, + ) -> None: ... + + +def registry(name: RegistryName) -> MutableSequence[CallbackTarget]: + import litellm + + return cast( + MutableSequence[CallbackTarget], getattr(litellm, _REGISTRY_ATTRIBUTES[name]) + ) # cast-ok: legacy module-level lists are untyped + + +def is_async_callable(callback: object) -> bool: + from litellm.litellm_core_utils.cached_imports import get_coroutine_checker + + return get_coroutine_checker().is_async_callable(callback) + + +def is_known_name(callback: str) -> bool: + import litellm + + known: Final = cast(Sequence[str], litellm._known_custom_logger_compatible_callbacks) # pyright: ignore[reportPrivateUsage] # cast-ok: registry list has no public typed accessor + return callback in known + + +def resolve_named_integration(callback: str) -> CustomLogger | None: + from litellm.litellm_core_utils import litellm_logging + + resolve: Final = cast( # cast-ok: legacy factory is untyped at its definition + Callable[..., CustomLogger | None], + litellm_logging._init_custom_logger_compatible_class, # pyright: ignore[reportPrivateUsage] # legacy factory + ) + return resolve(callback, internal_usage_cache=None, llm_router=None) + + +def async_success_registry_has_type(callback: object) -> bool: + return any(type(existing) is type(callback) for existing in registry("async_success")) + + +def bootstrap_pending() -> bool: + from litellm import utils + + return not utils.callback_list + + +def bootstrap(function_id: str | None) -> None: + from litellm import utils + from litellm.litellm_core_utils import cached_imports + + combined: Final = list({*registry("input"), *registry("success"), *registry("failure")}) + utils.callback_list = cast( + list[str], combined + ) # rebind-ok: legacy module global consumed by set_callbacks # cast-ok: legacy list annotation is narrower than its contents + set_callbacks: Final = cast(Callable[..., None], cached_imports.get_set_callbacks()) # pyright: ignore[reportUnknownMemberType] # cast-ok: cached import is untyped + set_callbacks(callback_list=combined, function_id=function_id) + + +def expand_named(callback: str, event: Literal["success", "failure"]) -> None: + from litellm import utils + + utils._add_custom_logger_callback_to_specific_event(callback, event) # pyright: ignore[reportPrivateUsage] # legacy expansion helper + + +def append_registry(name: RegistryName, callback: CallbackTarget) -> None: + import litellm + + manager: Final = cast(_CallbackManager, litellm.logging_callback_manager) # cast-ok: legacy manager is untyped + match name: + case "input" | "async_input": + registry(name).append(callback) + case "success": + manager.add_litellm_success_callback(callback) + case "async_success": + manager.add_litellm_async_success_callback(callback) + case "failure": + manager.add_litellm_failure_callback(callback) + case "async_failure": + manager.add_litellm_async_failure_callback(callback) + case "callbacks": + raise KeyError(name) + + +def remove_registry(name: RegistryName, callback: CallbackTarget) -> None: + target: Final = registry(name) + for index in range(len(target) - 1, -1, -1): + if target[index] is callback or target[index] == callback: + del target[index] + return + + +def logger_fn(callback: object) -> None: + from litellm import utils + + utils.user_logger_fn = callback # rebind-ok: legacy module global read by pre_call + + +def breadcrumb(kwargs: Mapping[str, object]) -> None: + from litellm import utils + + add_breadcrumb: Final = cast( + Callable[..., None] | None, utils.add_breadcrumb + ) # cast-ok: legacy sentry hook is untyped + if add_breadcrumb is None: + return + import litellm + from litellm.litellm_core_utils import core_helpers + + deep_copy: Final = cast( # cast-ok: legacy helper is untyped + Callable[[dict[str, object]], dict[str, object]], + core_helpers.safe_deep_copy, # pyright: ignore[reportUnknownMemberType] # legacy helper + ) + try: + copied: dict[str, object] = deep_copy(dict(kwargs)) + except Exception: # noqa: BLE001 # legacy breadcrumb falls back to the live mapping + copied = dict(kwargs) + hidden: Final = frozenset(("messages", "input", "prompt")) if litellm.turn_off_message_logging else frozenset[str]() + details: Final = {key: value for key, value in copied.items() if key not in hidden} + add_breadcrumb(category="litellm.llm_call", message=f"Keyword Args: {details}", level="info") + + +def prepare_environment() -> None: + from litellm import utils + + utils.custom_llm_setup() + + +def applied_guardrails(kwargs: Mapping[str, object]) -> list[str]: + from litellm.utils import get_applied_guardrails + + return get_applied_guardrails(dict(kwargs)) + + +def build_logging( + *, + call_type: str, + model: str | None, + kwargs: dict[str, object], + start_time: datetime.datetime, + asynchronous: bool, + dynamic_success: Sequence[CallbackTarget] | None, + dynamic_async_success: Sequence[CallbackTarget] | None, + dynamic_failure: Sequence[CallbackTarget] | None, + guardrails: Sequence[str], +) -> Logging: + from litellm.litellm_core_utils.cached_imports import get_litellm_logging_class + + function_id: Final = kwargs.get("id") + metadata: Final = kwargs.get("metadata") + trace_id: Final = kwargs.get("litellm_trace_id") + factory: Final = cast(_LoggingFactory, get_litellm_logging_class()) # cast-ok: legacy constructor is untyped + logger: Final = factory( + model=model, + messages="default-message-value", + stream=False, + litellm_call_id=str(kwargs["litellm_call_id"]), + litellm_trace_id=trace_id if isinstance(trace_id, str) else None, + function_id=function_id if isinstance(function_id, str) else "", + call_type=call_type, + start_time=start_time, + dynamic_success_callbacks=list(dynamic_success) if dynamic_success is not None else None, + dynamic_failure_callbacks=list(dynamic_failure) if dynamic_failure is not None else None, + dynamic_async_success_callbacks=list(dynamic_async_success) if dynamic_async_success is not None else None, + dynamic_async_failure_callbacks=None, + kwargs=kwargs, + applied_guardrails=list(guardrails), + supports_correlation_logging=asynchronous, + ) + litellm_metadata: Final = kwargs.get("litellm_metadata") + litellm_params: Final[dict[str, object]] = { + "api_base": "", + **({"metadata": kwargs["metadata"]} if "metadata" in kwargs else {}), + **( + { + "litellm_metadata": litellm_metadata, + **( + {} if metadata else {"metadata": dict(cast(Mapping[str, object], litellm_metadata))} + ), # cast-ok: isinstance narrows only to dict[Unknown, Unknown] + } + if isinstance(litellm_metadata, dict) + else {} + ), + } + update: Final = cast(_EnvironmentUpdater, logger.update_environment_variables) # cast-ok: legacy method is untyped + update( + model=model, + user="", + optional_params={}, + litellm_params=litellm_params, + stream_options=kwargs.get("stream_options"), + ) + return logger diff --git a/tests/test_litellm/rust_bridge/test_setup.py b/tests/test_litellm/rust_bridge/test_setup.py new file mode 100644 index 00000000000..ecfd68349c8 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_setup.py @@ -0,0 +1,169 @@ +import datetime +from collections.abc import Iterator +from typing import Final + +import pytest + +import litellm +from litellm import utils +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.rust_bridge.loader import native_bridge_available + +pytestmark = pytest.mark.skipif(not native_bridge_available(), reason="requires the Rust extension") + +REGISTRIES: Final = ( + "input_callback", + "_async_input_callback", + "success_callback", + "_async_success_callback", + "failure_callback", + "_async_failure_callback", + "callbacks", +) + + +@pytest.fixture +def clean_registries(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + for name in REGISTRIES: + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(utils, "callback_list", []) + monkeypatch.setattr(utils, "user_logger_fn", None) + monkeypatch.setenv("OPENMETER_API_KEY", "test") + monkeypatch.setenv("OPENMETER_API_ENDPOINT", "http://127.0.0.1:9") + yield + + +class SyncLogger(CustomLogger): + pass + + +class AsyncOnly(CustomLogger): + pass + + +def sync_fn(*args: object, **kwargs: object) -> None: + del args, kwargs + + +async def async_fn(*args: object, **kwargs: object) -> None: + del args, kwargs + + +def snapshot() -> dict[str, list[object]]: + return {name: list(getattr(litellm, name)) for name in REGISTRIES} + + +def run_legacy(kwargs: dict[str, object]) -> tuple[Logging, dict[str, object]]: + logger, prepared = utils.function_setup( + "ocr", + utils.Rules(), + datetime.datetime.now(), + is_async_call=False, + **{"litellm_call_id": "legacy", **kwargs}, + ) + assert isinstance(logger, Logging) + return logger, prepared + + +def run_native(kwargs: dict[str, object]) -> tuple[Logging, dict[str, object]]: + from litellm.rust_bridge import _native + + logger, prepared = _native._debug_setup( + "ocr", (), {"litellm_call_id": "native", **kwargs}, datetime.datetime.now(), False + ) + assert isinstance(logger, Logging) + return logger, prepared + + +@pytest.mark.parametrize( + "globals_before, kwargs", + [ + ({}, {}), + ({"callbacks": [SyncLogger()]}, {}), + ({"callbacks": [async_fn]}, {}), + ({}, {"callbacks": [SyncLogger(), sync_fn]}), + ({"success_callback": [async_fn, "openmeter", sync_fn]}, {}), + ({"failure_callback": [async_fn, sync_fn]}, {}), + ({"input_callback": [async_fn, sync_fn]}, {}), + ({}, {"success_callback": [sync_fn, async_fn, "s3", "dynamodb"], "failure_callback": [sync_fn]}), + ({"callbacks": [SyncLogger()]}, {"callbacks": [SyncLogger()], "success_callback": [async_fn]}), + ], + ids=[ + "empty", + "global-custom-logger", + "global-async-callable", + "dynamic-callbacks", + "success-safety-net", + "failure-safety-net", + "input-safety-net", + "per-call-success-failure-split", + "mixed", + ], +) +def test_native_setup_registry_side_effects_match_function_setup( + clean_registries: None, globals_before: dict[str, list[object]], kwargs: dict[str, object] +) -> None: + for name, values in globals_before.items(): + getattr(litellm, name).extend(values) + legacy_logger, legacy_kwargs = run_legacy( + {key: list(value) if isinstance(value, list) else value for key, value in kwargs.items()} + ) + legacy_snapshot: Final = snapshot() + legacy_bootstrap: Final = list(utils.callback_list or []) + + for name in REGISTRIES: + getattr(litellm, name).clear() + utils.callback_list = [] + for name, values in globals_before.items(): + getattr(litellm, name).extend(values) + native_logger, native_kwargs = run_native( + {key: list(value) if isinstance(value, list) else value for key, value in kwargs.items()} + ) + + assert snapshot() == legacy_snapshot + assert sorted(map(repr, utils.callback_list or [])) == sorted(map(repr, legacy_bootstrap)) + assert set(native_kwargs) - {"litellm_call_id"} == set(legacy_kwargs) - {"litellm_call_id"} + for attribute in ( + "dynamic_success_callbacks", + "dynamic_async_success_callbacks", + "dynamic_failure_callbacks", + "dynamic_async_failure_callbacks", + "call_type", + "stream", + "model", + ): + assert getattr(native_logger, attribute) == getattr(legacy_logger, attribute), attribute + + +def test_native_setup_honours_caller_supplied_logging_object(clean_registries: None) -> None: + class Supplied(Logging): + pass + + supplied: Final = Supplied( + model="mistral/mistral-ocr-latest", + messages=[], + stream=False, + call_type="ocr", + start_time=datetime.datetime.now(), + litellm_call_id="supplied", + function_id="", + ) + litellm.callbacks.append(SyncLogger()) + logger, kwargs = run_native({"litellm_logging_obj": supplied, "callbacks": [SyncLogger()]}) + assert logger is supplied + assert kwargs["callbacks"] is not None + assert litellm.success_callback == [] + + +def test_native_setup_records_logger_fn_and_metadata(clean_registries: None) -> None: + def logger_fn(details: object) -> None: + del details + + logger, kwargs = run_native( + {"model": "mistral/mistral-ocr-latest", "logger_fn": logger_fn, "metadata": {"source": "test"}} + ) + assert utils.user_logger_fn is logger_fn + assert logger.litellm_params["metadata"] == {"source": "test"} + assert logger.model == "mistral/mistral-ocr-latest" + assert kwargs["metadata"] == {"source": "test"} diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 2ca9e77db4f..04cc6eb41fc 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -967,26 +967,18 @@ async def test_terminal_registration_added_during_http_is_observed( @pytest.fixture def created_loggers(monkeypatch: pytest.MonkeyPatch) -> list[Logging]: - from litellm import utils + from litellm.rust_bridge import setup as native_setup - original_setup: Final = utils.function_setup + original_build: Final = native_setup.build_logging loggers: Final[list[Logging]] = [] - def setup( - call_type: str, - rules: utils.Rules, - start: datetime.datetime, - *args: object, - is_async_call: bool = True, - **kwargs: object, - ) -> tuple[Logging, dict[str, object]]: - logger, prepared = original_setup(call_type, rules, start, *args, is_async_call=is_async_call, **kwargs) - assert isinstance(logger, Logging) + def build_logging(**kwargs: object) -> Logging: + logger: Final = original_build(**kwargs) # pyright: ignore[reportArgumentType] # passthrough of the factory signature setattr(logger, "_defer_async_logging", True) loggers.append(logger) - return logger, prepared + return logger - monkeypatch.setattr(utils, "function_setup", setup) + monkeypatch.setattr(native_setup, "build_logging", build_logging) return loggers