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

wip

Browse source
This commit is contained in:
Yujong Lee 2026-09-08 08:10:45 -07:00
parent a459ac932f
commit 42b8cd80be
15 changed files with 1016 additions and 541 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

236
litellm-rust/crates/python-bridge/src/driver.rs
View file

@ -1,83 +1,197 @@
use std::ffi::CString;
use litellm_core::lifecycle::program::Operation;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use crate::errors::RustBridgeDriverError;
pub(crate) const ADDITIONAL_ARGS: &str = "additional_args";
pub(crate) const API_BASE: &str = "api_base";
pub(crate) const API_KEY: &str = "api_key";
pub(crate) const COMPLETE_INPUT_DICT: &str = "complete_input_dict";
pub(crate) const HEADERS: &str = "headers";
pub(crate) const INPUT: &str = "input";
const DRIVE: &str = r#"
def drive_sync(arguments):
host = Host(arguments, False)
while host.machine.complete() is None:
try:
_invoke(host.machine, host)
except Exception as error:
host.advance(1, error)
except BaseException as error:
host.advance(2, error)
else:
host.advance(0)
return host.result()
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct OperationBinding {
method: &'static str,
awaiting: bool,
}
async def drive_async(arguments):
host = Host(arguments, True)
while host.machine.complete() is None:
try:
awaiting, value = _invoke(host.machine, host)
if awaiting:
await value
except Exception as error:
host.advance(1, error)
except BaseException as error:
host.advance(2, error)
else:
host.advance(0)
return host.result()
"#;
pub(crate) fn compile<'py>(
py: Python<'py>,
fn operation_binding(
operation: Operation,
asynchronous: bool,
supports_pre_call: bool,
route: &str,
host: &str,
) -> PyResult<Bound<'py, PyModule>> {
let source = CString::new(format!("{host}\n{DRIVE}"))
.map_err(|_| RustBridgeDriverError::new_err("driver source contains a null byte"))?;
let filename = CString::new(format!("{route}_driver.py"))
.map_err(|_| RustBridgeDriverError::new_err("driver route name contains a null byte"))?;
let module_name = CString::new(format!("_{route}_driver"))
.map_err(|_| RustBridgeDriverError::new_err("driver route name contains a null byte"))?;
PyModule::from_code(py, &source, &filename, &module_name)
) -> PyResult<OperationBinding> {
let binding = match operation {
Operation::Setup => OperationBinding {
method: "setup",
awaiting: false,
},
Operation::DeploymentPre => OperationBinding {
method: "deployment_pre",
awaiting: true,
},
Operation::Prepare => OperationBinding {
method: "prepare",
awaiting: false,
},
Operation::PreCall if supports_pre_call => OperationBinding {
method: "pre_call",
awaiting: false,
},
Operation::PreCall => {
return Err(PyRuntimeError::new_err(format!(
"{route} lifecycle selected an unsupported pre-call operation"
)));
}
Operation::Send if asynchronous => OperationBinding {
method: "send",
awaiting: true,
},
Operation::Send => OperationBinding {
method: "send_sync",
awaiting: false,
},
Operation::DeploymentSuccess => OperationBinding {
method: "deployment_success",
awaiting: true,
},
Operation::DeploymentFailure => OperationBinding {
method: "deployment_failure",
awaiting: true,
},
Operation::SyncSuccess => OperationBinding {
method: "sync_success",
awaiting: false,
},
Operation::AsyncSuccess => OperationBinding {
method: "async_success",
awaiting: false,
},
Operation::SyncSuccessIfNeeded => OperationBinding {
method: "sync_success_if_needed",
awaiting: false,
},
Operation::SyncFailure => OperationBinding {
method: "sync_failure",
awaiting: false,
},
Operation::AsyncFailure => OperationBinding {
method: "async_failure",
awaiting: true,
},
Operation::Restore => OperationBinding {
method: "restore",
awaiting: false,
},
Operation::Complete(_) => {
return Err(PyRuntimeError::new_err(format!(
"{route} lifecycle is complete"
)));
}
};
Ok(binding)
}
pub(crate) fn invoke(
py: Python<'_>,
operation: Operation,
asynchronous: bool,
supports_pre_call: bool,
route: &str,
host: Py<PyAny>,
) -> PyResult<(bool, Py<PyAny>)> {
let binding = operation_binding(operation, asynchronous, supports_pre_call, route)?;
Ok((
binding.awaiting,
host.getattr(py, binding.method)?.call0(py)?,
))
}
#[cfg(test)]
mod tests {
use litellm_core::lifecycle::Outcome;
use super::*;
#[test]
fn null_byte_in_route_raises_driver_error() {
fn operation_bindings_cover_the_lifecycle_contract() {
Python::initialize();
Python::attach(|py| {
let error = compile(py, "invalid\0route", "class Host: pass")
.expect_err("route names containing null bytes should fail");
assert!(error.is_instance_of::<RustBridgeDriverError>(py));
assert_eq!(
error.to_string(),
"RustBridgeDriverError: driver route name contains a null byte"
);
Python::attach(|_| {
let cases = [
(Operation::Setup, false, false, "setup", false),
(
Operation::DeploymentPre,
false,
false,
"deployment_pre",
true,
),
(Operation::Prepare, false, false, "prepare", false),
(Operation::PreCall, false, true, "pre_call", false),
(Operation::Send, false, false, "send_sync", false),
(Operation::Send, true, false, "send", true),
(
Operation::DeploymentSuccess,
false,
false,
"deployment_success",
true,
),
(
Operation::DeploymentFailure,
false,
false,
"deployment_failure",
true,
),
(Operation::SyncSuccess, false, false, "sync_success", false),
(
Operation::AsyncSuccess,
false,
false,
"async_success",
false,
),
(
Operation::SyncSuccessIfNeeded,
false,
false,
"sync_success_if_needed",
false,
),
(Operation::SyncFailure, false, false, "sync_failure", false),
(Operation::AsyncFailure, false, false, "async_failure", true),
(Operation::Restore, false, false, "restore", false),
];
for (operation, asynchronous, pre_call, method, awaiting) in cases {
assert_eq!(
operation_binding(operation, asynchronous, pre_call, "test").unwrap(),
OperationBinding { method, awaiting },
);
}
});
}
#[test]
fn null_byte_in_source_raises_driver_error() {
fn invalid_operations_raise_route_specific_errors() {
Python::initialize();
Python::attach(|py| {
let error = compile(py, "test", "class Host:\0 pass")
.expect_err("driver source containing null bytes should fail");
assert!(error.is_instance_of::<RustBridgeDriverError>(py));
Python::attach(|_| {
let pre_call = operation_binding(Operation::PreCall, false, false, "messages")
.expect_err("unsupported pre-call should fail");
assert_eq!(
error.to_string(),
"RustBridgeDriverError: driver source contains a null byte"
pre_call.to_string(),
"RuntimeError: messages lifecycle selected an unsupported pre-call operation"
);
let complete = operation_binding(
Operation::Complete(Outcome::Success),
false,
false,
"messages",
)
.expect_err("complete lifecycle should fail");
assert_eq!(
complete.to_string(),
"RuntimeError: messages lifecycle is complete"
);
});
}

168
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -17,6 +17,7 @@ use pyo3::sync::PyOnceLock;
use pyo3::types::PyDict;
use serde_json::{Map, Value};
use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT};
use crate::errors::{RustBridgeDeclined, chat_completions_error_to_pyerr, core_error_to_pyerr};
use crate::marshal::optional_timeout;
@ -154,36 +155,15 @@ fn invoke(
let machine = machine.borrow(py);
(machine.machine.operation(), machine.asynchronous)
};
let (method, awaiting) = match operation {
Operation::Setup => ("setup", false),
Operation::DeploymentPre => ("deployment_pre", true),
Operation::Prepare => ("prepare", false),
Operation::PreCall => {
return Err(PyRuntimeError::new_err(
"chat completions lifecycle selected an unsupported pre-call operation",
));
}
Operation::Send if asynchronous => ("send", true),
Operation::Send => ("send_sync", false),
Operation::DeploymentSuccess => ("deployment_success", true),
Operation::DeploymentFailure => ("deployment_failure", true),
Operation::SyncSuccess => ("sync_success", false),
Operation::AsyncSuccess => ("async_success", false),
Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false),
Operation::SyncFailure => ("sync_failure", false),
Operation::AsyncFailure => ("async_failure", true),
Operation::Restore => ("restore", false),
Operation::Complete(_) => {
return Err(PyRuntimeError::new_err(
"chat completions lifecycle is complete",
));
}
};
Ok((awaiting, host.getattr(py, method)?.call0(py)?))
crate::driver::invoke(py, operation, asynchronous, false, "chat completions", host)
}
#[pyfunction]
fn prepare(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Py<ChatCompletionsState>> {
fn prepare(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
) -> PyResult<Py<ChatCompletionsState>> {
let bag = arguments.bind(py);
let admission = admission(bag)?;
let api_key = scalar(bag, "api_key")?;
@ -195,10 +175,6 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Py<ChatCompletions
.map(|value| value.extract::<f64>())
.transpose()?,
)?;
let logging = bag
.get_item("litellm_logging_obj")?
.filter(|value| !value.is_none())
.ok_or_else(|| PyRuntimeError::new_err("chat completions logging was not initialized"))?;
let complete_input = PyDict::new(py);
complete_input.set_item("model", &admission.model)?;
complete_input.set_item("messages", bag.get_item("messages")?)?;
@ -206,14 +182,16 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Py<ChatCompletions
complete_input.set_item(name, Pythonized(value))?;
}
let additional = PyDict::new(py);
additional.set_item("complete_input_dict", complete_input)?;
additional.set_item("api_base", bag.get_item("api_base")?)?;
additional.set_item("headers", bag.get_item("extra_headers")?)?;
additional.set_item(COMPLETE_INPUT_DICT, complete_input)?;
additional.set_item(API_BASE, bag.get_item("api_base")?)?;
additional.set_item(HEADERS, bag.get_item("extra_headers")?)?;
let kwargs = PyDict::new(py);
kwargs.set_item("input", bag.get_item("messages")?)?;
kwargs.set_item("api_key", bag.get_item("logging_api_key")?)?;
kwargs.set_item("additional_args", additional)?;
logging.call_method("pre_call", (), Some(&kwargs))?;
kwargs.set_item(INPUT, bag.get_item("messages")?)?;
kwargs.set_item(API_KEY, bag.get_item("logging_api_key")?)?;
kwargs.set_item(ADDITIONAL_ARGS, additional)?;
logging
.bind(py)
.call_method("pre_call", (), Some(&kwargs))?;
Py::new(
py,
ChatCompletionsState {
@ -348,13 +326,17 @@ fn validate_arguments(arguments: &Bound<'_, PyDict>) -> PyResult<()> {
#[pyfunction]
fn chat_completions(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
validate_arguments(arguments.bind(py))?;
driver(py)?.getattr("drive_sync")?.call1((arguments,))
runner(py)?
.getattr("_drive_sync")?
.call1((arguments, bindings(py)?))
}
#[pyfunction]
fn achat_completions(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
validate_arguments(arguments.bind(py))?;
driver(py)?.getattr("drive_async")?.call1((arguments,))
runner(py)?
.getattr("_drive_async")?
.call1((arguments, bindings(py)?))
}
#[pyfunction]
@ -376,94 +358,32 @@ fn chat_completions_decline(
.map(str::to_string)
}
fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static DRIVER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = DRIVER.get(py) {
fn runner(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static RUNNER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = RUNNER.get(py) {
return Ok(module.bind(py));
}
let module = crate::driver::compile(py, "chat_completions", HOST)?;
module.add("_Lifecycle", py.get_type::<ChatCompletionsLifecycle>())?;
module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?;
module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
module.add("_send", wrap_pyfunction!(send, &module)?)?;
module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add(
"_terminal_record",
wrap_pyfunction!(terminal_record, &module)?,
)?;
Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py))
let module = py.import("litellm.rust_bridge.chat_completions")?;
Ok(RUNNER.get_or_init(py, || module.unbind()).bind(py))
}
const HOST: &str = r#"
from datetime import datetime
from litellm import utils
from litellm.types.utils import CallTypes
from litellm.rust_bridge.chat_completions import build_model_response, initialize_logging, invoke_terminal
class Host:
def __init__(self, arguments, asynchronous):
self.machine = _Lifecycle(arguments, asynchronous, utils.is_internal_call.get())
self.arguments = arguments
self.current = arguments
self.asynchronous = asynchronous
self.logger = arguments.get('litellm_logging_obj')
self.state = None
self.response = None
self.error = None
self.start = datetime.now()
self.end = None
def setup(self):
self.logger = initialize_logging(self.arguments, self.asynchronous)
self.arguments['litellm_logging_obj'] = self.logger
async def deployment_pre(self):
modified = await utils.async_pre_call_deployment_hook(self.current, 'acompletion')
if modified is not None:
self.current = modified
self.current['litellm_logging_obj'] = self.logger
def prepare(self): self.state = _prepare(self.current)
def send_sync(self):
self.response = build_model_response(_send_sync(self.state), self.arguments['model_response'])
self.end = datetime.now()
async def send(self):
self.response = build_model_response(await _send(self.state), self.arguments['model_response'])
self.end = datetime.now()
async def deployment_success(self):
self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.acompletion)
async def deployment_failure(self):
await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'acompletion')
def terminal(self, action, value):
record = _terminal_record(self.state) if self.state is not None else None
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, record, value, self.start, self.end)
def sync_success(self): return self.terminal('sync_success', self.response)
def async_success(self): return self.terminal('async_success', self.response)
def sync_success_if_needed(self): return self.terminal('sync_success_if_needed', self.response)
def sync_failure(self): return self.terminal('sync_failure', self.error)
def async_failure(self): return self.terminal('async_failure', self.error)
def restore(self): utils._restore_correlation_context_if_supported(self.logger)
def advance(self, outcome, error=None):
if error is not None and self.end is None:
self.end = datetime.now()
if self.logger is None:
self.logger = self.arguments.get('litellm_logging_obj')
replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None)
if replace:
self.error = error
def result(self):
if self.machine.complete():
return self.response
raise self.error
"#;
fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static BINDINGS: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = BINDINGS.get(py) {
return Ok(module.bind(py));
}
let module = PyModule::new(py, "_chat_completions_bindings")?;
module.add("Lifecycle", py.get_type::<ChatCompletionsLifecycle>())?;
module.add("invoke", wrap_pyfunction!(invoke, &module)?)?;
module.add("prepare", wrap_pyfunction!(prepare, &module)?)?;
module.add("send", wrap_pyfunction!(send, &module)?)?;
module.add("send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add(
"terminal_record",
wrap_pyfunction!(terminal_record, &module)?,
)?;
Ok(BINDINGS.get_or_init(py, || module.unbind()).bind(py))
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
crate::routes::definition::add_function(module, wrap_pyfunction!(chat_completions, module)?)?;

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

@ -11,6 +11,7 @@ use pyo3::sync::PyOnceLock;
use pyo3::types::PyDict;
use serde_json::{Map, Value};
use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT};
use crate::errors::{RustUpstreamError, core_error_to_pyerr, messages_provider_error_to_pyerr};
use crate::marshal::optional_timeout;
@ -145,34 +146,15 @@ fn invoke(
let machine = machine.borrow(py);
(machine.machine.operation(), machine.asynchronous)
};
let (method, awaiting) = match operation {
Operation::Setup => ("setup", false),
Operation::DeploymentPre => ("deployment_pre", true),
Operation::Prepare => ("prepare", false),
Operation::PreCall => {
return Err(PyRuntimeError::new_err(
"messages lifecycle selected an unsupported pre-call operation",
));
}
Operation::Send if asynchronous => ("send", true),
Operation::Send => ("send_sync", false),
Operation::DeploymentSuccess => ("deployment_success", true),
Operation::DeploymentFailure => ("deployment_failure", true),
Operation::SyncSuccess => ("sync_success", false),
Operation::AsyncSuccess => ("async_success", false),
Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false),
Operation::SyncFailure => ("sync_failure", false),
Operation::AsyncFailure => ("async_failure", true),
Operation::Restore => ("restore", false),
Operation::Complete(_) => {
return Err(PyRuntimeError::new_err("messages lifecycle is complete"));
}
};
Ok((awaiting, host.getattr(py, method)?.call0(py)?))
crate::driver::invoke(py, operation, asynchronous, false, "messages", host)
}
#[pyfunction]
fn prepare(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Py<MessagesState>> {
fn prepare(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
) -> PyResult<Py<MessagesState>> {
let bag = arguments.bind(py);
let request = decode_request(py, bag)?;
let prepared = py
@ -185,23 +167,18 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Py<MessagesState>>
for (name, value) in &prepared.upstream_headers {
headers.set_item(name, value)?;
}
let logging = bag
.get_item("litellm_logging_obj")?
.filter(|value| !value.is_none())
.map(|value| value.unbind())
.ok_or_else(|| PyRuntimeError::new_err("messages logging was not initialized"))?;
let additional = PyDict::new(py);
additional.set_item(pyo3::intern!(py, "complete_input_dict"), &body)?;
additional.set_item(pyo3::intern!(py, "api_base"), &prepared.url)?;
additional.set_item(pyo3::intern!(py, "headers"), &headers)?;
additional.set_item(COMPLETE_INPUT_DICT, &body)?;
additional.set_item(API_BASE, &prepared.url)?;
additional.set_item(HEADERS, &headers)?;
let kwargs = PyDict::new(py);
let serialized = py.import("json")?.call_method1("dumps", (&body,))?;
let message = PyDict::new(py);
message.set_item("role", "user")?;
message.set_item("content", serialized)?;
kwargs.set_item("input", vec![message])?;
kwargs.set_item("api_key", "")?;
kwargs.set_item("additional_args", additional)?;
kwargs.set_item(INPUT, vec![message])?;
kwargs.set_item(API_KEY, "")?;
kwargs.set_item(ADDITIONAL_ARGS, additional)?;
logging
.bind(py)
.call_method(pyo3::intern!(py, "pre_call"), (), Some(&kwargs))?;
@ -296,123 +273,45 @@ fn committed_failure() -> PyResult<()> {
#[pyfunction]
fn messages(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
validate_arguments(arguments.bind(py))?;
driver(py)?.getattr("drive_sync")?.call1((arguments,))
runner(py)?
.getattr("_drive_sync")?
.call1((arguments, bindings(py)?))
}
#[pyfunction]
fn amessages(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
validate_arguments(arguments.bind(py))?;
driver(py)?.getattr("drive_async")?.call1((arguments,))
runner(py)?
.getattr("_drive_async")?
.call1((arguments, bindings(py)?))
}
fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static DRIVER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = DRIVER.get(py) {
fn runner(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static RUNNER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = RUNNER.get(py) {
return Ok(module.bind(py));
}
let module = crate::driver::compile(py, "messages", HOST)?;
module.add("_Lifecycle", py.get_type::<MessagesLifecycle>())?;
module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?;
module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
module.add("_send", wrap_pyfunction!(send, &module)?)?;
module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add(
"_committed_failure",
wrap_pyfunction!(committed_failure, &module)?,
)?;
Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py))
let module = py.import("litellm.rust_bridge.messages")?;
Ok(RUNNER.get_or_init(py, || module.unbind()).bind(py))
}
const HOST: &str = r#"
from datetime import datetime
from litellm import utils
from litellm.types.utils import CallTypes
from litellm.rust_bridge.messages import initialize_logging, invoke_terminal, retain_stream_response
class Host:
def __init__(self, arguments, asynchronous):
self.machine = _Lifecycle(asynchronous, utils.is_internal_call.get())
self.arguments = arguments
self.current = arguments
self.asynchronous = asynchronous
self.logger = arguments.get('litellm_logging_obj')
self.lifecycle_owned = self.logger is None
self.state = None
self.response = None
self.error = None
self.start = datetime.now()
self.end = None
self.streaming = False
def setup(self):
self.logger = initialize_logging(self.arguments, self.asynchronous)
self.arguments['litellm_logging_obj'] = self.logger
self.streaming = self.logger.stream is True
async def deployment_pre(self):
if not self.lifecycle_owned:
return
modified = await utils.async_pre_call_deployment_hook(self.current, 'anthropic_messages')
if modified is not None:
self.current = modified
self.current['litellm_logging_obj'] = self.logger
def prepare(self):
self.state = _prepare(self.current)
def send_sync(self):
self.response = _send_sync(self.state)
self.end = datetime.now()
async def send(self):
self.response = await _send(self.state)
self.end = datetime.now()
async def deployment_success(self):
if self.lifecycle_owned:
self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages)
if self.streaming:
self.response = retain_stream_response(
self.response,
(self.arguments, self.current, self.state),
self.logger,
self.start,
)
async def deployment_failure(self):
if self.lifecycle_owned:
await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'anthropic_messages')
def terminal(self, action, value):
if self.streaming or not self.lifecycle_owned:
return None
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, None, value, self.start, self.end)
def sync_success(self): return self.terminal('sync_success', self.response)
def async_success(self): return self.terminal('async_success', self.response)
def sync_success_if_needed(self): return self.terminal('sync_success_if_needed', self.response)
def sync_failure(self): return self.terminal('sync_failure', self.error)
def async_failure(self): return self.terminal('async_failure', self.error)
def restore(self):
if not self.streaming and self.lifecycle_owned:
utils._restore_correlation_context_if_supported(self.logger)
def advance(self, outcome, error=None):
if error is not None and self.end is None:
self.end = datetime.now()
if self.logger is None:
self.logger = self.arguments.get('litellm_logging_obj')
replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None)
if replace:
self.error = error
def result(self):
if self.machine.complete():
return self.response
if self.machine.failed_after_provider_response():
_committed_failure()
raise self.error
"#;
fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static BINDINGS: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = BINDINGS.get(py) {
return Ok(module.bind(py));
}
let module = PyModule::new(py, "_messages_bindings")?;
module.add("Lifecycle", py.get_type::<MessagesLifecycle>())?;
module.add("invoke", wrap_pyfunction!(invoke, &module)?)?;
module.add("prepare", wrap_pyfunction!(prepare, &module)?)?;
module.add("send", wrap_pyfunction!(send, &module)?)?;
module.add("send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add(
"committed_failure",
wrap_pyfunction!(committed_failure, &module)?,
)?;
Ok(BINDINGS.get_or_init(py, || module.unbind()).bind(py))
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
crate::routes::definition::add_function(module, wrap_pyfunction!(messages, module)?)?;

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

