scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

refactor(rust): route-neutral callback contract

Browse source 
Every legacy callback call from callbacks-legacy now goes through one typed
Python shim, litellm.rust_bridge.legacy_callbacks, the only Python module
the crate reaches. Before, the crate called Logging methods, litellm.utils
hooks, the logging worker, the executor and several litellm globals
directly, and its tests retyped those signatures by hand, so an outdated
fake could accept a call the real code rejects. python_contract.json lists
each shim function's parameters: a Python test pins it to the real
signatures and a Rust test pins it to the Rust enum.

The lifecycle contract changes to match the Python @client wrapper:
- the driver emits CallEvent::Started before begin, so every host sees one
  start time
- RequestContext carries the route-resolved api_key, so legacy pre_call and
  post_call receive it, and post_call's additional_args match the Python OCR
  path
- Passthrough and its re-aliasing are gone
- async deployment hooks always run, and the "no callbacks" shortcut that
  skipped the logging payload is removed, as in the Python path

The OCR api_key is a SecretValue from the wire request onward, so Debug
output upstream of the callback contract cannot leak it.

host-python's RouteHost now classifies native failures once through
classify, and host ops return HostOpError. The OCR route host keeps main's
public errors by sending both through the existing Python map_failure.
This commit is contained in:
Yujong Lee 2026-09-18 13:36:06 -07:00
parent bc6b540205
commit b4bfd92a2a
51 changed files with 1564 additions and 1297 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-rust/Cargo.lock generated
View file

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

4
litellm-rust/crates/auth/src/secret.rs
View file

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

14
litellm-rust/crates/callbacks-legacy/AGENTS.md
View file

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

4
litellm-rust/crates/callbacks-legacy/Cargo.toml
View file

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

98
litellm-rust/crates/callbacks-legacy/python_contract.json Normal file
View file

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

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

@ -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<Py<PyAny>>,
error: Option<Py<PyBaseException>>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
context: Option<RequestContext>,
asynchronous: bool,
internal: bool,
pending: Option<Pending>,
@ -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<bool> {
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<AdapterStep> {
fn prepare(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
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<AdapterStep> {
fn finalize(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
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<AdapterStep> {
fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
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<PyDict>,
started_at: f64,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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<WireRequest>,
context: &RequestContext,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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::<PyDict>()?;
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>(&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::<String>()?, value.extract::<String>()?)))
.collect::<PyResult<Vec<_>>>()?;
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<PyAny>,
timing: Timing,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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<PublicValue<'_>>,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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<Py<PyAny>>) -> PyResult<AdapterStep> {
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<LifecycleStep> {
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)
}
}

43
litellm-rust/crates/callbacks-legacy/src/call.rs
View file

@ -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<Option<Bound<'py, PyAny>>> {
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<H, M>(
@ -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();

242
litellm-rust/crates/callbacks-legacy/src/callbacks.rs
View file

@ -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<bool>;
/// `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<PyDict>>,
headers: Option<&Py<PyDict>>,
) -> 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<bool> {
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::<PyDict>()?,
&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::<PyDict>()?,
&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<PyDict>>,
headers: Option<&Py<PyDict>>,
) -> 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<PyAny>,
end: &Option<Py<PyAny>>,
) -> 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<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Option<Py<PyAny>>> {
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<PyAny>,
end: &Option<Py<PyAny>>,
) -> 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<PyAny>,
end: &Option<Py<PyAny>>,
) -> 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<Vec<String>> {
py.import("litellm.types.utils")?
.getattr("CustomPricingLiteLLMParams")?
.getattr("model_fields")?
.cast_into::<PyDict>()?
.keys()
.iter()
.map(|name| name.extract::<String>())
.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<bool> {
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()
}

155
litellm-rust/crates/callbacks-legacy/src/legacy_python.rs Normal file
View file

@ -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<Bound<'py, PyAny>>
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<Bound<'py, PyAny>>
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<Bound<'py, PyAny>>
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<Bound<'py, PyAny>>
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<String, serde_json::Value> =
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::<BTreeSet<_>>(), declared);
}
}

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

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

