scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-13 03:59:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

refactor(rust): classify legacy callbacks by layer

Browse source
This commit is contained in:
Yujong Lee 2026-10-06 15:42:05 -07:00
parent 2cee61626d
commit c26852c799
6 changed files with 239 additions and 113 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -44,7 +44,7 @@ struct LoggedRequest {
context: RequestContext,
}
pub struct LegacyLogging {
pub struct InferenceAdapter {
operation: LoggingOperation,
call: PublicCall,
logger: Option<PythonLogger>,
@ -67,7 +67,7 @@ fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
!error.is_instance_of::<PyException>(py)
}
impl LegacyLogging {
impl InferenceAdapter {
pub fn new(
py: Python<'_>,
operation: LoggingOperation,
@ -300,7 +300,7 @@ impl LegacyLogging {
}
}
impl LegacyLogging {
impl InferenceAdapter {
fn resume_begin(
&mut self,
py: Python<'_>,
@ -340,7 +340,7 @@ impl LegacyLogging {
}
}
impl LegacyLogging {
impl InferenceAdapter {
pub(crate) fn prepare_call(
&mut self,
py: Python<'_>,
@ -553,7 +553,7 @@ impl LegacyLogging {
}
}
impl PythonOwned for LegacyLogging {
impl PythonOwned for InferenceAdapter {
fn close(&mut self, py: Python<'_>) {
if let Some(logger) = self.logger.take()
&& let Err(error) = logger.restore_context(py)
@ -596,15 +596,15 @@ mod deployment_hooks_tests {
use pyo3::types::PyDict;
use rstest::rstest;
use super::LegacyLogging;
use super::InferenceAdapter;
use crate::test_support::{legacy_call, local, namespace, run};
fn resume<T>(
logging: &mut LegacyLogging,
step: HookStep<LegacyLogging, T>,
logging: &mut InferenceAdapter,
step: HookStep<InferenceAdapter, T>,
py: Python<'_>,
value: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<LegacyLogging, T>> {
) -> PyResult<HookStep<InferenceAdapter, T>> {
let HookStep::Await(_, continuation) = step else {
panic!("expected suspension")
};
@ -625,7 +625,7 @@ kwargs = {'logger': logger, 'document': document}
py: Python<'py>,
locals: &Bound<'py, PyDict>,
asynchronous: bool,
) -> (LegacyLogging, HookStep<LegacyLogging, Py<PyDict>>) {
) -> (InferenceAdapter, HookStep<InferenceAdapter, Py<PyDict>>) {
let mut logging = legacy_call(py, locals, asynchronous);
let kwargs = local(locals, "kwargs")
.cast_into::<PyDict>()
@ -637,7 +637,7 @@ kwargs = {'logger': logger, 'document': document}
fn arguments<'py>(
py: Python<'py>,
step: HookStep<LegacyLogging, Py<PyDict>>,
step: HookStep<InferenceAdapter, Py<PyDict>>,
) -> Bound<'py, PyDict> {
let HookStep::Ready(arguments) = step else {
panic!("expected the prepared arguments");
@ -645,7 +645,7 @@ kwargs = {'logger': logger, 'document': document}
arguments.into_bound(py)
}
fn awaits_deployment_hook<T>(step: &HookStep<LegacyLogging, T>) -> bool {
fn awaits_deployment_hook<T>(step: &HookStep<InferenceAdapter, T>) -> bool {
matches!(step, HookStep::Await(_, _))
}
@ -666,7 +666,7 @@ kwargs = {'logger': logger, 'document': document}
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, CALL);
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
operation,
..legacy_call(py, &locals, asynchronous)
};
@ -929,7 +929,7 @@ mod payload_tests {
use rstest::rstest;
use serde_json::{Map, Value, json};
use super::LegacyLogging;
use super::InferenceAdapter;
use crate::PythonLogger;
use crate::test_support::{legacy_call, local, namespace, run};
@ -999,7 +999,7 @@ check = lambda: None
locals.set_item(name, to_py(py, value).unwrap()).unwrap();
}
run(py, &locals, script);
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
@ -1028,12 +1028,12 @@ check = lambda: None
/// delivers it, so `pre_call` and `post_call` have both seen the retained payload.
fn send_and_receive<'a>(
py: Python<'_>,
logging: &'a mut LegacyLogging,
logging: &'a mut InferenceAdapter,
wire: WireRequest,
context: &RequestContext,
) -> (
&'a mut LegacyLogging,
HookStep<LegacyLogging, Box<WireRequest>>,
&'a mut InferenceAdapter,
HookStep<InferenceAdapter, Box<WireRequest>>,
) {
let step = logging
.before_provider_request(py, Box::new(wire), context)
@ -1073,11 +1073,11 @@ check = lambda: None
}
}
/// A Python object owning one `LegacyLogging`, so the interpreter's collector sees the
/// A Python object owning one `InferenceAdapter`, so the interpreter's collector sees the
/// edges the adapter reports and clears them the way the driver's `Execution` does.
#[pyclass(weakref)]
struct Retained {
logging: Option<LegacyLogging>,
logging: Option<InferenceAdapter>,
}
#[pymethods]
@ -1116,7 +1116,7 @@ kwargs['pages'] = original
prepared = {'pages': replacement}
",
);
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
operation,
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
@ -1150,7 +1150,7 @@ assert original == [0]
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
@ -1186,7 +1186,7 @@ assert reference() is None
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
@ -1638,7 +1638,7 @@ mod terminal_tests {
use pyo3::types::PyDict;
use rstest::rstest;
use super::LegacyLogging;
use super::InferenceAdapter;
use crate::PythonLogger;
use crate::test_support::{legacy_call, local, namespace, run};
@ -1647,8 +1647,8 @@ mod terminal_tests {
end_time: 1.0,
};
fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging {
LegacyLogging {
fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> InferenceAdapter {
InferenceAdapter {
logger: Some(PythonLogger::new(local(locals, "logger").unbind())),
..legacy_call(py, locals, asynchronous)
}
@ -1657,8 +1657,8 @@ mod terminal_tests {
fn succeed(
py: Python<'_>,
locals: &Bound<'_, PyDict>,
logging: &mut LegacyLogging,
) -> HookStep<LegacyLogging, ()> {
logging: &mut InferenceAdapter,
) -> HookStep<InferenceAdapter, ()> {
let response = local(locals, "response").unbind();
logging
.on_event(
@ -1674,8 +1674,8 @@ mod terminal_tests {
fn fail(
py: Python<'_>,
locals: &Bound<'_, PyDict>,
logging: &mut LegacyLogging,
) -> HookStep<LegacyLogging, ()> {
logging: &mut InferenceAdapter,
) -> HookStep<InferenceAdapter, ()> {
let failure = PyErr::from_value(local(locals, "failure"));
logging
.on_event(
@ -1769,7 +1769,7 @@ assert logger.calls[1][1] is response
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"first = b'first'\nlast = b'last'\nresponse = None\nhead = {'additional_headers': {'request-id': 'req_native'}}");
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
@ -1811,7 +1811,7 @@ assert chunks[1] is last
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"failure = ValueError('provider')");
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
internal: true,
..logged(py, &locals, asynchronous)
};
@ -1833,7 +1833,7 @@ assert chunks[1] is last
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
let mut logging = LegacyLogging {
let mut logging = InferenceAdapter {
internal: true,
..logged(py, &locals, true)
};

4
litellm-rust/crates/callbacks-legacy-python/src/lib.rs
View file

@ -14,11 +14,11 @@ mod deferred;
mod logger;
mod mapping;
mod python;
pub use adapter::LegacyLogging;
pub use adapter::InferenceAdapter;
pub use call::PublicCall;
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
pub use mapping::{CallBoundary, CallbackMapping, Dispatch, callback_mappings};
pub use mapping::{CallbackMapping, Dispatch, InferenceBoundary, callback_mappings};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LoggingOperation {

182
litellm-rust/crates/callbacks-legacy-python/src/mapping.rs
View file

@ -6,10 +6,10 @@ use litellm_host::{
use litellm_host_python::{HookStep, PythonCallEvent, PythonRuntime};
use pyo3::{prelude::*, types::PyDict};
use crate::LegacyLogging;
use crate::InferenceAdapter;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CallBoundary {
pub enum InferenceBoundary {
PrepareArguments,
BeforeProviderRequest,
AfterProviderResponse,
@ -18,13 +18,18 @@ pub enum CallBoundary {
Failed,
StreamOpened,
StreamChunk,
SucceededOrFailed,
NotFiredNativelyYet(&'static [&'static str]),
DeclarationOnly,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Dispatch {
Call(CallBoundary),
Python(&'static str),
DeclarationOnly,
Inference(InferenceBoundary),
Gateway(&'static str),
Router(&'static str),
Management(&'static str),
HandlerTrait(&'static str),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -34,7 +39,7 @@ pub struct CallbackMapping {
}
struct Binding<H> {
boundary: CallBoundary,
boundary: InferenceBoundary,
invoke: H,
callbacks: &'static [&'static str],
}
@ -43,82 +48,82 @@ impl<H> Binding<H> {
fn mappings(&self) -> impl Iterator<Item = CallbackMapping> {
self.callbacks.iter().map(|callback| CallbackMapping {
callback,
dispatch: Dispatch::Call(self.boundary),
dispatch: Dispatch::Inference(self.boundary),
})
}
}
type Step<T> = PyResult<HookStep<LegacyLogging, T>>;
type Prepare = fn(&mut LegacyLogging, Python<'_>, Py<PyDict>, f64) -> Step<Py<PyDict>>;
type Before =
fn(&mut LegacyLogging, Python<'_>, Box<WireRequest>, &RequestContext) -> Step<Box<WireRequest>>;
type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>;
type Transform = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
type Open = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
type Chunk = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
type Step<T> = PyResult<HookStep<InferenceAdapter, T>>;
type Prepare = fn(&mut InferenceAdapter, Python<'_>, Py<PyDict>, f64) -> Step<Py<PyDict>>;
type Before = fn(
&mut InferenceAdapter,
Python<'_>,
Box<WireRequest>,
&RequestContext,
) -> Step<Box<WireRequest>>;
type After = fn(&mut InferenceAdapter, Python<'_>, &RawResponse) -> Step<()>;
type Transform = fn(&mut InferenceAdapter, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
type Success = fn(&mut InferenceAdapter, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
type Failure = fn(&mut InferenceAdapter, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
type Open = fn(&mut InferenceAdapter, Python<'_>, &Py<PyAny>) -> PyResult<()>;
type Chunk = fn(&mut InferenceAdapter, Python<'_>, &Py<PyAny>) -> PyResult<()>;
const PREPARE: Binding<Prepare> = Binding {
boundary: CallBoundary::PrepareArguments,
invoke: LegacyLogging::prepare_call,
boundary: InferenceBoundary::PrepareArguments,
invoke: InferenceAdapter::prepare_call,
callbacks: &["async_pre_call_deployment_hook"],
};
const BEFORE: Binding<Before> = Binding {
boundary: CallBoundary::BeforeProviderRequest,
invoke: LegacyLogging::pre_call,
boundary: InferenceBoundary::BeforeProviderRequest,
invoke: InferenceAdapter::pre_call,
callbacks: &["log_pre_api_call", "log_input_event"],
};
const AFTER: Binding<After> = Binding {
boundary: CallBoundary::AfterProviderResponse,
invoke: LegacyLogging::post_call,
boundary: InferenceBoundary::AfterProviderResponse,
invoke: InferenceAdapter::post_call,
callbacks: &["log_post_api_call"],
};
const TRANSFORM: Binding<Transform> = Binding {
boundary: CallBoundary::TransformResponse,
invoke: LegacyLogging::transform_public_response,
boundary: InferenceBoundary::TransformResponse,
invoke: InferenceAdapter::transform_public_response,
callbacks: &["async_post_call_success_deployment_hook"],
};
const SUCCESS: Binding<Success> = Binding {
boundary: CallBoundary::Succeeded,
invoke: LegacyLogging::succeeded,
boundary: InferenceBoundary::Succeeded,
invoke: InferenceAdapter::succeeded,
callbacks: &[
"log_success_event",
"async_log_success_event",
"logging_hook",
"async_logging_hook",
"redact_standard_logging_payload_from_model_call_details",
"log_event",
"async_log_event",
],
};
const FAILURE: Binding<Failure> = Binding {
boundary: CallBoundary::Failed,
invoke: LegacyLogging::failed,
boundary: InferenceBoundary::Failed,
invoke: InferenceAdapter::failed,
callbacks: &[
"async_post_call_failure_deployment_hook",
"log_failure_event",
"async_log_failure_event",
"log_model_group_rate_limit_error",
"log_event",
"async_log_event",
],
};
const OPEN: Binding<Open> = Binding {
boundary: CallBoundary::StreamOpened,
invoke: LegacyLogging::stream_opened,
boundary: InferenceBoundary::StreamOpened,
invoke: InferenceAdapter::stream_opened,
callbacks: &[],
};
const CHUNK: Binding<Chunk> = Binding {
boundary: CallBoundary::StreamChunk,
invoke: LegacyLogging::stream_chunk,
boundary: InferenceBoundary::StreamChunk,
invoke: InferenceAdapter::stream_chunk,
callbacks: &[],
};
@ -132,10 +137,10 @@ pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
.chain(FAILURE.mappings())
.chain(OPEN.mappings())
.chain(CHUNK.mappings())
.chain(PYTHON_CALLBACKS.iter().copied())
.chain(OTHER_CALLBACKS.iter().copied())
}
impl CallHooks<PythonRuntime> for LegacyLogging {
impl CallHooks<PythonRuntime> for InferenceAdapter {
fn prepare_arguments(
&mut self,
py: Python<'_>,
@ -199,26 +204,27 @@ impl CallHooks<PythonRuntime> for LegacyLogging {
}
}
macro_rules! python_callbacks {
macro_rules! callbacks {
($($dispatch:expr => [$($callback:literal),* $(,)?]),* $(,)?) => {
const PYTHON_CALLBACKS: &[CallbackMapping] = &[
const OTHER_CALLBACKS: &[CallbackMapping] = &[
$($(CallbackMapping { callback: $callback, dispatch: $dispatch },)*)*
];
};
}
python_callbacks! {
Dispatch::Python("litellm.router") => [
callbacks! {
Dispatch::Inference(InferenceBoundary::SucceededOrFailed) => ["log_event", "async_log_event"],
Dispatch::Router("litellm.router") => [
"async_pre_routing_hook",
"async_filter_deployments",
"pre_call_check",
"async_pre_call_check",
],
Dispatch::Python("litellm.router_utils.fallback_event_handlers") => [
Dispatch::Router("litellm.router_utils.fallback_event_handlers") => [
"log_success_fallback_event",
"log_failure_fallback_event",
],
Dispatch::Python("litellm.proxy.utils") => [
Dispatch::Gateway("litellm.proxy.utils") => [
"async_pre_call_hook",
"async_post_call_response_headers_hook",
"async_post_call_failure_hook",
@ -227,62 +233,92 @@ python_callbacks! {
"async_post_call_streaming_hook",
"async_post_call_streaming_iterator_hook",
"async_filter_listed_models",
"apply_guardrail",
],
Dispatch::Python("litellm.litellm_core_utils.litellm_logging") => [
Dispatch::Gateway("litellm.litellm_core_utils.litellm_logging") => [
"async_post_mcp_tool_call_hook",
],
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&[
"litellm.litellm_core_utils.litellm_logging",
])) => [
"async_get_chat_completion_prompt",
"get_chat_completion_prompt",
"log_stream_event",
"async_log_stream_event",
"async_post_mcp_tool_call_hook",
],
Dispatch::Python("litellm.llms.anthropic.pass_through.messages.handler") => [
"async_pre_request_hook",
],
Dispatch::Python("litellm.litellm_core_utils.streaming_handler") => [
"async_post_call_streaming_deployment_hook",
],
Dispatch::Python("litellm.responses.streaming_iterator") => [
"async_post_call_streaming_deployment_hook",
],
Dispatch::Python("litellm.main") => [
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&[
"litellm.llms.anthropic.pass_through.messages.handler",
])) => ["async_pre_request_hook"],
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&[
"litellm.litellm_core_utils.streaming_handler",
"litellm.responses.streaming_iterator",
])) => ["async_post_call_streaming_deployment_hook"],
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&["litellm.main"])) => [
"translate_completion_input_params",
"translate_completion_output_params",
"translate_completion_output_params_streaming",
],
Dispatch::Python("litellm.integrations.argilla") => ["async_dataset_hook"],
Dispatch::Python("litellm.proxy.management_helpers.audit_logs") => ["async_log_audit_log_event"],
Dispatch::Python("litellm.llms.custom_httpx.llm_http_handler") => [
"async_should_run_agentic_loop",
"async_run_agentic_loop",
"async_build_agentic_loop_plan",
"async_post_agentic_loop_response_hook",
"async_agentic_loop_cleanup_hook",
"async_should_run_chat_completion_agentic_loop",
"async_run_chat_completion_agentic_loop",
"async_build_chat_completion_agentic_loop_plan",
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&["litellm.integrations.argilla"])) => [
"async_dataset_hook",
],
Dispatch::Python("litellm.litellm_core_utils.chat_completion_agentic_loop") => [
Dispatch::Management("litellm.proxy.management_helpers.audit_logs") => ["async_log_audit_log_event"],
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&[
"litellm.llms.custom_httpx.llm_http_handler",
"litellm.litellm_core_utils.chat_completion_agentic_loop",
])) => [
"async_should_run_agentic_loop",
"async_run_agentic_loop",
"async_build_agentic_loop_plan",
"async_post_agentic_loop_response_hook",
"async_agentic_loop_cleanup_hook",
],
Dispatch::Python("litellm.llms.openai.openai") => [
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&[
"litellm.llms.custom_httpx.llm_http_handler",
"litellm.llms.openai.openai",
])) => [
"async_should_run_chat_completion_agentic_loop",
"async_run_chat_completion_agentic_loop",
],
Dispatch::Python("litellm.proxy.spend_tracking.cold_storage_handler") => [
Dispatch::Inference(InferenceBoundary::NotFiredNativelyYet(&[
"litellm.llms.custom_httpx.llm_http_handler",
])) => ["async_build_chat_completion_agentic_loop_plan"],
Dispatch::Management("litellm.proxy.spend_tracking.cold_storage_handler") => [
"get_proxy_server_request_from_cold_storage_with_object_key",
],
Dispatch::Python("litellm.integrations.custom_logger") => [
Dispatch::HandlerTrait("litellm.integrations.custom_logger") => [
"truncate_standard_logging_payload_content",
"redacts_messages_itself",
"handle_callback_failure",
"get_callback_env_vars",
],
Dispatch::DeclarationOnly => [
Dispatch::Inference(InferenceBoundary::DeclarationOnly) => [
"async_log_pre_api_call",
"async_log_input_event",
],
Dispatch::HandlerTrait("litellm.integrations.custom_guardrail") => [
"add_standard_logging_guardrail_information_to_request_data",
"async_pre_call_hook_on_messages",
"filter_new_texts_for_session",
"get_config_model",
"get_disable_global_guardrail",
"get_guardrail_dynamic_request_body_params",
"get_guardrail_from_metadata",
"get_guardrails_messages_for_call_type",
"get_opted_out_global_guardrails_from_metadata",
"get_supported_event_hooks",
"handle_sensitive_data_detection",
"inject_advisory_message",
"mark_pre_call_hook_ran",
"mark_texts_scanned",
"mask_content_in_string",
"raise_passthrough_exception",
"raise_sensitive_data_route_exception",
"render_violation_message",
"should_route_on_sensitive_data",
"should_run_guardrail",
"structured_messages_cover_full_request",
"supports_scan_only_tool_results",
"update_in_memory_litellm_params",
"uses_apply_guardrail_interface",
],
}

6
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -3,7 +3,7 @@ use std::ffi::CStr;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};
use crate::{LegacyLogging, PublicCall};
use crate::{InferenceAdapter, PublicCall};
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
@ -178,7 +178,7 @@ pub(crate) fn legacy_call(
py: Python<'_>,
locals: &Bound<'_, PyDict>,
asynchronous: bool,
) -> LegacyLogging {
) -> InferenceAdapter {
let request = locals
.get_item("request")
.unwrap()
@ -189,5 +189,5 @@ pub(crate) fn legacy_call(
.map(|kwargs| kwargs.cast_into::<PyDict>().unwrap())
.unwrap_or_else(|| PyDict::new(py));
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous)
InferenceAdapter::new(py, crate::LoggingOperation::Ocr, call, asynchronous)
}

90
litellm-rust/crates/callbacks-legacy-python/tests/classification.rs Normal file
View file

@ -0,0 +1,90 @@
use std::{collections::BTreeSet, ffi::CString};
use litellm_callbacks_legacy_python::{Dispatch, callback_mappings};
use pyo3::{prelude::*, types::PyDict};
use rstest::rstest;
fn public_methods(source: &str, class: &str) -> BTreeSet<String> {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
locals.set_item("source", source).unwrap();
locals.set_item("class_name", class).unwrap();
py.run(
&CString::new(
"import ast\n\
tree = ast.parse(source)\n\
methods = {method.name for node in tree.body \
if isinstance(node, ast.ClassDef) and node.name == class_name \
for method in node.body if isinstance(method, (ast.FunctionDef, ast.AsyncFunctionDef)) \
and not method.name.startswith('_')}",
)
.unwrap(),
Some(&locals),
Some(&locals),
)
.unwrap();
locals
.get_item("methods")
.unwrap()
.unwrap()
.extract()
.unwrap()
})
}
fn unclassified(methods: &BTreeSet<String>) -> BTreeSet<String> {
let classified: BTreeSet<String> = callback_mappings()
.map(|entry| entry.callback.to_owned())
.collect();
methods.difference(&classified).cloned().collect()
}
#[rstest]
fn every_public_method_has_exactly_one_home() {
let methods: BTreeSet<String> = public_methods(
include_str!("../../../../litellm/integrations/custom_logger.py"),
"CustomLogger",
)
.union(&public_methods(
include_str!("../../../../litellm/integrations/custom_guardrail.py"),
"CustomGuardrail",
))
.cloned()
.collect();
assert!(
unclassified(&methods).is_empty(),
"{:?}",
unclassified(&methods)
);
let entries: Vec<_> = callback_mappings().collect();
assert_eq!(entries.len(), methods.len());
assert_eq!(
entries
.iter()
.map(|entry| entry.callback)
.collect::<BTreeSet<_>>()
.len(),
entries.len()
);
}
#[rstest]
fn a_new_public_method_requires_classification() {
let source = format!(
"{}\n def new_unclassified_hook(self): pass\n",
include_str!("../../../../litellm/integrations/custom_logger.py"),
);
assert_eq!(
unclassified(&public_methods(&source, "CustomLogger")),
BTreeSet::from(["new_unclassified_hook".to_owned()]),
);
}
#[rstest]
#[case::gateway(|dispatch| matches!(dispatch, Dispatch::Gateway(_)))]
#[case::router(|dispatch| matches!(dispatch, Dispatch::Router(_)))]
#[case::inference(|dispatch| matches!(dispatch, Dispatch::Inference(_)))]
fn each_request_layer_has_methods(#[case] belongs: fn(Dispatch) -> bool) {
assert!(callback_mappings().any(|entry| belongs(entry.dispatch)));
}

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

@ -9,7 +9,7 @@ pub(crate) mod token_counter;
pub(crate) mod traces;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_callbacks_legacy_python::{LegacyLogging, PublicCall};
use litellm_callbacks_legacy_python::{InferenceAdapter, PublicCall};
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
use litellm_host_python::{HookChain, PythonBinding, PythonCallHooks, PythonHostCalls};
use pyo3::{
@ -29,7 +29,7 @@ fn call_hooks(
let arguments = call.arguments(py);
Ok((
arguments,
LegacyLogging::new(py, operation, call, asynchronous),
InferenceAdapter::new(py, operation, call, asynchronous),
))
}