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

Merge pull request #41885 from BerriAI/litellm_rust_callback_contract

Browse source 
refactor(rust): formalize legacy callback contract
This commit is contained in:
yujonglee 2026-09-18 16:16:02 -07:00 • committed by GitHub
commit 018f640b30
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
96 changed files with 3654 additions and 1896 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

91
litellm-rust/Cargo.lock generated
View file

@ -2027,24 +2027,18 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-callbacks"
version = "0.1.0"
dependencies = [
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-callbacks-legacy"
version = "0.1.0"
dependencies = [
"litellm-callbacks",
"litellm-auth",
"litellm-host",
"litellm-host-python",
"proptest",
"pyo3",
"rstest",
"serde_json",
"strum",
]
[[package]]
@ -2056,8 +2050,8 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-callbacks",
"litellm-core-utils",
"litellm-host",
"litellm-llms",
"litellm-types",
"mime_guess",
@ -2110,12 +2104,22 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-host"
version = "0.1.0"
dependencies = [
"litellm-auth",
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-host-python"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-callbacks",
"litellm-host",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
@ -2139,9 +2143,9 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-callbacks",
"litellm-core-utils",
"litellm-framing",
"litellm-host",
"litellm-types",
"reqwest 0.12.28",
"rstest",
@ -2598,6 +2602,25 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "proptest"
version = "1.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744"
dependencies = [
"bit-set",
"bit-vec",
"bitflags",
"num-traits",
"rand 0.9.5",
"rand_chacha 0.9.0",
"rand_xorshift",
"regex-syntax",
"rusty-fork",
"tempfile",
"unarray",
]
[[package]]
name = "pyo3"
version = "0.29.2"
@ -2679,6 +2702,12 @@ dependencies = [
"serde",
]
[[package]]
name = "quick-error"
version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quinn"
version = "0.11.11"
@ -2842,6 +2871,15 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "rand_xorshift"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a"
dependencies = [
"rand_core 0.9.5",
]
[[package]]
name = "rayon"
version = "1.12.0"
@ -3241,6 +3279,18 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
[[package]]
name = "rusty-fork"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2"
dependencies = [
"fnv",
"quick-error",
"tempfile",
"wait-timeout",
]
[[package]]
name = "ryu"
version = "1.0.23"
@ -4096,6 +4146,12 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "unarray"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94"
[[package]]
name = "unicase"
version = "2.9.0"
@ -4209,6 +4265,15 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
[[package]]
name = "wait-timeout"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11"
dependencies = [
"libc",
]
[[package]]
name = "walkdir"
version = "2.5.0"

3
litellm-rust/Cargo.toml
View file

@ -10,7 +10,7 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-core = { path = "crates/core" }
litellm-callbacks = { path = "crates/callbacks" }
litellm-host = { path = "crates/host" }
litellm-callbacks-legacy = { path = "crates/callbacks-legacy" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
@ -26,6 +26,7 @@ litellm-token-counter = { path = "crates/token-counter" }
litellm-host-python = { path = "crates/host-python" }
bytes = "1"
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"

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-host`, `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

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

@ -7,10 +7,15 @@ repository.workspace = true
autotests = false
[dependencies]
litellm-callbacks.workspace = true
litellm-host.workspace = true
litellm-host-python.workspace = true
pyo3.workspace = true
strum.workspace = true
serde_json.workspace = true
[dev-dependencies]
litellm-auth.workspace = true
proptest.workspace = true
rstest.workspace = true
serde_json.workspace = true

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

@ -0,0 +1,118 @@
{
"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"
],
"stream_opened": [
"logger"
],
"stream_success": [
"logger",
"url_route",
"endpoint_type",
"request_body",
"chunks",
"start",
"end",
"first_chunk"
],
"stream_failure": [
"logger",
"endpoint_type",
"request_body",
"chunks",
"error"
]
}

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

@ -2,21 +2,26 @@
//! raises is answered with the same `Logging` calls, in the same order, as the Python
//! `@client` path makes them.
use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest};
use litellm_host::event::{
FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds,
};
use litellm_host_python::{
AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py,
LifecycleEvent, LifecycleStep, PythonLifecycle, from_py, missing_state, to_py,
};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
types::{PyDict, PyList},
};
use serde_json::Value;
use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call, prepare, setup,
finalize, is_internal_call,
legacy_python::Streaming,
prepare, setup,
};
/// What the legacy contract needs to know about the route it is logging.
@ -25,6 +30,22 @@ pub struct LegacySurface {
pub call_type: &'static str,
/// What `Logging.pre_call` is told the input was.
pub input_description: &'static str,
/// How a streamed response is billed; `None` for a route that never streams.
pub stream: Option<PassThroughStream>,
}
/// The pass-through billing a streamed response goes through once its chunks are in.
#[derive(Clone, Copy, Debug)]
pub struct PassThroughStream {
pub url_route: &'static str,
/// A value of Python's `EndpointType`.
pub endpoint_type: &'static str,
}
/// What the Messages stream iterator keeps for its end-of-stream billing.
struct DeliveredStream {
chunks: Py<PyList>,
first_chunk: Option<Py<PyAny>>,
}
enum Pending {
@ -44,6 +65,8 @@ pub struct LegacyLogging {
error: Option<Py<PyBaseException>>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
context: Option<RequestContext>,
stream: Option<DeliveredStream>,
asynchronous: bool,
internal: bool,
pending: Option<Pending>,
@ -77,6 +100,8 @@ impl LegacyLogging {
error: None,
body: None,
headers: None,
context: None,
stream: None,
asynchronous,
internal: false,
pending: None,
@ -85,8 +110,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 +120,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 +137,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 +170,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 {
@ -162,15 +185,72 @@ impl LegacyLogging {
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.surface.stream.ok_or_else(missing_state)?;
let billed = Streaming::Success.call(
py,
(
logger.object(py),
billing.url_route,
billing.endpoint_type,
&self.body,
&stream.chunks,
&self.start,
&self.end,
&stream.first_chunk,
),
);
match billed {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(logger.object(py)));
Ok(())
}
result => result.map(|_| ()),
}
}
/// A failure after the stream reached the caller bills the delivered chunks as
/// partial usage. The sync path has no loop to schedule that on, so it falls back to
/// the plain failure handler.
fn stream_failure(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
let (Some(logger), Some(error), Some(stream), Some(billing)) =
(&self.logger, &self.error, &self.stream, self.surface.stream)
else {
return Ok(LifecycleStep::Done);
};
if !self.asynchronous {
return self.dispatch_failure(py);
}
let scheduled = Streaming::Failure.call(
py,
(
logger.object(py),
billing.endpoint_type,
&self.body,
&stream.chunks,
error,
),
);
match scheduled {
Ok(awaitable) => {
self.pending = Some(Pending::AsyncFailure);
Ok(LifecycleStep::Await(awaitable.unbind()))
}
Err(failure) if is_cancellation(py, &failure) => Err(failure),
Err(_) => Ok(LifecycleStep::Done),
}
}
/// 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 +258,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 +292,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 +308,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)?;
}
}
@ -249,11 +327,11 @@ impl CallbackAdapter for LegacyLogging {
}
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 +340,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 +352,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,
@ -289,36 +367,50 @@ impl CallbackAdapter for LegacyLogging {
self.finalize(py)
}
fn emit(
&mut self,
py: Python<'_>,
event: &CallEvent,
public: Option<PublicValue<'_>>,
) -> PyResult<AdapterStep> {
match (event, public) {
(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)
fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult<LifecycleStep> {
match event {
LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done),
LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
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(),
self.headers.as_ref(),
)?;
Ok(LifecycleStep::Done)
}
(CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => {
LifecycleEvent::Succeeded { timing, response } => {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
self.dispatch_success(py)?;
Ok(AdapterStep::Done)
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
None => self.dispatch_success(py)?,
}
Ok(LifecycleStep::Done)
}
(CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => {
LifecycleEvent::Failed {
timing,
origin,
error,
} => {
self.end = Some(datetime(py, timing.end_time)?);
self.error = Some(error.clone_ref(py).into_value(py));
if *origin == FailureOrigin::Call
if self.stream.is_some() {
return self.stream_failure(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,
@ -327,11 +419,30 @@ impl CallbackAdapter for LegacyLogging {
}
self.dispatch_failure(py)
}
_ => Err(missing_state()),
}
}
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep> {
fn opened(&mut self, py: Python<'_>) -> PyResult<()> {
if self.surface.stream.is_none() {
return Err(missing_state());
}
Streaming::Opened.call(py, (self.logger()?.object(py),))?;
self.stream = Some(DeliveredStream {
chunks: PyList::empty(py).unbind(),
first_chunk: None,
});
Ok(())
}
fn delivered(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
if stream.first_chunk.is_none() {
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
}
stream.chunks.bind(py).append(chunk)
}
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 +456,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 +468,8 @@ impl CallbackAdapter for LegacyLogging {
error.write_unraisable(py, None);
}
self.body = None;
self.headers = None;
self.context = None;
self.stream = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
@ -369,8 +481,11 @@ impl CallbackAdapter for LegacyLogging {
visit.call(&self.end)?;
visit.call(&self.response)?;
visit.call(&self.error)?;
visit.call(&self.body)?;
visit.call(&self.headers)
if let Some(stream) = &self.stream {
visit.call(&stream.chunks)?;
visit.call(&stream.first_chunk)?;
}
visit.call(&self.body)
}
}

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

@ -3,8 +3,8 @@
//! lifetime. No other callback host has that obligation, which is why nothing outside
//! this crate holds them.
use litellm_callbacks::{machine::Machine, route::Route};
use litellm_host_python::{RouteHost, run_call};
use litellm_host::{machine::Machine, route::Route};
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();

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

@ -2,15 +2,14 @@
//! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls
//! duplication. All of it expires with the legacy callback contract.
use litellm_callbacks::event::{RequestContext, WireRequest};
use litellm_host::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,24 +20,23 @@ 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<()>;
@ -82,16 +80,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 +88,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 +114,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 +132,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 +141,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 +149,30 @@ 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 +182,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 +194,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 +206,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 +217,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 +230,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 +255,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()
}

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

@ -0,0 +1,183 @@
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),
Streaming(Streaming),
}
/// 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,
}
/// The Messages stream iterator's logging: the stream flag, the end-of-stream billing
/// from the delivered chunks, and the partial-usage failure path.
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
pub(crate) enum Streaming {
#[strum(serialize = "stream_opened")]
Opened,
#[strum(serialize = "stream_success")]
Success,
#[strum(serialize = "stream_failure")]
Failure,
}
impl LegacyPython {
fn name(self) -> &'static str {
match self {
Self::Wrapper(function) => function.into(),
Self::Logging(function) => function.into(),
Self::DeploymentHooks(function) => function.into(),
Self::Streaming(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 Streaming {
pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
where
A: pyo3::call::PyCallArgs<'py>,
{
LegacyPython::Streaming(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, Streaming, 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)),
)
.chain(
Streaming::VARIANTS
.iter()
.map(|&function| LegacyPython::Streaming(function)),
)
.map(LegacyPython::name)
.collect();
assert_eq!(called.len(), declared.len(), "a function is borrowed twice");
assert_eq!(called.into_iter().collect::<BTreeSet<_>>(), declared);
}
}

7
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)]
@ -20,8 +21,8 @@ mod preparation;
mod test_support;
pub(crate) use adapter::LegacyLogging;
pub use adapter::LegacySurface;
pub use call::{PublicCall, lookup, run_legacy_call};
pub use adapter::{LegacySurface, PassThroughStream};
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)]