91
litellm-rust/crates/callbacks-legacy/src/logger.rs
View file

@ -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<PyAny>,
bridge_owned: bool,
}
impl PythonLogger {
pub(crate) fn new(object: Py<PyAny>, bridge_owned: bool) -> Self {
Self {
object,
bridge_owned,
}
pub(crate) fn new(object: Py<PyAny>) -> 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<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
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<Self> {
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<PythonLogger> {
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<Py<PyDict>> {
@ -93,9 +65,8 @@ pub fn setup<'py>(
start: &Py<PyAny>,
asynchronous: bool,
) -> PyResult<SetupResult<'py>> {
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<PyAny>,
end: &Option<Py<PyAny>>,
) -> 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<bool> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("deployment_callbacks_needed")?
.call0()?
.extract()
}
pub fn before_call(
py: Python<'_>,
kwargs: &Py<PyDict>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
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<Py<PyAny>>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
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<PyBaseException>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
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::<Vec<String>>()
.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());
});
}
}

44
litellm-rust/crates/callbacks-legacy/src/preparation.rs
View file

@ -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<Bound<'py, PyDict>> {
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::<PyList>()?)
})?;
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<Bound<'py, PyList>>,
) -> 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::<PyList>()?;
let credentials = credential_list()?;
let names = credentials
.iter()
.map(|credential| CredentialEntry(credential).name())
.collect::<PyResult<Vec<_>>>()?;
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::<PyDict>()?,
|| {
Ok(locals
.get_item("credentials")?
.unwrap()
.cast_into::<PyList>()?)
},
)
}
@ -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();
}
});
}

18
litellm-rust/crates/callbacks-legacy/tests/deferred.rs
View file

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

18
litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs
View file

@ -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::<PyDict>()
@ -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();

113
litellm-rust/crates/callbacks-legacy/tests/payload.rs
View file

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

135
litellm-rust/crates/callbacks-legacy/tests/support.rs
View file

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

46
litellm-rust/crates/callbacks-legacy/tests/terminal.rs
View file

@ -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<String> = 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<String> = 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<String> = 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()));

1
litellm-rust/crates/callbacks/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
serde_json.workspace = true
[dev-dependencies]

77
litellm-rust/crates/callbacks/src/event.rs
View file

@ -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<String>,
}
/// 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<String>);
impl Passthrough {
pub fn unchanged(caller: &Map<String, Value>, 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<Item = &str> {
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<litellm_auth::SecretValue>,
}
#[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::<Vec<_>>(), expected);
}
}

36
litellm-rust/crates/callbacks/src/run.rs
View file

@ -11,6 +11,7 @@ where
H: Host<M::Route>,
{
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<Vec<f64>>);
impl Host<Unit> 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]);
}
}

13
litellm-rust/crates/core/src/ocr/handler.rs
View file

@ -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<String>,
api_key: Option<SecretValue>,
}
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<Error> for OcrCallHooks {
fn before_send(
&self,
wire: WireRequest,
passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
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))
}

4
litellm-rust/crates/core/src/ocr/mod.rs
View file

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

4
litellm-rust/crates/core/src/ocr/prepare.rs
View file

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

46
litellm-rust/crates/core/src/ocr/provider_config.rs
View file

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

18
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -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<String>,
pub api_key: Option<SecretValue>,
pub api_base: Option<String>,
pub extra_headers: Map<String, Value>,
pub timeout: Option<Duration>,
@ -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!(

4
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -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<D = Value> {
pub model: String,
pub document: D,
pub api_key: Option<String>,
pub api_key: Option<SecretValue>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,

2
litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
View file

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

22
litellm-rust/crates/core/tests/ocr.rs
View file

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

152
litellm-rust/crates/core/tests/ocr/document.rs Normal file
View file

@ -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<Value>,
}
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)
);
}

282
litellm-rust/crates/core/tests/ocr/passthrough.rs
View file

@ -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<String, Value>,
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<String, Value>,
result: Result<(), Error>,
before_send: Option<(WireRequest, RequestContext)>,
provider_body: Option<Value>,
}
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<String, Value> = 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)
);
}

