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

wip

Browse source
This commit is contained in:
Yujong Lee 2026-09-08 13:02:08 -07:00
parent f35ce6ddb8
commit eabd1c187d
24 changed files with 743 additions and 278 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -7,8 +7,8 @@ use super::client::http_client;
use super::prepare::prepare_provider_request;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
ResolvedChatCompletionsRequest,
ChatBodySnapshot, ChatCompletionsResponse, ChatEndpoint, ProviderChatCompletionsRequest,
ProviderChatResponseData, ResolvedChatCompletionsRequest, SettledChatRequest,
};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
@ -22,12 +22,32 @@ pub(super) async fn execute_chat_completions_provider_call(
))
})?;
let headers = signed_headers(&request, &body).await?;
execute_settled_request(
ChatBodySnapshot {
endpoint: ChatEndpoint {
model: request.model,
config: request.config,
url: request.url,
timeout: request.timeout,
},
body,
}
.settle_headers(headers),
)
.await
}
let mut request_builder = http_client().post(&request.url).body(body);
for (key, value) in &headers {
pub(super) async fn execute_settled_request(
request: SettledChatRequest,
) -> Result<ChatCompletionsResponse, Error> {
let endpoint = request.snapshot.endpoint;
let mut request_builder = http_client()
.post(&endpoint.url)
.body(request.snapshot.body);
for (key, value) in &request.headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
if let Some(duration) = endpoint.timeout {
request_builder = request_builder.timeout(duration);
}
@ -58,9 +78,9 @@ pub(super) async fn execute_chat_completions_provider_call(
let body: Value = serde_json::from_str(&text).map_err(|err| {
Error::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
})?;
request
endpoint
.config
.transform_response(&request.model, ProviderChatResponseData { body })
.transform_response(&endpoint.model, ProviderChatResponseData { body })
.map_err(as_response_error)
}

7
litellm-rust/crates/core/src/chat_completions/lifecycle.rs
View file

@ -67,7 +67,6 @@ impl LifecycleRoute for ChatCompletionsRoute {
program: CallProgram::new(ProgramOptions {
asynchronous: options.asynchronous,
internal_call: options.internal_call,
pre_call: false,
}),
}))
}
@ -142,7 +141,8 @@ mod tests {
false,
vec![
Operation::Setup,
Operation::Prepare,
Operation::BuildRequest,
Operation::PreCall,
Operation::Send,
Operation::SyncSuccess,
Operation::Restore,
@ -153,7 +153,8 @@ mod tests {
vec![
Operation::Setup,
Operation::DeploymentPre,
Operation::Prepare,
Operation::BuildRequest,
Operation::PreCall,
Operation::Send,
Operation::DeploymentSuccess,
Operation::AsyncSuccess,

17
litellm-rust/crates/core/src/chat_completions/mod.rs
View file

@ -20,6 +20,7 @@ pub mod types;
use serde_json::{Map, Value};
use handler::execute_chat_completions_provider_call;
pub use prepare::prepare_callback_request;
use prepare::{parse_messages, resolve_provider_config, resolve_request};
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
@ -39,9 +40,23 @@ pub async fn chat_completions(
pub async fn chat_completions_with_terminal(
request: ChatCompletionsRequest<'_>,
context: CallLifecycleContext,
) -> ExecutedCall<ChatCompletionsResponse, Error> {
with_terminal(chat_completions(request), context).await
}
pub async fn execute_settled_with_terminal(
request: types::SettledChatRequest,
context: CallLifecycleContext,
) -> ExecutedCall<ChatCompletionsResponse, Error> {
with_terminal(handler::execute_settled_request(request), context).await
}
async fn with_terminal(
call: impl std::future::Future<Output = Result<ChatCompletionsResponse, Error>>,
context: CallLifecycleContext,
) -> ExecutedCall<ChatCompletionsResponse, Error> {
let start_time = epoch_seconds();
match chat_completions(request).await {
match call.await {
Ok(response) => {
let usage = Usage {
prompt_tokens: response.usage.prompt_tokens,

51
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -143,3 +143,54 @@ pub(super) fn prepare_provider_request(
timeout: request.timeout,
})
}
pub async fn prepare_callback_request(
request: ChatCompletionsRequest<'_>,
) -> Result<super::types::ChatCallbackRequest, Error> {
use super::transformation::PreCallBody;
use super::types::{ChatBodySnapshot, ChatCallbackRequest, ChatEndpoint};
let prepared = prepare_provider_request(resolve_request(request)?)?;
let endpoint = ChatEndpoint {
model: prepared.model.clone(),
config: prepared.config,
url: prepared.url.clone(),
timeout: prepared.timeout,
};
match prepared.config.pre_call_body() {
PreCallBody::Live => {
let mut generated = prepared
.body
.as_object()
.cloned()
.ok_or_else(|| Error::InvalidRequest("chat body must be an object".into()))?;
let parameter_fields = prepared
.optional_params
.keys()
.filter(|name| generated.contains_key(*name))
.cloned()
.collect::<Vec<_>>();
for name in &parameter_fields {
generated.remove(name);
}
Ok(ChatCallbackRequest::Live {
endpoint,
generated,
parameter_fields,
headers: prepared.upstream_headers,
})
}
PreCallBody::Serialized => {
let logging_body = serde_json::to_string(&prepared.body).map_err(|error| {
Error::InvalidRequest(format!("could not encode chat request: {error}"))
})?;
let body = logging_body.as_bytes().to_vec();
let headers = super::handler::signed_headers(&prepared, &body).await?;
Ok(ChatCallbackRequest::Serialized {
snapshot: ChatBodySnapshot { endpoint, body },
logging_body,
headers,
})
}
}
}

10
litellm-rust/crates/core/src/chat_completions/transformation.rs
View file

@ -27,11 +27,21 @@ pub struct Unsupported(pub &'static str);
pub const STREAM_PARAM: &str = "stream";
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum PreCallBody {
Live,
Serialized,
}
/// Message fields that carry no meaning for the upstream body, so their
/// presence does not make a request untranslatable.
const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"];
pub trait ChatCompletionsProviderConfig: Sync {
fn pre_call_body(&self) -> PreCallBody {
PreCallBody::Live
}
fn complete_url(
&self,
api_base: Option<&str>,

62
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -52,6 +52,68 @@ pub struct ProviderChatRequestData {
pub body: Value,
}
pub struct ChatEndpoint {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) url: String,
pub(super) timeout: Option<Duration>,
}
pub struct ChatBodySnapshot {
pub(super) endpoint: ChatEndpoint,
pub(super) body: Vec<u8>,
}
pub struct SettledChatRequest {
pub(super) snapshot: ChatBodySnapshot,
pub(super) headers: Vec<(String, String)>,
}
pub enum ChatCallbackRequest {
Live {
endpoint: ChatEndpoint,
generated: Map<String, Value>,
parameter_fields: Vec<String>,
headers: Vec<(String, String)>,
},
Serialized {
snapshot: ChatBodySnapshot,
logging_body: String,
headers: Vec<(String, String)>,
},
}
impl ChatEndpoint {
pub fn capture_body(self, body: Value) -> Result<ChatBodySnapshot, crate::Error> {
let body = serde_json::to_vec(&body).map_err(|error| {
crate::Error::InvalidRequest(format!("could not encode chat request: {error}"))
})?;
Ok(ChatBodySnapshot {
endpoint: self,
body,
})
}
}
impl ChatBodySnapshot {
pub fn settle_headers(self, headers: Vec<(String, String)>) -> SettledChatRequest {
SettledChatRequest {
snapshot: self,
headers,
}
}
}
impl SettledChatRequest {
pub fn body(&self) -> &[u8] {
&self.snapshot.body
}
pub fn headers(&self) -> &[(String, String)] {
&self.headers
}
}
/// The raw provider response body handed back to a config for normalization.
pub struct ProviderChatResponseData {
pub body: Value,

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

@ -110,7 +110,6 @@ impl LifecycleRoute for OcrRoute {
program: CallProgram::new(ProgramOptions {
asynchronous: options.asynchronous,
internal_call: options.internal_call,
pre_call: true,
}),
identity: Identity {
requested_model: admission.model.clone(),

27
litellm-rust/crates/core/src/lifecycle/program.rs
View file

@ -7,7 +7,7 @@ use super::{
pub enum Operation {
Setup,
DeploymentPre,
Prepare,
BuildRequest,
PreCall,
Send,
DeploymentSuccess,
@ -53,7 +53,6 @@ pub struct Transition {
pub struct ProgramOptions {
pub asynchronous: bool,
pub internal_call: bool,
pub pre_call: bool,
}
#[derive(Debug)]
@ -125,14 +124,14 @@ impl CallProgram {
(DeploymentFailure, _) => failure,
(_, Outcome::Abort) => Restore,
(SyncFailure | AsyncFailure, Outcome::Failure) => Restore,
(Prepare | PreCall | Send, Outcome::Failure) if self.options.asynchronous => {
(BuildRequest | PreCall | Send, Outcome::Failure) if self.options.asynchronous => {
DeploymentFailure
}
(_, Outcome::Failure) => failure,
(Setup, Outcome::Success) if self.options.asynchronous => DeploymentPre,
(Setup | DeploymentPre, Outcome::Success) => Prepare,
(Prepare, Outcome::Success) if self.options.pre_call => PreCall,
(Prepare | PreCall, Outcome::Success) => Send,
(Setup | DeploymentPre, Outcome::Success) => BuildRequest,
(BuildRequest, Outcome::Success) => PreCall,
(PreCall, Outcome::Success) => Send,
(Send, Outcome::Success) if self.options.asynchronous => DeploymentSuccess,
(Send, Outcome::Success) => SyncSuccess,
(DeploymentSuccess, Outcome::Success) => {
@ -169,7 +168,8 @@ impl CallProgram {
pub fn actions_for(operation: Operation) -> &'static [ActionBinding] {
match operation {
Operation::Prepare | Operation::Send => &PROVIDER_ACTION,
Operation::BuildRequest => &REQUEST_BUILD_ACTION,
Operation::Send => &PROVIDER_ACTION,
Operation::PreCall => &PRE_CALL_ACTION,
Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => {
&FAILURE_ACTION
@ -180,6 +180,14 @@ pub fn actions_for(operation: Operation) -> &'static [ActionBinding] {
}
}
const REQUEST_BUILD_ACTION: [ActionBinding; 1] = [ActionBinding {
kind: ActionKind::RequestBuild,
delivery: Delivery::InlineDirect,
on_result: ResultPolicy::Replace,
on_error: FailurePolicy::Propagate,
owner: Owner::Core,
}];
const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding {
kind: ActionKind::ProviderCall,
delivery: Delivery::InlineAwaited,
@ -236,7 +244,6 @@ mod tests {
let mut before = CallProgram::new(ProgramOptions {
asynchronous: false,
internal_call: false,
pre_call: false,
});
let failure = before.advance(Outcome::Failure, observations()).unwrap();
assert_eq!(failure.commitment, Commitment::Replayable);
@ -245,10 +252,10 @@ mod tests {
let mut provider = CallProgram::new(ProgramOptions {
asynchronous: false,
internal_call: false,
pre_call: false,
});
provider.advance(Outcome::Success, observations()).unwrap();
provider.advance(Outcome::Success, observations()).unwrap();
provider.advance(Outcome::Success, observations()).unwrap();
let failure = provider.advance(Outcome::Failure, observations()).unwrap();
assert_eq!(failure.commitment, Commitment::ProviderStarted);
assert_eq!(failure.failure_stage, Some(FailureStage::ProviderCall));
@ -256,11 +263,11 @@ mod tests {
let mut after = CallProgram::new(ProgramOptions {
asynchronous: false,
internal_call: false,
pre_call: false,
});
after.advance(Outcome::Success, observations()).unwrap();
after.advance(Outcome::Success, observations()).unwrap();
after.advance(Outcome::Success, observations()).unwrap();
after.advance(Outcome::Success, observations()).unwrap();
let failure = after.advance(Outcome::Failure, observations()).unwrap();
assert_eq!(failure.commitment, Commitment::ResponseReceived);
assert_eq!(

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

@ -2,6 +2,7 @@ use serde::Serialize;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
pub enum ActionKind {
RequestBuild,
RequestPolicy,
ProviderCall,
Deployment,

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

@ -49,7 +49,6 @@ impl LifecycleRoute for MessagesRoute {
program: CallProgram::new(ProgramOptions {
asynchronous: options.asynchronous,
internal_call: options.internal_call,
pre_call: false,
}),
}))
}

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

@ -20,7 +20,7 @@ use std::sync::Arc;
use crate::lifecycle::StreamingCall;
pub use handler::execute_prepared_messages_provider_call;
pub use prepare::prepare_provider_request;
pub use prepare::{prepare_endpoint, prepare_provider_request};
use types::{AnthropicMessagesResponse, MessagesRequest};
pub async fn messages(request: MessagesRequest) -> Result<AnthropicMessagesResponse, Error> {

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

@ -3,12 +3,41 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use super::types::{MessagesRequest, ProviderMessagesRequest};
use super::types::{MessagesEndpoint, MessagesOptions, MessagesRequest, ProviderMessagesRequest};
use serde_json::{Map, Value};
pub fn prepare_provider_request(
request: MessagesRequest,
) -> Result<ProviderMessagesRequest, Error> {
let endpoint = prepare_endpoint(MessagesOptions {
model: request.model,
api_key: request.api_key,
api_base: request.api_base,
custom_llm_provider: request.custom_llm_provider,
extra_headers: request.extra_headers,
timeout: request.timeout,
})?;
let typed_request = serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
})?;
let transformed = endpoint.config.transform_request(typed_request)?;
let body = serde_json::to_value(transformed).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize Anthropic messages request: {err}"
))
})?;
Ok(ProviderMessagesRequest {
provider: endpoint.provider,
model: endpoint.model,
config: endpoint.config,
url: endpoint.url,
body,
upstream_headers: endpoint.headers,
timeout: endpoint.timeout,
})
}
pub fn prepare_endpoint(request: MessagesOptions) -> Result<MessagesEndpoint, Error> {
let provider_info =
get_custom_llm_provider(&request.model, request.custom_llm_provider.as_deref())
.or_else(|| {
@ -39,25 +68,14 @@ pub fn prepare_provider_request(
&env_lookup,
)?;
let typed_request = serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
})?;
let transformed = config.transform_request(typed_request)?;
let body = serde_json::to_value(transformed).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize Anthropic messages request: {err}"
))
})?;
let url = config.complete_url(request.api_base.as_deref(), &model, &env_lookup)?;
Ok(ProviderMessagesRequest {
Ok(MessagesEndpoint {
provider: provider.to_string(),
model,
config,
url,
body,
upstream_headers: headers,
headers,
timeout: request.timeout,
})
}

77
litellm-rust/crates/core/src/messages/types.rs
View file

@ -16,15 +16,82 @@ pub struct MessagesRequest {
}
pub struct ProviderMessagesRequest {
pub provider: String,
pub model: String,
pub(super) provider: String,
pub(super) model: String,
pub(super) config: &'static dyn AnthropicMessagesProviderConfig,
pub url: String,
pub body: Value,
pub upstream_headers: Vec<(String, String)>,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) timeout: Option<Duration>,
}
impl ProviderMessagesRequest {
pub fn body(&self) -> &Value {
&self.body
}
pub fn headers(&self) -> &[(String, String)] {
&self.upstream_headers
}
}
pub struct MessagesOptions {
pub model: String,
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>,
}
pub struct MessagesEndpoint {
pub(super) provider: String,
pub(super) model: String,
pub(super) config: &'static dyn AnthropicMessagesProviderConfig,
pub(super) url: String,
pub(super) headers: Vec<(String, String)>,
pub(super) timeout: Option<Duration>,
}
pub struct MessagesBodySnapshot(Value);
impl MessagesEndpoint {
pub fn url(&self) -> &str {
&self.url
}
pub fn headers(&self) -> &[(String, String)] {
&self.headers
}
pub fn capture_buffered_body(
&self,
mut body: Value,
) -> Result<MessagesBodySnapshot, crate::Error> {
let object = body.as_object_mut().ok_or_else(|| {
crate::Error::InvalidRequest("messages body must be an object".into())
})?;
object.remove("stream");
Ok(MessagesBodySnapshot(body))
}
pub fn settle(
self,
body: MessagesBodySnapshot,
headers: Vec<(String, String)>,
) -> ProviderMessagesRequest {
ProviderMessagesRequest {
provider: self.provider,
model: self.model,
config: self.config,
url: self.url,
body: body.0,
upstream_headers: headers,
timeout: self.timeout,
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum SystemPrompt {

4
litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs
View file

@ -105,6 +105,10 @@ fn has_blank_text(message: &ChatMessage) -> bool {
}
impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
fn pre_call_body(&self) -> crate::chat_completions::transformation::PreCallBody {
crate::chat_completions::transformation::PreCallBody::Serialized
}
fn complete_url(
&self,
api_base: Option<&str>,

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

@ -30,8 +30,8 @@ fn operation_binding(
method: "deployment_pre",
awaiting: true,
},
Operation::Prepare => OperationBinding {
method: "prepare",
Operation::BuildRequest => OperationBinding {
method: "build_request",
awaiting: false,
},
Operation::PreCall if supports_pre_call => OperationBinding {
@ -126,7 +126,13 @@ mod tests {
"deployment_pre",
true,
),
(Operation::Prepare, false, false, "prepare", false),
(
Operation::BuildRequest,
false,
false,
"build_request",
false,
),
(Operation::PreCall, false, true, "pre_call", false),
(Operation::Send, false, false, "send_sync", false),
(Operation::Send, true, false, "send", true),

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

@ -6,6 +6,7 @@ mod errors;
#[cfg(feature = "trace-parity")]
mod function_trace;
mod marshal;
mod retained;
mod routes;
use pyo3::prelude::*;

37
litellm-rust/crates/python-bridge/src/retained.rs Normal file
View file

@ -0,0 +1,37 @@
use pyo3::prelude::*;
use pyo3::pyclass::{PyTraverseError, PyVisit};
use pyo3::types::PyDict;
pub(crate) struct RequestRoots {
arguments: Py<PyDict>,
body: Py<PyAny>,
headers: Py<PyAny>,
}
impl RequestRoots {
pub(crate) fn new(arguments: Py<PyDict>, body: Py<PyAny>, headers: Py<PyAny>) -> Self {
Self {
arguments,
body,
headers,
}
}
pub(crate) fn arguments<'py>(&self, py: Python<'py>) -> Bound<'py, PyDict> {
self.arguments.clone_ref(py).into_bound(py)
}
pub(crate) fn body<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> {
self.body.clone_ref(py).into_bound(py)
}
pub(crate) fn headers<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> {
self.headers.clone_ref(py).into_bound(py)
}
pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.arguments)?;
visit.call(&self.body)?;
visit.call(&self.headers)
}
}

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

@ -4,7 +4,7 @@ use litellm_core::chat_completions::lifecycle::{
};
use litellm_core::chat_completions::types::ChatCompletionsRequest;
use litellm_core::chat_completions::{
chat_completions_decline_reason, chat_completions_with_terminal,
chat_completions_decline_reason, execute_settled_with_terminal, prepare_callback_request,
};
use litellm_core::lifecycle::{
CallLifecycleContext, ErrorDisposition, ExecutedCall, Lifecycle, Outcome, TerminalRecord,
@ -20,35 +20,32 @@ 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;
use crate::retained::RequestRoots;
#[pyclass]
struct ChatCompletionsState {
arguments: Option<Py<PyDict>>,
model: Option<String>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
timeout: Option<std::time::Duration>,
roots: Option<RequestRoots>,
pending: Option<PendingChatRequest>,
context: Option<CallLifecycleContext>,
terminal: Option<TerminalRecord>,
}
#[pymethods]
impl ChatCompletionsState {
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.arguments)?;
visit.call(&self.body)?;
visit.call(&self.headers)
if let Some(roots) = &self.roots {
roots.traverse(&visit)?;
}
Ok(())
}
fn __clear__(slf: &Bound<'_, Self>) {
let roots = {
let mut state = slf.borrow_mut();
(
state.arguments.take(),
state.body.take(),
state.headers.take(),
state.roots.take(),
state.pending.take(),
state.context.take(),
state.terminal.take(),
)
};
@ -164,33 +161,105 @@ fn invoke(
crate::driver::invoke(py, operation, asynchronous, false, "chat completions", host)
}
enum PendingChatRequest {
Live(litellm_core::chat_completions::types::ChatEndpoint),
Serialized(litellm_core::chat_completions::types::ChatBodySnapshot),
}
#[pyfunction]
fn prepare(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
) -> PyResult<Py<ChatCompletionsState>> {
use litellm_core::chat_completions::types::ChatCallbackRequest;
let bag = arguments.bind(py);
let admission = admission(bag)?;
let api_key = scalar(bag, "api_key")?;
let api_base = scalar(bag, "api_base")?;
let headers = bag
.get_item("extra_headers")?
.filter(|value| !value.is_none());
let timeout = optional_timeout(
bag.get_item("timeout_seconds")?
.filter(|value| !value.is_none())
.map(|value| value.extract::<f64>())
.transpose()?,
)?;
let complete_input = PyDict::new(py);
complete_input.set_item("model", &admission.model)?;
complete_input.set_item("messages", bag.get_item("messages")?)?;
if let Some(optional_params) = bag.get_item("optional_params")? {
complete_input.call_method1("update", (optional_params,))?;
let context = CallLifecycleContext::new(
"chat_completion",
&admission.model,
admission.custom_llm_provider.as_deref().unwrap_or_default(),
scalar(bag, "litellm_call_id")?.unwrap_or_default(),
);
let extra_headers = optional_map(bag, "extra_headers")?;
let prepared = run_sync_value(
py,
async move {
prepare_callback_request(ChatCompletionsRequest {
model: &admission.model,
messages: admission.messages,
optional_params: admission.optional_params,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: admission.custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
},
core_error_to_pyerr,
)?;
let (body, pending, header_values) = match prepared {
ChatCallbackRequest::Live {
endpoint,
generated,
parameter_fields,
headers,
} => {
let body = PyDict::new(py);
for (name, value) in generated {
body.set_item(name, to_py(py, &value)?)?;
}
if let Some(params) = bag.get_item("optional_params")? {
for name in parameter_fields {
body.set_item(&name, params.get_item(&name)?)?;
}
}
(body.into_any(), PendingChatRequest::Live(endpoint), headers)
}
ChatCallbackRequest::Serialized {
snapshot,
logging_body,
headers,
} => (
logging_body.into_pyobject(py)?.into_any(),
PendingChatRequest::Serialized(snapshot),
headers,
),
};
let headers = PyDict::new(py);
for (name, value) in header_values {
headers.set_item(name, value)?;
}
let headers = match &pending {
PendingChatRequest::Live(_) => {
match bag
.get_item("extra_headers")?
.filter(|value| !value.is_none())
{
Some(original) => {
original.call_method1("update", (&headers,))?;
original
}
None => headers.into_any(),
}
}
PendingChatRequest::Serialized(_) => py
.import("botocore.awsrequest")?
.getattr("HeadersDict")?
.call1((headers,))?,
};
let additional = PyDict::new(py);
additional.set_item(COMPLETE_INPUT_DICT, &complete_input)?;
additional.set_item(COMPLETE_INPUT_DICT, &body)?;
additional.set_item(API_BASE, bag.get_item("api_base")?)?;
additional.set_item(HEADERS, &headers)?;
let kwargs = PyDict::new(py);
@ -203,99 +272,60 @@ fn prepare(
Py::new(
py,
ChatCompletionsState {
arguments: Some(arguments),
model: Some(admission.model),
body: Some(complete_input.unbind()),
headers: headers.map(Bound::unbind),
api_key,
api_base,
custom_llm_provider: admission.custom_llm_provider,
timeout,
roots: Some(RequestRoots::new(
arguments,
body.unbind(),
headers.unbind(),
)),
pending: Some(pending),
context: Some(context),
terminal: None,
},
)
}
struct OwnedRequest {
model: String,
messages: Value,
optional_params: Map<String, Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Map<String, Value>>,
timeout: Option<std::time::Duration>,
call_id: String,
request: litellm_core::chat_completions::types::SettledChatRequest,
context: CallLifecycleContext,
}
fn take_request(py: Python<'_>, state: &Py<ChatCompletionsState>) -> PyResult<OwnedRequest> {
let (arguments, body, headers) = {
let state = state.borrow(py);
let arguments = state
.arguments
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("chat completions state was cleared"))?
.clone_ref(py);
let body = state
.body
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("chat completions body was cleared"))?
.clone_ref(py);
let headers = state.headers.as_ref().map(|value| value.clone_ref(py));
(arguments, body, headers)
};
let call_id = scalar(arguments.bind(py), "litellm_call_id")?.unwrap_or_default();
let messages = value(body.bind(py), "messages")?;
let optional_params = body
.bind(py)
.iter()
.filter_map(|(key, value)| match key.extract::<String>() {
Ok(key) if key == "model" || key == "messages" => None,
Ok(key) => Some(from_py(&value).map(|value| (key, value))),
Err(error) => Some(Err(error)),
})
.collect::<PyResult<Map<String, Value>>>()?;
let extra_headers = headers
.as_ref()
.map(|value| from_py(value.bind(py)))
.transpose()?;
let mut state = state.borrow_mut(py);
Ok(OwnedRequest {
model: state
.model
let (pending, context, body, headers) = {
let mut state = state.borrow_mut(py);
let pending = state.pending.take().ok_or_else(|| {
PyRuntimeError::new_err("chat completions request was already sent or cleared")
})?;
let context = state
.context
.take()
.ok_or_else(|| PyRuntimeError::new_err("chat completions request was already sent"))?,
messages,
optional_params,
api_key: state.api_key.take(),
api_base: state.api_base.take(),
custom_llm_provider: state.custom_llm_provider.take(),
extra_headers,
timeout: state.timeout.take(),
call_id,
.ok_or_else(|| PyRuntimeError::new_err("chat completions context was cleared"))?;
let roots = state
.roots
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("chat completions roots were cleared"))?;
(pending, context, roots.body(py), roots.headers(py))
};
let snapshot = match pending {
PendingChatRequest::Live(endpoint) => endpoint
.capture_body(from_py(&body)?)
.map_err(core_error_to_pyerr)?,
PendingChatRequest::Serialized(snapshot) => snapshot,
};
let headers = headers
.call_method0("items")?
.try_iter()?
.map(|item| item?.extract::<(String, String)>())
.collect::<PyResult<_>>()?;
Ok(OwnedRequest {
request: snapshot.settle_headers(headers),
context,
})
}
async fn execute(
request: OwnedRequest,
) -> ExecutedCall<litellm_core::chat_completions::types::ChatCompletionsResponse, Error> {
let provider = request.custom_llm_provider.clone().unwrap_or_default();
let context =
CallLifecycleContext::new("chat_completion", &request.model, provider, request.call_id);
chat_completions_with_terminal(
ChatCompletionsRequest {
model: &request.model,
messages: request.messages,
optional_params: request.optional_params,
api_key: request.api_key.as_deref(),
api_base: request.api_base.as_deref(),
custom_llm_provider: request.custom_llm_provider.as_deref(),
extra_headers: request.extra_headers,
timeout: request.timeout,
},
context,
)
.await
execute_settled_with_terminal(request.request, request.context).await
}
fn store_result(
@ -441,9 +471,8 @@ mod tests {
to_py(
py,
&(
request.messages,
request.optional_params,
request.extra_headers,
serde_json::from_slice::<Value>(request.request.body()).unwrap(),
request.request.headers(),
),
)
}
@ -476,10 +505,11 @@ class Logger:
self.body = view['complete_input_dict']
self.headers = view['headers']
assert kwargs['input'] is messages
assert self.body['messages'] is messages
assert self.body['stop'] is stops
assert self.body['messages'] is not messages
assert self.body['stop_sequences'] is stops
assert self.headers is headers
messages[0]['content'] = 'edited'
self.body['messages'][0]['content'][0]['text'] = 'body edit'
stops.append('second')
self.headers['x-hook'] = 'edited'
view['complete_input_dict'] = {'replacement': True}
@ -491,15 +521,15 @@ headers = {}
opaque = Opaque()
logger = Logger()
arguments = dict(model='claude-opus-5', messages=messages,
optional_params={'max_tokens': 16, 'stop': stops},
optional_params={'max_tokens': 16, 'stop_sequences': stops},
extra_headers=headers, api_key='test',
custom_llm_provider='anthropic', opaque=opaque,
litellm_logging_obj=logger)
state = native.prepare(arguments, logger)
wire_messages, wire_params, wire_headers = native.snapshot(state)
assert wire_messages[0]['content'] == 'edited'
assert wire_params['stop'] == ['first', 'second']
assert wire_headers == {'x-hook': 'edited'}
wire_body, wire_headers = native.snapshot(state)
assert wire_body['messages'][0]['content'][0]['text'] == 'body edit'
assert wire_body['stop_sequences'] == ['first', 'second']
assert dict(wire_headers)['x-hook'] == 'edited'
arguments['cycle'] = state
logger.body['cycle'] = state
headers['cycle'] = state

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

@ -1,9 +1,11 @@
use litellm_core::lifecycle::FailureStage;
use litellm_core::lifecycle::{ErrorDisposition, Lifecycle, Outcome};
use litellm_core::messages::lifecycle::{MessagesRoute, Observations, Operation, Options, machine};
use litellm_core::messages::types::{MessagesRequest, ProviderMessagesRequest};
use litellm_core::messages::{execute_prepared_messages_provider_call, prepare_provider_request};
use litellm_python_interop::{Pythonized, from_py, run_async_value, run_sync_value, to_py};
use litellm_core::messages::types::{
MessagesBodySnapshot, MessagesEndpoint, MessagesOptions, ProviderMessagesRequest,
};
use litellm_core::messages::{execute_prepared_messages_provider_call, prepare_endpoint};
use litellm_python_interop::{Pythonized, from_py, run_async_value, run_sync_value};
use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::pyclass::{PyTraverseError, PyVisit};
@ -14,32 +16,27 @@ 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;
use crate::retained::RequestRoots;
#[pyclass]
struct MessagesState {
arguments: Option<Py<PyDict>>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
prepared: Option<ProviderMessagesRequest>,
roots: Option<RequestRoots>,
prepared: Option<(MessagesEndpoint, MessagesBodySnapshot)>,
}
#[pymethods]
impl MessagesState {
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.arguments)?;
visit.call(&self.body)?;
visit.call(&self.headers)
if let Some(roots) = &self.roots {
roots.traverse(&visit)?;
}
Ok(())
}
fn __clear__(slf: &Bound<'_, Self>) {
let roots = {
let mut state = slf.borrow_mut();
(
state.arguments.take(),
state.body.take(),
state.headers.take(),
state.prepared.take(),
)
(state.roots.take(), state.prepared.take())
};
drop(roots);
}
@ -53,10 +50,7 @@ fn scalar(arguments: &Bound<'_, PyDict>, name: &str) -> PyResult<Option<String>>
.transpose()
}
fn decode_request(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<MessagesRequest> {
let body = arguments
.get_item(pyo3::intern!(py, "body"))?
.ok_or_else(|| PyValueError::new_err("messages requires body"))?;
fn decode_options(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<MessagesOptions> {
let timeout = optional_timeout(
arguments
.get_item(pyo3::intern!(py, "timeout_seconds"))?
@ -64,8 +58,7 @@ fn decode_request(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<Mes
.map(|value| value.extract::<f64>())
.transpose()?,
)?;
Ok(MessagesRequest {
body: from_py(&body)?,
Ok(MessagesOptions {
model: scalar(arguments, "model")?
.ok_or_else(|| PyValueError::new_err("messages requires model"))?,
api_key: scalar(arguments, "api_key")?,
@ -156,20 +149,24 @@ fn prepare(
logging: Py<PyAny>,
) -> PyResult<Py<MessagesState>> {
let bag = arguments.bind(py);
let request = decode_request(py, bag)?;
let prepared = py
.detach(|| prepare_provider_request(request))
let options = decode_options(py, bag)?;
let endpoint = py
.detach(|| prepare_endpoint(options))
.map_err(core_error_to_pyerr)?;
let body = to_py(py, &prepared.body)?
.into_bound(py)
let body = bag
.get_item("body")?
.ok_or_else(|| PyValueError::new_err("messages requires body"))?
.cast_into::<PyDict>()?;
let snapshot = endpoint
.capture_buffered_body(from_py(body.as_any())?)
.map_err(core_error_to_pyerr)?;
let headers = PyDict::new(py);
for (name, value) in &prepared.upstream_headers {
for (name, value) in endpoint.headers() {
headers.set_item(name, value)?;
}
let additional = PyDict::new(py);
additional.set_item(COMPLETE_INPUT_DICT, &body)?;
additional.set_item(API_BASE, &prepared.url)?;
additional.set_item(API_BASE, endpoint.url())?;
additional.set_item(HEADERS, &headers)?;
let kwargs = PyDict::new(py);
let serialized = py.import("json")?.call_method1("dumps", (&body,))?;
@ -185,39 +182,34 @@ fn prepare(
Py::new(
py,
MessagesState {
arguments: Some(arguments),
body: Some(body.unbind()),
headers: Some(headers.unbind()),
prepared: Some(prepared),
roots: Some(RequestRoots::new(
arguments,
body.unbind().into_any(),
headers.unbind().into_any(),
)),
prepared: Some((endpoint, snapshot)),
},
)
}
fn take_request(py: Python<'_>, state: &Py<MessagesState>) -> PyResult<ProviderMessagesRequest> {
let (mut prepared, body, headers) = {
let ((endpoint, snapshot), headers) = {
let mut state = state.borrow_mut(py);
let prepared = state.prepared.take().ok_or_else(|| {
PyRuntimeError::new_err("messages request was already sent or cleared")
})?;
let body = state
.body
let roots = state
.roots
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("messages body was cleared"))?
.clone_ref(py);
let headers = state
.headers
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("messages headers were cleared"))?
.clone_ref(py);
(prepared, body, headers)
.ok_or_else(|| PyRuntimeError::new_err("messages roots were cleared"))?;
(prepared, roots.headers(py))
};
prepared.body = from_py(body.bind(py).as_any())?;
prepared.upstream_headers = headers
.bind(py)
let headers = headers
.cast::<PyDict>()?
.iter()
.map(|(name, value)| Ok((name.extract()?, value.extract()?)))
.collect::<PyResult<_>>()?;
Ok(prepared)
Ok(endpoint.settle(snapshot, headers))
}
fn validate_arguments(arguments: &Bound<'_, PyDict>) -> PyResult<()> {
@ -328,6 +320,72 @@ pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
mod tests {
use super::*;
#[pyfunction]
fn snapshot(py: Python<'_>, state: Py<MessagesState>) -> PyResult<Py<PyAny>> {
let request = take_request(py, &state)?;
litellm_python_interop::to_py(py, &(request.body(), request.headers()))
}
#[test]
fn callback_aliases_survive_snapshot_and_roots_are_collectible() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "messages_test").unwrap();
module.add_function(wrap_pyfunction!(prepare, &module).unwrap()).unwrap();
module.add_function(wrap_pyfunction!(snapshot, &module).unwrap()).unwrap();
let globals = PyDict::new(py);
globals.set_item("native", module).unwrap();
py.run(c"
import gc
import weakref
class Opaque:
pass
class Logger:
def pre_call(self, **kwargs):
view = kwargs['additional_args']
self.body = view['complete_input_dict']
self.headers = view['headers']
assert self.body is body
assert self.body['messages'][0] is message
message['content'] = 'changed'
self.headers['x-hook'] = 'changed'
view['complete_input_dict'] = {'replacement': True}
view['headers'] = {'replacement': 'true'}
message = {'role': 'user', 'content': 'original'}
body = {'model': 'model', 'messages': [message], 'max_tokens': 16}
opaque = Opaque()
logger = Logger()
arguments = dict(model='model', body=body, api_key='test',
custom_llm_provider='anthropic', opaque=opaque,
litellm_logging_obj=logger)
state = native.prepare(arguments, logger)
wire_body, wire_headers = native.snapshot(state)
assert wire_body['messages'][0]['content'] == 'original'
assert body['messages'][0]['content'] == 'changed'
assert dict(wire_headers)['x-hook'] == 'changed'
try:
native.snapshot(state)
except RuntimeError:
pass
else:
raise AssertionError('request was sent twice')
arguments['cycle'] = state
logger.body['cycle'] = state
logger.headers['cycle'] = state
alive = weakref.ref(opaque)
del arguments, logger, opaque, body
gc.collect()
assert alive() is not None
del state
gc.collect()
assert alive() is None
", Some(&globals), Some(&globals)).unwrap();
});
}
#[test]
fn decode_state_rejects_invalid_timeouts_without_panicking() {
Python::initialize();
@ -337,7 +395,7 @@ mod tests {
arguments.set_item("model", "model").unwrap();
arguments.set_item("body", PyDict::new(py)).unwrap();
arguments.set_item("timeout_seconds", timeout).unwrap();
let error = match decode_request(py, &arguments) {
let error = match decode_options(py, &arguments) {
Ok(_) => panic!("invalid timeout should fail normally"),
Err(error) => error,
};

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

@ -21,13 +21,12 @@ 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 crate::retained::RequestRoots;
use litellm_python_interop::{run_async_value, run_sync_value};
#[pyclass]
struct OcrState {
arguments: Option<Py<PyDict>>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
roots: Option<RequestRoots>,
logging: Option<Py<PyAny>>,
pre_call: Option<Py<PyDict>>,
endpoint: Option<OcrEndpoint>,
@ -38,9 +37,9 @@ struct OcrState {
#[pymethods]
impl OcrState {
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.arguments)?;
visit.call(&self.body)?;
visit.call(&self.headers)?;
if let Some(roots) = &self.roots {
roots.traverse(&visit)?;
}
visit.call(&self.logging)?;
visit.call(&self.pre_call)
}
@ -49,9 +48,7 @@ impl OcrState {
let roots = {
let mut state = slf.borrow_mut();
(
state.arguments.take(),
state.body.take(),
state.headers.take(),
state.roots.take(),
state.logging.take(),
state.pre_call.take(),
state.endpoint.take(),
@ -342,17 +339,17 @@ fn prepare(
document_projection,
parameter_fields,
} = draft;
let body = to_py(py, &draft_body)?
.into_bound(py)
.cast_into::<PyDict>()?;
match document_projection {
OcrDocumentProjection::RetainedDocument => {
body.set_item(pyo3::intern!(py, "document"), &document)?
let body = PyDict::new(py);
for (name, value) in &draft_body {
match (name.as_str(), document_projection) {
("document", OcrDocumentProjection::RetainedDocument) => {
body.set_item(name, &document)?;
}
("document", OcrDocumentProjection::ShallowCopyDocument) => {
body.set_item(name, document.copy()?)?;
}
_ => body.set_item(name, to_py(py, value)?)?,
}
OcrDocumentProjection::ShallowCopyDocument => {
body.set_item(pyo3::intern!(py, "document"), document.copy()?)?
}
OcrDocumentProjection::Transformed => {}
}
let optional_params = PyDict::new(py);
for &name in parameter_fields {
@ -392,9 +389,11 @@ fn prepare(
Py::new(
py,
OcrState {
arguments: Some(arguments),
body: Some(body.unbind()),
headers: Some(headers.unbind()),
roots: Some(RequestRoots::new(
arguments,
body.unbind().into_any(),
headers.unbind().into_any(),
)),
logging: Some(logging),
pre_call: Some(pre_call.unbind()),
endpoint: Some(endpoint),
@ -435,24 +434,17 @@ fn request(py: Python<'_>, state: &Py<OcrState>) -> PyResult<OcrWireRequest> {
.endpoint
.take()
.ok_or_else(|| PyRuntimeError::new_err("OCR request was already sent or cleared"))?;
let body = state
.body
let roots = state
.roots
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("OCR body was cleared"))?
.clone_ref(py);
let headers = state
.headers
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("OCR headers were cleared"))?
.clone_ref(py);
.ok_or_else(|| PyRuntimeError::new_err("OCR roots were cleared"))?;
let body = roots.body(py);
let headers = roots.headers(py);
(endpoint, body, headers, state.asynchronous)
};
let model = endpoint.model().to_string();
let provider = endpoint.custom_llm_provider().to_string();
let request = endpoint.settle(
header_pairs(headers.bind(py))?,
from_py(body.bind(py).as_any())?,
);
let request = endpoint.settle(header_pairs(headers.cast::<PyDict>()?)?, from_py(&body)?);
Ok((request, model, provider, asynchronous))
}
@ -465,9 +457,13 @@ fn send(py: Python<'_>, state: Py<OcrState>) -> PyResult<Bound<'_, PyAny>> {
let call_id = Python::attach(|py| {
state
.borrow(py)
.arguments
.roots
.as_ref()
.and_then(|arguments| scalar(arguments.bind(py), "litellm_call_id").ok().flatten())
.and_then(|roots| {
scalar(&roots.arguments(py), "litellm_call_id")
.ok()
.flatten()
})
.unwrap_or_default()
});
let executed = run_async_value(
@ -527,9 +523,13 @@ fn send_sync(py: Python<'_>, state: Py<OcrState>) -> PyResult<Py<PyAny>> {
let error_provider = provider.clone();
let call_id = state
.borrow(py)
.arguments
.roots
.as_ref()
.and_then(|arguments| scalar(arguments.bind(py), "litellm_call_id").ok().flatten())
.and_then(|roots| {
scalar(&roots.arguments(py), "litellm_call_id")
.ok()
.flatten()
})
.unwrap_or_default();
let executed = run_sync_value(
py,
@ -754,9 +754,7 @@ mod tests {
let state = Py::new(
py,
OcrState {
arguments: None,
body: None,
headers: None,
roots: None,
logging: None,
pre_call: None,
endpoint: None,
@ -969,16 +967,12 @@ asyncio.run(exercise())
#[pyfunction]
fn snapshot(py: Python<'_>, state: Py<OcrState>) -> PyResult<Py<PyAny>> {
let state = state.borrow(py);
let headers = state
.headers
let roots = state
.roots
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("OCR headers were cleared"))?;
let body = state
.body
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("OCR body was cleared"))?;
let headers = header_pairs(headers.bind(py))?;
let body: serde_json::Value = from_py(body.bind(py).as_any())?;
.ok_or_else(|| PyRuntimeError::new_err("OCR roots were cleared"))?;
let headers = header_pairs(roots.headers(py).cast::<PyDict>()?)?;
let body: serde_json::Value = from_py(&roots.body(py))?;
to_py(py, &(headers, body))
}

3
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2442,7 +2442,6 @@ class BaseLLMHTTPHandler:
from litellm.rust_bridge import messages as rust_messages_bridge
upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"}
rust_response: Final = await rust_messages_bridge.amessages(
arguments=arguments,
request_arguments=request_arguments,
@ -2451,7 +2450,7 @@ class BaseLLMHTTPHandler:
messages=messages,
lifecycle_owner=rust_messages_bridge.LifecycleOwner.WRAPPER,
model=model,
body=upstream_body,
body=request_body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,

53
tests/test_litellm_rust/test_messages_callbacks.py
View file

@ -59,15 +59,18 @@ async def test_messages_pre_call_receives_expected_provider_request(messages_ser
"model": "claude-sonnet-4-5-20250929",
"messages": MESSAGES,
"max_tokens": 64,
"stream": False,
}
assert additional_args["headers"]["x-api-key"] == "test-key"
@pytest.mark.asyncio
@pytest.mark.parametrize("raise_after_edit", [False, True])
@pytest.mark.parametrize("native", [False, True])
async def test_messages_pre_call_edits_reach_later_callbacks_and_provider(
messages_server: RecordingServer, raise_after_edit: bool
messages_server: RecordingServer, raise_after_edit: bool, native: bool
) -> None:
litellm.rust(native)
observed: Final = []
class Edit(CustomLogger):
@ -85,7 +88,7 @@ async def test_messages_pre_call_edits_reach_later_callbacks_and_provider(
assert observed[0][0]["temperature"] == 0.25
assert observed[0][1]["x-audit-tag"] == "reviewed"
assert messages_server.requests[0].body["temperature"] == 0.25
assert "temperature" not in messages_server.requests[0].body
assert messages_server.requests[0].headers["x-audit-tag"] == "reviewed"
@ -129,6 +132,52 @@ async def test_messages_pre_call_state_reaches_terminal_callbacks(messages_serve
assert observed == [token]
@pytest.mark.asyncio
@pytest.mark.parametrize("native", [False, True])
@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"])
@pytest.mark.parametrize("raise_after_edit", [False, True])
async def test_messages_retained_aliases_preserve_identity_and_snapshot_timing(
messages_server: RecordingServer, native: bool, provider: str, raise_after_edit: bool
) -> None:
litellm.rust(native)
retained: Final = []
observed: Final = []
class Retain(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
body = request_body(kwargs)
message = kwargs["messages"][0]
assert body["messages"][0] is message
assert body["messages"][0]["content"] is message["content"]
assert body["messages"][0]["content"][0] is message["content"][0]
retained.append(message["content"][0])
class Edit(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
retained[0]["text"] = "changed through retained reference"
if raise_after_edit:
raise RuntimeError("export failed after mutation")
class Observe(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
block = request_body(kwargs)["messages"][0]["content"][0]
observed.append((block is retained[0], block["text"]))
response: Final = await litellm.anthropic.messages.acreate(
model=MESSAGES_MODEL.replace("anthropic/", f"{provider}/"),
messages=[{"role": "user", "content": [{"type": "text", "text": "original"}]}],
max_tokens=64,
api_key="test-key",
api_base=messages_server.base_url,
callbacks=[Retain(), Edit(), Observe()],
)
assert observed == [(True, "changed through retained reference")]
assert retained[0]["text"] == "changed through retained reference"
assert messages_server.requests[0].body["messages"][0]["content"][0]["text"] == "original"
if native:
assert response["_hidden_params"]["additional_headers"]["x-litellm-rust"] == "true"
@pytest.mark.asyncio
async def test_messages_callbacks_run_once(messages_server: RecordingServer) -> None:
recorder: Final = RecordingLogger()

29
tests/test_litellm_rust/test_ocr_callbacks.py
View file

@ -96,28 +96,43 @@ def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: Re
assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed"
def test_pre_call_nested_mutation_updates_retained_references(ocr_server: RecordingServer) -> None:
@pytest.mark.asyncio
@pytest.mark.parametrize("native", [False, True])
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_pre_call_nested_mutation_updates_retained_references(
ocr_server: RecordingServer, native: bool, asynchronous: bool
) -> None:
litellm.rust(native)
original: Final = dict(OCR_DOCUMENT)
replacement_url: Final = "data:application/pdf;base64,ZGVm"
retained: Final = []
class Retain(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
assert request_body(kwargs)["document"] is original
retained.append(request_body(kwargs)["document"])
class Edit(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
request_body(kwargs)["document"]["document_url"] = replacement_url
original["document_url"] = replacement_url
call_native_ocr(
ocr_server,
document=original,
callbacks=[Retain(), Edit()],
)
arguments: Final = {
"model": "mistral/mistral-ocr-latest",
"document": original,
"api_key": "test-key",
"api_base": ocr_server.base_url,
"callbacks": [Retain(), Edit()],
}
if asynchronous:
await litellm.aocr(**arguments)
else:
litellm.ocr(**arguments)
assert retained[0]["document_url"] == replacement_url
assert original["document_url"] == replacement_url
assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url
if native:
assert ocr_server.requests[0].headers["accept-encoding"] == "identity"
def test_pre_call_field_replacement_preserves_original_references(ocr_server: RecordingServer) -> None:

42
tests/test_litellm_rust/test_sdk_dispatch.py
View file

@ -71,13 +71,22 @@ def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingSer
@pytest.mark.parametrize("provider", ["anthropic", "bedrock"])
@pytest.mark.parametrize("rebind_logging_view", [False, True])
@pytest.mark.parametrize("status", [200, 429])
@pytest.mark.parametrize("native", [False, True])
async def test_chat_retains_callback_edits_through_public_dispatch(
recording_server: RecordingServer, asynchronous: bool, provider: str, rebind_logging_view: bool, status: int
recording_server: RecordingServer,
asynchronous: bool,
provider: str,
rebind_logging_view: bool,
status: int,
native: bool,
) -> None:
import threading
import json
litellm.rust(native)
from tests.test_litellm_rust.callback_recorder import RecordingLogger
from tests.test_litellm_rust.contracts import MESSAGES_RESPONSE, request_body, request_headers
from tests.test_litellm_rust.contracts import MESSAGES_RESPONSE
recording_server.default_response = ResponseSpec(
status=status,
@ -98,10 +107,15 @@ async def test_chat_retains_callback_edits_through_public_dispatch(
class EditingLogger(RecordingLogger):
def log_pre_api_call(self, model, messages, kwargs):
super().log_pre_api_call(model, messages, kwargs)
body = request_body(kwargs)
headers = request_headers(kwargs)
body["messages"][0]["content"] = "edited by callback"
body["max_tokens" if provider == "anthropic" else "maxTokens"] = 32
body = kwargs["additional_args"]["complete_input_dict"]
headers = kwargs["additional_args"]["headers"]
if provider == "anthropic":
assert isinstance(body, dict)
body["messages"][0]["content"][0]["text"] = "edited by callback"
body["max_tokens"] = 32
else:
assert isinstance(body, str)
assert json.loads(body)["messages"][0]["content"][0]["text"] == "original"
headers["x-retained-callback"] = "original"
observations.append((threading.current_thread(), body, headers))
if rebind_logging_view:
@ -139,12 +153,20 @@ async def test_chat_retains_callback_edits_through_public_dispatch(
event_name: Final = "async_log_success_event" if asynchronous else "log_success_event"
events: Final = await recorder.wait_for_async(event_name)
assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
if native:
assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
assert recorder.names.count("log_pre_api_call") == 1
assert recorder.names.count(event_name) == 1
assert observations[0][0] is caller_thread
if asynchronous and provider == "anthropic" and not native:
assert observations[0][0] is not caller_thread
else:
assert observations[0][0] is caller_thread
assert events[0].response is response
assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == "edited by callback"
assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == (
"edited by callback" if provider == "anthropic" else "original"
)
assert recording_server.requests[0].headers["x-retained-callback"] == "original"
body: Final = recording_server.requests[0].body
assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == 32
assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == (
32 if provider == "anthropic" else 64
)