diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 70b7ecb45fd..ff5ba6f8c5f 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -44,7 +44,7 @@ struct LoggedRequest { context: RequestContext, } -pub struct LegacyLogging { +pub struct InferenceAdapter { operation: LoggingOperation, call: PublicCall, logger: Option, @@ -67,7 +67,7 @@ fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool { !error.is_instance_of::(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( - logging: &mut LegacyLogging, - step: HookStep, + logging: &mut InferenceAdapter, + step: HookStep, py: Python<'_>, value: PyResult>, - ) -> PyResult> { + ) -> PyResult> { 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>) { + ) -> (InferenceAdapter, HookStep>) { let mut logging = legacy_call(py, locals, asynchronous); let kwargs = local(locals, "kwargs") .cast_into::() @@ -637,7 +637,7 @@ kwargs = {'logger': logger, 'document': document} fn arguments<'py>( py: Python<'py>, - step: HookStep>, + step: HookStep>, ) -> 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(step: &HookStep) -> bool { + fn awaits_deployment_hook(step: &HookStep) -> 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>, + &'a mut InferenceAdapter, + HookStep>, ) { 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, + logging: Option, } #[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 { + logging: &mut InferenceAdapter, + ) -> HookStep { 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 { + logging: &mut InferenceAdapter, + ) -> HookStep { 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) }; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 38c6b1aedbd..e200d7f14c2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -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 { diff --git a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs index a4aee4eb4da..a067b58db1b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs @@ -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 { - boundary: CallBoundary, + boundary: InferenceBoundary, invoke: H, callbacks: &'static [&'static str], } @@ -43,82 +48,82 @@ impl Binding { fn mappings(&self) -> impl Iterator { self.callbacks.iter().map(|callback| CallbackMapping { callback, - dispatch: Dispatch::Call(self.boundary), + dispatch: Dispatch::Inference(self.boundary), }) } } -type Step = PyResult>; -type Prepare = fn(&mut LegacyLogging, Python<'_>, Py, f64) -> Step>; -type Before = - fn(&mut LegacyLogging, Python<'_>, Box, &RequestContext) -> Step>; -type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>; -type Transform = fn(&mut LegacyLogging, Python<'_>, Py, Timing) -> Step>; -type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py) -> Step<()>; -type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>; -type Open = fn(&mut LegacyLogging, Python<'_>, &Py) -> PyResult<()>; -type Chunk = fn(&mut LegacyLogging, Python<'_>, &Py) -> PyResult<()>; +type Step = PyResult>; +type Prepare = fn(&mut InferenceAdapter, Python<'_>, Py, f64) -> Step>; +type Before = fn( + &mut InferenceAdapter, + Python<'_>, + Box, + &RequestContext, +) -> Step>; +type After = fn(&mut InferenceAdapter, Python<'_>, &RawResponse) -> Step<()>; +type Transform = fn(&mut InferenceAdapter, Python<'_>, Py, Timing) -> Step>; +type Success = fn(&mut InferenceAdapter, Python<'_>, Timing, &Py) -> Step<()>; +type Failure = fn(&mut InferenceAdapter, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>; +type Open = fn(&mut InferenceAdapter, Python<'_>, &Py) -> PyResult<()>; +type Chunk = fn(&mut InferenceAdapter, Python<'_>, &Py) -> PyResult<()>; const PREPARE: Binding = Binding { - boundary: CallBoundary::PrepareArguments, - invoke: LegacyLogging::prepare_call, + boundary: InferenceBoundary::PrepareArguments, + invoke: InferenceAdapter::prepare_call, callbacks: &["async_pre_call_deployment_hook"], }; const BEFORE: Binding = 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 = Binding { - boundary: CallBoundary::AfterProviderResponse, - invoke: LegacyLogging::post_call, + boundary: InferenceBoundary::AfterProviderResponse, + invoke: InferenceAdapter::post_call, callbacks: &["log_post_api_call"], }; const TRANSFORM: Binding = 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 = 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 = 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 = Binding { - boundary: CallBoundary::StreamOpened, - invoke: LegacyLogging::stream_opened, + boundary: InferenceBoundary::StreamOpened, + invoke: InferenceAdapter::stream_opened, callbacks: &[], }; const CHUNK: Binding = 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 { .chain(FAILURE.mappings()) .chain(OPEN.mappings()) .chain(CHUNK.mappings()) - .chain(PYTHON_CALLBACKS.iter().copied()) + .chain(OTHER_CALLBACKS.iter().copied()) } -impl CallHooks for LegacyLogging { +impl CallHooks for InferenceAdapter { fn prepare_arguments( &mut self, py: Python<'_>, @@ -199,26 +204,27 @@ impl CallHooks 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", + ], } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index f252a45b562..cfd8352c54a 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -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::().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) } diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/classification.rs b/litellm-rust/crates/callbacks-legacy-python/tests/classification.rs new file mode 100644 index 00000000000..7d8ea82dc9a --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy-python/tests/classification.rs @@ -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 { + 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) -> BTreeSet { + let classified: BTreeSet = 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 = 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::>() + .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))); +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 2380274001e..c6949067636 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -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), )) }