@ -19,6 +19,7 @@ use pyo3::pyclass::{PyTraverseError, PyVisit};
use pyo3::sync::PyOnceLock;
use pyo3::types::PyDict;
use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT};
use crate::errors::core_error_to_pyerr;
use litellm_python_interop::{run_async_value, run_sync_value};
@ -213,15 +214,13 @@ impl OcrLifecycle {
fn new(
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
logger: Option<&Bound<'_, PyAny>>,
asynchronous: bool,
internal_call: bool,
) -> PyResult<Self> {
let request = decode_request(py, arguments)?;
let logger = arguments
.get_item("litellm_logging_obj")?
.filter(|value| !value.is_none());
let identity = |name: &str| -> PyResult<Option<String>> {
if let Some(logger) = &logger {
if let Some(logger) = logger {
match logger.getattr(name) {
Ok(value) => {
if let Ok(value) = value.extract::<String>() {
@ -313,30 +312,16 @@ fn invoke(
let machine = machine.borrow(py);
(machine.machine.operation(), machine.asynchronous)
};
let (method, awaiting) = match operation {
Operation::Setup => ("setup", false),
Operation::DeploymentPre => ("deployment_pre", true),
Operation::Prepare => ("prepare", false),
Operation::PreCall => ("pre_call", false),
Operation::Send if asynchronous => ("send", true),
Operation::Send => ("send_sync", false),
Operation::DeploymentSuccess => ("deployment_success", true),
Operation::DeploymentFailure => ("deployment_failure", true),
Operation::SyncSuccess => ("sync_success", false),
Operation::AsyncSuccess => ("async_success", false),
Operation::SyncSuccessIfNeeded => ("sync_success_if_needed", false),
Operation::SyncFailure => ("sync_failure", false),
Operation::AsyncFailure => ("async_failure", true),
Operation::Restore => ("restore", false),
Operation::Complete(_) => return Err(PyRuntimeError::new_err("OCR lifecycle is complete")),
};
let value = host.getattr(py, method)?.call0(py)?;
Ok((awaiting, value))
crate::driver::invoke(py, operation, asynchronous, true, "OCR", host)
}
#[pyfunction]
#[pyo3(signature = (arguments, asynchronous=false))]
fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResult<Py<OcrState>> {
fn prepare(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
asynchronous: bool,
) -> PyResult<Py<OcrState>> {
let bag = arguments.bind(py);
let request = decode_request(py, bag)?;
let model = request.model.clone();
@ -380,10 +365,6 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResul
for (name, value) in &draft_headers {
headers.set_item(name, value)?;
}
let logging = py
.import("litellm.rust_bridge.ocr")?
.getattr("initialize_logging")?
.call1((bag, asynchronous))?;
let litellm_params = PyDict::new(py);
litellm_params.set_item("litellm_call_id", bag.get_item("litellm_call_id")?)?;
litellm_params.set_item(
@ -396,17 +377,18 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>, asynchronous: bool) -> PyResul
update.set_item("optional_params", optional_params)?;
update.set_item("litellm_params", litellm_params)?;
update.set_item("custom_llm_provider", endpoint.custom_llm_provider())?;
logging.call_method("update_from_kwargs", (), Some(&update))?;
logging
.bind(py)
.call_method("update_from_kwargs", (), Some(&update))?;
let additional_args = PyDict::new(py);
additional_args.set_item("complete_input_dict", &body)?;
additional_args.set_item(pyo3::intern!(py, "api_base"), endpoint.url())?;
additional_args.set_item(pyo3::intern!(py, "headers"), &headers)?;
additional_args.set_item(COMPLETE_INPUT_DICT, &body)?;
additional_args.set_item(API_BASE, endpoint.url())?;
additional_args.set_item(HEADERS, &headers)?;
let pre_call = PyDict::new(py);
pre_call.set_item("input", "OCR document processing")?;
pre_call.set_item("api_key", bag.get_item("api_key")?)?;
pre_call.set_item("additional_args", additional_args)?;
let logging = logging.unbind();
pre_call.set_item(INPUT, "OCR document processing")?;
pre_call.set_item(API_KEY, bag.get_item("api_key")?)?;
pre_call.set_item(ADDITIONAL_ARGS, additional_args)?;
Py::new(
py,
OcrState {
@ -593,127 +575,45 @@ fn terminal_record(py: Python<'_>, state: Py<OcrState>) -> PyResult<Py<PyAny>> {
#[pyfunction]
fn ocr(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
driver(py)?.getattr("drive_sync")?.call1((arguments,))
runner(py)?
.getattr("_drive_sync")?
.call1((arguments, bindings(py)?))
}
#[pyfunction]
fn aocr(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
driver(py)?.getattr("drive_async")?.call1((arguments,))
runner(py)?
.getattr("_drive_async")?
.call1((arguments, bindings(py)?))
}
// Compilation can re-enter through audit hooks; publish only a finished module.
fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static DRIVER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = DRIVER.get(py) {
fn runner(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static RUNNER: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = RUNNER.get(py) {
return Ok(module.bind(py));
}
let module = crate::driver::compile(
py,
"ocr",
"from datetime import datetime
from litellm import utils
from litellm.types.utils import CallTypes
from litellm.rust_bridge.ocr import initialize_logging, invoke_terminal
let module = py.import("litellm.rust_bridge.ocr")?;
Ok(RUNNER.get_or_init(py, || module.unbind()).bind(py))
}
class Host:
def __init__(self, arguments, asynchronous):
self.machine = _Lifecycle(arguments, asynchronous, utils.is_internal_call.get())
self.arguments = arguments
self.current = arguments
self.asynchronous = asynchronous
self.logger = arguments.get('litellm_logging_obj')
self.state = None
self.response = None
self.error = None
self.start = datetime.now()
self.end = None
def setup(self):
call_id, trace_id = self.machine.identity()
self.arguments['litellm_call_id'] = call_id
self.arguments['litellm_trace_id'] = trace_id
self.logger = initialize_logging(self.arguments, self.asynchronous)
self.arguments['litellm_logging_obj'] = self.logger
async def deployment_pre(self):
modified = await utils.async_pre_call_deployment_hook(self.current, 'aocr')
if modified is not None:
self.current = modified
self.current['litellm_logging_obj'] = self.logger
call_id, trace_id = self.machine.identity()
self.current['litellm_call_id'] = call_id
self.current['litellm_trace_id'] = trace_id
def prepare(self):
self.state = _prepare(self.current, self.asynchronous)
def pre_call(self):
_pre_call(self.state)
def send_sync(self):
self.response = _send_sync(self.state)
self.end = datetime.now()
async def send(self):
self.response = _finish(await _send(self.state))
self.end = datetime.now()
async def deployment_success(self):
self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aocr)
async def deployment_failure(self):
await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'aocr')
def terminal(self, action, value):
record = _terminal_record(self.state) if self.state is not None else None
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, record, value, self.start, self.end)
def sync_success(self):
return self.terminal('sync_success', self.response)
def async_success(self):
return self.terminal('async_success', self.response)
def sync_success_if_needed(self):
return self.terminal('sync_success_if_needed', self.response)
def sync_failure(self):
return self.terminal('sync_failure', self.error)
def async_failure(self):
return self.terminal('async_failure', self.error)
def restore(self):
utils._restore_correlation_context_if_supported(self.logger)
def advance(self, outcome, error=None):
if error is not None and self.end is None:
self.end = datetime.now()
if self.logger is None:
self.logger = self.arguments.get('litellm_logging_obj')
replace = self.machine.advance(outcome, self.logger is not None, self.current.get('fallbacks') is not None)
if replace:
self.error = error
def result(self):
if self.machine.complete():
return self.response
raise self.error
",
)?;
module.add("_Lifecycle", py.get_type::<OcrLifecycle>())?;
module.add("_invoke", wrap_pyfunction!(invoke, &module)?)?;
module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?;
module.add("_pre_call", wrap_pyfunction!(pre_call, &module)?)?;
module.add("_send", wrap_pyfunction!(send, &module)?)?;
module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add("_finish", wrap_pyfunction!(finish, &module)?)?;
fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
static BINDINGS: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
if let Some(module) = BINDINGS.get(py) {
return Ok(module.bind(py));
}
let module = PyModule::new(py, "_ocr_bindings")?;
module.add("Lifecycle", py.get_type::<OcrLifecycle>())?;
module.add("invoke", wrap_pyfunction!(invoke, &module)?)?;
module.add("prepare", wrap_pyfunction!(prepare, &module)?)?;
module.add("pre_call", wrap_pyfunction!(pre_call, &module)?)?;
module.add("send", wrap_pyfunction!(send, &module)?)?;
module.add("send_sync", wrap_pyfunction!(send_sync, &module)?)?;
module.add("finish", wrap_pyfunction!(finish, &module)?)?;
module.add(
"_terminal_record",
"terminal_record",
wrap_pyfunction!(terminal_record, &module)?,
)?;
Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py))
Ok(BINDINGS.get_or_init(py, || module.unbind()).bind(py))
}
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {

154
litellm/rust_bridge/_lifecycle.py
View file

@ -1,7 +1,157 @@
from __future__ import annotations
from collections.abc import Mapping
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime
from enum import IntEnum
from typing import TYPE_CHECKING, Final, Literal, Protocol
from pydantic import TypeAdapter
if TYPE_CHECKING:
from litellm.types.utils import CallTypes
LOGGING_OBJECT_KEY: Final = "litellm_logging_obj"
FALLBACKS_KEY: Final = "fallbacks"
TerminalAction = Literal[
"sync_success",
"async_success",
"sync_success_if_needed",
"sync_failure",
"async_failure",
]
_OPTIONAL_ARGUMENTS_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(
dict[str, object] | None
)
class NativeOutcome(IntEnum):
SUCCESS = 0
FAILURE = 1
ABORT = 2
class NativeLifecycle(Protocol):
def complete(self) -> bool | None: ...
def advance(self, outcome: int, logger_available: bool, has_fallbacks: bool) -> bool: ...
class NativeLifecycleBindings(Protocol):
invoke: Callable[[NativeLifecycle, object], tuple[bool, object]]
class LifecycleHost(Protocol):
@property
def machine(self) -> NativeLifecycle: ...
def invoke(self) -> tuple[bool, object]: ...
def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None: ...
def result(self) -> object: ...
class MutableLifecycleHost(LifecycleHost, Protocol):
arguments: dict[str, object]
current: dict[str, object]
logger: object | None
response: object
error: BaseException | None
end: datetime | None
def advance_host(host: MutableLifecycleHost, outcome: NativeOutcome, error: BaseException | None) -> None:
if error is not None and host.end is None:
host.end = datetime.now()
if host.logger is None:
host.logger = host.arguments.get(LOGGING_OBJECT_KEY)
replace: Final = host.machine.advance(
outcome,
host.logger is not None,
host.current.get(FALLBACKS_KEY) is not None,
)
if replace:
host.error = error
def host_result(host: MutableLifecycleHost) -> object:
if host.machine.complete():
return host.response
if host.error is None:
raise RuntimeError("native lifecycle failed without an error")
raise host.error
async def deployment_pre(arguments: dict[str, object], call_type: str) -> dict[str, object]:
from litellm import utils
modified: Final = _OPTIONAL_ARGUMENTS_ADAPTER.validate_python(
await utils.async_pre_call_deployment_hook(arguments, call_type)
)
return arguments if modified is None else modified
async def deployment_success(arguments: dict[str, object], response: object, call_type: CallTypes) -> object:
from litellm import utils
updated: object = await utils.async_post_call_success_deployment_hook( # pyright: ignore[reportUnknownMemberType] # legacy hook annotations expose an unknown return
arguments, response, call_type
)
return updated
async def deployment_failure(arguments: dict[str, object], error: BaseException | None, call_type: str) -> None:
from litellm import utils
if not isinstance(error, Exception):
raise RuntimeError("native lifecycle failure did not retain an exception")
await utils.async_post_call_failure_deployment_hook(arguments, error, call_type)
def restore_correlation_context(logger: object | None) -> None:
from litellm import utils
utils._restore_correlation_context_if_supported(logger) # pyright: ignore[reportPrivateUsage] # lifecycle cleanup has no public wrapper
def _invoke_sync(host: LifecycleHost) -> None:
host.invoke()
async def _invoke_async(host: LifecycleHost) -> None:
awaiting, value = host.invoke()
if not awaiting:
return
if not isinstance(value, Awaitable):
raise TypeError("native lifecycle operation did not return an awaitable")
await value
def drive_sync(host: LifecycleHost) -> object:
while host.machine.complete() is None:
try:
_invoke_sync(host)
except Exception as error:
host.advance(NativeOutcome.FAILURE, error)
except BaseException as error:
host.advance(NativeOutcome.ABORT, error)
else:
host.advance(NativeOutcome.SUCCESS)
return host.result()
async def drive_async(host: LifecycleHost) -> object:
while host.machine.complete() is None:
try:
await _invoke_async(host)
except Exception as error:
host.advance(NativeOutcome.FAILURE, error)
except BaseException as error:
host.advance(NativeOutcome.ABORT, error)
else:
host.advance(NativeOutcome.SUCCESS)
return host.result()
def initialize_logging(arguments: dict[str, object], asynchronous: bool, route: str) -> object:
@ -11,7 +161,7 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route:
def invoke_terminal(
action: str,
action: TerminalAction,
roots: object,
logger: object,
record: Mapping[str, object] | None,

141
litellm/rust_bridge/chat_completions.py
View file

@ -16,6 +16,7 @@ import inspect
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import (
TYPE_CHECKING,
Final,
@ -32,6 +33,22 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
convert_to_model_response_object,
)
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.rust_bridge._lifecycle import (
LOGGING_OBJECT_KEY,
NativeLifecycle,
NativeLifecycleBindings,
NativeOutcome,
TerminalAction,
advance_host,
deployment_failure,
deployment_pre,
deployment_success,
drive_async,
drive_sync,
host_result,
invoke_terminal,
restore_correlation_context,
)
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.timeouts import timeout_to_seconds
@ -579,3 +596,127 @@ def _arguments(
"timeout_seconds": timeout_to_seconds(timeout),
"logging_api_key": logging_api_key if logging_api_key is not None else api_key or "",
}
class _ChatCompletionsBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[dict[str, object], bool, bool], NativeLifecycle]
prepare: Callable[[dict[str, object], object], object]
send: Callable[[object], Awaitable[Mapping[str, object]]]
send_sync: Callable[[object], Mapping[str, object]]
terminal_record: Callable[[object], Mapping[str, object]]
class _ChatCompletionsHost:
def __init__(
self,
arguments: dict[str, object],
asynchronous: bool,
bindings: _ChatCompletionsBindings,
) -> None:
from litellm import utils
self.bindings: _ChatCompletionsBindings = bindings
self.machine: NativeLifecycle = bindings.Lifecycle(arguments, asynchronous, utils.is_internal_call.get())
self.arguments: dict[str, object] = arguments
self.current: dict[str, object] = arguments
self.asynchronous: bool = asynchronous
self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY)
self.state: object | None = None
self.response: object = None
self.error: BaseException | None = None
self.start: datetime = datetime.now()
self.end: datetime | None = None
def invoke(self) -> tuple[bool, object]:
return self.bindings.invoke(self.machine, self)
def setup(self) -> None:
self.logger = initialize_logging(self.arguments, self.asynchronous)
self.arguments[LOGGING_OBJECT_KEY] = self.logger
async def deployment_pre(self) -> None:
self.current = await deployment_pre(self.current, "acompletion")
self.current[LOGGING_OBJECT_KEY] = self.logger
def prepare(self) -> None:
if self.logger is None:
raise RuntimeError("chat completions logging was not initialized")
self.state = self.bindings.prepare(self.current, self.logger)
def send_sync(self) -> None:
model_response: Final = self.arguments["model_response"]
if not isinstance(model_response, ModelResponse):
raise TypeError("chat completions model_response must be a ModelResponse")
self.response = build_model_response(self.bindings.send_sync(self.state), model_response)
self.end = datetime.now()
async def send(self) -> None:
model_response: Final = self.arguments["model_response"]
if not isinstance(model_response, ModelResponse):
raise TypeError("chat completions model_response must be a ModelResponse")
self.response = build_model_response(await self.bindings.send(self.state), model_response)
self.end = datetime.now()
async def deployment_success(self) -> None:
from litellm.types.utils import CallTypes
self.response = await deployment_success(self.current, self.response, CallTypes.acompletion)
async def deployment_failure(self) -> None:
await deployment_failure(self.current, self.error, "acompletion")
def terminal(self, action: TerminalAction, value: object) -> object:
if self.logger is None or self.end is None:
raise RuntimeError("chat completions terminal state was not initialized")
record: Final = self.bindings.terminal_record(self.state) if self.state is not None else None
return invoke_terminal(
action,
(self.arguments, self.current, self.state),
self.logger,
record,
value,
self.start,
self.end,
)
def sync_success(self) -> object:
return self.terminal("sync_success", self.response)
def async_success(self) -> object:
return self.terminal("async_success", self.response)
def sync_success_if_needed(self) -> object:
return self.terminal("sync_success_if_needed", self.response)
def sync_failure(self) -> object:
return self.terminal("sync_failure", self.error)
def async_failure(self) -> object:
return self.terminal("async_failure", self.error)
def restore(self) -> None:
restore_correlation_context(self.logger)
def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None:
advance_host(self, outcome, error)
def result(self) -> object:
return host_result(self)
def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _ChatCompletionsBindings
) -> ModelResponse:
result: Final = drive_sync(_ChatCompletionsHost(arguments, False, bindings))
if not isinstance(result, ModelResponse):
raise TypeError("native chat completions driver returned an invalid response")
return result
async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _ChatCompletionsBindings
) -> ModelResponse:
result: Final = await drive_async(_ChatCompletionsHost(arguments, True, bindings))
if not isinstance(result, ModelResponse):
raise TypeError("native chat completions driver returned an invalid response")
return result

154
litellm/rust_bridge/messages.py
View file

@ -1,16 +1,31 @@
from __future__ import annotations
from collections.abc import Awaitable
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final, Protocol, cast
import httpx
from litellm.rust_bridge._lifecycle import (
LOGGING_OBJECT_KEY,
NativeLifecycle,
NativeLifecycleBindings,
NativeOutcome,
TerminalAction,
advance_host,
deployment_failure,
deployment_pre,
deployment_success,
drive_async,
drive_sync,
host_result,
invoke_terminal,
restore_correlation_context,
)
from litellm.rust_bridge._lifecycle import (
initialize_logging as initialize_lifecycle_logging,
)
from litellm.rust_bridge._lifecycle import invoke_terminal
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.llms.anthropic_messages.anthropic_response import (
@ -206,6 +221,141 @@ async def amessages(
)
class _MessagesLifecycle(NativeLifecycle, Protocol):
def failed_after_provider_response(self) -> bool: ...
class _MessagesBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[bool, bool], _MessagesLifecycle]
prepare: Callable[[dict[str, object], object], object]
send: Callable[[object], Awaitable[AnthropicMessagesResponse]]
send_sync: Callable[[object], AnthropicMessagesResponse]
committed_failure: Callable[[], None]
class _MessagesHost:
def __init__(self, arguments: dict[str, object], asynchronous: bool, bindings: _MessagesBindings) -> None:
from litellm import utils
self.bindings: _MessagesBindings = bindings
self.machine: _MessagesLifecycle = bindings.Lifecycle(asynchronous, utils.is_internal_call.get())
self.arguments: dict[str, object] = arguments
self.current: dict[str, object] = arguments
self.asynchronous: bool = asynchronous
self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY)
self.lifecycle_owned: bool = self.logger is None
self.state: object | None = None
self.response: object = None
self.error: BaseException | None = None
self.start: datetime = datetime.now()
self.end: datetime | None = None
self.streaming: bool = False
def invoke(self) -> tuple[bool, object]:
return self.bindings.invoke(self.machine, self)
def setup(self) -> None:
self.logger = initialize_logging(self.arguments, self.asynchronous)
self.arguments[LOGGING_OBJECT_KEY] = self.logger
self.streaming = getattr(self.logger, "stream", False) is True
async def deployment_pre(self) -> None:
if not self.lifecycle_owned:
return
self.current = await deployment_pre(self.current, "anthropic_messages")
self.current[LOGGING_OBJECT_KEY] = self.logger
def prepare(self) -> None:
if self.logger is None:
raise RuntimeError("messages logging was not initialized")
self.state = self.bindings.prepare(self.current, self.logger)
def send_sync(self) -> None:
self.response = self.bindings.send_sync(self.state)
self.end = datetime.now()
async def send(self) -> None:
self.response = await self.bindings.send(self.state)
self.end = datetime.now()
async def deployment_success(self) -> None:
from litellm.types.utils import CallTypes
if self.lifecycle_owned:
self.response = await deployment_success(self.current, self.response, CallTypes.aanthropic_messages)
if not self.streaming:
return
if self.logger is None:
raise RuntimeError("messages logging was not initialized")
self.response = retain_stream_response(
cast(AnthropicMessagesResponse, self.response),
(self.arguments, self.current, self.state),
cast(_MessagesLogging, self.logger),
self.start,
)
async def deployment_failure(self) -> None:
if not self.lifecycle_owned:
return
await deployment_failure(self.current, self.error, "anthropic_messages")
def terminal(self, action: TerminalAction, value: object) -> object:
if self.streaming or not self.lifecycle_owned:
return None
if self.logger is None or self.end is None:
raise RuntimeError("messages terminal state was not initialized")
return invoke_terminal(
action,
(self.arguments, self.current, self.state),
self.logger,
None,
value,
self.start,
self.end,
)
def sync_success(self) -> object:
return self.terminal("sync_success", self.response)
def async_success(self) -> object:
return self.terminal("async_success", self.response)
def sync_success_if_needed(self) -> object:
return self.terminal("sync_success_if_needed", self.response)
def sync_failure(self) -> object:
return self.terminal("sync_failure", self.error)
def async_failure(self) -> object:
return self.terminal("async_failure", self.error)
def restore(self) -> None:
if not self.streaming and self.lifecycle_owned:
restore_correlation_context(self.logger)
def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None:
advance_host(self, outcome, error)
def result(self) -> object:
if self.machine.complete():
return self.response
if self.machine.failed_after_provider_response():
self.bindings.committed_failure()
return host_result(self)
def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _MessagesBindings
) -> AnthropicMessagesResponse:
return cast(AnthropicMessagesResponse, drive_sync(_MessagesHost(arguments, False, bindings)))
async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _MessagesBindings
) -> AnthropicMessagesResponse:
return cast(AnthropicMessagesResponse, await drive_async(_MessagesHost(arguments, True, bindings)))
__all__ = (
"amessages",
"initialize_logging",

147
litellm/rust_bridge/ocr.py
View file

@ -3,13 +3,28 @@
from __future__ import annotations
import traceback
from collections.abc import Awaitable, Mapping
from collections.abc import Awaitable, Callable, Mapping
from contextvars import copy_context
from datetime import datetime
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
from uuid import uuid4
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge._lifecycle import (
LOGGING_OBJECT_KEY,
NativeLifecycle,
NativeLifecycleBindings,
NativeOutcome,
TerminalAction,
advance_host,
deployment_failure,
deployment_pre,
deployment_success,
drive_async,
drive_sync,
host_result,
restore_correlation_context,
)
from litellm.rust_bridge.bindings import NativeBinding
@ -63,7 +78,7 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route:
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.litellm_core_utils.litellm_logging import Logging, set_callbacks
supplied: Final = arguments.get("litellm_logging_obj")
supplied: Final = arguments.get(LOGGING_OBJECT_KEY)
if supplied is not None:
return supplied
callbacks: Final = tuple( # cast-ok: callback registry accepts heterogeneous legacy callback objects
@ -185,12 +200,12 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool, route:
cb for cb in dict.fromkeys(logger.dynamic_input_callbacks or ()) if cb not in litellm.input_callback
]
arguments["litellm_call_id"] = call_id
arguments["litellm_logging_obj"] = logger
arguments[LOGGING_OBJECT_KEY] = logger
return logger
def invoke_terminal(
action: str,
action: TerminalAction,
roots: object,
logger: object,
record: Mapping[str, object] | None,
@ -250,3 +265,127 @@ def invoke_terminal(
logging.failure_handler(exception, trace, start_time, end_time)
return None
return logging.async_failure_handler(exception, trace, start_time, end_time)
class _OcrLifecycle(NativeLifecycle, Protocol):
def identity(self) -> tuple[str, str | None]: ...
class _OcrBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[dict[str, object], object | None, bool, bool], _OcrLifecycle]
prepare: Callable[[dict[str, object], object, bool], object]
pre_call: Callable[[object], None]
send: Callable[[object], Awaitable[dict[str, object]]]
send_sync: Callable[[object], OCRResponse]
finish: Callable[[dict[str, object]], OCRResponse]
terminal_record: Callable[[object], Mapping[str, object]]
class _OcrHost:
def __init__(self, arguments: dict[str, object], asynchronous: bool, bindings: _OcrBindings) -> None:
from litellm import utils
self.bindings: _OcrBindings = bindings
self.arguments: dict[str, object] = arguments
self.current: dict[str, object] = arguments
self.asynchronous: bool = asynchronous
self.logger: object | None = arguments.get(LOGGING_OBJECT_KEY)
self.machine: _OcrLifecycle = bindings.Lifecycle(
arguments, self.logger, asynchronous, utils.is_internal_call.get()
)
self.state: object | None = None
self.response: object = None
self.error: BaseException | None = None
self.start: datetime = datetime.now()
self.end: datetime | None = None
def invoke(self) -> tuple[bool, object]:
return self.bindings.invoke(self.machine, self)
def setup(self) -> None:
call_id, trace_id = self.machine.identity()
self.arguments["litellm_call_id"] = call_id
self.arguments["litellm_trace_id"] = trace_id
self.logger = initialize_logging(self.arguments, self.asynchronous)
self.arguments[LOGGING_OBJECT_KEY] = self.logger
async def deployment_pre(self) -> None:
self.current = await deployment_pre(self.current, "aocr")
self.current[LOGGING_OBJECT_KEY] = self.logger
call_id, trace_id = self.machine.identity()
self.current["litellm_call_id"] = call_id
self.current["litellm_trace_id"] = trace_id
def prepare(self) -> None:
if self.logger is None:
raise RuntimeError("OCR logging was not initialized")
self.state = self.bindings.prepare(self.current, self.logger, self.asynchronous)
def pre_call(self) -> None:
self.bindings.pre_call(self.state)
def send_sync(self) -> None:
self.response = self.bindings.send_sync(self.state)
self.end = datetime.now()
async def send(self) -> None:
self.response = self.bindings.finish(await self.bindings.send(self.state))
self.end = datetime.now()
async def deployment_success(self) -> None:
from litellm.types.utils import CallTypes
self.response = await deployment_success(self.current, self.response, CallTypes.aocr)
async def deployment_failure(self) -> None:
await deployment_failure(self.current, self.error, "aocr")
def terminal(self, action: TerminalAction, value: object) -> object:
if self.logger is None or self.end is None:
raise RuntimeError("OCR terminal state was not initialized")
record: Final = self.bindings.terminal_record(self.state) if self.state is not None else None
return invoke_terminal(
action,
(self.arguments, self.current, self.state),
self.logger,
record,
value,
self.start,
self.end,
)
def sync_success(self) -> object:
return self.terminal("sync_success", self.response)
def async_success(self) -> object:
return self.terminal("async_success", self.response)
def sync_success_if_needed(self) -> object:
return self.terminal("sync_success_if_needed", self.response)
def sync_failure(self) -> object:
return self.terminal("sync_failure", self.error)
def async_failure(self) -> object:
return self.terminal("async_failure", self.error)
def restore(self) -> None:
restore_correlation_context(self.logger)
def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None:
advance_host(self, outcome, error)
def result(self) -> object:
return host_result(self)
def _drive_sync( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _OcrBindings
) -> OCRResponse:
return cast(OCRResponse, drive_sync(_OcrHost(arguments, False, bindings)))
async def _drive_async( # pyright: ignore[reportUnusedFunction] # called by the native extension
arguments: dict[str, object], bindings: _OcrBindings
) -> OCRResponse:
return cast(OCRResponse, await drive_async(_OcrHost(arguments, True, bindings)))

10
tests/test_litellm/integrations/test_prometheus_caller_identity.py
View file

@ -24,6 +24,7 @@ from litellm.types.integrations.prometheus import (
validate_prometheus_deployment_and_latency_caller_identity,
)
from litellm.types.utils import StandardLoggingPayload
from tests._prometheus_helpers import clear_prometheus_registry
TARGET_METRICS: Final[tuple[DEFINED_PROMETHEUS_METRICS, ...]] = cast(
tuple[DEFINED_PROMETHEUS_METRICS, ...],
@ -32,14 +33,9 @@ TARGET_METRICS: Final[tuple[DEFINED_PROMETHEUS_METRICS, ...]] = cast(
IDENTITY_MODES: Final = ("api_key_alias", "user_email", "both")
def _clear_prometheus_registry() -> None:
for collector in list(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage]
REGISTRY.unregister(collector)
@pytest.fixture(autouse=True)
def reset_prometheus_settings(monkeypatch: pytest.MonkeyPatch):
_clear_prometheus_registry()
clear_prometheus_registry()
monkeypatch.setattr(litellm, "prometheus_deployment_and_latency_caller_identity", "api_key_alias")
monkeypatch.setattr(litellm, "prometheus_metrics_config", None)
monkeypatch.setattr(litellm, "prometheus_exclude_metrics", None)
@ -47,7 +43,7 @@ def reset_prometheus_settings(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", [])
monkeypatch.setattr(litellm, "custom_prometheus_tags", [])
yield
_clear_prometheus_registry()
clear_prometheus_registry()
def _expected_identity_labels(baseline: list[str], mode: str) -> list[str]:

9
tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py
View file

@ -6,12 +6,7 @@ from prometheus_client import REGISTRY
import litellm
from litellm.integrations.prometheus import PrometheusLogger
def _clear_prometheus_registry() -> None:
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
from tests._prometheus_helpers import clear_prometheus_registry
def _create_prometheus_logger_with_custom_labels(monkeypatch: pytest.MonkeyPatch):
@ -20,7 +15,7 @@ def _create_prometheus_logger_with_custom_labels(monkeypatch: pytest.MonkeyPatch
"custom_prometheus_metadata_labels",
["metadata.department", "metadata.environment"],
)
_clear_prometheus_registry()
clear_prometheus_registry()
return PrometheusLogger()

19
tests/test_litellm/integrations/test_prometheus_labels.py
View file

@ -8,16 +8,7 @@ from litellm.types.integrations.prometheus import (
PrometheusMetricLabels,
UserAPIKeyLabelNames,
)
def _clear_prometheus_registry() -> None:
from prometheus_client import REGISTRY
for collector in list(REGISTRY._collector_to_names.keys()):
try:
REGISTRY.unregister(collector)
except Exception:
pass
from tests._prometheus_helpers import clear_prometheus_registry
def _collected_samples(metric_name: str):
@ -661,7 +652,7 @@ async def test_success_hook_emits_api_provider_value_on_token_metric():
"hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None},
}
_clear_prometheus_registry()
clear_prometheus_registry()
try:
logger = PrometheusLogger()
now = datetime.datetime.now()
@ -682,7 +673,7 @@ async def test_success_hook_emits_api_provider_value_on_token_metric():
f"{[s.labels.get('api_provider') for s in samples]}"
)
finally:
_clear_prometheus_registry()
clear_prometheus_registry()
@pytest.mark.asyncio
@ -698,7 +689,7 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric()
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import UserAPIKeyAuth
_clear_prometheus_registry()
clear_prometheus_registry()
try:
logger = PrometheusLogger()
await logger.async_post_call_failure_hook(
@ -713,7 +704,7 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric()
f"{[s.labels.get('api_provider') for s in samples]}"
)
finally:
_clear_prometheus_registry()
clear_prometheus_registry()
if __name__ == "__main__":

27
tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py
View file

@ -31,6 +31,7 @@ from litellm.types.integrations.prometheus import (
UserAPIKeyLabelNames,
UserAPIKeyLabelValues,
)
from tests._prometheus_helpers import clear_prometheus_registry
# ---------------------------------------------------------------------------
@ -482,16 +483,6 @@ KEY_AND_TEAM_RATE_LIMIT_METRICS = (
)
def _clear_prometheus_registry() -> None:
from prometheus_client import REGISTRY
for collector in list(REGISTRY._collector_to_names.keys()):
try:
REGISTRY.unregister(collector)
except Exception:
pass
def _collected_samples(metric_name: str) -> dict[tuple[tuple[str, str], ...], float]:
from prometheus_client import REGISTRY
@ -569,7 +560,7 @@ async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_head
and the window consumption as ``limit - remaining`` for each key / team
dimension, split by ``rate_limit_type``.
"""
_clear_prometheus_registry()
clear_prometheus_registry()
try:
await _run_success_event(
{
@ -624,7 +615,7 @@ async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_head
team_tokens: 40,
}
finally:
_clear_prometheus_registry()
clear_prometheus_registry()
@pytest.mark.asyncio
@ -634,7 +625,7 @@ async def test_should_emit_only_the_dimensions_the_limiter_enforced():
key/requests headers, so no tokens series and no team series may appear
(a phantom 0 or sys.maxsize series would misreport an unlimited dimension).
"""
_clear_prometheus_registry()
clear_prometheus_registry()
try:
await _run_success_event(
{
@ -653,7 +644,7 @@ async def test_should_emit_only_the_dimensions_the_limiter_enforced():
assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {}
assert _collected_samples("litellm_team_rate_limit_used_metric") == {}
finally:
_clear_prometheus_registry()
clear_prometheus_registry()
@pytest.mark.asyncio
@ -664,7 +655,7 @@ async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_
requests. The old allowed/used samples must disappear instead of keeping
a limit that no longer exists on the scrape.
"""
_clear_prometheus_registry()
clear_prometheus_registry()
try:
logger = PrometheusLogger()
await _run_success_event(
@ -698,7 +689,7 @@ async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_
assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {team_requests: 50}
assert _collected_samples("litellm_team_rate_limit_used_metric") == {team_requests: 4}
finally:
_clear_prometheus_registry()
clear_prometheus_registry()
@pytest.mark.asyncio
@ -715,11 +706,11 @@ async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_
async def test_should_emit_no_key_or_team_rate_limit_series_without_a_complete_int_pair(
additional_headers,
):
_clear_prometheus_registry()
clear_prometheus_registry()
try:
await _run_success_event(additional_headers)
for metric_name in KEY_AND_TEAM_RATE_LIMIT_METRICS:
assert _collected_samples(metric_name) == {}, metric_name
finally:
_clear_prometheus_registry()
clear_prometheus_registry()

15
tests/test_litellm/integrations/test_prometheus_service_tier_label.py
View file

@ -23,6 +23,7 @@ from litellm.types.integrations.prometheus import (
UserAPIKeyLabelNames,
UserAPIKeyLabelValues,
)
from tests._prometheus_helpers import clear_prometheus_registry
SERVICE_TIER_METRICS = [
"litellm_llm_api_latency_metric",
@ -32,16 +33,6 @@ SERVICE_TIER_METRICS = [
]
def _clear_prometheus_registry() -> None:
from prometheus_client import REGISTRY
for collector in list(REGISTRY._collector_to_names.keys()):
try:
REGISTRY.unregister(collector)
except Exception:
pass
def _collected_samples(metric_name: str):
from prometheus_client import REGISTRY
@ -211,7 +202,7 @@ async def test_success_event_emits_service_tier_on_latency_and_spend_metrics():
"end_time": now,
}
_clear_prometheus_registry()
clear_prometheus_registry()
try:
logger = PrometheusLogger()
await logger.async_log_success_event(kwargs, None, now, now)
@ -229,7 +220,7 @@ async def test_success_event_emits_service_tier_on_latency_and_spend_metrics():
f"{sorted({sample.labels.get('service_tier') for sample in samples})}"
)
finally:
_clear_prometheus_registry()
clear_prometheus_registry()
def test_allowlist_covers_every_modeled_service_tier():

91
tests/test_litellm/rust_bridge/test_lifecycle.py Normal file
View file

@ -0,0 +1,91 @@
from __future__ import annotations
import asyncio
from collections.abc import Callable
from typing import Final
import pytest
from litellm.rust_bridge._lifecycle import NativeOutcome, drive_async, drive_sync
class _Machine:
def __init__(self, replace: bool = False) -> None:
self.outcome: NativeOutcome | None = None
self.replace: Final = replace
def complete(self) -> bool | None:
if self.outcome is None:
return None
return self.outcome is NativeOutcome.SUCCESS
def advance(self, outcome: int, logger_available: bool, has_fallbacks: bool) -> bool:
self.outcome = NativeOutcome(outcome)
return self.replace
class _Host:
def __init__(self, invoke: Callable[[], tuple[bool, object]], replace: bool = False) -> None:
self.machine: Final = _Machine(replace)
self._invoke: Final = invoke
self.error: BaseException | None = None
def invoke(self) -> tuple[bool, object]:
return self._invoke()
def advance(self, outcome: NativeOutcome, error: BaseException | None = None) -> None:
if self.machine.advance(outcome, True, False):
self.error = error
def result(self) -> object:
if self.machine.complete():
return "complete"
if self.error is None:
raise RuntimeError("missing test error")
raise self.error
def test_drive_sync_returns_terminal_result() -> None:
host: Final = _Host(lambda: (False, None))
assert drive_sync(host) == "complete"
assert host.machine.outcome is NativeOutcome.SUCCESS
def test_drive_sync_replaces_an_ordinary_failure() -> None:
failure: Final = ValueError("failed")
host: Final = _Host(lambda: (_ for _ in ()).throw(failure), replace=True)
with pytest.raises(ValueError, match="failed") as raised:
drive_sync(host)
assert raised.value is failure
assert host.machine.outcome is NativeOutcome.FAILURE
def test_drive_sync_classifies_base_exception_as_abort() -> None:
class Abort(BaseException):
pass
failure: Final = Abort("aborted")
host: Final = _Host(lambda: (_ for _ in ()).throw(failure), replace=True)
with pytest.raises(Abort, match="aborted"):
drive_sync(host)
assert host.machine.outcome is NativeOutcome.ABORT
@pytest.mark.asyncio
async def test_drive_async_awaits_the_selected_operation() -> None:
completed: Final = asyncio.Event()
async def operation() -> None:
await asyncio.sleep(0)
completed.set()
host: Final = _Host(lambda: (True, operation()))
assert await drive_async(host) == "complete"
assert completed.is_set()
assert host.machine.outcome is NativeOutcome.SUCCESS

11
tests/test_litellm_rust/recording_server.py
View file

@ -3,6 +3,7 @@ import copy
import threading
import time
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass, field
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
@ -44,8 +45,8 @@ class RecordingServer:
self.responses.append(response)
@pytest.fixture
def recording_server() -> Iterator[RecordingServer]:
@contextmanager
def recording_service() -> Iterator[RecordingServer]:
requests: list[RecordedRequest] = []
responses: list[ResponseSpec] = []
@ -105,3 +106,9 @@ def recording_server() -> Iterator[RecordingServer]:
if recording_server.expected_requests is not None:
assert len(recording_server.requests) == recording_server.expected_requests
assert recording_server.responses == []
@pytest.fixture
def recording_server() -> Iterator[RecordingServer]:
with recording_service() as server:
yield server