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

refactor(python-bridge): classify callbacks natively instead of function_setup

Browse source 
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.
This commit is contained in:
Yujong Lee 2026-09-16 08:30:13 -07:00
parent 54be2dcae1
commit f8190bbe80
15 changed files with 1256 additions and 168 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

374
litellm-rust/crates/core/src/call_lifecycle/registration.rs Normal file
View file

@ -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<Entry>,
pub duplicate_type: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct RegistrationFacts {
pub candidates: Vec<Candidate>,
pub input: Vec<Entry>,
pub success: Vec<Entry>,
pub failure: Vec<Entry>,
pub async_success: Vec<CallbackId>,
pub async_failure: Vec<CallbackId>,
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<Entry>,
success: Vec<Entry>,
failure: Vec<Entry>,
async_success: HashSet<CallbackId>,
async_failure: HashSet<CallbackId>,
mutations: Vec<RegistryMutation>,
}
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<RegistryMutation> {
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
);
}
}

24
litellm-rust/crates/python-bridge/src/diagnostics.rs
View file

@ -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<Py<PyAny>> {
@ -9,6 +9,27 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
Ok(stats.into_any().unbind())
}
#[pyfunction]
fn _debug_setup(
py: Python<'_>,
call_type: String,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
start: Py<PyAny>,
asynchronous: bool,
) -> PyResult<(Py<PyAny>, Py<PyDict>)> {
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(())

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

@ -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<PyAny>);
@ -126,32 +126,6 @@ impl PythonLogger {
}
}
pub(super) struct SetupResult<'py>(Bound<'py, PyAny>);
impl SetupResult<'_> {
pub(super) fn logger(&self) -> PyResult<PythonLogger> {
self.0.getattr("logger")?.extract()
}
pub(super) fn kwargs(&self) -> PyResult<Py<PyDict>> {
Ok(self.0.getattr("kwargs")?.extract()?)
}
}
pub(super) fn setup<'py>(
py: Python<'py>,
call_type: &str,
args: &Py<PyTuple>,
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
asynchronous: bool,
) -> PyResult<SetupResult<'py>> {
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<Py<PyAny>>,
@ -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::<Vec<String>>()
.unwrap(),
["logger"]
);
assert!(
result
.kwargs()
.unwrap_err()
.is_instance_of::<PyTypeError>(py)
);
assert_eq!(
locals
.get_item("reads")
.unwrap()
.unwrap()
.extract::<Vec<String>>()
.unwrap(),
["logger", "kwargs"]
);
});
}
#[test]
fn logger_resolves_each_callback_at_invocation_and_preserves_arguments() {

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

@ -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<R: PythonRoute + 'static>(
}
}
pub(crate) fn debug_setup(
py: Python<'_>,
call_type: &'static str,
args: &Py<PyTuple>,
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
asynchronous: bool,
) -> PyResult<(Py<PyAny>, Py<PyDict>)> {
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<R: PythonRoute> PythonLifecycle<R> {
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(())
}

334
litellm-rust/crates/python-bridge/src/lifecycle/setup.rs Normal file
View file

@ -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<Bound<'py, PyAny>>,
}
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<Registration> {
if let Ok(name) = object.cast::<PyString>() {
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<Vec<Entry>> {
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<Bound<'py, PyAny>> {
setup.getattr("registry")?.call1((name,))
}
fn read_candidates<'py>(
setup: &Bound<'py, PyModule>,
targets: &mut Targets<'py>,
dynamic: Option<Bound<'py, PyAny>>,
) -> PyResult<Vec<Candidate>> {
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::<PyString>() {
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<Bound<'py, PyList>>,
async_success: Option<Bound<'py, PyList>>,
failure: Option<Bound<'py, PyList>>,
}
fn split_dynamic<'py>(
py: Python<'py>,
setup: &Bound<'py, PyModule>,
targets: &mut Targets<'py>,
kwargs: &Bound<'py, PyDict>,
) -> PyResult<DynamicLists<'py>> {
let success = match kwargs.get_item("success_callback")? {
Some(value) if value.is_instance_of::<PyList>() => {
let list = value.cast_into::<PyList>()?;
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::<PyString>()
.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::<PyList>() => {
kwargs.del_item("failure_callback")?;
Some(value.cast_into::<PyList>()?)
}
_ => 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<PyDict>,
}
pub(super) fn setup(
py: Python<'_>,
call_type: &str,
args: &Py<PyTuple>,
kwargs: &Py<PyDict>,
start: &Py<PyAny>,
asynchronous: bool,
) -> PyResult<Setup> {
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(),
})
}

15
litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs
View file

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

33
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -40,17 +40,28 @@ pub(super) struct PythonPayload {
impl PythonPayload {
fn from_request(py: Python<'_>, request: &OcrDuringCallRequest) -> PyResult<Self> {
let body = to_py(py, &request.body)?.into_bound(py).cast_into::<PyDict>()?;
let body = to_py(py, &request.body)?
.into_bound(py)
.cast_into::<PyDict>()?;
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<OcrDuringCallRequest> {
fn write_back(
&self,
py: Python<'_>,
mut request: OcrDuringCallRequest,
) -> PyResult<OcrDuringCallRequest> {
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::<String>()?, value.extract::<String>()?)))
.collect::<PyResult<Vec<_>>>()?;
@ -81,7 +92,9 @@ impl PythonOcrHost {
}
fn project(&mut self, py: Python<'_>) -> PyResult<OcrHostResult> {
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<OcrPostCallRequest> {
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)?))
}

17
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

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

19
litellm-rust/crates/python-bridge/src/routes/ocr/project.rs
View file

@ -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<Result<FileDocumentInput, litellm_core::ocr::Error>> {
fn project_document(
document: &Bound<'_, PyAny>,
) -> PyResult<Result<FileDocumentInput, litellm_core::ocr::Error>> {
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"));
});
}

8
litellm/rust_bridge/_native.pyi
View file

@ -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]]: ...

30
litellm/rust_bridge/lifecycle.py
View file

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

262
litellm/rust_bridge/setup.py Normal file
View file

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

169
tests/test_litellm/rust_bridge/test_setup.py Normal file
View file

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

20
tests/test_litellm_rust/ocr/test_lifecycle.py
View file

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