10
litellm-rust/crates/core/tests/ocr/support.rs
View file

@ -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<Error> for NoHooks {
fn before_send(
&self,
wire: WireRequest,
_passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
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,

5
litellm-rust/crates/host-python/AGENTS.md
View file

@ -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<T>` for retention; binding/unbinding does not copy payloads
- Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized<T>`

55
litellm-rust/crates/host-python/src/adapter.rs
View file

@ -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<PyAny>),
Arguments(Py<PyDict>),
Wire(Box<WireRequest>),
@ -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<PyDict>,
started_at: f64,
) -> PyResult<AdapterStep>;
) -> PyResult<LifecycleStep>;
fn before_send(
&mut self,
py: Python<'_>,
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<AdapterStep>;
) -> PyResult<LifecycleStep>;
fn after_success(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<AdapterStep>;
) -> PyResult<LifecycleStep>;
fn emit(
&mut self,
py: Python<'_>,
event: &CallEvent,
public: Option<PublicValue<'_>>,
) -> PyResult<AdapterStep>;
) -> PyResult<LifecycleStep>;
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep>;
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<LifecycleStep>;
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<E> {
Native(E),
Python(PyErr),
}
/// `arguments` is the keyword view the callback adapter's `begin` produced, not the
impl<E> From<PyErr> for HostOpError<E> {
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<Error: std::fmt::Display>;
/// The public exception a native failure maps to, kept as a value until the driver
/// raises it.
type Failure: Into<PyErr>;
/// `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: <Self::Route as Route>::Op,
) -> PyResult<<Self::Route as Route>::OpResult>;
) -> Result<<Self::Route as Route>::OpResult, HostOpError<<Self::Route as Route>::Error>>;
fn complete(
&mut self,
@ -92,12 +111,14 @@ pub trait RouteHost: Send + Sync {
response: <Self::Route as Route>::Response,
) -> PyResult<Py<PyAny>>;
fn native_error(error: <Self::Route as Route>::Error) -> PyErr;
fn classify(
&self,
py: Python<'_>,
error: <Self::Route as Route>::Error,
) -> PyResult<Self::Failure>;
fn host_error(error: &PyErr) -> <Self::Route as Route>::Error;
fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult<PyErr>;
fn close(&mut self, py: Python<'_>);
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;

51
litellm-rust/crates/host-python/src/argument.rs Normal file
View file

@ -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<Option<Bound<'py, PyAny>>> {
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::<PyDict>().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());
});
}
}

364
litellm-rust/crates/host-python/src/driver.rs
View file