64
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::event::{FailureOrigin, Timing};
use litellm_host_python::{LifecycleEvent, LifecycleStep, 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]
@ -97,6 +97,43 @@ assert checked is prepared
});
}
#[rstest]
#[case::synchronous(false)]
#[case::asynchronous(true)]
fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object(
#[case] asynchronous: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
opaque = object()
hooked = []
logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs}
kwargs = {'logger': logger, 'vendor_extension': opaque}
",
);
let (mut logging, step) = begin(py, &locals, asynchronous);
let step = match step {
LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(),
step => step,
};
locals.set_item("prepared", arguments(py, step)).unwrap();
locals.set_item("asynchronous", asynchronous).unwrap();
run(
py,
&locals,
c"
assert prepared['vendor_extension'] is opaque
[checked] = [value for name, value in logger.calls if name == 'check_limits']
assert checked['vendor_extension'] is opaque
assert hooked == ([opaque] if asynchronous else []), hooked
",
);
});
}
#[test]
fn response_returned_by_the_post_call_hook_is_finalized_and_returned() {
Python::initialize();
@ -121,7 +158,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")));
@ -180,13 +217,12 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle
.resume(py, Ok(local(&locals, "kwargs").unbind()))
.unwrap();
let failure = PyErr::from_value(local(&locals, "failure"));
let failed = CallEvent::Failed {
let failed = LifecycleEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Call,
error: &failure,
};
let step = logging
.emit(py, &failed, Some(PublicValue::Error(&failure)))
.unwrap();
let step = logging.emit(py, failed).unwrap();
assert!(awaits_deployment_hook(&step));
let hook_result = if cancelled {
Err(CancelledError::new_err("cancelled"))
@ -195,7 +231,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 +273,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();

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

@ -1,10 +1,12 @@
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_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest};
use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py};
use proptest::prelude::*;
use pyo3::prelude::*;
use rstest::rstest;
use serde_json::{Value, json};
use serde_json::{Map, Value, json};
use super::LegacyLogging;
use crate::PythonLogger;
@ -23,20 +25,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,33 +46,47 @@ 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 {
before_send_bound(&[], script, optional_params, body, secret_fields)
}
/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs.
fn before_send_bound(
bindings: &[(&str, &Value)],
script: &CStr,
optional_params: Value,
body: Value,
secret_fields: &[&str],
) -> WireRequest {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
for &(name, value) in bindings {
locals.set_item(name, to_py(py, value).unwrap()).unwrap();
}
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(),
@ -86,17 +94,17 @@ fn before_send_with_secrets(
body,
};
let step = logging.before_send(py, Box::new(wire), &context).unwrap();
let raw = CallEvent::ResponseReceived {
let raw = MachineEvent::ResponseReceived {
raw: RawResponse {
body: "raw response".into(),
},
};
assert!(matches!(
logging.emit(py, &raw, None).unwrap(),
AdapterStep::Done
logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(),
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 +137,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 +153,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 +171,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 +179,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 +211,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 +224,6 @@ def on_pre_call(args):
args['headers']['x-callback'] = 'edited'
",
json!({}),
json!({}),
);
assert_eq!(
wire.headers,
@ -288,7 +306,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 +320,6 @@ def on_pre_call(args):
retained['x-retained'] = 'sent'
",
json!({}),
json!({}),
);
assert_eq!(
wire.headers,
@ -314,52 +331,193 @@ 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_and_headers_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'],
'headers': logger.pre['headers'],
}, 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})
);
}
/// What one pre-call callback does to the payload it is handed.
#[derive(Clone, Debug)]
enum Edit {
Nothing,
Set(String, Value),
Remove(String),
Rebind(Value),
RebindThenSetRetained(String, Value),
}
impl Edit {
fn script(&self) -> Value {
match self {
Self::Nothing => json!({"kind": "nothing"}),
Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}),
Self::Remove(key) => json!({"kind": "remove", "key": key}),
Self::Rebind(value) => json!({"kind": "rebind", "value": value}),
Self::RebindThenSetRetained(key, value) => {
json!({"kind": "rebind_then_set_retained", "key": key, "value": value})
}
}
}
/// The legacy contract: the provider is sent the body object `pre_call` received, as
/// the callback left it. Rebinding the envelope's key points the envelope elsewhere and
/// leaves that object alone.
fn sent(&self, body: &Map<String, Value>) -> Value {
let mut sent = body.clone();
match self {
Self::Nothing | Self::Rebind(_) => {}
Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => {
sent.insert(key.clone(), value.clone());
}
Self::Remove(key) => {
sent.remove(key);
}
}
Value::Object(sent)
}
}
/// How the caller's keyword for a body key relates to what the route sends under it.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Caller {
PassedUnchanged,
RewrittenByTheRoute,
NotPassed,
}
const MODEL: &CStr = c"
aliased = {}
def on_pre_call(args):
body = args['complete_input_dict']
aliased.update({name: body[name] is kwargs[name] for name in unchanged})
kind = edit['kind']
if kind == 'set':
body[edit['key']] = edit['value']
elif kind == 'remove':
body.pop(edit['key'], None)
elif kind == 'rebind':
args['complete_input_dict'] = edit['value']
elif kind == 'rebind_then_set_retained':
args['complete_input_dict'] = {}
body[edit['key']] = edit['value']
def check():
assert aliased == {name: True for name in unchanged}, aliased
assert logger.names() == ['pre_call', 'post_call'], logger.calls
";
fn json_value() -> impl Strategy<Value = Value> {
let leaf = prop_oneof![
Just(Value::Null),
any::<bool>().prop_map(Value::from),
any::<i64>().prop_map(Value::from),
any::<f64>()
.prop_filter("JSON has no NaN or infinity", |number| number.is_finite())
.prop_map(Value::from),
".{0,8}".prop_map(Value::from),
];
leaf.prop_recursive(3, 24, 4, |inner| {
prop_oneof![
prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from),
prop::collection::btree_map(key(), inner, 0..4)
.prop_map(|fields| Value::Object(fields.into_iter().collect())),
]
})
}
fn key() -> impl Strategy<Value = String> {
"[a-z]{1,6}"
}
fn caller() -> impl Strategy<Value = Caller> {
prop_oneof![
Just(Caller::PassedUnchanged),
Just(Caller::RewrittenByTheRoute),
Just(Caller::NotPassed),
]
}
fn edit() -> impl Strategy<Value = Edit> {
prop_oneof![
Just(Edit::Nothing),
(key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)),
key().prop_map(Edit::Remove),
json_value().prop_map(Edit::Rebind),
(key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)),
]
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(128))]
/// For any body, any caller keywords and any callback edit: every keyword the route
/// sends unchanged reaches `pre_call` as the caller's own object, and the provider is
/// sent exactly what the model says, so a callback that edits nothing changes nothing.
#[test]
fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it(
fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5),
edit in edit(),
) {
let body: Map<String, Value> = fields
.iter()
.map(|(name, (value, _))| (name.clone(), value.clone()))
.collect();
let kwargs: Map<String, Value> = fields
.iter()
.filter_map(|(name, (value, caller))| match caller {
Caller::PassedUnchanged => Some((name.clone(), value.clone())),
Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))),
Caller::NotPassed => None,
})
.collect();
let unchanged: Value = fields
.iter()
.filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged)
.map(|(name, _)| Value::from(name.clone()))
.collect();
let wire = before_send_bound(
&[
("kwargs", &Value::Object(kwargs)),
("unchanged", &unchanged),
("edit", &edit.script()),
],
MODEL,
json!({}),
Value::Object(body.clone()),
&[],
);
prop_assert_eq!(wire.body, edit.sent(&body));
prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]);
}
}

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

@ -5,64 +5,94 @@ 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),
'stream_opened': lambda logger: logger.record('stream_opened', None),
'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error),
}
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 +107,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 +122,6 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.needed = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
@ -147,6 +162,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
@ -181,6 +197,7 @@ pub(crate) fn legacy_call(
LegacySurface {
call_type: "test",
input_description: "test input",
stream: None,
},
call,
asynchronous,

58
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::event::{FailureOrigin, Timing};
use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle};
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
@ -19,52 +19,52 @@ 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(
py,
&CallEvent::Succeeded { timing: TIMING },
Some(PublicValue::Response(&response)),
LifecycleEvent::Succeeded {
timing: TIMING,
response: &response,
},
)
.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(
py,
&CallEvent::Failed {
LifecycleEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Host,
error: &failure,
},
Some(PublicValue::Error(&failure)),
)
.unwrap()
}
#[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 +76,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 +109,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 +160,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 +176,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 +189,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 +227,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 +265,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()));

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

@ -1,135 +0,0 @@
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::{Map, Value};
/// Seconds since the Unix epoch, on one clock for every host.
pub fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or(0.0)
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Timing {
pub start_time: f64,
pub end_time: f64,
}
/// The provider request as it is about to leave, offered to the host for rewriting.
#[derive(Clone, Debug, PartialEq)]
pub struct WireRequest {
pub url: String,
pub headers: Vec<(String, String)>,
pub body: Value,
}
/// What the route knows about the request it is sending, for a host that logs it. The
/// route owns these facts; a host reads them beside the wire request and never rewrites
/// them.
#[derive(Clone, Debug, PartialEq)]
pub struct RequestContext {
pub model: String,
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)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RawResponse {
pub body: String,
}
/// Whether a failure surfaced inside the call, including a host op the call asked for,
/// or in a host step around it (preparing the arguments, finalizing the response).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FailureOrigin {
Call,
Host,
}
#[derive(Clone, Debug, PartialEq)]
pub enum CallEvent {
ResponseReceived {
raw: RawResponse,
},
Succeeded {
timing: Timing,
},
Failed {
timing: Timing,
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);
}
}

2
litellm-rust/crates/core/Cargo.toml
View file

@ -9,7 +9,7 @@ autotests = false
[dependencies]
litellm-types.workspace = true
litellm-core-utils.workspace = true
litellm-callbacks.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
base64.workspace = true

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

@ -2,7 +2,6 @@ pub mod audio_transcription;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod machine;
pub mod messages;
pub mod ocr;
pub mod responses;

128
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -1,88 +1,54 @@
use litellm_llms::custom_httpx::http_handler::http_request;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use std::time::Duration;
use super::{
Error, client::http_client, common_utils::truncate_error_body,
prepare::prepare_provider_request,
use litellm_llms::{
base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
custom_httpx::{http_handler::http_request, transport::Error as TransportError},
};
use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, messages::types::MessagesRequest};
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use serde_json::Value;
pub(super) async fn execute_messages_provider_call(
request: MessagesRequest<'_>,
use super::{Error, client::http_client, common_utils::truncate_error_body};
pub(super) fn network(error: reqwest::Error) -> Error {
Error::Transport(TransportError::Network(error.to_string()))
}
pub(super) async fn send(
url: &str,
headers: &[(String, String)],
body: &Value,
timeout: Option<Duration>,
) -> Result<reqwest::Response, Error> {
let builder = headers.iter().fold(
http_client().post(url).json(body),
|builder, (key, value)| builder.header(key, value),
);
let builder = match timeout {
Some(duration) => builder.timeout(duration),
None => builder,
};
http_request(builder).await.map_err(network)
}
pub(super) async fn provider_error(response: reqwest::Response) -> Error {
let status = response.status().as_u16();
match response.text().await {
Ok(text) => Error::Transport(TransportError::Http {
status,
body: truncate_error_body(&text),
}),
Err(error) => network(error),
}
}
pub(super) fn decode_response(
config: &dyn BaseAnthropicMessagesConfig,
model: &str,
text: &str,
) -> Result<AnthropicMessagesResponse, Error> {
let request = prepare_provider_request(request)?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder).await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
let status = response.status();
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
if !status.is_success() {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
let response = serde_json::from_str(&text)
let response = serde_json::from_str(text)
.map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?;
request
.config
.transform_anthropic_messages_response(&request.model, response)
config
.transform_anthropic_messages_response(model, response)
.map_err(Error::from)
}
pub(super) async fn execute_messages_provider_stream(
request: MessagesRequest<'_>,
) -> Result<reqwest::Response, Error> {
let request = prepare_provider_request(request)?;
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::Unsupported("streaming messages for this provider"));
}
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder).await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
let status = response.status();
if !status.is_success() {
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
Ok(response)
}

36
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -1,11 +1,8 @@
//! The Anthropic Messages call, the Rust equivalent of Python's
//! `litellm.messages()`.
//!
//! [`messages`] is the top-level entrypoint: give it a model, a body, and
//! credentials, and it resolves the provider, transforms the request, calls the
//! provider, and returns a typed non-streaming response. [`messages_stream`]
//! is the streaming variant; it hands the raw upstream response back so a host
//! can splice the event stream to its own caller.
//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs
//! it in process for a caller that already holds the request and wants the message.
mod error;
pub mod types;
@ -14,17 +11,34 @@ mod client;
mod common_utils;
mod handler;
mod prepare;
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
pub mod route;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine};
use serde_json::Value;
use crate::messages::types::MessagesRequest;
pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {
execute_messages_provider_call(request).await
}
pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {
execute_messages_provider_stream(request).await
let Value::Object(body) = request.body else {
return Err(Error::InvalidRequest(
"messages body must be an object".into(),
));
};
let call = MessagesCall {
model: request.model.into(),
body,
api_key: request.api_key.map(Into::into),
api_base: request.api_base.map(Into::into),
custom_llm_provider: request.custom_llm_provider.map(Into::into),
extra_headers: request.extra_headers,
timeout: request.timeout,
};
match litellm_host::run::run(messages_machine(), &LocalMessagesHost::new(call)).await? {
MessagesOutput::Message(message) => Ok(*message),
MessagesOutput::Streamed => Err(Error::Unsupported(
"streamed responses need a streaming host",
)),
}
}
#[cfg(test)]

