diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c359ca19986..0265e0adbc2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2031,6 +2031,7 @@ dependencies = [ name = "litellm-callbacks" version = "0.1.0" dependencies = [ + "litellm-auth", "rstest", "serde_json", "tokio", @@ -2040,11 +2041,13 @@ dependencies = [ name = "litellm-callbacks-legacy" version = "0.1.0" dependencies = [ + "litellm-auth", "litellm-callbacks", "litellm-host-python", "pyo3", "rstest", "serde_json", + "strum", ] [[package]] diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth/src/secret.rs index 3ecb0a835ee..a07fe3eaad9 100644 --- a/litellm-rust/crates/auth/src/secret.rs +++ b/litellm-rust/crates/auth/src/secret.rs @@ -1,6 +1,8 @@ +use serde::Deserialize; use veil::Redact; -#[derive(Redact, Clone)] +#[derive(Redact, Clone, Deserialize)] +#[serde(transparent)] pub struct SecretValue(#[redact(with = "[REDACTED]")] String); impl SecretValue { diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy/AGENTS.md index e4762d3037a..e184e2fb415 100644 --- a/litellm-rust/crates/callbacks-legacy/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy/AGENTS.md @@ -1,15 +1,17 @@ - Target invariants, not completion claims - Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits) - - The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call + - The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call +- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`) + - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it + - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy -- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it - - A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case - - A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run +- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does + - Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's - Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - - Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None` + - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields` + - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view - Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index 96c9c9ed560..3cf9382f857 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -9,8 +9,12 @@ autotests = false [dependencies] litellm-callbacks.workspace = true litellm-host-python.workspace = true + pyo3.workspace = true +strum.workspace = true +serde_json.workspace = true [dev-dependencies] +litellm-auth.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json new file mode 100644 index 00000000000..840c0abfa45 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -0,0 +1,98 @@ +{ + "setup": [ + "call_type", + "args", + "kwargs", + "start_time", + "asynchronous" + ], + "check_limits": [ + "kwargs" + ], + "finalize": [ + "response", + "logger", + "kwargs", + "start_time", + "end_time" + ], + "update_logging": [ + "logger", + "kwargs", + "model", + "optional_params", + "litellm_params", + "custom_llm_provider" + ], + "pre_call": [ + "logger", + "input", + "api_key", + "additional_args" + ], + "post_call": [ + "logger", + "original_response", + "api_key", + "additional_args" + ], + "defers_async_logging": [ + "logger" + ], + "defer_success": [ + "logger", + "pending" + ], + "sync_success_for_async_call": [ + "logger", + "response", + "start", + "end" + ], + "failure_handler": [ + "logger", + "error", + "start", + "end", + "asynchronous" + ], + "submit_success": [ + "logger", + "response", + "start", + "end" + ], + "async_success_handler": [ + "logger", + "response", + "start", + "end" + ], + "enqueue_logging": [ + "coroutine" + ], + "restore_context": [ + "logger" + ], + "custom_pricing_fields": [], + "is_internal_call": [], + "credential_list": [], + "warn_unknown_credential": [ + "name", + "loaded" + ], + "before_deployment_call": [ + "kwargs", + "call_type" + ], + "after_deployment_success": [ + "kwargs", + "response", + "call_type" + ], + "after_deployment_failure": [ + "kwargs", + "error", + "call_type" + ] +} diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index df346506094..7204207cd62 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -4,7 +4,7 @@ use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest}; use litellm_host_python::{ - AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py, + LifecycleStep, PublicValue, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, @@ -12,6 +12,7 @@ use pyo3::{ prelude::*, types::PyDict, }; +use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, @@ -43,7 +44,7 @@ pub struct LegacyLogging { response: Option>, error: Option>, body: Option>, - headers: Option>, + context: Option, asynchronous: bool, internal: bool, pending: Option, @@ -76,7 +77,7 @@ impl LegacyLogging { response: None, error: None, body: None, - headers: None, + context: None, asynchronous, internal: false, pending: None, @@ -85,8 +86,8 @@ impl LegacyLogging { /// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never /// runs them. - fn deployment_hooks(&self, py: Python<'_>) -> PyResult { - Ok(self.asynchronous && DeploymentHooks::needed(py)?) + fn runs_deployment_hooks(&self) -> bool { + self.asynchronous } fn logger(&self) -> PyResult<&PythonLogger> { @@ -95,13 +96,13 @@ impl LegacyLogging { }) } - fn prepare(&mut self, py: Python<'_>) -> PyResult { + fn prepare(&mut self, py: Python<'_>) -> PyResult { let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind(); self.call.set_kwargs(prepared); - Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py))) + Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py))) } - fn finalize(&mut self, py: Python<'_>) -> PyResult { + fn finalize(&mut self, py: Python<'_>) -> PyResult { finalize( py, &self.response, @@ -112,7 +113,7 @@ impl LegacyLogging { )?; self.response .as_ref() - .map(|response| AdapterStep::Response(response.clone_ref(py))) + .map(|response| LifecycleStep::Response(response.clone_ref(py))) .ok_or_else(missing_state) } @@ -145,9 +146,7 @@ impl LegacyLogging { .get_item("fallbacks")? .is_none_or(|value| value.is_none()) { - if !logger.callbacks_needed(py, "async_success")? { - logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; - } else if logger.defers_async_logging(py) { + if logger.defers_async_logging(py) { let pending = Py::new( py, PendingLogging { @@ -165,12 +164,12 @@ impl LegacyLogging { /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. - fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { + fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { let (Some(logger), Some(error)) = (&self.logger, &self.error) else { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); }; if self.asynchronous && self.internal { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false) && is_cancellation(py, &failure) @@ -178,27 +177,27 @@ impl LegacyLogging { return Err(failure); } if !self.asynchronous { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } match logger.failure(py, error, &self.start, &self.end, true) { Ok(Some(awaitable)) => { self.pending = Some(Pending::AsyncFailure); - Ok(AdapterStep::Await(awaitable)) + Ok(LifecycleStep::Await(awaitable)) } - Ok(None) => Ok(AdapterStep::Done), + Ok(None) => Ok(LifecycleStep::Done), Err(failure) if is_cancellation(py, &failure) => Err(failure), - Err(_) => Ok(AdapterStep::Done), + Err(_) => Ok(LifecycleStep::Done), } } } -impl CallbackAdapter for LegacyLogging { +impl PythonLifecycle for LegacyLogging { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult { + ) -> PyResult { self.call.set_kwargs(arguments); self.start = datetime(py, started_at)?; self.internal = is_internal_call(py)?; @@ -212,9 +211,9 @@ impl CallbackAdapter for LegacyLogging { )?; self.logger = Some(result.logger()?); self.call.set_kwargs(result.kwargs()?); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPreCall); - return Ok(AdapterStep::Await(DeploymentHooks::before_call( + return Ok(LifecycleStep::Await(DeploymentHooks::before_call( py, self.call.kwargs(), self.surface.call_type, @@ -228,18 +227,16 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult { + ) -> PyResult { let logger = self.logger()?; logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; - if !logger.callbacks_needed(py, "payload")? { - logger.record_api_call_start(py)?; - return Ok(AdapterStep::Wire(wire)); - } let body = to_py(py, &wire.body)? .into_bound(py) .cast_into::()?; - for name in context.passthrough_fields.iter() { - if let Some(value) = self.call.lookup(py, name)? { + for (name, sent) in wire.body.as_object().into_iter().flatten() { + if let Some(value) = self.call.lookup(py, name)? + && from_py::(&value).is_ok_and(|caller| caller == *sent) + { body.set_item(name, value)?; } } @@ -248,12 +245,11 @@ impl CallbackAdapter for LegacyLogging { headers.set_item(name, value)?; } self.body = Some(body.clone().unbind()); - self.headers = Some(headers.clone().unbind()); - let api_key = self.call.lookup(py, "api_key")?; + self.context = Some(context.clone()); self.logger()?.pre_call( py, self.surface.input_description, - api_key.as_ref(), + context.api_key.as_ref().map(|api_key| api_key.expose()), &body, &headers, &wire.url, @@ -262,7 +258,7 @@ impl CallbackAdapter for LegacyLogging { .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { body: from_py(&body)?, headers, ..*wire @@ -274,12 +270,12 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult { + ) -> PyResult { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPostCall); - return Ok(AdapterStep::Await(DeploymentHooks::after_success( + return Ok(LifecycleStep::Await(DeploymentHooks::after_success( py, self.call.kwargs(), &self.response, @@ -294,31 +290,35 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult { + ) -> PyResult { match (event, public) { + (CallEvent::Started { .. }, _) => Ok(LifecycleStep::Done), (CallEvent::ResponseReceived { raw }, _) => { - let logger = self.logger()?; - if logger.callbacks_needed(py, "payload")? { - logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?; - } - Ok(AdapterStep::Done) + let api_key = self + .context + .as_ref() + .and_then(|context| context.api_key.as_ref()) + .map(|api_key| api_key.expose()); + self.logger()? + .post_call(py, &raw.body, api_key, self.body.as_ref())?; + Ok(LifecycleStep::Done) } (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); self.dispatch_success(py)?; - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); if *origin == FailureOrigin::Call && self.logger.is_some() - && self.deployment_hooks(py)? + && self.runs_deployment_hooks() { let error = self.error.as_ref().ok_or_else(missing_state)?; self.pending = Some(Pending::DeploymentFailure); - return Ok(AdapterStep::Await(DeploymentHooks::after_failure( + return Ok(LifecycleStep::Await(DeploymentHooks::after_failure( py, self.call.kwargs(), error, @@ -331,7 +331,7 @@ impl CallbackAdapter for LegacyLogging { } } - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { self.call @@ -345,7 +345,7 @@ impl CallbackAdapter for LegacyLogging { Pending::DeploymentFailure => self.dispatch_failure(py), Pending::AsyncFailure => match result { Err(failure) if is_cancellation(py, &failure) => Err(failure), - _ => Ok(AdapterStep::Done), + _ => Ok(LifecycleStep::Done), }, } } @@ -357,7 +357,7 @@ impl CallbackAdapter for LegacyLogging { error.write_unraisable(py, None); } self.body = None; - self.headers = None; + self.context = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -369,8 +369,7 @@ impl CallbackAdapter for LegacyLogging { visit.call(&self.end)?; visit.call(&self.response)?; visit.call(&self.error)?; - visit.call(&self.body)?; - visit.call(&self.headers) + visit.call(&self.body) } } diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy/src/call.rs index 59090ee8d60..bd1e2525d3d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy/src/call.rs @@ -4,7 +4,7 @@ //! this crate holds them. use litellm_callbacks::{machine::Machine, route::Route}; -use litellm_host_python::{RouteHost, run_call}; +use litellm_host_python::{RouteHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -63,21 +63,6 @@ impl PublicCall { } } -/// The caller's own object for a public argument, as every legacy reader resolves it: the -/// keyword if given, even an explicit `None`, else the bound request's attribute. A route -/// host projecting from the prepared keyword view uses the same rule, so the callbacks -/// and the provider see one object per argument. -pub fn lookup<'py>( - kwargs: &Bound<'py, PyDict>, - request: &Bound<'py, PyAny>, - name: &str, -) -> PyResult>> { - if let Some(value) = kwargs.get_item(name)? { - return Ok(Some(value)); - } - request.getattr_opt(name) -} - /// Runs one native call under the legacy `Logging` contract: the route host projects from /// the keyword view the contract prepares, and the contract observes the call. pub fn run_legacy_call( @@ -121,32 +106,6 @@ mod tests { (call, locals) } - #[test] - fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { - Python::initialize(); - Python::attach(|py| { - let (call, locals) = capture( - py, - c" -key = object() -document = {'type': 'document_url'} -class Request: - api_key = 'from-request' - api_base = 'from-request' - document = document -request = Request() -kwargs = {'api_key': key, 'api_base': None} -", - ); - let key = locals.get_item("key").unwrap().unwrap(); - let document = locals.get_item("document").unwrap().unwrap(); - assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key)); - assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none()); - assert!(call.lookup(py, "document").unwrap().unwrap().is(&document)); - assert!(call.lookup(py, "model").unwrap().is_none()); - }); - } - #[test] fn capture_copies_the_keyword_dict_without_copying_its_values() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index aa586013e75..9464f1d6612 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -6,11 +6,10 @@ use litellm_callbacks::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use crate::legacy_python::{Logging, Wrapper}; use crate::logger::PythonLogger; pub trait LegacyCallbacks { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult; - /// `Logging.update_from_kwargs`: what the logger is told about the request it is /// about to see, with consumed credentials redacted. fn update_from_kwargs( @@ -21,26 +20,24 @@ pub trait LegacyCallbacks { context: &RequestContext, ) -> PyResult<()>; - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>; - - /// `Logging.pre_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.pre_call`. fn pre_call( &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, ) -> PyResult<()>; - /// `Logging.post_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.post_call`. fn post_call( &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, - headers: Option<&Py>, ) -> PyResult<()>; fn defers_async_logging(&self, py: Python<'_>) -> bool; @@ -82,16 +79,6 @@ pub trait LegacyCallbacks { } impl LegacyCallbacks for PythonLogger { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { - if !self.bridge_owned() { - return Ok(true); - } - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("callbacks_needed")? - .call1((self.object(py), phase))? - .extract() - } - fn update_from_kwargs( &self, py: Python<'_>, @@ -100,18 +87,13 @@ impl LegacyCallbacks for PythonLogger { context: &RequestContext, ) -> PyResult<()> { let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect(); - let update = PyDict::new(py); - update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?; - update.set_item("model", &context.model)?; - update.set_item( - "optional_params", - redact( - py, - &to_py(py, &context.optional_params)? - .into_bound(py) - .cast_into::()?, - &secret_fields, - )?, + let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?; + let optional_params = redact( + py, + &to_py(py, &context.optional_params)? + .into_bound(py) + .cast_into::()?, + &secret_fields, )?; let params = PyDict::new(py); params.set_item( @@ -131,15 +113,17 @@ impl LegacyCallbacks for PythonLogger { params.set_item(name, value)?; } } - update.set_item("litellm_params", params)?; - update.set_item("custom_llm_provider", &context.custom_llm_provider)?; - self.object(py) - .call_method("update_from_kwargs", (), Some(&update))?; - Ok(()) - } - - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> { - self.object(py).call_method0("record_api_call_start_time")?; + Logging::Update.call( + py, + ( + self.object(py), + redacted_kwargs, + &context.model, + optional_params, + params, + &context.custom_llm_provider, + ), + )?; Ok(()) } @@ -147,7 +131,7 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, @@ -156,17 +140,7 @@ impl LegacyCallbacks for PythonLogger { additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; additional.set_item("api_base", url)?; - let kwargs = PyDict::new(py); - kwargs.set_item("input", input)?; - kwargs.set_item("api_key", api_key)?; - kwargs.set_item("additional_args", &additional)?; - if self.callbacks_needed(py, "input")? { - self.object(py).call_method("pre_call", (), Some(&kwargs))?; - } else { - self.object(py) - .call_method("_pre_call", (), Some(&kwargs))?; - self.record_api_call_start(py)?; - } + Logging::PreCall.call(py, (self.object(py), input, api_key, &additional))?; Ok(()) } @@ -174,37 +148,28 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, - headers: Option<&Py>, ) -> PyResult<()> { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; - additional.set_item("headers", headers)?; - if self.callbacks_needed(py, "input")? { - let kwargs = PyDict::new(py); - kwargs.set_item("original_response", original_response)?; - kwargs.set_item("additional_args", &additional)?; - self.object(py) - .call_method("post_call", (), Some(&kwargs))?; - } else { - let response = py - .import("json")? - .call_method1("dumps", (original_response,))?; - self.object(py).call_method1( - "record_post_call", - (response, py.None(), py.None(), additional), - )?; - } + Logging::PostCall.call( + py, + (self.object(py), original_response, api_key, &additional), + )?; Ok(()) } + fn defers_async_logging(&self, py: Python<'_>) -> bool { - self.object(py) - .getattr("_defer_async_logging") - .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + Logging::DefersAsync + .call(py, (self.object(py),)) + .and_then(|value| value.extract()) + .unwrap_or(false) } fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - self.object(py).setattr("_native_pending_logging", pending) + Logging::DeferSuccess.call(py, (self.object(py), pending))?; + Ok(()) } fn sync_success_for_async_call( @@ -214,13 +179,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success_async")? { - return Ok(()); - } - self.object(py).call_method1( - "handle_sync_success_callbacks_for_async_calls", - (response, start, end), - )?; + Logging::SyncSuccessForAsyncCall.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -232,34 +191,11 @@ impl LegacyCallbacks for PythonLogger { end: &Option>, asynchronous: bool, ) -> PyResult>> { - if !self.callbacks_needed( - py, - if asynchronous { - "async_failure" - } else { - "sync_failure" - }, - )? { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("failure_bookkeeping")? - .call1((self.object(py), error, start, end, asynchronous))?; - return Ok(None); - } - let trace = py - .import("traceback")? - .getattr("format_exception")? - .call1((error,))?; - let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?; - let value = self.object(py).call_method1( - if asynchronous { - "async_failure_handler" - } else { - "failure_handler" - }, - (error, trace, start, end), - )?; + let value = + Logging::FailureHandler.call(py, (self.object(py), error, start, end, asynchronous))?; Ok(asynchronous.then(|| value.unbind())) } + fn submit_success( &self, py: Python<'_>, @@ -267,22 +203,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success")? { - return self.success_bookkeeping(py, response, start, end, false); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - py.import("litellm.litellm_core_utils.litellm_logging")? - .getattr("executor")? - .call_method1( - "submit", - ( - context.getattr("run")?, - self.object(py).getattr("success_handler")?, - response, - start, - end, - ), - )?; + Logging::SubmitSuccess.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -293,18 +214,9 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "async_success")? { - return self.success_bookkeeping(py, response, start, end, true); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - let worker = py - .import("litellm.litellm_core_utils.logging_worker")? - .getattr("GLOBAL_LOGGING_WORKER")? - .getattr("ensure_initialized_and_enqueue")?; - let coroutine = self - .object(py) - .call_method1("async_success_handler", (response, start, end))?; - let enqueue = context.call_method1("run", (worker, &coroutine)); + let coroutine = + Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") { @@ -315,14 +227,7 @@ impl LegacyCallbacks for PythonLogger { } fn custom_pricing_fields(py: Python<'_>) -> PyResult> { - py.import("litellm.types.utils")? - .getattr("CustomPricingLiteLLMParams")? - .getattr("model_fields")? - .cast_into::()? - .keys() - .iter() - .map(|name| name.extract::()) - .collect() + Logging::CustomPricingFields.call(py, ())?.extract() } fn redact( @@ -347,58 +252,5 @@ fn redact( /// Proxy-internal calls skip the legacy success fan-out. pub fn is_internal_call(py: Python<'_>) -> PyResult { - py.import("litellm._internal_context")? - .getattr("is_internal_call")? - .call_method0("get")? - .extract() -} - -#[cfg(test)] -mod tests { - use pyo3::types::PyDict; - - use super::*; - - fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger { - let locals = PyDict::new(py); - py.run( - c" -import sys -import types -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): - sys.modules.setdefault(name, types.ModuleType(name)) -legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -class Logger: - needed = {'input': False} -logger = Logger() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - PythonLogger::new( - locals.get_item("logger").unwrap().unwrap().unbind(), - bridge_owned, - ) - } - - #[test] - fn a_caller_owned_logger_is_observed_in_full() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, false); - assert!(logger.callbacks_needed(py, "input").unwrap()); - }); - } - - #[test] - fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, true); - assert!(!logger.callbacks_needed(py, "input").unwrap()); - assert!(logger.callbacks_needed(py, "payload").unwrap()); - }); - } + Wrapper::IsInternalCall.call(py, ())?.extract() } diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs new file mode 100644 index 00000000000..a924775070c --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs @@ -0,0 +1,155 @@ +use pyo3::prelude::*; +use strum::{IntoStaticStr, VariantArray}; + +const MODULE: &str = "litellm.rust_bridge.legacy_callbacks"; + +/// Every litellm Python internal the native call still borrows, grouped by the subsystem it +/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today +/// (span tracking, the standard logging payload, spend, callback fan-out) keeps working. +/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a +/// user's own callback is not borrowing and does not belong here. +/// +/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and +/// `python_contract.json` pins each function's parameters on both sides. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LegacyPython { + Wrapper(Wrapper), + Logging(Logging), + DeploymentHooks(DeploymentHooks), +} + +/// The `@client` wrapper around the call: `function_setup`, limits, credentials, +/// response metadata and the correlation context. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Wrapper { + #[strum(serialize = "setup")] + Setup, + #[strum(serialize = "check_limits")] + CheckLimits, + #[strum(serialize = "credential_list")] + CredentialList, + #[strum(serialize = "warn_unknown_credential")] + WarnUnknownCredential, + #[strum(serialize = "is_internal_call")] + IsInternalCall, + #[strum(serialize = "finalize")] + Finalize, + #[strum(serialize = "restore_context")] + RestoreContext, +} + +/// litellm's `Logging` object and the sync and async callback fan-out behind it. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Logging { + #[strum(serialize = "custom_pricing_fields")] + CustomPricingFields, + #[strum(serialize = "update_logging")] + Update, + #[strum(serialize = "pre_call")] + PreCall, + #[strum(serialize = "post_call")] + PostCall, + #[strum(serialize = "defers_async_logging")] + DefersAsync, + #[strum(serialize = "defer_success")] + DeferSuccess, + #[strum(serialize = "sync_success_for_async_call")] + SyncSuccessForAsyncCall, + #[strum(serialize = "submit_success")] + SubmitSuccess, + #[strum(serialize = "async_success_handler")] + AsyncSuccessHandler, + #[strum(serialize = "enqueue_logging")] + Enqueue, + #[strum(serialize = "failure_handler")] + FailureHandler, +} + +/// The `litellm.utils` fan-outs that run every callback's deployment hook. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum DeploymentHooks { + #[strum(serialize = "before_deployment_call")] + BeforeDeploymentCall, + #[strum(serialize = "after_deployment_success")] + AfterDeploymentSuccess, + #[strum(serialize = "after_deployment_failure")] + AfterDeploymentFailure, +} + +impl LegacyPython { + fn name(self) -> &'static str { + match self { + Self::Wrapper(function) => function.into(), + Self::Logging(function) => function.into(), + Self::DeploymentHooks(function) => function.into(), + } + } + + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + py.import(MODULE)?.getattr(self.name())?.call1(args) + } +} + +impl Wrapper { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Wrapper(self).call(py, args) + } +} + +impl Logging { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Logging(self).call(py, args) + } +} + +impl DeploymentHooks { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::DeploymentHooks(self).call(py, args) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use strum::VariantArray; + + use super::{DeploymentHooks, LegacyPython, Logging, Wrapper}; + use crate::test_support::PYTHON_CONTRACT; + + #[test] + fn every_borrowed_function_is_in_the_python_contract() { + let contract: serde_json::Map = + serde_json::from_str(PYTHON_CONTRACT).unwrap(); + let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); + let called: Vec<&str> = Wrapper::VARIANTS + .iter() + .map(|&function| LegacyPython::Wrapper(function)) + .chain( + Logging::VARIANTS + .iter() + .map(|&function| LegacyPython::Logging(function)), + ) + .chain( + DeploymentHooks::VARIANTS + .iter() + .map(|&function| LegacyPython::DeploymentHooks(function)), + ) + .map(LegacyPython::name) + .collect(); + assert_eq!(called.len(), declared.len(), "a function is borrowed twice"); + assert_eq!(called.into_iter().collect::>(), declared); + } +} diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy/src/lib.rs index 06783ac255d..42ffd545e2b 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy/src/lib.rs @@ -2,7 +2,7 @@ //! sync and async callback registries it fans out to, the deployment hooks, the deferred //! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name //! inheritance, budget and retry-count limits). All of it sits behind one -//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and +//! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and //! core never learn which Python object is on the other end. //! //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`] @@ -13,6 +13,7 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod legacy_python; mod logger; mod preparation; #[cfg(test)] @@ -21,7 +22,7 @@ mod test_support; pub(crate) use adapter::LegacyLogging; pub use adapter::LegacySurface; -pub use call::{PublicCall, lookup, run_legacy_call}; +pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; diff --git a/litellm-rust/crates/callbacks-legacy/src/logger.rs b/litellm-rust/crates/callbacks-legacy/src/logger.rs index a0e525000b8..061941f05b9 100644 --- a/litellm-rust/crates/callbacks-legacy/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy/src/logger.rs @@ -5,34 +5,25 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -/// The `Logging` instance one call fans out through, and who owns it. A logger the caller -/// handed in is observed in full, because the caller reads it after the call; one this -/// crate built through `function_setup` is elided wherever no registry needs it. +use crate::legacy_python::{self, Wrapper}; + +/// The `Logging` instance one call fans out through. pub struct PythonLogger { object: Py, - bridge_owned: bool, } impl PythonLogger { - pub(crate) fn new(object: Py, bridge_owned: bool) -> Self { - Self { - object, - bridge_owned, - } + pub(crate) fn new(object: Py) -> Self { + Self { object } } pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> { self.object.bind(py) } - pub(crate) fn bridge_owned(&self) -> bool { - self.bridge_owned - } - pub fn clone_ref(&self, py: Python<'_>) -> Self { Self { object: self.object.clone_ref(py), - bridge_owned: self.bridge_owned, } } @@ -40,34 +31,17 @@ impl PythonLogger { visit.call(&self.object) } - pub fn success_bookkeeping( - &self, - py: Python<'_>, - response: &Option>, - start: &Py, - end: &Option>, - asynchronous: bool, - ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("success_bookkeeping")? - .call1((self.object(py), response, start, end, asynchronous))?; - Ok(()) - } - pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> { - py.import("litellm.utils")? - .getattr("_restore_correlation_context_if_supported")? - .call1((self.object(py),))?; + Wrapper::RestoreContext.call(py, (self.object(py),))?; Ok(()) } } -/// A bare Python object was not obtained from `setup`, so it is caller-owned. impl FromPyObject<'_, '_> for PythonLogger { type Error = PyErr; fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult { - Ok(Self::new(object.to_owned().unbind(), false)) + Ok(Self::new(object.to_owned().unbind())) } } @@ -75,9 +49,7 @@ pub struct SetupResult<'py>(Bound<'py, PyAny>); impl SetupResult<'_> { pub fn logger(&self) -> PyResult { - let object = self.0.getattr("logger")?.unbind(); - let bridge_owned = self.0.getattr("bridge_owned")?.extract()?; - Ok(PythonLogger::new(object, bridge_owned)) + Ok(PythonLogger::new(self.0.getattr("logger")?.unbind())) } pub fn kwargs(&self) -> PyResult> { @@ -93,9 +65,8 @@ pub fn setup<'py>( start: &Py, asynchronous: bool, ) -> PyResult> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("setup")? - .call1((call_type, args, kwargs, start, asynchronous)) + Wrapper::Setup + .call(py, (call_type, args, kwargs, start, asynchronous)) .map(SetupResult) } @@ -107,30 +78,20 @@ pub fn finalize( start: &Py, end: &Option>, ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("finalize")? - .call1((response, logger.object(py), kwargs, start, end))?; + Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?; Ok(()) } pub struct DeploymentHooks; impl DeploymentHooks { - pub fn needed(py: Python<'_>) -> PyResult { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("deployment_callbacks_needed")? - .call0()? - .extract() - } - pub fn before_call( py: Python<'_>, kwargs: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_pre_call_deployment_hook")? - .call1((kwargs, call_type)) + legacy_python::DeploymentHooks::BeforeDeploymentCall + .call(py, (kwargs, call_type)) .map(Bound::unbind) } @@ -140,9 +101,8 @@ impl DeploymentHooks { response: &Option>, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_success_deployment_hook")? - .call1((kwargs, response, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentSuccess + .call(py, (kwargs, response, call_type)) .map(Bound::unbind) } @@ -152,9 +112,8 @@ impl DeploymentHooks { error: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_failure_deployment_hook")? - .call1((kwargs, error, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentFailure + .call(py, (kwargs, error, call_type)) .map(Bound::unbind) } } @@ -185,10 +144,6 @@ class Setup: reads.append('logger') return logger @property - def bridge_owned(self): - reads.append('bridge_owned') - return True - @property def kwargs(self): reads.append('kwargs') return [] @@ -206,7 +161,6 @@ result = Setup() .object(py) .is(locals.get_item("logger").unwrap().unwrap()) ); - assert!(logger.bridge_owned()); assert!( result .kwargs() @@ -220,17 +174,8 @@ result = Setup() .unwrap() .extract::>() .unwrap(), - ["logger", "bridge_owned", "kwargs"] + ["logger", "kwargs"] ); }); } - - #[test] - fn a_logger_extracted_from_a_bare_object_is_caller_owned() { - Python::initialize(); - Python::attach(|py| { - let logger: PythonLogger = py.None().into_bound(py).extract().unwrap(); - assert!(!logger.bridge_owned()); - }); - } } diff --git a/litellm-rust/crates/callbacks-legacy/src/preparation.rs b/litellm-rust/crates/callbacks-legacy/src/preparation.rs index 981b1702f2e..fa1ff9acd4d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/preparation.rs +++ b/litellm-rust/crates/callbacks-legacy/src/preparation.rs @@ -3,6 +3,8 @@ use pyo3::{ types::{PyDict, PyList}, }; +use crate::legacy_python::Wrapper; + struct CredentialEntry<'py>(Bound<'py, PyAny>); impl<'py> CredentialEntry<'py> { @@ -22,18 +24,19 @@ pub fn prepare<'py>( ) -> PyResult> { let arguments = kwargs.copy()?; arguments.set_item("litellm_logging_obj", logger.object(py))?; - let litellm = py.import("litellm")?; - inherit_credentials(py, &litellm, &arguments)?; - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("check_limits")? - .call1((&arguments,))?; + inherit_credentials(py, &arguments, || { + Ok(Wrapper::CredentialList + .call(py, ())? + .cast_into::()?) + })?; + Wrapper::CheckLimits.call(py, (&arguments,))?; Ok(arguments) } -fn inherit_credentials( - py: Python<'_>, - litellm: &Bound<'_, PyModule>, - arguments: &Bound<'_, PyDict>, +fn inherit_credentials<'py>( + py: Python<'py>, + arguments: &Bound<'py, PyDict>, + credential_list: impl FnOnce() -> PyResult>, ) -> PyResult<()> { let Some(requested) = arguments .get_item("litellm_credential_name")? @@ -45,16 +48,13 @@ fn inherit_credentials( return Ok(()); } let requested: String = requested.extract()?; - let credentials = litellm.getattr("credential_list")?.cast_into::()?; + let credentials = credential_list()?; let names = credentials .iter() .map(|credential| CredentialEntry(credential).name()) .collect::>>()?; let Some(index) = names.iter().position(|name| *name == requested) else { - py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1( - "warning", - ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()), - )?; + Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?; return Ok(()); }; let selected = CredentialEntry(credentials.get_item(index)?); @@ -80,19 +80,19 @@ mod tests { } fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> { - let litellm = PyModule::new(py, "credential_host")?; - litellm.setattr( - "credential_list", - locals.get_item("credentials").unwrap().unwrap(), - )?; inherit_credentials( py, - &litellm, &locals .get_item("arguments") .unwrap() .unwrap() .cast_into::()?, + || { + Ok(locals + .get_item("credentials")? + .unwrap() + .cast_into::()?) + }, ) } @@ -304,11 +304,11 @@ arguments = {'litellm_credential_name': 'ocr-test'} fn falsy_credential_names_return_before_loading_credentials() { Python::initialize(); Python::attach(|py| { - let litellm = PyModule::new(py, "credential_host").unwrap(); for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] { let arguments = PyDict::new(py); arguments.set_item("litellm_credential_name", name).unwrap(); - inherit_credentials(py, &litellm, &arguments).unwrap(); + inherit_credentials(py, &arguments, || panic!("credentials must not be loaded")) + .unwrap(); } }); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs index 3daea8840d8..289ea1b2e7f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs @@ -16,7 +16,7 @@ fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { py, PendingLogging { pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind(), true), + logger: PythonLogger::new(local(&locals, "logger").unbind()), response: Some(local(&locals, "response").unbind()), start: py.None(), end: Some(py.None()), @@ -79,22 +79,6 @@ assert logger.calls == [], logger.calls }); } -#[test] -fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c"logger.needed = {'async_success': False}"); - run( - py, - &locals, - c" -pending.release(True) -assert logger.calls == [('success_bookkeeping', True)], logger.calls -", - ); - }); -} - #[rstest] #[case::ordinary_error(c"RuntimeError('queue full')", false)] #[case::cancellation(c"asyncio.CancelledError()", true)] diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 3ceda4441a7..ea3510de17e 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -24,7 +24,7 @@ fn begin<'py>( py: Python<'py>, locals: &Bound<'py, PyDict>, asynchronous: bool, -) -> (LegacyLogging, AdapterStep) { +) -> (LegacyLogging, LifecycleStep) { let mut logging = legacy_call(py, locals, asynchronous); let kwargs = local(locals, "kwargs") .cast_into::() @@ -34,15 +34,15 @@ fn begin<'py>( (logging, step) } -fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> { - let AdapterStep::Arguments(arguments) = step else { +fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { panic!("expected the prepared arguments"); }; arguments.into_bound(py) } -fn awaits_deployment_hook(step: &AdapterStep) -> bool { - matches!(step, AdapterStep::Await(_)) +fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) } #[rstest] @@ -121,7 +121,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs} let step = logging .resume(py, Ok(local(&locals, "replacement").unbind())) .unwrap(); - let AdapterStep::Response(returned) = step else { + let LifecycleStep::Response(returned) = step else { panic!("expected the finalized response"); }; assert!(returned.bind(py).is(local(&locals, "replacement"))); @@ -195,7 +195,7 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle }; assert!(matches!( logging.resume(py, hook_result).unwrap(), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); run( py, @@ -237,7 +237,7 @@ kwargs = {'logger': logger} .unwrap() .unbind(); let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), + LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), step => Ok(step), }); let error = result.err().unwrap(); diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 480bedf8548..68c0a2b1e15 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,7 +1,8 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{AdapterStep, CallbackAdapter}; +use litellm_auth::SecretValue; +use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleStep, PythonLifecycle}; use pyo3::prelude::*; use rstest::rstest; use serde_json::{Value, json}; @@ -23,20 +24,12 @@ class PayloadLogger(StubLogger): def pre_call(self, input, api_key, additional_args): self.record('pre_call', None) self.pre = additional_args + self.pre_api_key = api_key on_pre_call(additional_args) - def _pre_call(self, input, api_key, additional_args): - self.record('_pre_call', None) - - def record_api_call_start_time(self): - self.record('record_api_call_start_time', None) - - def post_call(self, original_response, additional_args): + def post_call(self, original_response, api_key, additional_args): self.record('post_call', None) - self.post = (original_response, additional_args) - - def record_post_call(self, response, *rest): - self.record('record_post_call', response) + self.post = (original_response, api_key, additional_args) request = Request() kwargs = {} @@ -52,16 +45,16 @@ fn document(source: &str) -> Value { json!({"type": "document_url", "document_url": source}) } -fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest { - before_send_with_secrets(script, caller, body, &[]) +fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) } -/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the -/// Python objects `script` binds, then delivers the provider's raw response the way the +/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with +/// the Python objects `script` binds, then delivers the provider's raw response the way the /// driver does and runs the script's `check()`. fn before_send_with_secrets( script: &CStr, - caller: Value, + optional_params: Value, body: Value, secret_fields: &[&str], ) -> WireRequest { @@ -70,15 +63,15 @@ fn before_send_with_secrets( let locals = namespace(py, PAYLOAD_LOGGER); run(py, &locals, script); let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), ..legacy_call(py, &locals, false) }; let context = RequestContext { model: "model".into(), custom_llm_provider: "provider".into(), - optional_params: caller.clone(), - passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body), + optional_params, secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), }; let wire = WireRequest { url: "https://provider.invalid/ocr".into(), @@ -93,10 +86,10 @@ fn before_send_with_secrets( }; assert!(matches!( logging.emit(py, &raw, None).unwrap(), - AdapterStep::Done + LifecycleStep::Done )); run(py, &locals, c"check()"); - let AdapterStep::Wire(wire) = step else { + let LifecycleStep::Wire(wire) = step else { panic!("before_send did not hand back the wire request"); }; *wire @@ -129,11 +122,7 @@ def check(): ")] fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send( - script, - json!({"document": document(DOCUMENT), "pages": [0]}), - body.clone(), - ); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); } @@ -149,7 +138,6 @@ def check(): assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' ", json!({"document": document(DOCUMENT)}), - json!({"document": document(DOCUMENT)}), ); assert_eq!(wire.body["document"], document(EDITED)); } @@ -168,7 +156,6 @@ def check(): assert observed == [False], observed assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} ", - json!({"document": document("https://example.invalid/scan.pdf")}), json!({"document": document(DOCUMENT)}), ); assert_eq!( @@ -177,6 +164,23 @@ def check(): ); } +#[test] +fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); +} + #[rstest] #[case::body( c" @@ -192,7 +196,7 @@ def on_pre_call(args): )] fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({}), body.clone()); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); } @@ -205,7 +209,6 @@ def on_pre_call(args): args['headers']['x-callback'] = 'edited' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -288,7 +291,7 @@ def on_pre_call(args): )] fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({"document": document(DOCUMENT)}), body); + let wire = before_send(script, body); assert_eq!(wire.body, expected); } @@ -302,7 +305,6 @@ def on_pre_call(args): retained['x-retained'] = 'sent' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -314,52 +316,33 @@ def on_pre_call(args): } #[test] -fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() { +fn post_call_receives_the_raw_response_the_route_key_and_the_body_pre_call_saw() { before_send( c" def check(): - original_response, additional_args = logger.post + original_response, api_key, additional_args = logger.post assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == {'complete_input_dict': logger.pre['complete_input_dict']}, additional_args assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] - assert additional_args['headers'] is logger.pre['headers'] ", - json!({}), json!({"document": document(DOCUMENT)}), ); } -#[rstest] -#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])] -#[case::no_input_callback( - c"{'input': False}", - &["_pre_call", "record_api_call_start_time", "record_post_call"] -)] -#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])] -fn payload_callbacks_run_only_for_the_phases_someone_listens_to( - #[case] needed: &CStr, - #[case] expected_calls: &[&str], -) { - let script = std::ffi::CString::new(format!( - " -logger.needed = {needed} +#[test] +fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" def on_pre_call(args): args['complete_input_dict']['include_image_base64'] = True def check(): - assert logger.names() == {expected_calls:?}, logger.calls + assert logger.names() == ['pre_call', 'post_call'], logger.calls ", - needed = needed.to_str().unwrap(), - expected_calls = expected_calls, - )) - .unwrap(); - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(&script, json!({}), body.clone()); - let edited = json!({"document": document(DOCUMENT), "include_image_base64": true}); + json!({"document": document(DOCUMENT)}), + ); assert_eq!( wire.body, - if expected_calls.contains(&"pre_call") { - edited - } else { - body - } + json!({"document": document(DOCUMENT), "include_image_base64": true}) ); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 1663e11963e..444655ea77b 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -5,64 +5,89 @@ use pyo3::types::{PyDict, PyTuple}; use crate::{LegacyLogging, LegacySurface, PublicCall}; -/// Stand-ins for every litellm function the legacy contract calls. Tests share one -/// interpreter and run concurrently, so each stub is installed idempotently and forwards to -/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// The parameters of every `legacy_callbacks` function, as the real module declares them. +/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python +/// signatures, and [`namespace`] binds every fake call against it. +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + +/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests +/// share one interpreter and run concurrently, so each fake is installed idempotently and +/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// Every fake is bound against the contract first, so a call the real module would reject +/// fails here too. const STUBS: &CStr = c" import contextvars +import inspect +import json import sys +import traceback import types -for name in ( - 'litellm', - 'litellm.utils', - 'litellm.types', - 'litellm.types.utils', - 'litellm._internal_context', - 'litellm.litellm_core_utils', - 'litellm.litellm_core_utils.logging_worker', - 'litellm.litellm_core_utils.litellm_logging', - 'litellm.rust_bridge', - 'litellm.rust_bridge.legacy_callbacks', -): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): sys.modules.setdefault(name, types.ModuleType(name)) legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - bridge_owned=True, -) -legacy.deployment_callbacks_needed = lambda: True -legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments) -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record( - 'success_bookkeeping', asynchronous -) -legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record( - 'failure_bookkeeping', asynchronous -) -legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response) +CONTRACT = json.loads(python_contract) -utils = sys.modules['litellm.utils'] -utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook( - 'pre', kwargs, call_type -) -utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[ - 'logger' -].hook('success', response, call_type) -utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[ - 'logger' -].hook('failure', error, call_type) -utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None) -internal = sys.modules['litellm._internal_context'] -if not hasattr(internal, 'is_internal_call'): - internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False) +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) -sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type( - 'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}} -) + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) unraisable = sys.modules.setdefault( @@ -77,20 +102,6 @@ def unraisable_from(owner): return [error for source, error in unraisable.events if source is owner] -class Worker: - def ensure_initialized_and_enqueue(self, coroutine): - return coroutine.enqueue() - - -class Executor: - def submit(self, run, handler, *args): - handler.__self__.record('submit', args) - - -sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker() -sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor() - - class StubCoroutine: def __init__(self, logger): self.logger = logger @@ -106,7 +117,6 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] - self.needed = {} self.hooks = {} self.on_enqueue = lambda coroutine: None @@ -147,6 +157,7 @@ logger = StubLogger() /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); py.run(script, Some(&locals), Some(&locals)).unwrap(); locals diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 9b9d29108f6..3094d7b88d2 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host_python::{LifecycleStep, PublicValue, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -19,12 +19,16 @@ const TIMING: Timing = Timing { fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), ..legacy_call(py, locals, asynchronous) } } -fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, +) -> LifecycleStep { let response = local(locals, "response").unbind(); logging .emit( @@ -35,7 +39,7 @@ fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLoggi .unwrap() } -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { let failure = PyErr::from_value(local(locals, "failure")); logging .emit( @@ -51,20 +55,14 @@ fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) #[rstest] #[case::sync_listened(false, c"", &["submit"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])] #[case::async_listened( true, c"", &["async_success_handler", "enqueued", "sync_success_for_async_call"] )] -#[case::async_unlistened( - true, - c"logger.needed = {'async_success': False, 'sync_success_async': False}", - &["success_bookkeeping"] -)] #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_only_the_callbacks_that_listen( +fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -76,7 +74,7 @@ fn success_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); let names: Vec = local(&locals, "logger") .call_method0("names") @@ -109,7 +107,10 @@ fn internal_calls_skip_failure_callbacks_only_when_asynchronous( internal: true, ..logged(py, &locals, asynchronous) }; - assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done)); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -157,7 +158,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); assert!( logging @@ -173,14 +174,8 @@ logger = FailingLogger() #[rstest] #[case::sync_listened(false, c"", &["failure_handler"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])] #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -#[case::async_unlistened( - true, - c"logger.needed = {'sync_failure': False, 'async_failure': False}", - &["failure_bookkeeping", "failure_bookkeeping"] -)] -fn failure_reaches_only_the_callbacks_that_listen( +fn failure_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -192,7 +187,10 @@ fn failure_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); let step = fail(py, &locals, &mut logging); let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -227,7 +225,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( fail(py, &locals, &mut logging), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); assert!( logging @@ -265,7 +263,7 @@ fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( }; let expected = result.as_ref().err().map(|error| error.value(py).clone()); match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, AdapterStep::Done)), + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), Err(propagated) => { assert!(!done); assert!(propagated.value(py).is(expected.unwrap())); diff --git a/litellm-rust/crates/callbacks/Cargo.toml b/litellm-rust/crates/callbacks/Cargo.toml index 4b966271478..a68ebc26a8d 100644 --- a/litellm-rust/crates/callbacks/Cargo.toml +++ b/litellm-rust/crates/callbacks/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-auth.workspace = true serde_json.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs index e6f88fd9709..3bf12e553b9 100644 --- a/litellm-rust/crates/callbacks/src/event.rs +++ b/litellm-rust/crates/callbacks/src/event.rs @@ -1,6 +1,6 @@ use std::time::{SystemTime, UNIX_EPOCH}; -use serde_json::{Map, Value}; +use serde_json::Value; /// Seconds since the Unix epoch, on one clock for every host. pub fn epoch_seconds() -> f64 { @@ -33,34 +33,10 @@ pub struct RequestContext { pub custom_llm_provider: String, /// The route's parameters before the provider transformation. pub optional_params: Value, - pub passthrough_fields: Passthrough, /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, -} - -/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to -/// build one is to compare the two, so a route cannot name a key it rewrote. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Passthrough(Vec); - -impl Passthrough { - pub fn unchanged(caller: &Map, body: &Value) -> Self { - Self( - caller - .iter() - .filter(|(name, value)| body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.clone()) - .collect(), - ) - } - - pub fn iter(&self) -> impl Iterator { - self.0.iter().map(String::as_str) - } - - pub fn contains(&self, name: &str) -> bool { - self.0.iter().any(|field| field == name) - } + /// The credential the route resolved for the provider call. + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -78,6 +54,9 @@ pub enum FailureOrigin { #[derive(Clone, Debug, PartialEq)] pub enum CallEvent { + Started { + start_time: f64, + }, ResponseReceived { raw: RawResponse, }, @@ -89,47 +68,3 @@ pub enum CallEvent { origin: FailureOrigin, }, } - -#[cfg(test)] -mod tests { - use rstest::rstest; - use serde_json::json; - - use super::*; - - #[rstest] - #[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])] - #[case::unchanged_nested_object( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}), - &["document"] - )] - #[case::rewritten_value( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}), - &[] - )] - #[case::dropped_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - &[] - )] - #[case::added_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}), - &[] - )] - #[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])] - #[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])] - #[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])] - fn passthrough_is_exactly_the_callers_unchanged_keys( - #[case] caller: Value, - #[case] body: Value, - #[case] expected: &[&str], - ) { - let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body); - assert_eq!(passthrough.iter().collect::>(), expected); - } -} diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/callbacks/src/run.rs index 57bf134f345..5accaa2e25e 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/callbacks/src/run.rs @@ -11,6 +11,7 @@ where H: Host, { let start_time = epoch_seconds(); + let _ = host.emit(&CallEvent::Started { start_time }).await; let mut result = None; let outcome = loop { let step = match machine.resume(result.take()).await { @@ -102,6 +103,7 @@ mod tests { async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { self.seen.lock().unwrap().push(match event { + CallEvent::Started { .. } => "started".into(), CallEvent::Succeeded { .. } => "succeeded".into(), CallEvent::Failed { .. } => "failed".into(), other => format!("{other:?}"), @@ -124,7 +126,7 @@ mod tests { assert_eq!(outcome, Ok(())); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "succeeded"] + ["started", "route:project", "route:send", "succeeded"] ); } @@ -133,7 +135,7 @@ mod tests { let host = Recording::default(); let outcome = run(scripted(&[], Err("boom")), &host).await; assert_eq!(outcome, Err("boom")); - assert_eq!(*host.seen.lock().unwrap(), ["failed"]); + assert_eq!(*host.seen.lock().unwrap(), ["started", "failed"]); let host = Recording { fail: Some("send"), @@ -143,7 +145,35 @@ mod tests { assert_eq!(outcome, Err("host failed")); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "failed"] + ["started", "route:project", "route:send", "failed"] ); } + + struct StartTimes(Mutex>); + + impl Host for StartTimes { + async fn route(&self, _: &'static str) -> Result<(), &'static str> { + Ok(()) + } + + async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { + if let CallEvent::Started { start_time } + | CallEvent::Succeeded { + timing: Timing { start_time, .. }, + } = event + { + self.0.lock().unwrap().push(*start_time); + } + Err("observer failed") + } + } + + #[tokio::test] + async fn started_opens_the_call_at_the_terminal_start_time_and_cannot_fail_it() { + let host = StartTimes(Mutex::default()); + assert_eq!(run(scripted(&["project"], Ok(())), &host).await, Ok(())); + let times = host.0.lock().unwrap(); + assert_eq!(times.len(), 2); + assert_eq!(times[0], times[1]); + } } diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 33cb8a8d32a..a6af190eb91 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,5 +1,6 @@ use futures_util::future::BoxFuture; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; +use litellm_auth::SecretValue; +use litellm_callbacks::event::{CallEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -36,6 +37,7 @@ pub(crate) struct OcrCallHooks { custom_llm_provider: &'static str, optional_params: Value, secret_fields: Vec, + api_key: Option, } impl OcrCallHooks { @@ -51,22 +53,19 @@ impl OcrCallHooks { .filter(|name| is_secret_param(name)) .cloned() .collect(), + api_key: request.connection.api_key.clone(), } } } impl CallHooks for OcrCallHooks { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { let context = RequestContext { model: self.model.clone(), custom_llm_provider: self.custom_llm_provider.into(), optional_params: self.optional_params.clone(), - passthrough_fields, secret_fields: self.secret_fields.clone(), + api_key: self.api_key.clone(), }; Box::pin(self.host.before_send(wire, context)) } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index e7f77acc3f8..c977f721a70 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -21,8 +21,8 @@ mod cohere_tests; #[path = "../../tests/deepseek_ocr.rs"] mod deepseek_tests; #[cfg(test)] -#[path = "../../tests/ocr/passthrough.rs"] -mod passthrough_tests; +#[path = "../../tests/ocr/document.rs"] +mod document_tests; #[cfg(test)] #[path = "../../tests/reducto_ocr.rs"] mod reducto_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 24c3f43e2b4..8ac038290b7 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,4 +1,4 @@ -use litellm_auth::{InputSource, Sourced}; +use litellm_auth::{InputSource, SecretValue, Sourced}; use litellm_llms::base_llm::ocr::transformation::{ OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env, }; @@ -22,7 +22,7 @@ pub(crate) fn prepare_request( .config .get_api_key_env_var() .and_then(credential_env) - .map(|value| Sourced::new(value, InputSource::Environment)) + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index d12b8cfee95..14b34ea4564 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -277,12 +277,18 @@ mod tests { #[test] fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Request, @@ -292,7 +298,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("dynamic-key") ); assert_eq!( @@ -318,22 +324,31 @@ mod tests { fn empty_or_missing_dynamic_credentials_preserve_explicit_values( #[case] dynamic_value: Option<&str>, ) { - let dynamic = + let dynamic_key = dynamic_value.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Environment, + ) + }); + let dynamic_base = dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: dynamic.clone(), - dynamic_api_base: dynamic, + dynamic_api_key: dynamic_key, + dynamic_api_base: dynamic_base, }); assert_eq!( connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("explicit-key") ); assert_eq!( @@ -356,11 +371,18 @@ mod tests { ) { let connection = OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params( OcrCredentialInputs { - api_key: explicit_key - .map(|value| Sourced::new(value.into(), InputSource::Deployment)), + api_key: explicit_key.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Deployment, + ) + }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Deployment, @@ -371,7 +393,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), explicit_key.map(|_| "dynamic-key") ); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 75202ed52a5..6316088dec8 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, @@ -56,7 +56,7 @@ pub struct OcrFileContent { /// credentials, and per-field provenance in `input_sources`. #[derive(Clone, Debug, Default)] pub struct OcrConnectionInputs { - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub extra_headers: Map, pub timeout: Option, @@ -237,6 +237,16 @@ mod tests { .unwrap() } + #[test] + fn connection_inputs_debug_hides_the_api_key() { + let inputs = OcrConnectionInputs { + api_key: Some(SecretValue::new("caller-api-key")), + ..OcrConnectionInputs::default() + }; + + assert!(!format!("{inputs:?}").contains("caller-api-key")); + } + #[test] fn from_inputs_applies_connection_overrides_with_field_sources() { let request = LiteLLMOcrRequest::from_inputs( @@ -245,7 +255,7 @@ mod tests { None, Default::default(), OcrConnectionInputs { - api_key: Some(" key ".into()), + api_key: Some(SecretValue::new(" key ")), api_base: Some("".into()), extra_headers: json!({"x-a": "1"}).as_object().unwrap().clone(), timeout: Some(Duration::from_secs(7)), @@ -259,7 +269,7 @@ mod tests { .unwrap(); let api_key = request.credentials.api_key.as_ref().unwrap(); - assert_eq!(api_key.clone().into_value(), "key"); + assert_eq!(api_key.value().expose(), "key"); assert_eq!(api_key.source(), InputSource::Request); assert!(request.credentials.api_base.is_none()); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 29345e38885..b9c60f57e3c 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::InputSource; +use litellm_auth::{InputSource, SecretValue}; use litellm_llms::base_llm::ocr::{ error::Error, transformation::{OcrDocument, decode_request_value}, @@ -44,7 +44,7 @@ pub fn consumed_optional_param_names( pub struct OcrWireRequest { pub model: String, pub document: D, - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, pub extra_headers: Option>, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 01a4e5efb3b..1fc4d6c2b9e 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -69,7 +69,7 @@ async fn rejects_invalid_pages_features_and_format( let result = decode_request(OcrWireRequest { model: "azure_ai/doc-intelligence/prebuilt-read".into(), document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: Some(base), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 1f591d74d5d..ca2e14a7f0d 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -81,7 +81,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { let request = OcrWireRequest { model: "mistral/model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: None, extra_headers: None, @@ -97,7 +97,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { decode_request(OcrWireRequest { model: "model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: Some("unknown".into()), extra_headers: None, @@ -194,6 +194,7 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { + CallEvent::Started { .. } => "started", CallEvent::ResponseReceived { .. } => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", @@ -235,7 +236,7 @@ async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { } #[tokio::test] -async fn before_send_context_names_passthrough_fields_and_secrets() { +async fn before_send_context_names_the_route_and_its_secrets() { let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; let observed = Arc::new(Mutex::new(None)); let captured = observed.clone(); @@ -254,8 +255,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { assert_eq!(context.custom_llm_provider, "mistral"); assert_eq!(context.model, "model"); assert_eq!(wire.body["pages"], json!([0])); - assert!(context.passthrough_fields.contains("pages")); - assert!(context.passthrough_fields.contains("document")); assert!(context.secret_fields.is_empty()); assert_eq!(context.optional_params["req_format"], "native"); @@ -279,7 +278,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { perform_ocr_with(host).await.unwrap(); server.await.unwrap(); let context = observed.lock().unwrap().take().unwrap(); - assert!(!context.passthrough_fields.contains("document")); assert_eq!(context.secret_fields, ["client_secret"]); } @@ -296,7 +294,7 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { server.await.unwrap(); assert_eq!( *events.lock().unwrap(), - ["before_send", "response", "success"] + ["started", "before_send", "response", "success"] ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -311,7 +309,10 @@ async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { ); let error = perform_ocr_with(host).await.unwrap_err(); assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); } #[tokio::test] @@ -330,7 +331,10 @@ async fn upstream_failure_emits_one_terminal_failure() { ); assert!(perform_ocr_with(host).await.is_err()); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs new file mode 100644 index 00000000000..5e10ce3ad2a --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/document.rs @@ -0,0 +1,152 @@ +use litellm_callbacks::event::WireRequest; +use litellm_llms::base_llm::ocr::error::Error; +use rstest::rstest; +use serde_json::{Value, json}; + +use super::test_support::{ + MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, + wire_request_with_document, +}; +use crate::ocr::route::LocalOcrHost; + +#[derive(Clone, Copy, Debug)] +enum Route { + Mistral, + AzureAi, + VertexMistral, + AzureCohereParse, + Cohere, +} + +impl Route { + fn model(self) -> &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } +} + +/// What the host does to the wire request in `before_send`. +#[derive(Clone, Copy, Debug)] +enum Host { + Detached, + ReplacesDocument, +} + +const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + +impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } +} + +struct Sent { + result: Result<(), Error>, + provider_body: Option, +} + +async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } +} + +fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) +} + +#[rstest] +#[case::azure_ai(Route::AzureAi)] +#[case::vertex_mistral(Route::VertexMistral)] +#[case::azure_cohere_parse(Route::AzureCohereParse)] +#[tokio::test] +async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); +} + +#[rstest] +#[tokio::test] +async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, +) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/passthrough.rs b/litellm-rust/crates/core/tests/ocr/passthrough.rs deleted file mode 100644 index 0273b48664d..00000000000 --- a/litellm-rust/crates/core/tests/ocr/passthrough.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::{ - collections::BTreeSet, - sync::{Arc, Mutex}, -}; - -use litellm_callbacks::event::{RequestContext, WireRequest}; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use rstest_reuse::{self, apply, template}; -use serde_json::{Map, Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -#[derive(Clone, Copy, Debug)] -enum Source { - Inline, - Remote, - RemoteWithExtraField, -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - /// What `litellm-callbacks-legacy` does before `pre_call`: every passthrough body key - /// is replaced by the caller's own value. - Realiasing, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send( - self, - caller: &Map, - wire: WireRequest, - context: &RequestContext, - ) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::Realiasing => { - let aliased = context - .passthrough_fields - .contains(&name) - .then(|| caller.get(&name).cloned()) - .flatten() - .unwrap_or(value); - (name, aliased) - } - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - caller: Map, - result: Result<(), Error>, - before_send: Option<(WireRequest, RequestContext)>, - provider_body: Option, -} - -fn caller_document(route: Route, source: Source, document_base: &str) -> Value { - let document_type = route.document_type(); - let remote = format!("{document_base}/scan.png"); - match source { - Source::Inline => { - json!({"type": document_type, document_type: "data:image/png;base64,YWJj"}) - } - Source::Remote => json!({"type": document_type, document_type: remote}), - Source::RemoteWithExtraField => { - json!({"type": document_type, document_type: remote, "document_name": "scan.png"}) - } - } -} - -async fn send(route: Route, source: Source, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document = caller_document(route, source, document_base); - let caller: Map = route - .options() - .as_object() - .unwrap() - .clone() - .into_iter() - .chain([("document".to_string(), document.clone())]) - .collect(); - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host_caller = caller.clone(); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(host.before_send(&host_caller, wire, context)) - }); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - let before_send = observed.lock().unwrap().take(); - Sent { - caller, - result, - before_send, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[template] -#[rstest] -fn every_route_and_source( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, - #[values(Source::Inline, Source::Remote, Source::RemoteWithExtraField)] source: Source, -) { -} - -#[template] -#[rstest] -fn every_route( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { -} - -#[template] -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -fn inlining_routes(#[case] route: Route) {} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn passthrough_fields_are_exactly_the_caller_values_sent_unchanged( - route: Route, - source: Source, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, source, Host::Detached, &document_base).await; - sent.result.unwrap(); - let (wire, context) = sent.before_send.unwrap(); - let passthrough: BTreeSet<&str> = context.passthrough_fields.iter().collect(); - let unchanged: BTreeSet<&str> = sent - .caller - .iter() - .filter(|(name, value)| wire.body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.as_str()) - .collect(); - assert_eq!( - passthrough, - unchanged, - "body: {:#}\ncaller: {:#}", - wire.body, - Value::Object(sent.caller.clone()) - ); -} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn realiasing_leaves_the_provider_request_unchanged(route: Route, source: Source) { - let (document_base, _documents) = document_server().await; - let detached = send(route, source, Host::Detached, &document_base).await; - let realiased = send(route, source, Host::Realiasing, &document_base).await; - detached.result.unwrap(); - realiased.result.unwrap(); - assert_eq!(realiased.provider_body, detached.provider_body); -} - -#[apply(inlining_routes)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document( - route: Route, - #[values(Host::Detached, Host::Realiasing)] host: Host, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Source::Remote, host, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[apply(every_route)] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider(route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send( - route, - Source::Remote, - Host::ReplacesDocument, - &document_base, - ) - .await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index f3adf27cfa6..44313d5f552 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_callbacks::event::WireRequest; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::{CallHooks, OcrClient}, @@ -23,11 +23,7 @@ use crate::ocr::{ pub(crate) struct NoHooks; impl CallHooks for NoHooks { - fn before_send( - &self, - wire: WireRequest, - _passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { Box::pin(async move { Ok(wire) }) } @@ -70,7 +66,7 @@ pub(crate) fn wire_request_with_document( decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some("test-key".into()), + api_key: Some(litellm_auth::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index a3fdd2340b3..903ccb06c77 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,9 +1,10 @@ - Target invariants; implementation and runtime validation may lag these rules -- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `CallbackAdapter`/`RouteHost` traits +- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - - A failure that surfaces inside the call, including a host op the call asked for, is mapped through the route's `map_failure`; a failure in `begin` or `after_success` is raised as is + - A native failure, including one a host op returns as `HostOpError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is + - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index f1bc3142a25..4aa7a2163ca 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -11,7 +11,7 @@ pub fn missing_state() -> PyErr { /// What an adapter step produced: either the value the driver asked for, or a Python /// awaitable the driver hands back to the caller's task before asking again. -pub enum AdapterStep { +pub enum LifecycleStep { Await(Py), Arguments(Py), Wire(Box), @@ -28,55 +28,74 @@ pub enum PublicValue<'a> { /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in /// order: `begin` before the machine starts, `before_send` and `emit` while it runs, /// `after_success` and one terminal `emit` after it completes. Whenever a step returns -/// [`AdapterStep::Await`], the driver awaits it in the caller's task and continues the +/// [`LifecycleStep::Await`], the driver awaits it in the caller's task and continues the /// same step through `resume`. /// /// A step that fails with an ordinary exception fails the call with that exception, /// except on a terminal event, where the adapter is expected to report and swallow its /// own errors. An exception that is not a `PyException`, such as a cancellation, ends /// the call without further dispatch. -pub trait CallbackAdapter: Send + Sync { +pub trait PythonLifecycle: Send + Sync { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult; + ) -> PyResult; fn before_send( &mut self, py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult; + ) -> PyResult; fn after_success( &mut self, py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult; + ) -> PyResult; fn emit( &mut self, py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult; + ) -> PyResult; - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } -/// The Python side of one route: answers the route's own operations, builds the public -/// response and maps failures to public exceptions. -pub trait RouteHost: Send + Sync { - type Route: Route; +/// Why a route operation the host answered did not produce a result: the route's own code +/// rejected it, which the route classifies like any other native failure, or Python code +/// raised, which reaches the caller as it was raised. +#[derive(Debug)] +pub enum HostOpError { + Native(E), + Python(PyErr), +} - /// `arguments` is the keyword view the callback adapter's `begin` produced, not the +impl From for HostOpError { + fn from(error: PyErr) -> Self { + Self::Python(error) + } +} + +/// The Python side of one route: answers the route's own operations, builds the public +/// response and classifies native failures into public exceptions. +pub trait RouteHost: Send + Sync { + type Route: Route; + + /// The public exception a native failure maps to, kept as a value until the driver + /// raises it. + type Failure: Into; + + /// `arguments` is the keyword view the lifecycle's `begin` produced, not the /// caller's own dict. A route host that projects from it inherits whatever that /// adapter rewrote. fn invoke( @@ -84,7 +103,7 @@ pub trait RouteHost: Send + Sync { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: ::Op, - ) -> PyResult<::OpResult>; + ) -> Result<::OpResult, HostOpError<::Error>>; fn complete( &mut self, @@ -92,12 +111,14 @@ pub trait RouteHost: Send + Sync { response: ::Response, ) -> PyResult>; - fn native_error(error: ::Error) -> PyErr; + fn classify( + &self, + py: Python<'_>, + error: ::Error, + ) -> PyResult; fn host_error(error: &PyErr) -> ::Error; - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult; - fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; diff --git a/litellm-rust/crates/host-python/src/argument.rs b/litellm-rust/crates/host-python/src/argument.rs new file mode 100644 index 00000000000..34e07cdfbd5 --- /dev/null +++ b/litellm-rust/crates/host-python/src/argument.rs @@ -0,0 +1,51 @@ +use pyo3::{prelude::*, types::PyDict}; + +/// The caller's own object for a public argument: the keyword if given, even an explicit +/// `None`, else the bound request's attribute. Every reader of a public Python call uses +/// this rule, so the callbacks and the provider see one object per argument. +pub fn lookup<'py>( + kwargs: &Bound<'py, PyDict>, + request: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + if let Some(value) = kwargs.get_item(name)? { + return Ok(Some(value)); + } + request.getattr_opt(name) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { + crate::initialize_python(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +key = object() +document = {'type': 'document_url'} +class Request: + api_key = 'from-request' + api_base = 'from-request' + document = document +request = Request() +kwargs = {'api_key': key, 'api_base': None} +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let item = |name: &str| locals.get_item(name).unwrap().unwrap(); + let kwargs = item("kwargs").cast_into::().unwrap(); + let request = item("request"); + let find = |name: &str| lookup(&kwargs, &request, name).unwrap(); + assert!(find("api_key").unwrap().is(item("key"))); + assert!(find("api_base").unwrap().is_none()); + assert!(find("document").unwrap().is(item("document"))); + assert!(find("model").is_none()); + }); + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 8bda13b44d0..7895d803cb5 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -12,7 +12,9 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use tokio::sync::Mutex; -use crate::adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +use crate::adapter::{ + HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, +}; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -43,6 +45,7 @@ enum Stage { #[derive(Clone, Copy)] enum Expect { + Started, Arguments, Wire, Emitted, @@ -66,7 +69,7 @@ where M: Machine> + 'static, { route: H, - adapter: Box, + adapter: Box, machine: Option>>>, arguments: Option>, started_at: f64, @@ -84,7 +87,7 @@ pub fn run_call( py: Python<'_>, machine: M, route: H, - adapter: Box, + adapter: Box, arguments: Py, asynchronous: bool, ) -> PyResult> @@ -146,9 +149,11 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let arguments = self.arguments.take().ok_or_else(missing_state)?; - match self.adapter.begin(py, arguments, self.started_at) { - Ok(step) => self.on_adapter(py, step, Expect::Arguments), + let started = CallEvent::Started { + start_time: self.started_at, + }; + match self.adapter.emit(py, &started, None) { + Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } } @@ -170,27 +175,28 @@ where fn on_adapter( &mut self, py: Python<'_>, - step: AdapterStep, + step: LifecycleStep, expect: Expect, ) -> PyResult { match (expect, step) { - (_, AdapterStep::Await(awaitable)) => { + (_, LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(expect)); Ok(ExecutionStep::Await(awaitable)) } - (Expect::Arguments, AdapterStep::Arguments(arguments)) => { + (Expect::Started, LifecycleStep::Done) => self.begin(py), + (Expect::Arguments, LifecycleStep::Arguments(arguments)) => { self.arguments = Some(arguments); self.stage = Stage::Call; self.resume_machine(py, None) } - (Expect::Wire, AdapterStep::Wire(wire)) => { + (Expect::Wire, LifecycleStep::Wire(wire)) => { self.resume_machine(py, Some(Ok(HostResult::BeforeSend(wire)))) } - (Expect::Emitted, AdapterStep::Done) => { + (Expect::Emitted, LifecycleStep::Done) => { self.resume_machine(py, Some(Ok(HostResult::Emitted))) } - (Expect::Response, AdapterStep::Response(response)) => self.succeeded(py, response), - (Expect::Terminal, AdapterStep::Done) => match &self.stage { + (Expect::Response, LifecycleStep::Response(response)) => self.succeeded(py, response), + (Expect::Terminal, LifecycleStep::Done) => match &self.stage { Stage::Succeeded(response) => Ok(ExecutionStep::Return(response.clone_ref(py))), Stage::Failed(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())), _ => Err(missing_state()), @@ -199,6 +205,14 @@ where } } + fn begin(&mut self, py: Python<'_>) -> PyResult { + let arguments = self.arguments.take().ok_or_else(missing_state)?; + match self.adapter.begin(py, arguments, self.started_at) { + Ok(step) => self.on_adapter(py, step, Expect::Arguments), + Err(error) => self.adapter_failed(py, error), + } + } + fn adapter_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { match self.stage { Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), @@ -248,14 +262,20 @@ where let answer = match op { HostOp::Route(op) => { let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; - self.route - .invoke(py, arguments.bind(py), op) - .map(HostResult::Route) + match self.route.invoke(py, arguments.bind(py), op) { + Ok(result) => Ok(HostResult::Route(result)), + Err(HostOpError::Native(error)) => { + return self + .resume_core(py, Some(Err(HostFailure::Error(error)))) + .map(Next::Continue); + } + Err(HostOpError::Python(error)) => Err(error), + } } HostOp::BeforeSend { wire, context } => { match self.adapter.before_send(py, wire, &context) { - Ok(AdapterStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Wire)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -264,8 +284,8 @@ where } } HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { - Ok(AdapterStep::Done) => Ok(HostResult::Emitted), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -360,11 +380,29 @@ where self.ended_at.get_or_insert_with(epoch_seconds); let error = match self.interrupted.take() { Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), - None => H::native_error(error), + None => self.classified(py, error), }; self.failure(py, error, FailureOrigin::Call) } + /// The route's public exception for a native failure. When classification itself + /// fails, that failure is raised with the native error's text as its `__context__`. + fn classified(&self, py: Python<'_>, error: ErrorOf) -> PyErr { + let native = error.to_string(); + let classifier_error = match self.route.classify(py, error) { + Ok(failure) => return failure.into(), + Err(classifier_error) => classifier_error, + }; + let attached = classifier_error.value(py).setattr( + "__context__", + PyRuntimeError::new_err(native).into_value(py), + ); + match attached { + Ok(()) => classifier_error, + Err(error) => error, + } + } + fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { let event = CallEvent::Succeeded { timing: self.timing(), @@ -386,18 +424,14 @@ where if is_cancellation(py, &error) { return Err(error); } - let public = match origin { - FailureOrigin::Call => self.route.map_failure(py, &error).unwrap_or(error), - FailureOrigin::Host => error, - }; let event = CallEvent::Failed { timing: self.timing(), origin, }; let step = self .adapter - .emit(py, &event, Some(PublicValue::Error(&public)))?; - self.stage = Stage::Failed(public.into_value(py)); + .emit(py, &event, Some(PublicValue::Error(&error)))?; + self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -489,6 +523,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri #[derive(Clone, Debug, PartialEq, Eq)] struct Error(String); + impl std::fmt::Display for Error { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } + } + struct Synthetic; impl Route for Synthetic { @@ -518,8 +558,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri model: "model".into(), custom_llm_provider: "provider".into(), optional_params: serde_json::json!({}), - passthrough_fields: Default::default(), secret_fields: Vec::new(), + api_key: None, } } @@ -566,25 +606,46 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } + #[derive(Clone, Copy)] + enum OpScript { + Answer, + RaisePython, + RejectNatively, + } + struct SyntheticHost { log: Log, - fail_op: bool, + op: OpScript, + classifier_fails: bool, + } + + /// The fake route's public exception, kept as a value so a test sees what `classify` + /// produced before the driver raises it. + #[derive(Debug, PartialEq, Eq)] + struct Classified(String); + + impl From for PyErr { + fn from(classified: Classified) -> Self { + PyValueError::new_err(format!("classified: {}", classified.0)) + } } impl RouteHost for SyntheticHost { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, _: Python<'_>, arguments: &Bound<'_, PyDict>, op: &'static str, - ) -> PyResult { + ) -> Result> { self.log.push(format!("route:{op}")); - if self.fail_op { - return Err(PyValueError::new_err("op failed")); + match self.op { + OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), + OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), + OpScript::RejectNatively => Err(HostOpError::Native(Error("op rejected".into()))), } - Ok(format!("{op}:{}", arguments.len())) } fn complete(&mut self, py: Python<'_>, response: String) -> PyResult> { @@ -594,22 +655,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unbind()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.log.push(format!("classify:{error}")); + if self.classifier_fails { + return Err(pyo3::exceptions::PyTypeError::new_err("classifier failed")); + } + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - self.log.push("map_failure"); - Ok(PyValueError::new_err(format!( - "mapped: {}", - error.value(py) - ))) - } - fn close(&mut self, _: Python<'_>) { self.log.push("route.close"); } @@ -632,13 +689,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri script: AdapterScript, } - impl CallbackAdapter for SyntheticAdapter { - fn begin(&mut self, _: Python<'_>, arguments: Py, _: f64) -> PyResult { + impl PythonLifecycle for SyntheticAdapter { + fn begin( + &mut self, + _: Python<'_>, + arguments: Py, + _: f64, + ) -> PyResult { self.log.push("begin"); if matches!(self.script, AdapterScript::FailBegin) { return Err(PyValueError::new_err("begin failed")); } - Ok(AdapterStep::Arguments(arguments)) + Ok(LifecycleStep::Arguments(arguments)) } fn before_send( @@ -646,9 +708,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri _: Python<'_>, wire: Box, _: &RequestContext, - ) -> PyResult { + ) -> PyResult { self.log.push("before_send"); - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { url: "rewritten".into(), ..*wire }))) @@ -659,17 +721,17 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, response: Py, _: Timing, - ) -> PyResult { + ) -> PyResult { self.log.push("after_success"); match self.script { - AdapterScript::ReplaceResponse => Ok(AdapterStep::Response( + AdapterScript::ReplaceResponse => Ok(LifecycleStep::Response( "replaced".into_pyobject(py)?.into_any().unbind(), )), AdapterScript::FailAfterSuccess => { Err(PyValueError::new_err("after_success failed")) } AdapterScript::Plain | AdapterScript::FailBegin => { - Ok(AdapterStep::Response(response)) + Ok(LifecycleStep::Response(response)) } } } @@ -679,8 +741,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, event: &CallEvent, public: Option>, - ) -> PyResult { + ) -> PyResult { self.log.push(match (event, public) { + (CallEvent::Started { .. }, None) => "started".into(), (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { format!("succeeded:{}", value.bind(py)) @@ -690,10 +753,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } _ => "unexpected".into(), }); - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } - fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -709,15 +772,31 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn run_scripted( py: Python<'_>, machine: ScriptedMachine, - fail_op: bool, + op: OpScript, script: AdapterScript, asynchronous: bool, ) -> (PyResult>, Vec) { - let log = Log::default(); - let route = SyntheticHost { - log: Log(log.0.clone()), - fail_op, - }; + run_hosted( + py, + machine, + SyntheticHost { + log: Log::default(), + op, + classifier_fails: false, + }, + script, + asynchronous, + ) + } + + fn run_hosted( + py: Python<'_>, + machine: ScriptedMachine, + route: SyntheticHost, + script: AdapterScript, + asynchronous: bool, + ) -> (PyResult>, Vec) { + let log = Log(route.log.0.clone()); let adapter = SyntheticAdapter { log: Log(log.0.clone()), script, @@ -777,7 +856,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::Plain, asynchronous, ); @@ -785,6 +864,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "route:project", "before_send", @@ -800,28 +880,75 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + fn failing_machine() -> ScriptedMachine { + ScriptedMachine { + ops: vec![HostOp::Route("project")], + outcome: Some(Err(Error("provider exploded".into()))), + answers: Vec::new(), + } + } + #[test] - fn machine_failures_are_mapped_and_dispatched_once_as_call_failures() { + fn a_native_failure_is_classified_once_and_reported_classified() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let machine = ScriptedMachine { - ops: vec![HostOp::Route("project")], - outcome: Some(Err(Error("provider exploded".into()))), - answers: Vec::new(), - }; - let (result, log) = run_scripted(py, machine, false, AdapterScript::Plain, false); - let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: provider exploded"); + install_lifecycle_module(py); + for asynchronous in [false, true] { + let (result, log) = run_scripted( + py, + failing_machine(), + OpScript::Answer, + AdapterScript::Plain, + asynchronous, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classified: provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classified: provider exploded", + "adapter.close", + "route.close", + ] + ); + } + }); + } + + #[test] + fn a_native_rejection_from_a_host_operation_is_classified_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RejectNatively, + AdapterScript::Plain, + false, + ); + assert_eq!( + result.unwrap_err().value(py).to_string(), + "classified: op rejected" + ); assert_eq!( log, [ + "started", "begin", "route:project", - "map_failure", - "failed:Call:mapped: provider exploded", + "classify:op rejected", + "failed:Call:classified: op rejected", "adapter.close", "route.close", ] @@ -830,18 +957,72 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } #[test] - fn host_operation_failures_interrupt_the_call_and_keep_the_python_exception() { + fn a_python_exception_from_a_host_operation_is_reported_as_raised() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let (result, log) = - run_scripted(py, success_machine(), true, AdapterScript::Plain, false); + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RaisePython, + AdapterScript::Plain, + false, + ); let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: op failed"); - assert!(!log.contains(&"before_send".to_string())); - assert!(log.contains(&"failed:Call:mapped: op failed".to_string())); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "op failed"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "failed:Call:op failed", + "adapter.close", + "route.close", + ] + ); + }); + } + + #[test] + fn a_failing_classifier_surfaces_with_the_native_error_as_context() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_hosted( + py, + failing_machine(), + SyntheticHost { + log: Log::default(), + op: OpScript::Answer, + classifier_fails: true, + }, + AdapterScript::Plain, + false, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classifier failed"); + let context = error.value(py).getattr("__context__").unwrap(); + assert!(context.is_instance_of::()); + assert_eq!(context.str().unwrap().to_string(), "provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classifier failed", + "adapter.close", + "route.close", + ] + ); }); } @@ -855,7 +1036,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailBegin, false, ); @@ -864,6 +1045,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "failed:Host:begin failed", "adapter.close", @@ -885,7 +1067,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::ReplaceResponse, asynchronous, ); @@ -908,7 +1090,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailAfterSuccess, asynchronous, ); @@ -938,12 +1120,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri struct Cancelling(Log); impl RouteHost for Cancelling { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, py: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, - ) -> PyResult { + ) -> Result> { self.0.push("route"); Err(PyErr::from_value( py.import("asyncio") @@ -952,21 +1135,19 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unwrap() .call0() .unwrap(), - )) + ) + .into()) } fn complete(&mut self, _: Python<'_>, _: String) -> PyResult> { Err(missing_state()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.0.push("classify"); + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, _: Python<'_>, _: &PyErr) -> PyResult { - self.0.push("map_failure"); - Err(missing_state()) - } fn close(&mut self, _: Python<'_>) {} fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { Ok(()) @@ -988,7 +1169,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) .unwrap_err(); assert!(!error.is_instance_of::(py)); - assert_eq!(log.entries(), ["begin", "route", "adapter.close"]); + assert_eq!( + log.entries(), + ["started", "begin", "route", "adapter.close"] + ); }); } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index bb0b5b1c3b1..8889b1513db 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,9 +1,10 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and //! asyncio glue, and the driver that runs a native [`Machine`](litellm_callbacks::machine::Machine) -//! against a Python route host and a callback adapter. Everything here is Python-specific by +//! against a Python route host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. mod adapter; +mod argument; mod callable; mod driver; mod execution; @@ -11,7 +12,10 @@ mod gil; mod handle; mod marshal; -pub use adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +pub use adapter::{ + HostOpError, LifecycleStep, PublicValue, PythonLifecycle, RouteHost, missing_state, +}; +pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value}; diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 87945bf8785..2e398d0287e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -150,7 +150,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { api_key: inputs.api_key.and_then(|key| { inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(Some(key)) }), api_base: inputs.api_base.and_then(|base| { @@ -592,12 +592,17 @@ impl AzureDocumentIntelligenceOcrConfig { )?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::super::common_utils::validate_destination(connection, key.source())?; return Ok( @@ -796,7 +801,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2012f740173..7ef051e8986 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -142,12 +142,17 @@ impl AzureAiOcrConfig { super::common_utils::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::common_utils::validate_destination(connection, key.source())?; return Ok(bearer_headers(connection, key.value())); @@ -196,7 +201,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -288,7 +293,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e6fe5d9556d..8321dcfb4ce 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, future::Future, time::Duration}; -use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, @@ -90,21 +90,22 @@ pub enum OcrResponseFormat { #[derive(Clone, Default)] pub struct OcrCredentialInputs { - pub api_key: Option>, - pub dynamic_api_key: Option>, + pub api_key: Option>, + pub dynamic_api_key: Option>, pub api_base: Option>, pub dynamic_api_base: Option>, } impl OcrCredentialInputs { pub fn new( - api_key: Option, + api_key: Option, api_key_source: InputSource, api_base: Option, api_base_source: InputSource, ) -> Self { Self { - api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)), + api_key: nonblank(api_key.as_ref().map(|key| key.expose().to_string())) + .map(|value| Sourced::new(SecretValue::new(value), api_key_source)), dynamic_api_key: None, api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)), dynamic_api_base: None, @@ -159,7 +160,7 @@ fn nonblank(value: Option) -> Option { #[derive(Clone)] pub struct OcrConnection { - pub api_key: Option, + pub api_key: Option, pub api_key_source: InputSource, pub api_base: Option, pub api_base_source: InputSource, @@ -209,7 +210,7 @@ impl Default for OcrConnection { #[derive(Clone, Default)] pub struct ResolvedOcrCredentials { - pub api_key: Option>, + pub api_key: Option>, pub api_base: Option>, } @@ -428,7 +429,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { ResolvedOcrCredentials { api_key: inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(inputs.api_key), api_base: inputs .dynamic_api_base diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index f353c22d8c4..2528c967f41 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -179,8 +179,8 @@ impl CohereParseConfig { } let key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -718,7 +718,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index e93ddee3c50..7635bdd3d04 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -3,9 +3,9 @@ use std::{sync::OnceLock, time::Duration}; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_callbacks::event::WireRequest; use serde::{Serialize, de::DeserializeOwned}; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::{ base_llm::ocr::{ @@ -26,11 +26,7 @@ use crate::{ /// The route's view of one call, handed to provider code that has to reach the /// caller's hooks mid-flight (guardrails on the outgoing body, raw response events). pub trait CallHooks: Send + Sync { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result>; + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result>; fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), E>>; } @@ -231,9 +227,8 @@ pub async fn transform_request_body( config.get_supported_ocr_params(&request.model), )?; config.validate_request_body(&composed)?; - let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed); let changed = hooks - .before_send(wire_request(url, headers, composed), passthrough_fields) + .before_send(wire_request(url, headers, composed)) .await?; if !changed.body.is_object() { return Err(Error::RequestField { @@ -252,21 +247,6 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq } } -fn caller_inputs(request: &PreparedOcrRequest) -> Result, Error> { - let document = request - .caller_document - .then(|| serde_json::to_value(&request.document)) - .transpose() - .map_err(|_| Error::RequestField { - path: "document".into(), - })?; - let params: Map = request.optional_params.clone().into(); - Ok(params - .into_iter() - .chain(document.map(|document| ("document".to_string(), document))) - .collect()) -} - pub fn build_http_request( client: &OcrClient, request: &PreparedOcrRequest, @@ -294,9 +274,7 @@ pub async fn guardrail_document( let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField { path: "document".into(), })?; - let changed = hooks - .before_send(wire_request(url, headers, body), Passthrough::default()) - .await?; + let changed = hooks.before_send(wire_request(url, headers, body)).await?; let document = decode_request_value(changed.body, "guardrail.document")?; Ok((document, changed.headers)) } diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index c2038d0552d..9028f09c5ab 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -135,8 +135,8 @@ impl MistralOcrConfig { } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -212,7 +212,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(str::to_string), + api_key: api_key.map(litellm_auth::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index ca2bae9c3bb..ec876fafb8f 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -442,8 +442,8 @@ fn resolve_headers( } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -629,7 +629,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some("passed-key".into()), + api_key: Some(litellm_auth::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -639,7 +639,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index ea0bcf3d08c..c2cb23d0010 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -134,7 +134,10 @@ impl VertexAiOcrConfig { .vertex_auth() .validate_environment( connection.extra_headers.clone(), - connection.api_key.as_deref(), + connection + .api_key + .as_ref() + .map(litellm_auth::SecretValue::expose), config, &credential_env, ) diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 9932594e2f5..5dccfb4aca8 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,7 +1,7 @@ - Target invariants, not completion claims; these supersede older conflicting bridge guidance - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, `passthrough_fields` re-aliasing) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy + - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy` owns `Logging` dispatch policy - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 9dc891a91d6..77212e3d38e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,9 +1,9 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{RouteHost, missing_state, to_py}; +use litellm_host_python::{HostOpError, RouteHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ - exceptions::PyBaseException, + exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -57,12 +57,8 @@ impl OcrRouteHost { .ok_or_else(missing_state)? .acquire(py) } -} -impl RouteHost for OcrRouteHost { - type Route = Ocr; - - fn invoke( + fn answer( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, @@ -88,6 +84,40 @@ impl RouteHost for OcrRouteHost { } } + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let provider = match &self.data { + OcrHostData::Projected(handles) => handles.provider, + _ => "", + }; + let mapped = py + .import("litellm.rust_bridge.ocr.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), provider))) + .and_then(|mapped| mapped.extract::>().map_err(PyErr::from)); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for OcrRouteHost { + type Route = Ocr; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: OcrOp, + ) -> Result> { + self.answer(py, arguments, op) + .map_err(|error| HostOpError::Python(self.map_failure(py, error))) + } + fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { py.import("litellm.rust_bridge.ocr.route_host")? .getattr("response")? @@ -95,27 +125,14 @@ impl RouteHost for OcrRouteHost { .map(Bound::unbind) } - fn native_error(error: Error) -> PyErr { - ocr_error_to_pyerr(error) + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } fn host_error(error: &PyErr) -> Error { Error::InvalidRequest(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - let provider = match &self.data { - OcrHostData::Projected(handles) => handles.provider, - _ => "", - }; - let mapped: Py = py - .import("litellm.rust_bridge.ocr.route_host")? - .getattr("map_failure")? - .call1((error.value(py), self.request.bind(py), provider))? - .extract()?; - Ok(PyErr::from_value(mapped.into_bound(py).into_any())) - } - fn close(&mut self, _: Python<'_>) { self.data = OcrHostData::Released; } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 7ffa129f85c..5dd2aa804b8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,3 +1,4 @@ +use litellm_auth::SecretValue; use litellm_core::ocr::{ types::{LiteLLMOcrRequest, OcrDocumentInput}, wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, @@ -31,7 +32,7 @@ struct OcrArguments<'a, 'py> { impl<'py> OcrArguments<'_, 'py> { fn lookup(&self, name: &str) -> PyResult> { - litellm_callbacks_legacy::lookup(self.kwargs, self.request, name)? + litellm_host_python::lookup(self.kwargs, self.request, name)? .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) } @@ -47,8 +48,11 @@ impl<'py> OcrArguments<'_, 'py> { self.lookup("document") } - fn api_key(&self) -> PyResult> { - self.lookup("api_key")?.extract() + fn api_key(&self) -> PyResult> { + Ok(self + .lookup("api_key")? + .extract::>()? + .map(SecretValue::new)) } fn api_base(&self) -> PyResult> { diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4a9a65b1485..7248c2f3590 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1265,11 +1265,6 @@ class Logging(LiteLLMLoggingBaseClass): additional_args.get("api_base", "") ) - def record_api_call_start_time(self) -> None: - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] - def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API try: @@ -1334,7 +1329,15 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.record_api_call_start_time() + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + # Set-once first provider-handoff instant. api_call_start_time + # is overwritten on every retry, so it can't measure one-time + # preprocessing; pinning the first attempt excludes retry loops + # + backoff. Logging object only — must NOT go into + # litellm_params["metadata"] (caller request metadata, typed + # Dict[str, str], echoed downstream; a datetime breaks it). + if self.model_call_details.get("first_api_call_start_time") is None: + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1468,21 +1471,16 @@ class Logging(LiteLLMLoggingBaseClass): """ return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) - def record_post_call( - self, original_response: object, input: object, api_key: object, additional_args: dict[str, object] - ) -> None: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" - def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: - self.record_post_call(original_response, input, api_key, additional_args) + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" attr: Literal["warning", "debug"] if self.litellm_request_debug: @@ -2177,7 +2175,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, - build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2202,9 +2199,6 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) - if not build_logging_payload: - return - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2266,7 +2260,6 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, - build_logging_payload: bool = True, ): try: if start_time is None: @@ -2304,7 +2297,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, - build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3328,9 +3320,7 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn( - self, exception, traceback_exception, start_time=None, end_time=None, build_logging_payload: bool = True - ): + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time if end_time is None: @@ -3365,9 +3355,6 @@ class Logging(LiteLLMLoggingBaseClass): metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) - if not build_logging_payload: - return start_time, end_time - ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index e05d9368fa8..406dc55cfea 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -6,24 +6,22 @@ registries it fans out to. It expires with that contract. from __future__ import annotations +import contextvars import datetime -import os +import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Coroutine, Mapping from dataclasses import dataclass from typing import ( TYPE_CHECKING, Final, - Literal, Protocol, - TypeAlias, cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations ) -from typing_extensions import assert_never - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CredentialItem class MetadataUpdater(Protocol): @@ -42,7 +40,6 @@ class MetadataUpdater(Protocol): class CallSetup: logger: Logging kwargs: dict[str, object] - bridge_owned: bool def setup( @@ -61,9 +58,9 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments, bridge_owned=False) + return CallSetup(supplied, arguments) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared, bridge_owned=True) + return CallSetup(logger, prepared) def check_limits(kwargs: Mapping[str, object]) -> None: @@ -93,87 +90,219 @@ def finalize( update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) -def deployment_callbacks_needed() -> bool: - import litellm - from litellm.integrations.custom_logger import CustomLogger +class LoggingSurface(Protocol): + def update_from_kwargs( + self, + kwargs: dict[str, object], + litellm_params: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + model: str | None = None, + user: str | None = None, + **additional_params: object, + ) -> None: ... - return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) + def pre_call( + self, input: object, api_key: object, model: object = None, additional_args: dict[str, object] = ... + ) -> object: ... + + def post_call( + self, + original_response: object, + input: object = None, + api_key: object = None, + additional_args: dict[str, object] = ..., + ) -> object: ... + + def handle_sync_success_callbacks_for_async_calls( + self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: object = None + ) -> None: ... + + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: ... + + def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> Coroutine[object, object, None]: ... + + def success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> None: ... + + def async_success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> Coroutine[object, object, None]: ... -Phase: TypeAlias = Literal[ - "input", "sync_success", "sync_success_async", "async_success", "sync_failure", "async_failure", "payload" -] +if TYPE_CHECKING: + _LOGGING_CONFORMS: type[LoggingSurface] = Logging -def callbacks_needed(logger: Logging, phase: Phase) -> bool: - import litellm - from litellm._logging import ( - _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging +class LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... + + +class DeploymentHook(Protocol): + def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... + + +class DeploymentSuccessHook(Protocol): + def __call__(self, request_data: dict[str, object], response: object, call_type: object) -> Awaitable[object]: ... + + +class DeploymentFailureHook(Protocol): + def __call__(self, request_data: Mapping[str, object], exception: Exception, call_type: str) -> Awaitable[None]: ... + + +def update_logging( + logger: LoggingSurface, + kwargs: dict[str, object], + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + custom_llm_provider: str, +) -> None: + logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, ) - if ( - _is_debugging_on() - or getattr(logger, "litellm_request_debug", False) - or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") - ): - return True - input_needed: Final = bool( - litellm.input_callback - or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_input_callbacks - or callable(getattr(logger, "logger_fn", None)) - or logger.log_raw_request_response - or litellm.log_raw_request_response + +def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: + logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) + + +def post_call( + logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] +) -> None: + logger.post_call(original_response=original_response, api_key=api_key, additional_args=additional_args) + + +def defers_async_logging(logger: LoggingSurface) -> bool: + return bool(getattr(logger, "_defer_async_logging", False)) + + +def defer_success(logger: LoggingSurface, pending: object) -> None: + setattr(logger, "_native_pending_logging", pending) + + +def sync_success_for_async_call( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> None: + logger.handle_sync_success_callbacks_for_async_calls(result=response, start_time=start, end_time=end) + + +def failure_handler( + logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> Coroutine[object, object, None] | None: + trace: Final = "".join(traceback.format_exception(error)) + if asynchronous: + return logger.async_failure_handler(error, trace, start, end) + logger.failure_handler(error, trace, start, end) + return None + + +def submit_success(logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime) -> None: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, logger.success_handler, response, start, end) + + +def async_success_handler( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> Coroutine[object, object, None]: + return logger.async_success_handler(response, start, end) + + +def enqueue_logging(coroutine: Coroutine[object, object, None]) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + worker: Final = cast( # cast-ok: bounded adapter for the untyped logging worker + LoggingWorker, GLOBAL_LOGGING_WORKER ) - match phase: - case "input": - return input_needed - case "sync_success": - return bool(litellm.success_callback or logger.dynamic_success_callbacks) - case "sync_success_async": - return bool( - (litellm.success_callback or logger.dynamic_success_callbacks) - and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks - ) - case "async_success": - return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "sync_failure": - return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) - case "async_failure": - return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "payload": - return bool( - input_needed - or litellm.success_callback - or litellm.failure_callback - or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_success_callbacks - or logger.dynamic_async_success_callbacks - or logger.dynamic_failure_callbacks - or logger.dynamic_async_failure_callbacks - ) - case _: - assert_never(phase) + contextvars.copy_context().run(worker.ensure_initialized_and_enqueue, coroutine) -def success_bookkeeping( - logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_success" if asynchronous else "sync_success" - if logger.should_run_logging(phase): - logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload - result=response, start_time=start, end_time=end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def restore_context(logger: LoggingSurface) -> None: + from litellm.utils import ( + _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # the @client wrapper restores the same correlation context + ) + + _restore_correlation_context_if_supported(logger) -def failure_bookkeeping( - logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_failure" if asynchronous else "sync_failure" - if logger.should_run_logging(phase): - logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload - error, "", start, end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def custom_pricing_fields() -> tuple[str, ...]: + from litellm.types.utils import CustomPricingLiteLLMParams + + return tuple(CustomPricingLiteLLMParams.model_fields) + + +def is_internal_call() -> bool: + from litellm._internal_context import is_internal_call as internal + + return internal.get() + + +def credential_list() -> list[CredentialItem]: + import litellm + + return litellm.credential_list + + +def warn_unknown_credential(name: str, loaded: int) -> None: + from litellm._logging import verbose_logger + + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + name, + loaded, + ) + + +def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitable[object]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentHook, utils.async_pre_call_deployment_hook + ) + return hook(kwargs, call_type) + + +def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]: + from litellm import utils + from litellm.types.utils import CallTypes + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentSuccessHook, utils.async_post_call_success_deployment_hook + ) + return hook(kwargs, response, CallTypes(call_type)) + + +def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_type: str) -> Awaitable[None]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentFailureHook, utils.async_post_call_failure_deployment_hook + ) + return hook(kwargs, error, call_type) diff --git a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py index a4474c85230..a0906c7c5be 100644 --- a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py +++ b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py @@ -1,12 +1,16 @@ import datetime +import inspect from collections.abc import Mapping +from pathlib import Path from types import MappingProxyType from typing import Final import pytest +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.rust_bridge import legacy_callbacks as legacy from litellm.rust_bridge.legacy_callbacks import check_limits, setup _OCR_KWARGS: Final = MappingProxyType( @@ -56,13 +60,12 @@ def _supplied_logger() -> Logging: ) -def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: +def test_setup_reuses_a_supplied_logger() -> None: supplied: Final = _supplied_logger() result: Final = setup( "aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True ) assert result.logger is supplied - assert result.bridge_owned is False @pytest.mark.parametrize( @@ -73,7 +76,15 @@ def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: ], ids=["ocr", "embedding"], ) -def test_setup_owns_every_logger_it_builds(call_type: str, kwargs: Mapping[str, object]) -> None: +def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Mapping[str, object]) -> None: result: Final = setup(call_type, (), kwargs, datetime.datetime.now(), asynchronous=True) - assert result.bridge_owned is True assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] + + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy/python_contract.json" + + +def test_the_rust_contract_matches_the_shim_signatures() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 085ea4a14c0..5fca927bea3 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -4,7 +4,7 @@ import gc import json import threading import weakref -from collections.abc import Coroutine +from collections.abc import Awaitable, Callable, Coroutine from contextvars import ContextVar from typing import Final @@ -400,7 +400,7 @@ async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: Rec @pytest.mark.asyncio @pytest.mark.parametrize("failure", [False, True]) -async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( +async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper( ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: bool, @@ -414,8 +414,12 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( submissions = 0 enqueues = 0 - def deployment(self, *args: object, **kwargs: object) -> None: - self.deployments += 1 + def counting(self, hook: Callable[..., Awaitable[object]]) -> Callable[..., Awaitable[object]]: + async def counted(*args: object, **kwargs: object) -> object: + self.deployments += 1 + return await hook(*args, **kwargs) + + return counted def submit(self, *args: object, **kwargs: object) -> None: self.submissions += 1 @@ -430,7 +434,7 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( "async_post_call_success_deployment_hook", "async_post_call_failure_deployment_hook", ): - monkeypatch.setattr(utils, name, probe.deployment) + monkeypatch.setattr(utils, name, probe.counting(getattr(utils, name))) monkeypatch.setattr(litellm_logging, "executor", probe) monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) if failure: @@ -447,17 +451,16 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( assert response._hidden_params["response_cost"] is not None assert response._hidden_params["_response_ms"] > 0 assert trace_id_var.get() == "callback-free-parent" - assert probe.deployments == probe.submissions == probe.enqueues == 0 + assert probe.deployments == 2 + assert probe.submissions == probe.enqueues == 0 assert len(created_loggers) == 1 logger: Final = created_loggers[0] - assert not hasattr(logger, "_native_pending_logging") - assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] - assert "standard_logging_object" not in logger.model_call_details - assert ( - "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None - ) - assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) - assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) + if failure: + assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] + assert logger.model_call_details["response_cost"] == 0 + else: + assert getattr(logger, "_native_pending_logging", None) is not None + assert "end_time" not in logger.model_call_details @pytest.mark.asyncio