@ -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<Route = H::Route, Complete = ResponseOf<H>> + 'static,
{
route: H,
adapter: Box<dyn CallbackAdapter>,
adapter: Box<dyn PythonLifecycle>,
machine: Option<Arc<Mutex<MachineState<M>>>>,
arguments: Option<Py<PyDict>>,
started_at: f64,
@ -84,7 +87,7 @@ pub fn run_call<H, M>(
py: Python<'_>,
machine: M,
route: H,
adapter: Box<dyn CallbackAdapter>,
adapter: Box<dyn PythonLifecycle>,
arguments: Py<PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
@ -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<ExecutionStep> {
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<ExecutionStep> {
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<ExecutionStep> {
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<H>) -> 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<PyAny>) -> PyResult<ExecutionStep> {
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<Classified> 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<String> {
) -> Result<String, HostOpError<Error>> {
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<Py<PyAny>> {
@ -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<Classified> {
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<PyErr> {
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<PyDict>, _: f64) -> PyResult<AdapterStep> {
impl PythonLifecycle for SyntheticAdapter {
fn begin(
&mut self,
_: Python<'_>,
arguments: Py<PyDict>,
_: f64,
) -> PyResult<LifecycleStep> {
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<WireRequest>,
_: &RequestContext,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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<PyAny>,
_: Timing,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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<PublicValue<'_>>,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
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<Py<PyAny>>) -> PyResult<AdapterStep> {
fn resume(&mut self, _: Python<'_>, _: PyResult<Py<PyAny>>) -> PyResult<LifecycleStep> {
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<Py<PyAny>>, Vec<String>) {
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<Py<PyAny>>, Vec<String>) {
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::<PyValueError>(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::<PyValueError>(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::<pyo3::exceptions::PyTypeError>(py));
assert_eq!(error.value(py).to_string(), "classifier failed");
let context = error.value(py).getattr("__context__").unwrap();
assert!(context.is_instance_of::<PyRuntimeError>());
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<String> {
) -> Result<String, HostOpError<Error>> {
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<Py<PyAny>> {
Err(missing_state())
}
fn native_error(error: Error) -> PyErr {
PyValueError::new_err(error.0)
fn classify(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
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<PyErr> {
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::<pyo3::exceptions::PyException>(py));
assert_eq!(log.entries(), ["begin", "route", "adapter.close"]);
assert_eq!(
log.entries(),
["started", "begin", "route", "adapter.close"]
);
});
}

8
litellm-rust/crates/host-python/src/lib.rs
View file

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

21
litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs
View file

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

21
litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs
View file

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

17
litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs
View file

@ -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<Sourced<String>>,
pub dynamic_api_key: Option<Sourced<String>>,
pub api_key: Option<Sourced<SecretValue>>,
pub dynamic_api_key: Option<Sourced<SecretValue>>,
pub api_base: Option<Sourced<String>>,
pub dynamic_api_base: Option<Sourced<String>>,
}
impl OcrCredentialInputs {
pub fn new(
api_key: Option<String>,
api_key: Option<SecretValue>,
api_key_source: InputSource,
api_base: Option<String>,
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<String>) -> Option<String> {
#[derive(Clone)]
pub struct OcrConnection {
pub api_key: Option<String>,
pub api_key: Option<SecretValue>,
pub api_key_source: InputSource,
pub api_base: Option<String>,
pub api_base_source: InputSource,
@ -209,7 +210,7 @@ impl Default for OcrConnection {
#[derive(Clone, Default)]
pub struct ResolvedOcrCredentials {
pub api_key: Option<Sourced<String>>,
pub api_key: Option<Sourced<SecretValue>>,
pub api_base: Option<Sourced<String>>,
}
@ -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

6
litellm-rust/crates/llms/src/cohere/ocr/transformation.rs
View file

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

32
litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs
View file

@ -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<E>: Send + Sync {
fn before_send(
&self,
wire: WireRequest,
passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, E>>;
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, E>>;
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), E>>;
}
@ -231,9 +227,8 @@ pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
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<Map<String, Value>, Error> {
let document = request
.caller_document
.then(|| serde_json::to_value(&request.document))
.transpose()
.map_err(|_| Error::RequestField {
path: "document".into(),
})?;
let params: Map<String, Value> = request.optional_params.clone().into();
Ok(params
.into_iter()
.chain(document.map(|document| ("document".to_string(), document)))
.collect())
}
pub fn build_http_request<B: Serialize>(
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))
}

6
litellm-rust/crates/llms/src/mistral/ocr/transformation.rs
View file

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

8
litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

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

5
litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs
View file

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

2
litellm-rust/crates/python-bridge/AGENTS.md
View file

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

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

@ -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::<PyException>(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::<Py<PyBaseException>>().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<OcrOpResult, HostOpError<Error>> {
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<PyAny>> {
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<PyErr> {
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<PyErr> {
let provider = match &self.data {
OcrHostData::Projected(handles) => handles.provider,
_ => "",
};
let mapped: Py<PyBaseException> = 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;
}

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

@ -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<Bound<'py, PyAny>> {
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<Option<String>> {
self.lookup("api_key")?.extract()
fn api_key(&self) -> PyResult<Option<SecretValue>> {
Ok(self
.lookup("api_key")?
.extract::<Option<String>>()?
.map(SecretValue::new))
}
fn api_base(&self) -> PyResult<Option<String>> {

43
litellm/litellm_core_utils/litellm_logging.py
View file

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

291
litellm/rust_bridge/legacy_callbacks.py
View file

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

19
tests/test_litellm/rust_bridge/test_legacy_callbacks.py
View file

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

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

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