11
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -2,6 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_l
use litellm_llms::base_llm::anthropic_messages::transformation::{
BaseAnthropicMessagesConfig, MessagesAuthStrategy,
};
use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest;
use serde_json::{Map, Value};
use super::{
@ -37,10 +38,14 @@ pub(super) fn prepare_provider_request(
let headers =
validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?;
let typed_request = serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
let typed_request: AnthropicMessagesRequest =
serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
})?;
let transformed = config.transform_anthropic_messages_request(AnthropicMessagesRequest {
model: model.clone(),
..typed_request
})?;
let transformed = config.transform_anthropic_messages_request(typed_request)?;
let body = serde_json::to_value(transformed).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize Anthropic messages request: {err}"

196
litellm-rust/crates/core/src/messages/route.rs Normal file
View file

@ -0,0 +1,196 @@
use std::{sync::Mutex, time::Duration};
use bytes::Bytes;
use litellm_auth::SecretValue;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
host::{Demand, Host},
machine::{HostChannel, MachineFault, RouteMachine},
route::Route,
};
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use serde_json::{Map, Value};
use super::{
Error,
common_utils::messages_provider_config,
handler::{decode_response, network, provider_error, send},
prepare::prepare_provider_request,
types::MessagesRequest,
};
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MessagesOp {
ProjectRequest,
}
pub enum MessagesOpResult {
Request(Box<MessagesCall>),
}
/// The caller's request as the host projects it.
pub struct MessagesCall {
pub model: String,
pub body: Map<String, Value>,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
}
impl MessagesCall {
fn streams(&self) -> bool {
self.body.get("stream").and_then(Value::as_bool) == Some(true)
}
}
pub enum MessagesOutput {
Message(Box<AnthropicMessagesResponse>),
/// Every chunk already reached the host through `Deliver`.
Streamed,
}
pub struct Messages;
impl Route for Messages {
type Response = MessagesOutput;
type Error = Error;
type Op = MessagesOp;
type OpResult = MessagesOpResult;
type Chunk = Bytes;
type StreamHead = ();
}
impl From<MachineFault> for Error {
fn from(fault: MachineFault) -> Self {
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "messages host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("messages {message}"),
MachineFault::Mismatch => "invalid messages host operation result".into(),
})
}
}
pub type MessagesHost = HostChannel<Messages>;
pub type MessagesMachine = RouteMachine<Messages>;
/// Whether this route serves the request, decided before any callback runs so a host
/// can still run its own path.
pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool {
let provider = get_custom_llm_provider(model, custom_llm_provider)
.map(|resolved| resolved.custom_llm_provider)
.or(custom_llm_provider);
match provider {
Some(ANTHROPIC_MESSAGES_PROVIDER) => true,
Some(provider) => !stream && messages_provider_config(provider).is_some(),
None => false,
}
}
/// The in-process host for a request already in hand. It answers projection once and
/// observes nothing.
pub struct LocalMessagesHost {
call: Mutex<Option<MessagesCall>>,
}
impl LocalMessagesHost {
pub fn new(call: MessagesCall) -> Self {
Self {
call: Mutex::new(Some(call)),
}
}
}
impl Host<Messages> for LocalMessagesHost {
async fn route(&self, op: MessagesOp) -> Result<MessagesOpResult, Error> {
match op {
MessagesOp::ProjectRequest => self
.call
.lock()
.unwrap_or_else(|error| error.into_inner())
.take()
.map(|call| MessagesOpResult::Request(Box::new(call)))
.ok_or_else(|| {
Error::InvalidRequest("messages request was already projected".into())
}),
}
}
}
pub fn messages_machine() -> MessagesMachine {
RouteMachine::new(|host| Box::pin(execute(host)))
}
async fn execute(host: MessagesHost) -> Result<MessagesOutput, Error> {
let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?;
let stream = call.streams();
let request = prepare_provider_request(MessagesRequest {
model: &call.model,
body: Value::Object(call.body.clone()),
api_key: call.api_key.as_deref(),
api_base: call.api_base.as_deref(),
custom_llm_provider: call.custom_llm_provider.as_deref(),
extra_headers: call.extra_headers.clone(),
timeout: call.timeout,
})?;
if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::Unsupported("streaming messages for this provider"));
}
let context = RequestContext {
model: request.model.clone(),
custom_llm_provider: request.provider.clone(),
optional_params: Value::Object(
call.body
.iter()
.filter(|(name, _)| !matches!(name.as_str(), "model" | "messages"))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
),
secret_fields: Vec::new(),
api_key: call.api_key.clone().map(SecretValue::new),
};
let wire = host
.before_send(
WireRequest {
url: request.url,
headers: request.upstream_headers,
body: request.body,
},
context,
)
.await?;
let response = send(&wire.url, &wire.headers, &wire.body, request.timeout).await?;
if !response.status().is_success() {
return Err(provider_error(response).await);
}
if stream {
return relay(&host, response).await;
}
let text = response.text().await.map_err(network)?;
host.emit(MachineEvent::ResponseReceived {
raw: RawResponse { body: text.clone() },
})
.await?;
decode_response(request.config, &request.model, &text)
.map(|message| MessagesOutput::Message(Box::new(message)))
}
/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading
/// ends the upstream read, and the call completes with what it delivered.
async fn relay(
host: &MessagesHost,
mut response: reqwest::Response,
) -> Result<MessagesOutput, Error> {
if host.open(()).await? == Demand::Detached {
return Ok(MessagesOutput::Streamed);
}
while let Some(chunk) = response.chunk().await.map_err(network)? {
if host.deliver(chunk).await? == Demand::Detached {
break;
}
}
Ok(MessagesOutput::Streamed)
}

2
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -12,7 +12,7 @@ pub async fn perform(
client: &OcrClient,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {

15
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_host::event::{MachineEvent, 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,28 +53,25 @@ 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))
}
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(self.host.emit(CallEvent::ResponseReceived {
Box::pin(self.host.emit(MachineEvent::ResponseReceived {
raw: RawResponse {
body: String::from_utf8_lossy(body).into_owned(),
},

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

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

@ -1,8 +1,9 @@
use std::sync::{Arc, Mutex};
use litellm_auth::ResolvedCredential;
use litellm_callbacks::{
use litellm_host::{
event::{CallEvent, RequestContext, WireRequest},
machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute},
route::Route,
};
use litellm_llms::{
@ -11,10 +12,7 @@ use litellm_llms::{
};
use super::handler::perform_ocr_request;
use crate::{
machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute},
ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest},
};
use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrOp {
@ -39,6 +37,8 @@ impl Route for Ocr {
type Error = Error;
type Op = OcrOp;
type OpResult = OcrOpResult;
type Chunk = std::convert::Infallible;
type StreamHead = std::convert::Infallible;
}
impl TokenRoute for Ocr {
@ -54,16 +54,6 @@ impl TokenRoute for Ocr {
}
}
impl From<MachineFault> for Error {
fn from(fault: MachineFault) -> Self {
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "OCR host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("OCR {message}"),
MachineFault::Mismatch => "invalid OCR host operation result".into(),
})
}
}
pub type OcrHost = HostChannel<Ocr>;
pub type OcrMachine = RouteMachine<Ocr>;
@ -173,7 +163,7 @@ impl LocalOcrHost {
}
}
impl litellm_callbacks::host::Host<Ocr> for LocalOcrHost {
impl litellm_host::host::Host<Ocr> for LocalOcrHost {
async fn route(&self, op: OcrOp) -> Result<OcrOpResult, Error> {
match op {
OcrOp::ProjectRequest => self

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

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

@ -1,4 +1,4 @@
use litellm_callbacks::event::CallEvent;
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_llms::base_llm::ocr::error::Error;
use rstest::rstest;
use serde_json::{Value, json};
@ -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,
@ -263,7 +263,7 @@ async fn accepted_response_emits_response_received_before_polling() {
json!({}),
))
.with_observer(move |event| {
let CallEvent::ResponseReceived { raw } = event else {
let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else {
return;
};
match request_count.lock().unwrap().len() {
@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() {
mod transformation {
use std::sync::{Arc, Mutex};
use litellm_callbacks::event::CallEvent;
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_llms::base_llm::ocr::transformation::OcrDocument;
use serde_json::{Value, json};
@ -646,7 +646,7 @@ mod transformation {
json!({}),
))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
observed
.lock()
.unwrap()

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

@ -1,7 +1,7 @@
use std::sync::{Arc, Mutex};
use litellm_callbacks::{
event::{CallEvent, WireRequest},
use litellm_host::{
event::{CallEvent, MachineEvent, WireRequest},
host::{Host, HostOp, HostResult},
machine::{HostFailure, Machine, MachineStep},
};
@ -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,7 +194,8 @@ async fn facade_uses_the_injected_http_client() {
fn event_name(event: &CallEvent) -> &'static str {
match event {
CallEvent::ResponseReceived { .. } => "response",
CallEvent::Started { .. } => "started",
CallEvent::Machine(MachineEvent::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);
}
@ -371,6 +375,7 @@ async fn drive_until(
intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire)))
}
HostOp::Emit(event) => {
let event = CallEvent::Machine(event);
ops.push(event_name(&event));
host.emit(&event)
.await
@ -414,7 +419,7 @@ async fn invalid_provider_response_emits_response_received_before_normalization_
let observed = responses_received.clone();
let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer(
move |event| {
if let CallEvent::ResponseReceived { raw } = event {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
observed.lock().unwrap().push(raw.body.clone());
}
},
@ -815,7 +820,7 @@ impl Host<crate::ocr::route::Ocr> for CallerTokenHost {
async fn before_send(
&self,
wire: WireRequest,
_: &litellm_callbacks::event::RequestContext,
_: &litellm_host::event::RequestContext,
) -> Result<WireRequest, OcrError> {
let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization");
let authorization = wire
@ -850,7 +855,7 @@ async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_
trace: Mutex::new(Vec::new()),
};
litellm_callbacks::run::run(ocr_machine(ocr_client()), &host)
litellm_host::run::run(ocr_machine(ocr_client()), &host)
.await
.unwrap();
server.await.unwrap();

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

@ -0,0 +1,152 @@
use litellm_host::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)
);
}

12
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_host::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) })
}
@ -49,7 +45,7 @@ pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcr
}
pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
litellm_callbacks::run::run(ocr_machine(ocr_client()), &host).await
litellm_host::run::run(ocr_machine(ocr_client()), &host).await
}
pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
@ -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,

8
litellm-rust/crates/core/tests/reducto_ocr.rs
View file

@ -1,4 +1,4 @@
use litellm_callbacks::event::{CallEvent, WireRequest};
use litellm_host::event::{CallEvent, MachineEvent, WireRequest};
use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument};
use rstest::rstest;
use serde_json::{Value, json};
@ -139,7 +139,7 @@ async fn response_received_stays_after_reducto_upload_and_parse() {
let request_count = seen.clone();
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer(
move |event| {
if let CallEvent::ResponseReceived { raw } = event {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
assert_eq!(request_count.lock().unwrap().len(), 2);
assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
}
@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() {
}
mod transformation {
use litellm_callbacks::event::{CallEvent, WireRequest};
use litellm_host::event::{CallEvent, MachineEvent, WireRequest};
use litellm_llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext},
reducto::ocr::transformation::*,
@ -506,7 +506,7 @@ mod transformation {
let request_count = seen.clone();
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
assert_eq!(request_count.lock().unwrap().len(), 2);
assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
}

7
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
- No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features
- 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-host`: 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 `InvokeError::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>`

2
litellm-rust/crates/host-python/Cargo.toml
View file

@ -7,7 +7,7 @@ repository.workspace = true
[dependencies]
futures-util.workspace = true
litellm-callbacks.workspace = true
litellm-host.workspace = true
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
pythonize.workspace = true

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

@ -1,5 +1,5 @@
use litellm_callbacks::event::{CallEvent, RequestContext, Timing, WireRequest};
use litellm_callbacks::route::Route;
use litellm_host::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest};
use litellm_host::route::Route;
use pyo3::exceptions::PyRuntimeError;
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
@ -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>),
@ -19,64 +19,97 @@ pub enum AdapterStep {
Done,
}
/// The host-typed value the driver attaches to a terminal event.
pub enum PublicValue<'a> {
Response(&'a Py<PyAny>),
Error(&'a PyErr),
/// What a lifecycle observes: the driver's start, the machine's own events, and one
/// terminal event carrying the public value the caller receives.
pub enum LifecycleEvent<'a> {
Started {
start_time: f64,
},
Machine(&'a MachineEvent),
Succeeded {
timing: Timing,
response: &'a Py<PyAny>,
},
Failed {
timing: Timing,
origin: FailureOrigin,
error: &'a PyErr,
},
}
/// 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>;
fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult<LifecycleStep>;
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep>;
/// The call streams and its stream was handed to the caller. The caller is not
/// inside an await here, so this step and `delivered` cannot suspend.
fn opened(&mut self, py: Python<'_>) -> PyResult<()>;
/// One chunk of an open stream is about to reach the caller.
fn delivered(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()>;
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 InvokeError<E> {
Native(E),
Python(PyErr),
}
/// `arguments` is the keyword view the callback adapter's `begin` produced, not the
impl<E> From<PyErr> for InvokeError<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 +117,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, InvokeError<<Self::Route as Route>::Error>>;
fn complete(
&mut self,
@ -92,12 +125,21 @@ pub trait RouteHost: Send + Sync {
response: <Self::Route as Route>::Response,
) -> PyResult<Py<PyAny>>;
fn native_error(error: <Self::Route as Route>::Error) -> PyErr;
/// One streamed chunk as the caller receives it.
fn chunk(
&mut self,
py: Python<'_>,
chunk: <Self::Route as Route>::Chunk,
) -> PyResult<Py<PyAny>>;
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());
});
}
}

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

@ -2,17 +2,19 @@ use std::sync::Arc;
use std::task::Poll;
use futures_util::future::{AbortHandle, Abortable};
use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds};
use litellm_callbacks::host::{HostOp, HostResult, HostStep};
use litellm_callbacks::machine::{HostFailure, Machine, MachineStep};
use litellm_callbacks::route::Route;
use litellm_host::event::{FailureOrigin, Timing, epoch_seconds};
use litellm_host::host::{Demand, HostOp, HostResult, HostStep};
use litellm_host::machine::{HostFailure, Machine, MachineStep};
use litellm_host::route::Route;
use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError};
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use tokio::sync::Mutex;
use crate::adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state};
use crate::adapter::{
InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state,
};
use crate::execution::{poll_async_value, run_async_value, run_sync_value};
use crate::handle::{Execution, ExecutionBody, ExecutionStep};
@ -36,6 +38,7 @@ struct MachineState<M: Machine> {
enum Stage {
Begin,
Call,
Streaming,
AfterSuccess,
Succeeded(Py<PyAny>),
Failed(Py<PyBaseException>),
@ -43,6 +46,7 @@ enum Stage {
#[derive(Clone, Copy)]
enum Expect {
Started,
Arguments,
Wire,
Emitted,
@ -53,6 +57,8 @@ enum Expect {
enum Pending {
Native,
Adapter(Expect),
/// The stream handed to the caller waits for its next read or its close.
Consumer,
}
enum Next<H: RouteHost> {
@ -66,7 +72,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 +90,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>>
@ -118,7 +124,14 @@ where
}
match driver.resume(None)? {
ExecutionStep::Return(value) => Ok(value),
ExecutionStep::Await(_) => Err(PyRuntimeError::new_err("sync call suspended")),
ExecutionStep::Open => py
.import("litellm.rust_bridge.lifecycle")?
.getattr("SyncStream")?
.call1((Py::new(py, Execution::suspended(driver))?,))
.map(Bound::unbind),
ExecutionStep::Await(_) | ExecutionStep::Yield(_) => {
Err(PyRuntimeError::new_err("sync call suspended"))
}
}
}
@ -146,9 +159,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 = LifecycleEvent::Started {
start_time: self.started_at,
};
match self.adapter.emit(py, started) {
Ok(step) => self.on_adapter(py, step, Expect::Started),
Err(error) => self.adapter_failed(py, error),
}
}
@ -157,6 +172,14 @@ where
self.run_steps(py, HostStep::Ready(result))
}
(Some(Pending::Native), Some(Err(error))) => self.interrupt(py, error),
(Some(Pending::Consumer), Some(read)) => {
let demand = if read.is_ok() {
Demand::More
} else {
Demand::Detached
};
self.resume_machine(py, Some(Ok(HostResult::Demand(demand))))
}
(Some(Pending::Adapter(expect)), Some(result)) => {
match self.adapter.resume(py, result) {
Ok(step) => self.on_adapter(py, step, expect),
@ -170,27 +193,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,10 +223,18 @@ 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),
Stage::Call => self.interrupt(py, error),
Stage::Call | Stage::Streaming => self.interrupt(py, error),
Stage::Succeeded(_) | Stage::Failed(_) => Err(error),
}
}
@ -248,14 +280,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(InvokeError::Native(error)) => {
return self
.resume_core(py, Some(Err(HostFailure::Error(error))))
.map(Next::Continue);
}
Err(InvokeError::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)));
}
@ -263,9 +301,11 @@ where
Err(error) => Err(error),
}
}
HostOp::Emit(event) => match self.adapter.emit(py, &event, None) {
Ok(AdapterStep::Done) => Ok(HostResult::Emitted),
Ok(AdapterStep::Await(awaitable)) => {
HostOp::Open(_) => return self.opened(py).map(Next::Return),
HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return),
HostOp::Emit(event) => match self.adapter.emit(py, LifecycleEvent::Machine(&event)) {
Ok(LifecycleStep::Done) => Ok(HostResult::Emitted),
Ok(LifecycleStep::Await(awaitable)) => {
self.pending = Some(Pending::Adapter(Expect::Emitted));
return Ok(Next::Return(ExecutionStep::Await(awaitable)));
}
@ -279,6 +319,35 @@ where
}
}
fn opened(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
self.stage = Stage::Streaming;
match self.adapter.opened(py) {
Ok(()) => {
self.pending = Some(Pending::Consumer);
Ok(ExecutionStep::Open)
}
Err(error) => self.interrupt(py, error),
}
}
fn delivered(
&mut self,
py: Python<'_>,
chunk: <RouteOf<H> as Route>::Chunk,
) -> PyResult<ExecutionStep> {
let chunk = match self.route.chunk(py, chunk) {
Ok(chunk) => chunk,
Err(error) => return self.interrupt(py, error),
};
match self.adapter.delivered(py, &chunk) {
Ok(()) => {
self.pending = Some(Pending::Consumer);
Ok(ExecutionStep::Yield(chunk))
}
Err(error) => self.interrupt(py, error),
}
}
fn interrupt(&mut self, py: Python<'_>, error: PyErr) -> PyResult<ExecutionStep> {
let cancelled = is_cancellation(py, &error);
let native = H::host_error(&error);
@ -349,6 +418,9 @@ where
Ok(public) => public,
Err(error) => return self.failure(py, error, FailureOrigin::Call),
};
if let Stage::Streaming = self.stage {
return self.succeeded(py, public);
}
self.stage = Stage::AfterSuccess;
match self.adapter.after_success(py, public, self.timing()) {
Ok(step) => self.on_adapter(py, step, Expect::Response),
@ -360,18 +432,35 @@ 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)
}
fn succeeded(&mut self, py: Python<'_>, response: Py<PyAny>) -> PyResult<ExecutionStep> {
let event = CallEvent::Succeeded {
timing: self.timing(),
/// 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 step = self
.adapter
.emit(py, &event, Some(PublicValue::Response(&response)))?;
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 = LifecycleEvent::Succeeded {
timing: self.timing(),
response: &response,
};
let step = self.adapter.emit(py, event)?;
self.stage = Stage::Succeeded(response);
self.on_adapter(py, step, Expect::Terminal)
}
@ -386,18 +475,13 @@ 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 {
let event = LifecycleEvent::Failed {
timing: self.timing(),
origin,
error: &error,
};
let step = self
.adapter
.emit(py, &event, Some(PublicValue::Error(&public)))?;
self.stage = Stage::Failed(public.into_value(py));
let step = self.adapter.emit(py, event)?;
self.stage = Stage::Failed(error.into_value(py));
self.on_adapter(py, step, Expect::Terminal)
}
@ -450,8 +534,8 @@ where
mod tests {
use std::sync::{Arc, Mutex};
use litellm_callbacks::event::{RequestContext, WireRequest};
use litellm_callbacks::machine::{Interrupted, Step};
use litellm_host::event::{MachineEvent, RequestContext, WireRequest};
use litellm_host::machine::{Interrupted, Step};
use pyo3::exceptions::{PyBaseException, PyValueError};
use pyo3::types::PyDict;
@ -489,6 +573,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 {
@ -496,6 +586,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
type Error = Error;
type Op = &'static str;
type OpResult = String;
type Chunk = std::convert::Infallible;
type StreamHead = std::convert::Infallible;
}
/// Yields the scripted ops in order, then completes or fails as scripted.
@ -518,8 +610,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,
}
}
@ -534,6 +626,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
HostResult::Route(value) => value,
HostResult::BeforeSend(wire) => wire.url,
HostResult::Emitted => "emitted".into(),
HostResult::Demand(demand) => format!("{demand:?}"),
});
}
if !self.ops.is_empty() {
@ -566,25 +659,50 @@ 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, InvokeError<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(InvokeError::Native(Error("op rejected".into()))),
}
Ok(format!("{op}:{}", arguments.len()))
}
fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
match chunk {}
}
fn complete(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
@ -594,22 +712,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 +746,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 +765,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,41 +778,48 @@ 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))
}
}
}
fn emit(
&mut self,
py: Python<'_>,
event: &CallEvent,
public: Option<PublicValue<'_>>,
) -> PyResult<AdapterStep> {
self.log.push(match (event, public) {
(CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body),
(CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => {
format!("succeeded:{}", value.bind(py))
fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult<LifecycleStep> {
self.log.push(match event {
LifecycleEvent::Started { .. } => "started".into(),
LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
format!("response:{}", raw.body)
}
(CallEvent::Failed { origin, .. }, Some(PublicValue::Error(error))) => {
LifecycleEvent::Succeeded { response, .. } => {
format!("succeeded:{}", response.bind(py))
}
LifecycleEvent::Failed { origin, error, .. } => {
format!("failed:{origin:?}:{}", error.value(py))
}
_ => "unexpected".into(),
});
Ok(AdapterStep::Done)
Ok(LifecycleStep::Done)
}
fn resume(&mut self, _: Python<'_>, _: PyResult<Py<PyAny>>) -> PyResult<AdapterStep> {
fn opened(&mut self, _: Python<'_>) -> PyResult<()> {
self.log.push("opened");
Ok(())
}
fn delivered(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("delivered");
Ok(())
}
fn resume(&mut self, _: Python<'_>, _: PyResult<Py<PyAny>>) -> PyResult<LifecycleStep> {
Err(missing_state())
}
@ -709,15 +835,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,
@ -756,8 +898,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
wire: Box::new(wire()),
context: Box::new(context()),
},
HostOp::Emit(CallEvent::ResponseReceived {
raw: litellm_callbacks::event::RawResponse { body: "raw".into() },
HostOp::Emit(MachineEvent::ResponseReceived {
raw: litellm_host::event::RawResponse { body: "raw".into() },
}),
],
outcome: Some(Ok("done".into())),
@ -777,7 +919,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 +927,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
assert_eq!(
log,
[
"started",
"begin",
"route:project",
"before_send",
@ -800,28 +943,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 +1020,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 +1099,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 +1108,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 +1130,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 +1153,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 +1183,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, InvokeError<Error>> {
self.0.push("route");
Err(PyErr::from_value(
py.import("asyncio")
@ -952,21 +1198,26 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
.unwrap()
.call0()
.unwrap(),
))
)
.into())
}
fn chunk(
&mut self,
_: Python<'_>,
chunk: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match chunk {}
}
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 +1239,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"]
);
});
}

13
litellm-rust/crates/host-python/src/handle.rs
View file

@ -8,6 +8,10 @@ use pyo3::prelude::*;
pub enum ExecutionStep {
Return(Py<PyAny>),
Await(Py<PyAny>),
/// The call streams: the caller gets a stream over this execution, which stays
/// suspended until the stream asks for a chunk.
Open,
Yield(Py<PyAny>),
}
pub trait ExecutionBody: Send + Sync {
@ -34,6 +38,13 @@ impl Execution {
}
}
/// An execution already started elsewhere and now waiting for its next input.
pub fn suspended(body: impl ExecutionBody + 'static) -> Self {
Self {
state: ExecutionState::Suspended(Box::new(body)),
}
}
fn advance(
slf: &Bound<'_, Self>,
py: Python<'_>,
@ -64,6 +75,8 @@ impl Execution {
let step = body.resume(result)?;
let (tag, value, suspended) = match step {
ExecutionStep::Await(value) => ("Await", value, true),
ExecutionStep::Open => ("Open", py.None(), true),
ExecutionStep::Yield(value) => ("Yield", value, true),
ExecutionStep::Return(value) => ("Complete", value, false),
};
let step = py

10
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
//! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine)
//! 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::{
InvokeError, LifecycleEvent, LifecycleStep, 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};

5
litellm-rust/crates/callbacks/Cargo.toml → litellm-rust/crates/host/Cargo.toml
View file

@ -1,13 +1,14 @@
[package]
name = "litellm-callbacks"
name = "litellm-host"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros"] }

76
litellm-rust/crates/host/src/event.rs Normal file
View file

@ -0,0 +1,76 @@
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::Value;
/// Seconds since the Unix epoch, on one clock for every host.
pub fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or(0.0)
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Timing {
pub start_time: f64,
pub end_time: f64,
}
/// The provider request as it is about to leave, offered to the host for rewriting.
#[derive(Clone, Debug, PartialEq)]
pub struct WireRequest {
pub url: String,
pub headers: Vec<(String, String)>,
pub body: Value,
}
/// What the route knows about the request it is sending, for a host that logs it. The
/// route owns these facts; a host reads them beside the wire request and never rewrites
/// them.
#[derive(Clone, Debug, PartialEq)]
pub struct RequestContext {
pub model: String,
pub custom_llm_provider: String,
/// The route's parameters before the provider transformation.
pub optional_params: Value,
/// Optional-param names that carry credentials and must be redacted when logged.
pub secret_fields: Vec<String>,
/// The credential the route resolved for the provider call.
pub api_key: Option<litellm_auth::SecretValue>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RawResponse {
pub body: String,
}
/// Whether a failure surfaced inside the call, including a host op the call asked for,
/// or in a host step around it (preparing the arguments, finalizing the response).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FailureOrigin {
Call,
Host,
}
/// What a machine reports while it runs.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum MachineEvent {
ResponseReceived { raw: RawResponse },
}
/// What an in-process host observes: the machine's own events between the driver's
/// start and terminal ones.
#[derive(Clone, Debug, PartialEq)]
pub enum CallEvent {
Started {
start_time: f64,
},
Machine(MachineEvent),
Succeeded {
timing: Timing,
},
Failed {
timing: Timing,
origin: FailureOrigin,
},
}

25
litellm-rust/crates/callbacks/src/host.rs → litellm-rust/crates/host/src/host.rs
View file

@ -1,6 +1,6 @@
use std::future::Future;
use crate::event::{CallEvent, RequestContext, WireRequest};
use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest};
use crate::route::Route;
/// One suspension point of a native call, performed by the host.
@ -10,13 +10,26 @@ pub enum HostOp<R: Route> {
wire: Box<WireRequest>,
context: Box<RequestContext>,
},
Emit(CallEvent),
Emit(MachineEvent),
/// The response streams: the host hands the caller a stream and answers once the
/// caller asks for the first chunk or goes away.
Open(R::StreamHead),
/// The next chunk of an open stream, answered once the caller asks for the one after.
Deliver(R::Chunk),
}
pub enum HostResult<R: Route> {
Route(R::OpResult),
BeforeSend(Box<WireRequest>),
Emitted,
Demand(Demand),
}
/// Whether the caller of a streamed call still reads it.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Demand {
More,
Detached,
}
/// A host answer that is either available now or arrives once the host's own
@ -42,4 +55,12 @@ pub trait Host<R: Route>: Send + Sync {
fn emit(&self, _event: &CallEvent) -> impl Future<Output = Result<(), R::Error>> + Send {
async { Ok(()) }
}
fn open(&self, _head: R::StreamHead) -> impl Future<Output = Result<Demand, R::Error>> + Send {
async { Ok(Demand::More) }
}
fn deliver(&self, _chunk: R::Chunk) -> impl Future<Output = Result<Demand, R::Error>> + Send {
async { Ok(Demand::More) }
}
}

2
litellm-rust/crates/callbacks/src/lib.rs → litellm-rust/crates/host/src/lib.rs
View file

@ -1,7 +1,7 @@
//! The contract between a native call and the host runtime that drives it.
//!
//! A host is whatever sits on the far side of the language boundary: CPython today,
//! another runtime later. Core implements [`machine::Machine`] per route and never learns
//! another runtime later. Core runs each route on a [`machine::RouteMachine`] and never learns
//! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers
//! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent.

5
litellm-rust/crates/core/src/machine/auth.rs → litellm-rust/crates/host/src/machine/auth.rs
View file

@ -1,9 +1,8 @@
use std::sync::Arc;
use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
use litellm_callbacks::route::Route;
use super::{HostChannel, MachineFault};
use crate::route::Route;
use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
/// A route whose host can mint credentials on the call's behalf.
pub trait TokenRoute: Route {

6
litellm-rust/crates/callbacks/src/machine.rs → litellm-rust/crates/host/src/machine/mod.rs
View file

@ -1,6 +1,12 @@
mod auth;
mod route_machine;
use std::future::Future;
use std::pin::Pin;
pub use auth::{HostTokenProvider, TokenRoute};
pub use route_machine::{ExecuteFuture, HostChannel, MachineFault, RouteMachine};
use crate::host::{HostOp, HostResult};
use crate::route::Route;

31
litellm-rust/crates/core/src/machine/mod.rs → litellm-rust/crates/host/src/machine/route_machine.rs
View file

@ -2,18 +2,16 @@
//! place, and turns the host operations that future requests into [`Machine`] steps. No
//! task is spawned; dropping the machine drops the in-flight call.
mod auth;
use std::{future::Future, pin::Pin};
pub use auth::{HostTokenProvider, TokenRoute};
use litellm_callbacks::{
event::{CallEvent, RequestContext, WireRequest},
host::{HostOp, HostResult},
machine::{HostFailure, Interrupted, Machine, MachineStep, Step},
use tokio::sync::{mpsc, oneshot};
use super::{HostFailure, Interrupted, Machine, MachineStep, Step};
use crate::{
event::{MachineEvent, RequestContext, WireRequest},
host::{Demand, HostOp, HostResult},
route::Route,
};
use tokio::sync::{mpsc, oneshot};
/// The machine's own failures, distinct from anything the provider call reports.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -82,12 +80,27 @@ where
}
}
pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> {
pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> {
match self.invoke(HostOp::Emit(event)).await? {
HostResult::Emitted => Ok(()),
_ => Err(MachineFault::Mismatch.into()),
}
}
pub async fn open(&self, head: R::StreamHead) -> Result<Demand, R::Error> {
self.demand(HostOp::Open(head)).await
}
pub async fn deliver(&self, chunk: R::Chunk) -> Result<Demand, R::Error> {
self.demand(HostOp::Deliver(chunk)).await
}
async fn demand(&self, op: HostOp<R>) -> Result<Demand, R::Error> {
match self.invoke(op).await? {
HostResult::Demand(demand) => Ok(demand),
_ => Err(MachineFault::Mismatch.into()),
}
}
}
enum Execution<R: Route> {

5
litellm-rust/crates/callbacks/src/route.rs → litellm-rust/crates/host/src/route.rs
View file

@ -6,4 +6,9 @@ pub trait Route: Send + Sync + 'static {
type Error: Clone + Send + Sync + 'static;
type Op: Send + 'static;
type OpResult: Send + 'static;
/// One piece of a streamed response, handed to the caller as it arrives. A route
/// that never streams uses `Infallible`.
type Chunk: Send + 'static;
/// What the route knows once a streamed response starts, before its first chunk.
type StreamHead: Send + 'static;
}

45
litellm-rust/crates/callbacks/src/run.rs → litellm-rust/crates/host/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 {
@ -24,7 +25,12 @@ where
.before_send(*wire, &context)
.await
.map(|wire| HostResult::BeforeSend(Box::new(wire))),
HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted),
HostOp::Emit(event) => host
.emit(&CallEvent::Machine(event))
.await
.map(|()| HostResult::Emitted),
HostOp::Open(head) => host.open(head).await.map(HostResult::Demand),
HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand),
};
match answer {
Ok(answer) => result = Some(answer),
@ -60,6 +66,8 @@ mod tests {
type Error = &'static str;
type Op = &'static str;
type OpResult = ();
type Chunk = std::convert::Infallible;
type StreamHead = std::convert::Infallible;
}
struct Scripted {
@ -102,6 +110,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 +133,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 +142,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 +152,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]);
}
}

2
litellm-rust/crates/llms/Cargo.toml
View file

@ -15,7 +15,7 @@ litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-callbacks.workspace = true
litellm-host.workspace = true
litellm-framing.workspace = true
base64.workspace = true
bytes.workspace = true

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,

11
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -102,6 +102,17 @@ pub enum Error {
Headers(#[from] crate::custom_httpx::http_handler::HeaderError),
}
impl From<litellm_host::machine::MachineFault> for Error {
fn from(fault: litellm_host::machine::MachineFault) -> Self {
use litellm_host::machine::MachineFault;
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "OCR host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("OCR {message}"),
MachineFault::Mismatch => "invalid OCR host operation result".into(),
})
}
}
impl From<litellm_core_utils::call_arguments::ArgumentError> for Error {
fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self {
Self::RequestField {

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_host::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

16
litellm-rust/crates/python-bridge/src/marshal.rs
View file

@ -18,10 +18,6 @@ pub(crate) struct RouteOptions {
pub(crate) timeout: Option<Duration>,
}
pub(crate) fn body_argument(value: &Bound<'_, PyAny>) -> PyResult<Map<String, Value>> {
required_object("body", from_py_argument(value)?)
}
pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult<Vec<Value>> {
match from_py_argument(value)? {
Value::Array(values) => Ok(values),
@ -192,18 +188,6 @@ mod tests {
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
);
let body = py
.eval(
c"{'model': 'claude', 'metadata': {'user': '1'}}",
None,
None,
)
.unwrap();
assert_eq!(
Value::Object(body_argument(&body).unwrap()),
json!({"model": "claude", "metadata": {"user": "1"}})
);
let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap();
assert_eq!(
optional_params_argument(&params).unwrap(),

88
litellm-rust/crates/python-bridge/src/routes/messages.rs
View file

@ -1,88 +0,0 @@
use litellm_core::messages::{Error, messages as run_messages, types::MessagesRequest};
use litellm_host_python::{run_async, run_sync};
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use pyo3::prelude::*;
use serde_json::{Map, Value};
use crate::{
errors::messages_error_to_pyerr,
marshal::{RouteOptions, body_argument, extra_headers_argument, optional_timeout},
};
async fn execute(
body: Map<String, Value>,
options: RouteOptions,
) -> Result<AnthropicMessagesResponse, Error> {
let RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout,
} = options;
run_messages(MessagesRequest {
model: &model,
body: Value::Object(body),
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
}
#[pyfunction]
#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[expect(
clippy::too_many_arguments,
reason = "one parameter per Python keyword"
)]
pub(crate) fn messages(
py: Python<'_>,
model: String,
#[pyo3(from_py_with = body_argument)] body: Map<String, Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let options = RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout: optional_timeout(timeout_seconds),
};
run_sync(py, execute(body, options), messages_error_to_pyerr)
}
#[pyfunction]
#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[expect(
clippy::too_many_arguments,
reason = "one parameter per Python keyword"
)]
pub(crate) fn amessages<'py>(
py: Python<'py>,
model: String,
#[pyo3(from_py_with = body_argument)] body: Map<String, Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option<Map<String, Value>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let options = RouteOptions {
model,
api_key,
api_base,
custom_llm_provider,
extra_headers,
timeout: optional_timeout(timeout_seconds),
};
run_async(py, execute(body, options), messages_error_to_pyerr)
}

186
litellm-rust/crates/python-bridge/src/routes/messages/host.rs Normal file
View file

@ -0,0 +1,186 @@
use bytes::Bytes;
use litellm_core::messages::{
Error,
route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput},
};
use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py};
use litellm_llms::custom_httpx::transport::Error as TransportError;
use pyo3::{
exceptions::{PyException, PyValueError},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyBytes, PyDict},
};
use serde_json::{Map, Value};
use crate::{
errors::{RustUpstreamError, messages_error_to_pyerr},
marshal::{optional_timeout, python_timeout_seconds},
};
/// The Anthropic Messages body fields a caller may pass besides `model` and `messages`,
/// as `AnthropicMessagesRequestOptionalParams` declares them.
const BODY_FIELDS: [&str; 20] = [
"max_tokens",
"metadata",
"stop_sequences",
"stream",
"system",
"temperature",
"thinking",
"tool_choice",
"tools",
"top_k",
"inference_geo",
"top_p",
"mcp_servers",
"context_management",
"container",
"output_format",
"speed",
"output_config",
"cache_control",
"reasoning_effort",
];
/// The Python side of the Messages route: projects the prepared arguments and builds the
/// public response, chunks and exceptions.
pub(super) struct MessagesRouteHost {
request: Py<PyAny>,
}
impl MessagesRouteHost {
pub(super) fn new(request: Py<PyAny>) -> Self {
Self { request }
}
fn project(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<MessagesCall> {
let request = self.request.bind(py);
let argument = |name: &str| -> PyResult<Option<Bound<'_, PyAny>>> {
Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none()))
};
let string = |name: &str| -> PyResult<Option<String>> {
argument(name)?.map(|value| value.extract()).transpose()
};
let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?;
let messages =
argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?;
let fields = BODY_FIELDS
.iter()
.filter_map(|name| match argument(name) {
Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))),
Ok(None) => None,
Err(error) => Some(Err(error)),
})
.collect::<PyResult<Vec<(String, Value)>>>()?;
let body = [
("model".to_string(), Value::String(model.clone())),
("messages".to_string(), from_py(&messages)?),
]
.into_iter()
.chain(fields)
.collect::<Map<String, Value>>();
let timeout = argument("timeout")?
.map(|value| python_timeout_seconds(py, value.unbind()))
.transpose()?
.flatten();
Ok(MessagesCall {
model,
body,
api_key: string("api_key")?,
api_base: string("api_base")?,
custom_llm_provider: string("custom_llm_provider")?,
extra_headers: argument("extra_headers")?
.map(|value| from_py(&value))
.transpose()?,
timeout: optional_timeout(timeout),
})
}
fn provider(&self, py: Python<'_>) -> String {
self.request
.bind(py)
.getattr("custom_llm_provider")
.and_then(|value| value.extract::<Option<String>>())
.ok()
.flatten()
.unwrap_or_else(|| "anthropic".into())
}
fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr {
if !error.is_instance_of::<PyException>(py) {
return error;
}
let mapped = py
.import("litellm.rust_bridge.messages.route_host")
.and_then(|module| module.getattr("map_failure"))
.and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py))))
.and_then(|mapped| {
mapped
.extract::<Py<pyo3::exceptions::PyBaseException>>()
.map_err(PyErr::from)
});
match mapped {
Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()),
Err(_) => error,
}
}
}
impl RouteHost for MessagesRouteHost {
type Route = Messages;
type Failure = PyErr;
fn invoke(
&mut self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
op: MessagesOp,
) -> Result<MessagesOpResult, InvokeError<Error>> {
match op {
MessagesOp::ProjectRequest => self
.project(py, arguments)
.map(|call| MessagesOpResult::Request(Box::new(call)))
.map_err(|error| InvokeError::Python(self.map_failure(py, error))),
}
}
fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult<Py<PyAny>> {
match response {
MessagesOutput::Message(message) => py
.import("litellm.rust_bridge.messages.route_host")?
.getattr("response")?
.call1((to_py(py, message.as_ref())?,))
.map(Bound::unbind),
MessagesOutput::Streamed => Ok(py.None()),
}
}
fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult<Py<PyAny>> {
Ok(PyBytes::new(py, &chunk).into_any().unbind())
}
fn classify(&self, py: Python<'_>, error: Error) -> PyResult<PyErr> {
let native = match error {
Error::Transport(TransportError::Http { status, body }) => {
let error = RustUpstreamError::new_err((status, body));
error
.value(py)
.setattr("headers", Vec::<(String, String)>::new())?;
error
}
other => messages_error_to_pyerr(other),
};
Ok(self.map_failure(py, native))
}
fn host_error(error: &PyErr) -> Error {
Error::InvalidRequest(error.to_string())
}
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.request)
}
}

68
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs Normal file
View file

@ -0,0 +1,68 @@
mod host;
use host::MessagesRouteHost;
use litellm_callbacks_legacy::{LegacySurface, PassThroughStream, PublicCall, run_legacy_call};
use litellm_core::messages::route::{messages_machine, supports};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use crate::errors::RustBridgeDeclined;
const SURFACE: LegacySurface = LegacySurface {
call_type: "anthropic_messages",
input_description: "Messages",
stream: Some(PassThroughStream {
url_route: "/v1/messages",
endpoint_type: "anthropic",
}),
};
fn run_messages(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let model: String = request.getattr("model")?.extract()?;
let provider: Option<String> = request.getattr("custom_llm_provider")?.extract()?;
let stream = request
.getattr("stream")?
.extract::<Option<bool>>()?
.unwrap_or(false);
if !supports(&model, provider.as_deref(), stream) {
return Err(RustBridgeDeclined::new_err(
"the Rust Messages route does not serve this provider",
));
}
run_legacy_call(
py,
SURFACE,
PublicCall::capture(&request, &args, &kwargs)?,
messages_machine(),
MessagesRouteHost::new(request.unbind()),
asynchronous,
)
}
#[pyfunction]
pub(crate) fn messages(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_messages(py, request, args, kwargs, false)
}
#[pyfunction]
pub(crate) fn amessages(
py: Python<'_>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
run_messages(py, request, args, kwargs, true)
}

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

@ -22,11 +22,6 @@ mod tests {
"atranscription",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"messages",
"amessages",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
(
"chat_completions",
"achat_completions",
@ -113,25 +108,6 @@ value = Broken()
);
assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string());
let invalid_body = PyList::empty(py);
let sync_messages_error = module
.getattr("messages")
.and_then(|function| function.call1(("model", &invalid_body)))
.expect_err("sync Messages should reject a non-dict body");
let async_messages_error = module
.getattr("amessages")
.and_then(|function| function.call1(("model", &invalid_body)))
.expect_err("async Messages should reject a non-dict body");
assert_eq!(
sync_messages_error.to_string(),
"ValueError: body must be a dict"
);
assert_eq!(
async_messages_error.to_string(),
sync_messages_error.to_string()
);
let invalid_headers = PyList::empty(py);
let kwargs = PyDict::new(py);
kwargs
@ -193,13 +169,6 @@ value = Broken()
headers_kwargs
.set_item("extra_headers", &invalid)
.expect("kwargs should accept extra_headers");
let invalid_body = PyList::empty(py);
let error = module
.getattr("messages")
.and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs)))
.expect_err("body should be validated before headers");
assert_eq!(error.to_string(), "ValueError: body must be a dict");
let invalid_payload =
PyModule::new(py, "invalid_payload").expect("invalid payload should be created");
let error = module

65
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::{InvokeError, 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, InvokeError<Error>> {
self.answer(py, arguments, op)
.map_err(|error| InvokeError::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,18 @@ impl RouteHost for OcrRouteHost {
.map(Bound::unbind)
}
fn native_error(error: Error) -> PyErr {
ocr_error_to_pyerr(error)
fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
match chunk {}
}
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;
}

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

@ -15,6 +15,7 @@ use pyo3::{
const SURFACE: LegacySurface = LegacySurface {
call_type: "ocr",
input_description: "OCR document processing",
stream: None,
};
const ASYNC_SURFACE: LegacySurface = LegacySurface {

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(

28
litellm/rust_bridge/_native.pyi
View file

@ -1,9 +1,11 @@
from asyncio import Future
from collections.abc import Coroutine, Mapping, Sequence
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from typing import Never, final
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
@ -39,23 +41,15 @@ def atranscription(
timeout_seconds: float | None = None,
) -> Future[dict[str, object]]: ...
def messages(
model: str,
body: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, object] | None = None,
timeout_seconds: float | None = None,
) -> dict[str, object]: ...
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
) -> AnthropicMessagesResponse | Iterator[bytes]: ...
def amessages(
model: str,
body: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, object] | None = None,
timeout_seconds: float | None = None,
) -> Future[dict[str, object]]: ...
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: dict[str, object],
) -> Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]]: ...
def chat_completions_decline(
model: str,
messages: Sequence[object],

1
litellm/rust_bridge/catalog.py
View file

@ -59,6 +59,7 @@ Rules: TypeAlias = tuple[Rule, ...]
RULES: Final[Rules] = (
Rule(Route.OCR, Rollout.RUST_OPT_OUT),
Rule(Route.MESSAGES, Rollout.RUST_OPT_IN),
Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})),
)

43
litellm/rust_bridge/failures.py
View file

@ -5,8 +5,37 @@ from __future__ import annotations
from collections.abc import Mapping
from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper
import httpx
import openai
from pydantic import TypeAdapter, ValidationError
import litellm
_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str])
_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]])
class UpstreamFailure(Exception):
def __init__(self, response: httpx.Response, cause: Exception) -> None:
super().__init__(str(cause))
self.message: Final = str(cause)
self.response: Final = response
self.status_code: Final = response.status_code
self.__cause__ = cause
def _upstream_failure(error: Exception, api_base: str | None) -> Exception:
try:
status, body = _UPSTREAM_ARGS.validate_python(error.args)
headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None))
except ValidationError:
return error
http_request: Final = httpx.Request("POST", api_base or "https://docs.litellm.ai/docs")
return UpstreamFailure(
httpx.Response(status, content=body.encode(), headers=headers, request=http_request),
error,
)
class ExceptionMapper(Protocol):
def __call__(
@ -35,3 +64,17 @@ def map_failure(error: Exception, model: str, request_provider: str, kwargs: Map
except Exception as public_error:
public_error.__context__ = error
return public_error
def map_native_failure(
error: Exception, model: str, request_provider: str, kwargs: Mapping[str, object], api_base: str | None = None
) -> Exception:
"""`map_failure`, reading a native `(status, body)` provider failure as the HTTP response it was."""
original: Final = _upstream_failure(error, api_base)
public_error: Final = map_failure(original, model, request_provider, kwargs)
if isinstance(original, UpstreamFailure) and public_error.__context__ is original:
public_error.__context__ = error
if isinstance(public_error, openai.APIStatusError):
public_error.response = original.response
public_error.status_code = original.status_code
return public_error

389
litellm/rust_bridge/legacy_callbacks.py
View file

@ -6,24 +6,23 @@ registries it fans out to. It expires with that contract.
from __future__ import annotations
import asyncio
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 +41,6 @@ class MetadataUpdater(Protocol):
class CallSetup:
logger: Logging
kwargs: dict[str, object]
bridge_owned: bool
def setup(
@ -61,19 +59,23 @@ 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:
import litellm
from litellm import (
BudgetExceededError,
_current_cost, # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor
max_budget,
num_retries_per_request,
)
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor
if litellm.max_budget and current_cost > litellm.max_budget:
raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget)
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
if max_budget and _current_cost > max_budget:
raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget)
if max_retries_per_request_hit(kwargs, num_retries_per_request):
raise RuntimeError("Max retries per request hit!")
@ -93,87 +95,304 @@ 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 StreamingLogBuilder(Protocol):
def __call__(
self,
*,
litellm_logging_obj: Logging,
passthrough_success_handler_obj: object,
url_route: str,
request_body: dict[str, object],
endpoint_type: object,
start_time: datetime.datetime,
raw_bytes: list[bytes],
end_time: datetime.datetime,
) -> Coroutine[object, object, 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
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 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]:
from litellm import credential_list as credentials
return credentials
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)
def stream_opened(logger: Logging) -> None:
logger.stream = True
logger.model_call_details["stream"] = True
def stream_success(
logger: Logging,
url_route: str,
endpoint_type: str,
request_body: dict[str, object],
chunks: list[bytes],
start: datetime.datetime,
end: datetime.datetime,
first_chunk: datetime.datetime | None,
) -> 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)
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
)
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
if first_chunk is not None:
logger.completion_start_time = first_chunk
logger.model_call_details["completion_start_time"] = first_chunk
build: Final = cast( # cast-ok: bounded adapter for the untyped pass-through logging builder
StreamingLogBuilder,
PassThroughStreamingHandler._route_streaming_logging_to_handler, # pyright: ignore[reportPrivateUsage] # the Messages stream iterator bills through the same builder
)
coroutine: Final = build(
litellm_logging_obj=logger,
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
url_route=url_route,
request_body=request_body,
endpoint_type=EndpointType(endpoint_type),
start_time=start,
raw_bytes=chunks,
end_time=end,
)
if getattr(logger, "_on_deferred_stream_complete", None) is not None:
logger._deferred_stream_complete_args = (coroutine,) # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release reads this slot
return
try:
asyncio.get_running_loop()
except RuntimeError:
from litellm.litellm_core_utils.litellm_logging import executor
executor.submit(contextvars.copy_context().run, asyncio.run, coroutine)
return
enqueue_logging(coroutine)
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 stream_failure(
logger: Logging,
endpoint_type: str,
request_body: dict[str, object],
chunks: list[bytes],
error: Exception,
) -> Coroutine[object, object, None]:
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
return PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=logger,
endpoint_type=EndpointType(endpoint_type),
request_body=request_body,
raw_bytes=chunks,
exception=error,
)

137
litellm/rust_bridge/lifecycle.py
View file

@ -1,8 +1,8 @@
from __future__ import annotations
from collections.abc import Awaitable
from collections.abc import AsyncIterator, Awaitable, Iterator
from dataclasses import dataclass
from typing import Protocol
from typing import Final, Protocol
@dataclass(frozen=True, slots=True)
@ -15,28 +15,133 @@ class Complete:
value: object
@dataclass(frozen=True, slots=True)
class Open:
value: None
@dataclass(frozen=True, slots=True)
class Yield:
value: object
Settled = Complete | Open | Yield
Step = Await | Settled
class Execution(Protocol):
def start(self) -> Await | Complete: ...
def start(self) -> Step: ...
def resume_value(self, value: object) -> Await | Complete: ...
def resume_value(self, value: object) -> Step: ...
def resume_error(self, error: BaseException) -> Await | Complete: ...
def resume_error(self, error: BaseException) -> Step: ...
def close(self) -> None: ...
class StreamClosed(Exception):
"""Tells a streaming execution that its caller stopped reading."""
async def _settle(execution: Execution, step: Step) -> Settled:
while isinstance(step, Await):
try:
value = await step.awaitable # rebind-ok: each selected await produces the next protocol input
except GeneratorExit:
raise
except BaseException as error:
step = execution.resume_error(error) # rebind-ok: advance the execution protocol
else:
step = execution.resume_value(value) # rebind-ok: advance the execution protocol
return step
def _settled(step: Step) -> Settled:
if isinstance(step, Await):
raise RuntimeError("sync call suspended")
return step
async def drive(execution: Execution) -> object:
handed_off = False # rebind-ok: set once the execution belongs to the returned stream
try:
step = execution.start() # rebind-ok: the execution protocol advances after each selected await
while isinstance(step, Await):
try:
value = await step.awaitable # rebind-ok: each selected await produces the next protocol input
except GeneratorExit:
raise
except BaseException as error:
step = execution.resume_error(error) # rebind-ok: advance the execution protocol
else:
step = execution.resume_value(value) # rebind-ok: advance the execution protocol
step: Final = await _settle(execution, execution.start())
if isinstance(step, Open):
handed_off = True
return Stream(execution)
return step.value
finally:
execution.close()
if not handed_off:
execution.close()
class Stream(AsyncIterator[object]):
"""A streamed native call: each read resumes the execution until its next chunk."""
def __init__(self, execution: Execution) -> None:
self._execution: Final = execution
self._done = False
def __aiter__(self) -> Stream:
return self
async def __anext__(self) -> object:
if self._done:
raise StopAsyncIteration
try:
step: Final = await _settle(self._execution, self._execution.resume_value(None))
except BaseException:
self._finish()
raise
if isinstance(step, Yield):
return step.value
self._finish()
raise StopAsyncIteration
async def aclose(self) -> None:
if self._done:
return
try:
await _settle(self._execution, self._execution.resume_error(StreamClosed()))
finally:
self._finish()
def _finish(self) -> None:
self._done = True
self._execution.close()
class SyncStream(Iterator[object]):
"""The sync form of `Stream`; its execution never suspends on an awaitable."""
def __init__(self, execution: Execution) -> None:
self._execution: Final = execution
self._done = False
def __iter__(self) -> SyncStream:
return self
def __next__(self) -> object:
if self._done:
raise StopIteration
try:
step: Final = _settled(self._execution.resume_value(None))
except BaseException:
self._finish()
raise
if isinstance(step, Yield):
return step.value
self._finish()
raise StopIteration
def close(self) -> None:
if self._done:
return
try:
_settled(self._execution.resume_error(StreamClosed()))
finally:
self._finish()
def _finish(self) -> None:
self._done = True
self._execution.close()

10
litellm/rust_bridge/messages/entrypoints.py
View file

@ -1,6 +1,6 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
@ -26,7 +26,7 @@ class NativeMessages(Protocol):
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse: ...
) -> AnthropicMessagesResponse | Iterator[bytes]: ...
class NativeAmessages(Protocol):
@ -35,7 +35,7 @@ class NativeAmessages(Protocol):
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[AnthropicMessagesResponse]: ...
) -> Awaitable[AnthropicMessagesResponse | AsyncIterator[bytes]]: ...
def _messages_binding(value: object) -> NativeMessages | None:
@ -50,5 +50,5 @@ def _amessages_binding(value: object) -> NativeAmessages | None:
return cast("NativeAmessages", value) # cast-ok: callable validated at the native binding boundary
NATIVE_MESSAGES: Final = NativeBinding("anthropic_messages_handler", validate=_messages_binding)
NATIVE_AMESSAGES: Final = NativeBinding("anthropic_messages", validate=_amessages_binding)
NATIVE_MESSAGES: Final = NativeBinding("messages", validate=_messages_binding)
NATIVE_AMESSAGES: Final = NativeBinding("amessages", validate=_amessages_binding)

2
litellm/rust_bridge/messages/route_host.py
View file

@ -20,4 +20,4 @@ def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]:
def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))
return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base)

40
litellm/rust_bridge/ocr/route_host.py
View file

@ -4,40 +4,17 @@ from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import httpx
import openai
from pydantic import TypeAdapter, ValidationError
from pydantic import TypeAdapter
import litellm
from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse
from litellm.rust_bridge import failures
from litellm.rust_bridge.failures import UpstreamFailure
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
__all__ = ("UpstreamFailure", "arguments", "map_failure", "response")
_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object])
_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str])
_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]])
class UpstreamFailure(Exception):
def __init__(self, response: httpx.Response, cause: Exception) -> None:
super().__init__(str(cause))
self.message: Final = str(cause)
self.response: Final = response
self.status_code: Final = response.status_code
self.__cause__ = cause
def _upstream_failure(error: Exception, request: LiteLLMOcrRequest) -> Exception:
try:
status, body = _UPSTREAM_ARGS.validate_python(error.args)
headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None))
except ValidationError:
return error
http_request: Final = httpx.Request("POST", request.api_base or "https://docs.litellm.ai/docs")
return UpstreamFailure(
httpx.Response(status, content=body.encode(), headers=headers, request=http_request),
error,
)
def response(value: Mapping[str, object]) -> OCRResponse:
@ -61,11 +38,4 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider:
model=request.model.removeprefix(f"{request_provider}/"),
llm_provider=request_provider,
)
original: Final = _upstream_failure(error, request)
public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request))
if isinstance(original, UpstreamFailure) and public_error.__context__ is original:
public_error.__context__ = error
if isinstance(public_error, openai.APIStatusError):
public_error.response = original.response
public_error.status_code = original.status_code
return public_error
return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base)

31
tests/test_litellm/rust_bridge/native_route_wheel_test.py
View file

@ -73,7 +73,7 @@ def assert_native_request(
headers: HTTPMessage,
body: object,
) -> None:
if route not in {"transcription", "messages", "chat_completions"}:
if route not in {"transcription", "chat_completions"}:
raise AssertionError(f"unexpected route marker: {route!r}")
if outcome not in {"success", "429", "hang"}:
raise AssertionError(f"unexpected outcome marker: {outcome!r}")
@ -89,10 +89,6 @@ def assert_native_request(
assert path == "/v1/messages"
assert headers.get("x-api-key") == "sk-native"
assert body["model"] == "claude-sonnet-4-5"
if route == "messages":
assert body["max_tokens"] == 16
assert body["messages"][0]["content"] == "hello-from-messages"
return
assert body["max_tokens"] == 17
assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}]
@ -132,17 +128,6 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
"language": "en",
},
}
if route == "messages":
return common | {
"model": "claude-sonnet-4-5",
"body": {
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello-from-messages"}],
},
"api_key": "sk-native",
"custom_llm_provider": "anthropic",
}
if route == "chat_completions":
return common | {
"model": "anthropic/claude-sonnet-4-5",
@ -165,8 +150,6 @@ def assert_success(route: str, response: object) -> None:
def success_value(route: str, response: dict[object, object]) -> object:
if route == "transcription":
return response["text"]
if route == "messages":
return response["content"][0]["text"]
return response["choices"][0]["message"]["content"]
@ -181,7 +164,7 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
def exercise_sync(native: object, api_base: str) -> None:
for route in ("transcription", "messages", "chat_completions"):
for route in ("transcription", "chat_completions"):
function: Final = getattr(native, route)
assert_success(route, function(**route_kwargs(route, api_base, "success")))
try:
@ -193,7 +176,7 @@ def exercise_sync(native: object, api_base: str) -> None:
async def exercise_async(native: object, api_base: str) -> None:
for route in ("transcription", "messages", "chat_completions"):
for route in ("transcription", "chat_completions"):
function: Final = getattr(native, f"a{route}")
assert_success(route, await function(**route_kwargs(route, api_base, "success")))
try:
@ -206,11 +189,11 @@ async def exercise_async(native: object, api_base: str) -> None:
async def exercise_async_concurrency(native: object, api_base: str) -> None:
responses: Final = await asyncio.wait_for(
asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))),
asyncio.gather(*(native.achat_completions(**route_kwargs("chat_completions", api_base, "success")) for _ in range(32))),
timeout=15,
)
for response in responses:
assert_success("messages", response)
assert_success("chat_completions", response)
def exercise_routes(native_path: Path, api_base: str) -> object:
@ -223,8 +206,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object:
def exercise_signal(native: object, api_base: str) -> int:
try:
native.messages(
**route_kwargs("messages", api_base, "hang"),
native.chat_completions(
**route_kwargs("chat_completions", api_base, "hang"),
)
except KeyboardInterrupt:
sys.stdout.write("KeyboardInterrupt\n")

4
tests/test_litellm/rust_bridge/test_bindings.py
View file

@ -43,8 +43,8 @@ def test_binding_validates_native_attribute(
ROUTE_BINDINGS: Final = (
("completion", chat_completions.NATIVE_COMPLETION),
("acompletion", chat_completions.NATIVE_ACOMPLETION),
("anthropic_messages_handler", messages.NATIVE_MESSAGES),
("anthropic_messages", messages.NATIVE_AMESSAGES),
("messages", messages.NATIVE_MESSAGES),
("amessages", messages.NATIVE_AMESSAGES),
("responses", responses.NATIVE_RESPONSES),
("aresponses", responses.NATIVE_ARESPONSES),
("ocr", ocr.NATIVE_OCR),

4
tests/test_litellm/rust_bridge/test_catalog.py
View file

@ -40,6 +40,10 @@ def test_shipped_decisions(
enabled: Final = environment == "1" if environment is not None else process is not False
assert catalog.rollout(context) is Rollout.RUST_OPT_OUT
assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON)
elif route is Route.MESSAGES:
enabled: Final = environment == "1" if environment is not None else process is True
assert catalog.rollout(context) is Rollout.RUST_OPT_IN
assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON)
elif route is Route.TRANSCRIPTION and provider == "bedrock":
assert catalog.rollout(context) is Rollout.RUST_REQUIRED
assert catalog.decision(context) is Decision.RUST_REQUIRED

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}

1
tests/test_litellm/rust_bridge/test_runtime.py
View file

@ -157,7 +157,6 @@ def test_context_outside_rule_stays_on_python() -> None:
(
Context(Route.CHAT_COMPLETIONS, provider="anthropic"),
Context(Route.CHAT_COMPLETIONS, provider="bedrock"),
Context(Route.MESSAGES, provider="anthropic"),
Context(Route.RESPONSES, provider="openai"),
Context(Route.TRANSCRIPTION, provider="openai"),
),

63
tests/test_litellm_rust/conftest.py
View file

@ -1,10 +1,9 @@
import asyncio
import os
from collections.abc import AsyncIterator, Generator, Iterator
from collections.abc import AsyncIterator, Generator
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack, contextmanager
from types import ModuleType
from typing import Final, cast
from contextlib import ExitStack
from typing import Final
import pytest
import pytest_asyncio
@ -18,62 +17,20 @@ from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivate
_parse_env_bool,
)
from tests.test_litellm_rust.support.callback_recorder import drain_logging
from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound
from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
CALLBACK_ATTRIBUTES: Final = (
"callbacks",
"input_callback",
"success_callback",
"failure_callback",
"_async_input_callback",
"_async_success_callback",
"_async_failure_callback",
)
def _list_attribute(container: ModuleType, attribute: str) -> list[object]:
value: Final = getattr(container, attribute)
if not isinstance(value, list):
raise AssertionError(f"{container.__name__}.{attribute} is not a list")
return cast(list[object], value)
@contextmanager
def _isolated_list(container: ModuleType, attribute: str) -> Iterator[None]:
source: Final = _list_attribute(container, attribute)
original: Final = list(source)
source.clear() # mutable-ok: test isolation mutates global registries by design
try:
yield
finally:
source.clear()
source.extend(original)
setattr(container, attribute, source)
@contextmanager
def _rebound(container: object, attribute: str, value: object) -> Iterator[None]:
original: Final[object] = getattr(container, attribute)
setattr(container, attribute, value)
try:
yield
finally:
setattr(container, attribute, original)
@pytest_asyncio.fixture(autouse=True, loop_scope="function")
async def isolate_ocr_test_state() -> AsyncIterator[None]:
with ExitStack() as stack:
for attribute in CALLBACK_ATTRIBUTES:
stack.enter_context(_isolated_list(litellm, attribute))
stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor
stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry
stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache
stack.enter_context(_rebound(_CONFIGURATION, "override", None))
stack.enter_context(isolated_callback_registries())
stack.enter_context(rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache
stack.enter_context(rebound(_CONFIGURATION, "override", None))
executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging")
stack.enter_context(_rebound(litellm_logging, "executor", executor))
stack.enter_context(_rebound(utils, "executor", executor))
stack.enter_context(_rebound(thread_pool_executor, "executor", executor))
stack.enter_context(rebound(litellm_logging, "executor", executor))
stack.enter_context(rebound(utils, "executor", executor))
stack.enter_context(rebound(thread_pool_executor, "executor", executor))
try:
yield
finally:

0
tests/test_litellm_rust/messages/__init__.py Normal file
View file

175
tests/test_litellm_rust/messages/test_callbacks.py Normal file
View file

@ -0,0 +1,175 @@
from collections.abc import AsyncIterator, Iterator
from typing import Final
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.requests import (
MESSAGES,
MESSAGES_EVENTS,
MESSAGES_MODEL,
MESSAGES_RESPONSE,
request_body,
)
pytestmark = pytest.mark.requires_rust_extension
STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS)
@pytest.fixture
def messages_server(recording_server: RecordingServer) -> RecordingServer:
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
return recording_server
def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
return {
"model": MESSAGES_MODEL,
"messages": [dict(message) for message in MESSAGES],
"max_tokens": 64,
"api_key": "test-key",
"api_base": server.base_url,
**kwargs,
}
def assert_served_natively(server: RecordingServer) -> None:
assert len(server.requests) == 1
assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx")
@pytest.mark.asyncio
async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response(
messages_server: RecordingServer,
) -> None:
recorder: Final = RecordingLogger()
response: Final = await litellm.anthropic.messages.acreate(
**arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success")
)
assert_served_natively(messages_server)
assert response["content"] == MESSAGES_RESPONSE["content"]
sent: Final = messages_server.requests[0]
assert sent.path == "/v1/messages"
assert sent.body == {"model": "claude-sonnet-5", "messages": list(MESSAGES), "max_tokens": 64, "stream": False}
pre_call: Final = recorder.wait_for("log_pre_api_call")
assert request_body(pre_call[0].kwargs) == sent.body
success: Final = await recorder.wait_for_async("async_log_success_event")
assert len(success) == 1
assert success[0].call_type == "anthropic_messages"
assert success[0].kwargs["litellm_call_id"] == "messages-success"
assert success[0].response.choices[0].message.content == "Hello from native Messages"
@pytest.mark.asyncio
async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None:
class Edit(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
request_body(kwargs)["temperature"] = 0.25
await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()]))
assert messages_server.requests[0].body["temperature"] == 0.25
@pytest.mark.asyncio
async def test_native_messages_provider_error_reaches_caller_and_failure_callbacks_as_one_public_error(
messages_server: RecordingServer,
) -> None:
messages_server.enqueue(
ResponseSpec(body={"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}, status=400)
)
observed: Final = []
class Observe(CustomLogger):
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
observed.append(("sync", kwargs["exception"]))
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
observed.append(("async", kwargs["exception"]))
with pytest.raises(litellm.BadRequestError) as raised:
await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()]))
assert_served_natively(messages_server)
assert [phase for phase, _ in observed] == ["sync", "async"]
assert all(error is raised.value for _, error in observed)
def sse_payload() -> bytes:
return b"".join(STREAM.payloads())
@pytest.mark.asyncio
async def test_native_messages_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
messages_server: RecordingServer,
) -> None:
messages_server.enqueue(STREAM)
recorder: Final = RecordingLogger()
stream: Final = await litellm.anthropic.messages.acreate(
**arguments(messages_server, stream=True, callbacks=[recorder])
)
assert isinstance(stream, AsyncIterator)
first: Final = await anext(stream)
await drain_logging()
assert "async_log_success_event" not in recorder.names
rest: Final = [chunk async for chunk in stream]
assert first + b"".join(rest) == sse_payload()
assert_served_natively(messages_server)
assert messages_server.requests[0].body["stream"] is True
success: Final = await recorder.wait_for_async("async_log_success_event")
assert len(success) == 1
assert success[0].kwargs["stream"] is True
assert success[0].kwargs["completion_start_time"] is not None
assert "log_failure_event" not in recorder.names
@pytest.mark.asyncio
async def test_native_messages_stream_closed_early_logs_success_once_for_what_was_delivered(
messages_server: RecordingServer,
) -> None:
messages_server.enqueue(STREAM)
recorder: Final = RecordingLogger()
stream: Final = await litellm.anthropic.messages.acreate(
**arguments(messages_server, stream=True, callbacks=[recorder])
)
assert isinstance(stream, AsyncIterator)
await anext(stream)
await stream.aclose()
success: Final = await recorder.wait_for_async("async_log_success_event")
assert len(success) == 1
with pytest.raises(StopAsyncIteration):
await anext(stream)
def test_native_sync_messages_stream_relays_provider_events_and_logs_success_once(
messages_server: RecordingServer,
) -> None:
messages_server.enqueue(STREAM)
recorder: Final = RecordingLogger()
stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder]))
assert isinstance(stream, Iterator)
assert b"".join(stream) == sse_payload()
assert_served_natively(messages_server)
assert len(recorder.wait_for("async_log_success_event")) == 1
def test_native_sync_messages_returns_the_provider_message(messages_server: RecordingServer) -> None:
recorder: Final = RecordingLogger()
response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder]))
assert_served_natively(messages_server)
assert response["content"] == MESSAGES_RESPONSE["content"]
assert len(recorder.wait_for("log_success_event")) == 1

162
tests/test_litellm_rust/ocr/test_callbacks.py
View file

@ -1,24 +1,31 @@
import asyncio
import copy
import gc
import queue
import threading
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import pytest
from hypothesis import HealthCheck, given, settings
from hypothesis import strategies as st
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.isolation import isolated_callback_registries
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.requests import (
OCR_DOCUMENT,
OCR_RESPONSE,
call_native,
call_native_aocr,
call_native_ocr,
request_body,
request_headers,
)
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
pytestmark = pytest.mark.requires_rust_extension
@ -123,9 +130,7 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_
"callbacks": [Retain(), Edit()],
}
response: Final = (
await call_native_aocr(ocr_server, **arguments)
if asynchronous
else call_native_ocr(ocr_server, **arguments)
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
)
assert aliases == [True]
@ -291,6 +296,153 @@ def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registere
assert "log_failure_event" not in recorder.names
JSON_SCALARS: Final = (
st.none()
| st.booleans()
| st.integers(min_value=-(2**63), max_value=2**63 - 1)
| st.floats(allow_nan=False, allow_infinity=False)
| st.text(max_size=8)
)
JSON_VALUES: Final = st.recursive(
JSON_SCALARS,
lambda children: st.lists(children, max_size=3) | st.dictionaries(st.text(max_size=6), children, max_size=3),
max_leaves=8,
)
class ApplyEdits(CustomLogger):
def __init__(self, edits: Mapping[str, object]) -> None:
super().__init__()
self.edits: Final = edits
def log_pre_api_call(self, model, messages, kwargs):
request_body(kwargs).update(copy.deepcopy(dict(self.edits)))
@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
@given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3))
def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it(
ocr_server: RecordingServer, edits: dict[str, object]
) -> None:
ocr_server.expected_requests = None
with isolated_callback_registries():
call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))])
assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits}
@pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"])
def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None:
recorder: Final = RecordingLogger()
call_native_ocr_with_callbacks(ocr_server, [recorder])
[event] = recorder.wait_for(hook)
assert event.loop is None
assert (event.thread is threading.current_thread()) == (hook == "log_pre_api_call")
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact(
ocr_server: RecordingServer, asynchronous: bool
) -> None:
retained: Final = []
class Retain(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
retained.append((kwargs, request_body(kwargs), request_headers(kwargs)))
await call_native(ocr_server, asynchronous, callbacks=[Retain()])
await drain_logging()
gc.collect()
[(details, body, headers)] = retained
assert body == ocr_server.requests[0].body
assert headers
assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items())
assert details["additional_args"]["complete_input_dict"] is body
assert details["additional_args"]["headers"] is headers
@pytest.mark.asyncio
@pytest.mark.parametrize("family", ["sync", "async"])
async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None:
queued: Final = []
finished: Final = threading.Event()
def queue_payload(kwargs: dict[str, object]) -> None:
queued.append(kwargs["standard_logging_object"])
def strip_payload(kwargs: dict[str, object]) -> None:
payload: Final = kwargs["standard_logging_object"]
assert isinstance(payload, dict)
payload["stripped-by-a-later-callback"] = True
finished.set()
class QueuePayload(CustomLogger):
if family == "sync":
def log_success_event(self, kwargs, response_obj, start_time, end_time):
queue_payload(kwargs)
else:
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
queue_payload(kwargs)
class StripPayload(CustomLogger):
if family == "sync":
def log_success_event(self, kwargs, response_obj, start_time, end_time):
strip_payload(kwargs)
else:
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
strip_payload(kwargs)
await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()])
await drain_logging()
assert await asyncio.to_thread(finished.wait, 10)
assert [payload["stripped-by-a-later-callback"] for payload in queued] == [True]
@pytest.mark.asyncio
async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks(
ocr_server: RecordingServer,
) -> None:
token: Final = object()
observed: Final = []
class Blocked(Exception):
pass
class Block(CustomLogger):
async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token
raise Blocked("blocked after the provider answered")
def log_success_event(self, kwargs, response_obj, start_time, end_time):
observed.append(("success", None, None))
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
observed.append(("sync", kwargs.get("blocked-by"), kwargs["exception"]))
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
observed.append(("async", kwargs.get("blocked-by"), kwargs["exception"]))
litellm.callbacks.append(Block())
with pytest.raises(Blocked) as raised:
await call_native_aocr(ocr_server)
await drain_logging()
assert observed == [("sync", token, raised.value), ("async", token, raised.value)]
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context(

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

58
tests/test_litellm_rust/support/isolation.py Normal file
View file

@ -0,0 +1,58 @@
from collections.abc import Generator
from contextlib import ExitStack, contextmanager
from types import ModuleType
from typing import Final, cast
import litellm
from litellm import utils
from litellm.litellm_core_utils import litellm_logging
CALLBACK_ATTRIBUTES: Final = (
"callbacks",
"input_callback",
"success_callback",
"failure_callback",
"_async_input_callback",
"_async_success_callback",
"_async_failure_callback",
)
def _list_attribute(container: ModuleType, attribute: str) -> list[object]:
value: Final = getattr(container, attribute)
if not isinstance(value, list):
raise AssertionError(f"{container.__name__}.{attribute} is not a list")
return cast(list[object], value)
@contextmanager
def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]:
source: Final = _list_attribute(container, attribute)
original: Final = list(source)
source.clear() # mutable-ok: test isolation mutates global registries by design
try:
yield
finally:
source.clear()
source.extend(original)
setattr(container, attribute, source)
@contextmanager
def rebound(container: object, attribute: str, value: object) -> Generator[None]:
original: Final[object] = getattr(container, attribute)
setattr(container, attribute, value)
try:
yield
finally:
setattr(container, attribute, original)
@contextmanager
def isolated_callback_registries() -> Generator[None]:
with ExitStack() as stack:
for attribute in CALLBACK_ATTRIBUTES:
stack.enter_context(_isolated_list(litellm, attribute))
stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor
stack.enter_context(rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry
yield

16
tests/test_litellm_rust/support/recording_server.py
View file

@ -25,6 +25,12 @@ class ResponseSpec:
status: int = 200
headers: dict[str, str] = field(default_factory=dict)
delay: float = 0
events: tuple[tuple[str, object], ...] = ()
def payloads(self) -> tuple[bytes, ...]:
if not self.events:
return (json.dumps(self.body).encode(),)
return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events)
@dataclass
@ -73,15 +79,17 @@ def recording_service() -> Iterator[RecordingServer]:
response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response)
if response.delay:
time.sleep(response.delay)
payload: Final = json.dumps(response.body).encode()
payloads: Final = response.payloads()
self.send_response(response.status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.send_header("Content-Type", "text/event-stream" if response.events else "application/json")
self.send_header("Content-Length", str(sum(len(payload) for payload in payloads)))
for name, value in response.headers.items():
self.send_header(name, value)
self.end_headers()
try:
self.wfile.write(payload)
for payload in payloads:
self.wfile.write(payload)
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass

35
tests/test_litellm_rust/support/requests.py
View file

@ -12,6 +12,41 @@ OCR_RESPONSE: Final = {
"usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
}
MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5"
MESSAGES: Final = ({"role": "user", "content": "Hello"},)
MESSAGES_RESPONSE: Final = {
"id": "msg_native",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [{"type": "text", "text": "Hello from native Messages"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 5, "output_tokens": 4},
}
MESSAGES_EVENTS: Final = (
("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Hello from native Messages"},
},
),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 4},
},
),
("message_stop", {"type": "message_stop"}),
)
def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
return {