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

refactor(rust): align request lifecycle naming

Browse source
This commit is contained in:
Yujong Lee 2026-09-08 13:13:54 -07:00
parent eabd1c187d
commit 277a71cd30
25 changed files with 437 additions and 304 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -4,7 +4,7 @@ use crate::error::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::prepare::prepare_provider_request;
use super::request::build_provider_request;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatBodySnapshot, ChatCompletionsResponse, ChatEndpoint, ProviderChatCompletionsRequest,
@ -15,7 +15,7 @@ use super::types::{
pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
) -> Result<ChatCompletionsResponse, Error> {
let request = prepare_provider_request(request)?;
let request = build_provider_request(request)?;
let body = serde_json::to_vec(&request.body).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize chat completions request: {err}"
@ -89,7 +89,7 @@ pub(super) async fn execute_settled_request(
///
/// A config reports the same variants on either side of the call: a missing
/// field or an unsupported block can mean "this request cannot be translated"
/// during prepare and "this response cannot be normalized" here. Only the
/// while building and "this response cannot be normalized" here. Only the
/// second kind has already been billed, and a host that keeps a reference
/// implementation must not retry those, so collapse them to one variant that
/// can only mean the provider was already called.

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

@ -12,7 +12,7 @@ mod common_utils;
pub mod conversation;
pub(crate) mod handler;
pub mod lifecycle;
mod prepare;
pub mod request;
pub mod response_utils;
pub mod transformation;
pub mod types;
@ -20,8 +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 request::{parse_messages, resolve_provider_config, resolve_request};
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
use crate::integrations::custom_logger::CallbackTiming;

34
litellm-rust/crates/core/src/chat_completions/prepare.rs → litellm-rust/crates/core/src/chat_completions/request.rs
View file

@ -116,7 +116,7 @@ fn validate_environment(
Ok((headers, auth))
}
pub(super) fn prepare_provider_request(
pub(super) fn build_provider_request(
request: ResolvedChatCompletionsRequest<'_>,
) -> Result<ProviderChatCompletionsRequest, Error> {
let (headers, auth) = validate_environment(&request, &request.model, request.config)?;
@ -144,27 +144,27 @@ pub(super) fn prepare_provider_request(
})
}
pub async fn prepare_callback_request(
pub async fn build_pre_call_request(
request: ChatCompletionsRequest<'_>,
) -> Result<super::types::ChatCallbackRequest, Error> {
) -> Result<super::types::ChatPreCallRequest, Error> {
use super::transformation::PreCallBody;
use super::types::{ChatBodySnapshot, ChatCallbackRequest, ChatEndpoint};
use super::types::{ChatBodySnapshot, ChatEndpoint, ChatPreCallRequest};
let prepared = prepare_provider_request(resolve_request(request)?)?;
let built = build_provider_request(resolve_request(request)?)?;
let endpoint = ChatEndpoint {
model: prepared.model.clone(),
config: prepared.config,
url: prepared.url.clone(),
timeout: prepared.timeout,
model: built.model.clone(),
config: built.config,
url: built.url.clone(),
timeout: built.timeout,
};
match prepared.config.pre_call_body() {
match built.config.pre_call_body() {
PreCallBody::Live => {
let mut generated = prepared
let mut generated = built
.body
.as_object()
.cloned()
.ok_or_else(|| Error::InvalidRequest("chat body must be an object".into()))?;
let parameter_fields = prepared
let parameter_fields = built
.optional_params
.keys()
.filter(|name| generated.contains_key(*name))
@ -173,20 +173,20 @@ pub async fn prepare_callback_request(
for name in &parameter_fields {
generated.remove(name);
}
Ok(ChatCallbackRequest::Live {
Ok(ChatPreCallRequest::Live {
endpoint,
generated,
parameter_fields,
headers: prepared.upstream_headers,
headers: built.upstream_headers,
})
}
PreCallBody::Serialized => {
let logging_body = serde_json::to_string(&prepared.body).map_err(|error| {
let logging_body = serde_json::to_string(&built.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 {
let headers = super::handler::signed_headers(&built, &body).await?;
Ok(ChatPreCallRequest::Serialized {
snapshot: ChatBodySnapshot { endpoint, body },
logging_body,
headers,

114
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -2,14 +2,14 @@ use serde_json::{Map, Value, json};
use crate::error::Error;
use super::prepare::{prepare_provider_request, resolve_request};
use super::request::{build_provider_request, resolve_request};
use super::transformation::ChatCompletionsAuth;
use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
fn prepare_chat_completions_call(
fn build_chat_completions_request(
request: ChatCompletionsRequest<'_>,
) -> Result<ProviderChatCompletionsRequest, Error> {
prepare_provider_request(resolve_request(request)?)
build_provider_request(resolve_request(request)?)
}
fn request<'a>(
@ -36,59 +36,59 @@ fn request<'a>(
/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers
/// carry resolved credentials), so unwrap the failure case by hand.
fn decline(request: ChatCompletionsRequest<'_>) -> Error {
match prepare_chat_completions_call(request) {
match build_chat_completions_request(request) {
Err(error) => error,
Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url),
Ok(built) => panic!("expected a decline, built a call to {}", built.url),
}
}
#[test]
fn resolves_the_provider_from_the_model_prefix() {
let prepared = prepare_chat_completions_call(request(
let built = build_chat_completions_request(request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
.expect("builds");
assert_eq!(built.model, "claude-sonnet-4-5");
assert_eq!(built.url, "https://api.anthropic.com/v1/messages");
assert_eq!(built.body["model"], json!("claude-sonnet-4-5"));
}
#[test]
fn strips_an_explicit_provider_prefix_from_the_model() {
let prepared = prepare_chat_completions_call(request(
let built = build_chat_completions_request(request(
"anthropic/claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
.expect("builds");
assert_eq!(built.model, "claude-sonnet-4-5");
}
#[test]
fn adds_the_auth_and_default_headers() {
let prepared = prepare_chat_completions_call(request(
let built = build_chat_completions_request(request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
.expect("builds");
assert!(
prepared
built
.upstream_headers
.contains(&("x-api-key".to_string(), "sk-test".to_string()))
);
assert!(
prepared
built
.upstream_headers
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
);
assert!(matches!(
prepared.auth,
built.auth,
ChatCompletionsAuth::Header {
name: "x-api-key",
..
@ -111,13 +111,13 @@ fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
"X-Api-Key".to_string(),
json!("sk-caller"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let keys: Vec<_> = prepared
let built = build_chat_completions_request(call).expect("builds");
let keys: Vec<_> = built
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
assert_eq!(keys.len(), 1, "got {:?}", built.upstream_headers);
assert_eq!(keys[0].1, "sk-test");
}
@ -139,17 +139,17 @@ fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let built = build_chat_completions_request(call).expect("builds");
assert!(
!prepared
!built
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
"the resolved key must not be applied over an OAuth bearer, got {:?}",
prepared.upstream_headers
built.upstream_headers
);
assert!(
prepared
built
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
@ -172,22 +172,22 @@ fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
("Authorization".to_string(), json!("Bearer unrelated")),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let keys: Vec<_> = prepared
let built = build_chat_completions_request(call).expect("builds");
let keys: Vec<_> = built
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
assert_eq!(keys.len(), 1, "got {:?}", built.upstream_headers);
assert_eq!(keys[0].1, "sk-test");
assert!(
prepared
built
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer unrelated"),
"the unrelated authorization must survive, got {:?}",
prepared.upstream_headers
built.upstream_headers
);
}
@ -272,7 +272,7 @@ fn rejects_non_string_extra_headers() {
#[cfg(feature = "bedrock-auth")]
#[test]
fn prepares_a_bedrock_call_without_resolving_credentials() {
fn builds_a_bedrock_call_without_resolving_credentials() {
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
@ -280,26 +280,26 @@ fn prepares_a_bedrock_call_without_resolving_credentials() {
json!({"maxTokens": 16}),
);
call.api_key = None;
let prepared = prepare_chat_completions_call(call).expect("prepares");
let built = build_chat_completions_request(call).expect("builds");
assert_eq!(
prepared.url,
built.url,
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
assert_eq!(
prepared.auth,
built.auth,
ChatCompletionsAuth::AwsSigV4 {
region: "us-east-1".to_string()
}
);
// SigV4 signs the serialized body, so prepare must not have added an
// SigV4 signs the serialized body, so build must not have added an
// Authorization header; the handler does it.
assert!(
!prepared
!built
.upstream_headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
);
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
assert_eq!(built.body["inferenceConfig"], json!({"maxTokens": 16}));
}
#[cfg(feature = "bedrock-auth")]
@ -324,8 +324,8 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
"x-request-id".to_string(),
json!("abc-123"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
let built = build_chat_completions_request(call).expect("builds");
let signed = super::handler::signed_headers(&built, br#"{"a":1}"#)
.await
.expect("signs");
@ -375,8 +375,8 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
);
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
let built = build_chat_completions_request(call).expect("builds");
let error = super::handler::signed_headers(&built, br#"{"a":1}"#)
.await
.expect_err("{forwarded} should decline instead of being signed");
assert!(
@ -403,8 +403,8 @@ fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
"Authorization".to_string(),
json!("Bearer caller-supplied"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let authorizations: Vec<_> = prepared
let built = build_chat_completions_request(call).expect("builds");
let authorizations: Vec<_> = built
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
@ -436,16 +436,16 @@ fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
"authorization".to_string(),
json!("Bearer sk-ant-oat01-forwarded"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let keys: Vec<_> = prepared
let built = build_chat_completions_request(call).expect("builds");
let keys: Vec<_> = built
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.map(|(_, value)| value.as_str())
.collect();
assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers);
assert!(keys.is_empty(), "got {:?}", built.upstream_headers);
assert!(
prepared
built
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
@ -459,26 +459,26 @@ fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
// The configured bearer identity has its own account and quota boundary,
// so a request carrying one must not be signed as whatever principal the
// host's AWS credentials resolve to.
let prepared = prepare_chat_completions_call(request(
let built = build_chat_completions_request(request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
))
.expect("prepares");
.expect("builds");
assert_eq!(
prepared.auth,
built.auth,
ChatCompletionsAuth::Bearer {
token: "sk-test".to_string()
}
);
assert!(
prepared
built
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-test"),
"prepare did not carry the bearer token"
"build did not carry the bearer token"
);
}
@ -496,7 +496,7 @@ fn decline_reason(
}
#[test]
fn the_gate_accepts_what_prepare_accepts() {
fn the_gate_accepts_what_the_builder_accepts() {
assert_eq!(
decline_reason(
"anthropic/claude-sonnet-4-5",
@ -553,8 +553,8 @@ fn the_gate_declines_without_resolving_credentials_or_calling_out() {
}
#[test]
fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
// A gate that accepts what prepare then declines would make the host emit
fn the_gate_agrees_with_the_builder_on_every_case_it_accepts() {
// A gate that accepts what build then declines would make the host emit
// its pre-call logging on a path that falls back, so pin the agreement.
for (messages, params) in [
(
@ -580,13 +580,13 @@ fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
None,
"gate declined {messages}"
);
prepare_chat_completions_call(request(
build_chat_completions_request(request(
"anthropic/claude-sonnet-4-5",
None,
messages.clone(),
params,
))
.unwrap_or_else(|error| panic!("prepare declined {messages}: {error}"));
.unwrap_or_else(|error| panic!("build declined {messages}: {error}"));
}
}

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

@ -7,7 +7,7 @@ use super::types::{
};
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
/// the request builder; SigV4 needs the serialized body, so the handler signs it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ChatCompletionsAuth {
Header { name: &'static str, value: String },

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

@ -69,7 +69,7 @@ pub struct SettledChatRequest {
pub(super) headers: Vec<(String, String)>,
}
pub enum ChatCallbackRequest {
pub enum ChatPreCallRequest {
Live {
endpoint: ChatEndpoint,
generated: Map<String, Value>,

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

@ -1,5 +1,5 @@
use crate::Error;
use crate::ocr::{OcrAdmissionRequest, prepare};
use crate::ocr::{OcrAdmissionRequest, request};
use super::program::{CallProgram, ProgramOptions, actions_for};
use super::{ActionBinding, LifecycleRoute, Outcome};
@ -96,7 +96,7 @@ impl LifecycleRoute for OcrRoute {
admission: &Self::Admission,
options: Self::Options,
) -> Result<Result<Self::State, Self::Decline>, Self::Error> {
match prepare::admission_capabilities(admission) {
match request::admission_capabilities(admission) {
Err(Error::Unsupported(reason)) => return Ok(Err(Decline(reason))),
Err(error) => return Err(error),
Ok(()) => {}
@ -221,14 +221,14 @@ mod tests {
for (asynchronous, expected) in [
(
false,
vec![Setup, Prepare, PreCall, Send, SyncSuccess, Restore],
vec![Setup, BuildRequest, PreCall, Send, SyncSuccess, Restore],
),
(
true,
vec![
Setup,
DeploymentPre,
Prepare,
BuildRequest,
PreCall,
Send,
DeploymentSuccess,
@ -259,7 +259,7 @@ mod tests {
vec![
Setup,
DeploymentPre,
Prepare,
BuildRequest,
PreCall,
Send,
DeploymentSuccess,
@ -267,7 +267,7 @@ mod tests {
SyncSuccessIfNeeded,
]
} else {
vec![Setup, Prepare, PreCall, Send, SyncSuccess]
vec![Setup, BuildRequest, PreCall, Send, SyncSuccess]
};
for stage in stages {
for outcome in [Outcome::Failure, Outcome::Abort] {
@ -277,7 +277,7 @@ mod tests {
assert_eq!(transition.error, ErrorDisposition::Replace);
let expected = if outcome == Outcome::Abort {
Restore
} else if asynchronous && matches!(stage, Prepare | PreCall | Send) {
} else if asynchronous && matches!(stage, BuildRequest | PreCall | Send) {
DeploymentFailure
} else {
SyncFailure

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

@ -275,4 +275,32 @@ mod tests {
Some(FailureStage::AfterProviderResponse)
);
}
#[test]
fn request_build_and_pre_call_failures_are_replayable() {
for success_count in [1, 2] {
let mut program = CallProgram::new(ProgramOptions {
asynchronous: false,
internal_call: false,
});
for _ in 0..success_count {
program.advance(Outcome::Success, observations()).unwrap();
}
let failure = program.advance(Outcome::Failure, observations()).unwrap();
assert_eq!(failure.commitment, Commitment::Replayable);
assert_eq!(failure.failure_stage, Some(FailureStage::BeforeProvider));
}
}
#[test]
fn request_build_has_its_own_action_kind() {
assert_eq!(
actions_for(Operation::BuildRequest)[0].kind,
ActionKind::RequestBuild
);
assert_eq!(
actions_for(Operation::Send)[0].kind,
ActionKind::ProviderCall
);
}
}

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

@ -5,18 +5,18 @@ use crate::lifecycle::{StreamingMetadata, StreamingSource};
use super::client::http_client;
use super::common_utils::truncate_error_body;
use super::prepare::prepare_provider_request;
use super::request::build_provider_request;
use super::types::{AnthropicMessagesResponse, MessagesRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn execute_messages_provider_call(
request: MessagesRequest,
) -> Result<AnthropicMessagesResponse, Error> {
let request = prepare_provider_request(request)?;
execute_prepared_messages_provider_call(request).await
let request = build_provider_request(request)?;
execute_provider_messages_request(request).await
}
pub async fn execute_prepared_messages_provider_call(
pub async fn execute_provider_messages_request(
request: super::types::ProviderMessagesRequest,
) -> Result<AnthropicMessagesResponse, Error> {
let mut request_builder = http_client().post(&request.url).json(&request.body);
@ -52,7 +52,7 @@ pub async fn execute_prepared_messages_provider_call(
pub(super) async fn execute_messages_provider_stream(
request: MessagesRequest,
) -> Result<StreamingSource, Error> {
let request = prepare_provider_request(request)?;
let request = build_provider_request(request)?;
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::InvalidRequest(
"streaming messages is not supported for this provider".to_string(),

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

@ -83,6 +83,61 @@ impl Lifecycle<MessagesRoute> {
}
}
#[cfg(test)]
mod program_tests {
use super::*;
#[test]
fn sync_and_async_sequences_build_then_run_pre_call() {
for (asynchronous, expected) in [
(
false,
vec![
Operation::Setup,
Operation::BuildRequest,
Operation::PreCall,
Operation::Send,
Operation::SyncSuccess,
Operation::Restore,
],
),
(
true,
vec![
Operation::Setup,
Operation::DeploymentPre,
Operation::BuildRequest,
Operation::PreCall,
Operation::Send,
Operation::DeploymentSuccess,
Operation::AsyncSuccess,
Operation::SyncSuccessIfNeeded,
Operation::Restore,
],
),
] {
let mut machine = machine(Options {
asynchronous,
..Options::default()
})
.unwrap();
for operation in expected {
assert_eq!(machine.operation(), operation);
machine
.advance(
Outcome::Success,
Observations {
logger_available: true,
has_fallbacks: false,
},
)
.unwrap();
}
assert_eq!(machine.operation(), Operation::Complete(Outcome::Success));
}
}
}
pub trait MessagesServices:
RequestPolicy<MessagesRequest, MessagesRequest> + TerminalDispatcher + Clock
{

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

@ -12,15 +12,14 @@ mod client;
mod common_utils;
mod handler;
pub mod lifecycle;
mod prepare;
pub mod request;
pub mod transformation;
pub mod types;
use std::sync::Arc;
use crate::lifecycle::StreamingCall;
pub use handler::execute_prepared_messages_provider_call;
pub use prepare::{prepare_endpoint, prepare_provider_request};
pub use handler::execute_provider_messages_request;
use types::{AnthropicMessagesResponse, MessagesRequest};
pub async fn messages(request: MessagesRequest) -> Result<AnthropicMessagesResponse, Error> {

8
litellm-rust/crates/core/src/messages/prepare.rs → litellm-rust/crates/core/src/messages/request.rs
View file

@ -6,10 +6,8 @@ use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrateg
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 {
pub fn build_provider_request(request: MessagesRequest) -> Result<ProviderMessagesRequest, Error> {
let endpoint = build_endpoint(MessagesOptions {
model: request.model,
api_key: request.api_key,
api_base: request.api_base,
@ -37,7 +35,7 @@ pub fn prepare_provider_request(
})
}
pub fn prepare_endpoint(request: MessagesOptions) -> Result<MessagesEndpoint, Error> {
pub fn build_endpoint(request: MessagesOptions) -> Result<MessagesEndpoint, Error> {
let provider_info =
get_custom_llm_provider(&request.model, request.custom_llm_provider.as_deref())
.or_else(|| {

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

@ -1,4 +1,4 @@
pub mod prepare;
pub mod request;
pub mod transformation;
pub mod types;
@ -11,7 +11,7 @@ use crate::error::json_type_name;
use crate::http_utils::{buffered_post, has_header};
pub use types::{
OcrAdmissionRequest, OcrDraft, OcrEndpoint, OcrResponseData, OcrTransportRequest,
OcrAdmissionRequest, OcrEndpoint, OcrPreCallRequest, OcrResponseData, OcrTransportRequest,
OcrTransportResponse, SettledOcrRequest,
};
use types::{OcrDocument, OcrDocumentProjection};
@ -136,8 +136,8 @@ pub(crate) async fn send<S: OcrTransport>(
headers,
body,
} = request;
let config = prepare::provider_config(&endpoint.custom_llm_provider, &endpoint.model)?;
prepare::validate_capabilities(config)?;
let config = request::provider_config(&endpoint.custom_llm_provider, &endpoint.model)?;
request::validate_capabilities(config)?;
let object = body.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&body),

25
litellm-rust/crates/core/src/ocr/prepare.rs → litellm-rust/crates/core/src/ocr/request.rs
View file

@ -9,7 +9,7 @@ use crate::providers::vertex_ai::ocr::transformation as vertex_ai;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::transformation::{OcrProviderConfig, OcrResponseHandling};
use super::types::{OcrAdmissionRequest, OcrDraft, OcrEndpoint};
use super::types::{OcrAdmissionRequest, OcrEndpoint, OcrPreCallRequest};
fn request_config(
request: &OcrAdmissionRequest,
@ -80,7 +80,7 @@ fn check_admission_capabilities(
Ok(())
}
pub fn prepare(request: OcrAdmissionRequest) -> Result<OcrDraft, Error> {
pub fn build_pre_call_request(request: OcrAdmissionRequest) -> Result<OcrPreCallRequest, Error> {
let (provider, config) = request_config(&request)?;
let env_lookup = |key: &str| std::env::var(key).ok();
let headers = config
@ -130,7 +130,7 @@ pub fn prepare(request: OcrAdmissionRequest) -> Result<OcrDraft, Error> {
let Value::Object(body) = template.data else {
return Err(Error::Unsupported("non-object OCR request template"));
};
Ok(OcrDraft {
Ok(OcrPreCallRequest {
endpoint: OcrEndpoint {
model: provider.model.to_string(),
custom_llm_provider: provider.custom_llm_provider.to_string(),
@ -217,16 +217,20 @@ mod tests {
assert!(
matches!(check_admission_capabilities(&request, &|_| None), Err(Error::Unsupported(message)) if message == expected)
);
assert!(
matches!(prepare(request), Err(Error::Unsupported(message)) if message == expected)
);
assert!(matches!(
build_pre_call_request(request),
Err(Error::Unsupported(message)) if message == expected
));
}
#[test]
fn admission_leaves_url_preparation_and_revalidation_until_prepare() {
fn admission_leaves_url_building_and_revalidation_until_request_build() {
let request = request();
assert!(check_admission_capabilities(&request, &|_| None).is_ok());
assert!(matches!(prepare(request), Err(Error::InvalidRequest(_))));
assert!(matches!(
build_pre_call_request(request),
Err(Error::InvalidRequest(_))
));
let mut request = self::request();
assert!(check_admission_capabilities(&request, &|_| None).is_ok());
@ -236,7 +240,10 @@ mod tests {
check_admission_capabilities(&request, &|_| None),
Err(Error::InvalidRequest(_))
));
assert!(matches!(prepare(request), Err(Error::InvalidRequest(_))));
assert!(matches!(
build_pre_call_request(request),
Err(Error::InvalidRequest(_))
));
}
#[test]

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

@ -59,7 +59,7 @@ pub enum OcrDocumentProjection {
Transformed,
}
pub struct OcrDraft {
pub struct OcrPreCallRequest {
pub endpoint: OcrEndpoint,
pub headers: Vec<(String, String)>,
pub body: Map<String, Value>,

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

@ -6,9 +6,9 @@ use std::time::Duration;
use litellm_core::Error;
use litellm_core::lifecycle::CallLifecycleContext;
use litellm_core::ocr::prepare::prepare;
use litellm_core::ocr::request::build_pre_call_request;
use litellm_core::ocr::types::{OcrDocument, OcrDocumentProjection};
use litellm_core::ocr::{DefaultOcrServices, OcrAdmissionRequest as OcrRequest, OcrDraft};
use litellm_core::ocr::{DefaultOcrServices, OcrAdmissionRequest as OcrRequest, OcrPreCallRequest};
use serde_json::{Value, json};
fn request() -> OcrRequest {
@ -30,8 +30,8 @@ fn request() -> OcrRequest {
}
}
fn body(prepared: &OcrDraft) -> Value {
let mut body = prepared.body.clone();
fn body(built: &OcrPreCallRequest) -> Value {
let mut body = built.body.clone();
body.insert(
"document".into(),
json!({"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"}),
@ -42,15 +42,15 @@ fn body(prepared: &OcrDraft) -> Value {
}
async fn ocr(
prepared: OcrDraft,
built: OcrPreCallRequest,
headers: Vec<(String, String)>,
body: Value,
) -> Result<litellm_core::ocr::OcrResponseData, Error> {
let model = prepared.endpoint.model().to_string();
let provider = prepared.endpoint.custom_llm_provider().to_string();
let model = built.endpoint.model().to_string();
let provider = built.endpoint.custom_llm_provider().to_string();
let response = litellm_core::ocr::ocr(
&DefaultOcrServices,
prepared.endpoint.settle(headers, body),
built.endpoint.settle(headers, body),
Default::default(),
CallLifecycleContext::new("ocr", model, provider, "test-call"),
)
@ -60,38 +60,38 @@ async fn ocr(
}
#[test]
fn prepares_provider_template_auth_and_url() {
let prepared = prepare(OcrRequest {
fn builds_provider_template_auth_and_url() {
let built = build_pre_call_request(OcrRequest {
api_base: Some(" https://ocr.example/v1/ ".into()),
extra_headers: vec![("X-Request-Id".into(), "request-1".into())],
request_format: Some("litellm".into()),
..request()
})
.unwrap();
assert_eq!(prepared.endpoint.model(), "mistral-ocr-latest");
assert_eq!(prepared.endpoint.custom_llm_provider(), "mistral");
assert_eq!(prepared.endpoint.url(), "https://ocr.example/v1/ocr");
assert_eq!(prepared.endpoint.timeout_seconds(), 2.0);
assert_eq!(built.endpoint.model(), "mistral-ocr-latest");
assert_eq!(built.endpoint.custom_llm_provider(), "mistral");
assert_eq!(built.endpoint.url(), "https://ocr.example/v1/ocr");
assert_eq!(built.endpoint.timeout_seconds(), 2.0);
assert_eq!(
prepared.document_projection,
built.document_projection,
OcrDocumentProjection::RetainedDocument
);
assert_eq!(
Value::Object(prepared.body),
Value::Object(built.body),
json!({"model": "mistral-ocr-latest", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"}})
);
assert_eq!(
prepared.parameter_fields,
built.parameter_fields,
litellm_core::providers::mistral::ocr::transformation::supported_ocr_params()
);
assert_eq!(
prepared.headers,
built.headers,
vec![
("Authorization".into(), "Bearer test-key".into()),
("X-Request-Id".into(), "request-1".into()),
]
);
let explicit = prepare(OcrRequest {
let explicit = build_pre_call_request(OcrRequest {
model: "mistral-ocr-latest".into(),
custom_llm_provider: Some("mistral".into()),
api_key: None,
@ -108,9 +108,9 @@ fn prepares_provider_template_auth_and_url() {
}
#[test]
fn prepare_environment_credentials() {
fn build_request_environment_credentials() {
if let Ok(case) = std::env::var("LITELLM_OCR_ENV_TEST") {
let result = prepare(OcrRequest {
let result = build_pre_call_request(OcrRequest {
api_key: Some(" ".into()),
..request()
});
@ -120,7 +120,7 @@ fn prepare_environment_credentials() {
vec![("Authorization".into(), "Bearer env-key".into())]
);
assert_eq!(
prepare(request()).unwrap().headers,
build_pre_call_request(request()).unwrap().headers,
vec![("Authorization".into(), "Bearer test-key".into())]
);
} else {
@ -135,7 +135,7 @@ fn prepare_environment_credentials() {
] {
let mut command = Command::new(std::env::current_exe().unwrap());
command
.args(["--exact", "prepare_environment_credentials"])
.args(["--exact", "build_request_environment_credentials"])
.env("LITELLM_OCR_ENV_TEST", case)
.env_remove("MISTRAL_API_KEY");
if let Some(key) = key {
@ -146,9 +146,9 @@ fn prepare_environment_credentials() {
}
#[test]
fn prepare_rejects_unsupported_providers_formats_and_invalid_metadata() {
fn build_request_rejects_unsupported_providers_formats_and_invalid_metadata() {
for (provider, capability) in [("reducto", "OCR provider"), ("openai", "OCR provider")] {
let result = prepare(OcrRequest {
let result = build_pre_call_request(OcrRequest {
custom_llm_provider: Some(provider.into()),
api_key: None,
api_base: Some("not a URL".into()),
@ -158,7 +158,7 @@ fn prepare_rejects_unsupported_providers_formats_and_invalid_metadata() {
}
for format in ["native", "json", "", "LiteLLM"] {
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
request_format: Some(format.into()),
..request()
}),
@ -167,7 +167,7 @@ fn prepare_rejects_unsupported_providers_formats_and_invalid_metadata() {
}
for timeout_seconds in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::MAX] {
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
timeout_seconds,
..request()
}),
@ -175,14 +175,14 @@ fn prepare_rejects_unsupported_providers_formats_and_invalid_metadata() {
));
}
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
model: "mistral-ocr-latest".into(),
..request()
}),
Err(Error::InvalidProvider(_))
));
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
api_base: Some("file:///secret".into()),
..request()
}),
@ -197,7 +197,7 @@ fn cloud_capabilities_fail_only_when_required() {
"vertex_ai/mistral-ocr-latest",
] {
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
model: model.into(),
api_base: Some("http://127.0.0.1:1".into()),
document: OcrDocument::ImageUrl {
@ -215,7 +215,7 @@ fn cloud_capabilities_fail_only_when_required() {
"azure_ai/documentintelligence/prebuilt-layout",
] {
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
model: model.into(),
api_base: Some("http://127.0.0.1:1".into()),
..request()
@ -231,7 +231,7 @@ fn cloud_capabilities_fail_only_when_required() {
"cohere/parse-v5.0",
] {
assert!(matches!(
prepare(OcrRequest {
build_pre_call_request(OcrRequest {
model: model.into(),
..request()
}),
@ -263,19 +263,19 @@ fn cloud_credentials_use_native_keys_headers_or_narrow_acquisition_stub() {
api_base: Some("http://127.0.0.1:1".into()),
..request()
};
let result = prepare(make_request());
let result = build_pre_call_request(make_request());
if std::env::var("LITELLM_CLOUD_OCR_ENV_TEST").unwrap() == "present" {
assert_eq!(result.unwrap().headers, vec![(header.into(), key.into())]);
} else {
assert!(matches!(result, Err(Error::Unsupported(message)) if message == operation));
}
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
extra_headers: vec![("aUtHoRiZaTiOn".into(), "Bearer supplied".into())],
..make_request()
})
.unwrap();
assert_eq!(
prepared.headers,
built.headers,
vec![("aUtHoRiZaTiOn".into(), "Bearer supplied".into())]
);
}
@ -332,20 +332,20 @@ async fn cloud_providers_use_existing_auth_urls_and_request_response_transforms(
json!({"pages": [{"index": 0, "markdown": "proof"}], "usage_info": {"pages_processed": 1}})
};
let (base, handle) = server(200, "", &response.to_string(), Duration::ZERO);
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
model: model.into(),
api_base: Some(base),
..request()
})
.unwrap();
let mut body = Value::Object(prepared.body.clone());
let mut body = Value::Object(built.body.clone());
body[if deepseek {
"temperature"
} else {
"include_image_base64"
}] = if deepseek { json!(0.1) } else { json!(true) };
assert_eq!(
prepared.headers,
built.headers,
vec![(
if auth_name == "api-key" {
"Api-Key"
@ -356,8 +356,8 @@ async fn cloud_providers_use_existing_auth_urls_and_request_response_transforms(
auth_value.into()
)]
);
let headers = prepared.headers.clone();
let response = ocr(prepared, headers, body.clone()).await.unwrap();
let headers = built.headers.clone();
let response = ocr(built, headers, body.clone()).await.unwrap();
let (headers, sent) = handle.join().unwrap();
assert!(headers.starts_with(&format!("POST {path} HTTP/1.1\r\n")));
assert!(headers.contains(&format!("{auth_name}: {auth_value}\r\n")));
@ -455,22 +455,19 @@ async fn posts_filled_body_and_normalizes_provider_response() {
"private_provider_field": "not forwarded"
});
let (base, handle) = server(200, "", &response_json.to_string(), Duration::ZERO);
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
api_base: Some(base),
..request()
})
.unwrap();
let body = body(&prepared);
let headers = prepared
let body = body(&built);
let headers = built
.headers
.iter()
.cloned()
.chain([("X-Retained".into(), "header".into())])
.collect();
let response = ocr(prepared, headers, body.clone())
.await
.unwrap()
.into_json();
let response = ocr(built, headers, body.clone()).await.unwrap().into_json();
let (headers, sent_body) = handle.join().unwrap();
let headers = headers.to_ascii_lowercase();
assert!(headers.starts_with("post /v1/ocr http/1.1\r\n"));
@ -495,15 +492,15 @@ async fn posts_filled_body_and_normalizes_provider_response() {
#[tokio::test]
async fn preserves_callback_body_and_header_changes() {
let (base, handle) = server(200, "", "{}", Duration::ZERO);
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
api_base: Some(base),
..request()
})
.unwrap();
let mut body = body(&prepared);
let mut body = body(&built);
body["model"] = json!("callback-model");
body["custom_provider_field"] = json!({"nested": [1, true, null]});
let headers = prepared
let headers = built
.headers
.iter()
.cloned()
@ -512,7 +509,7 @@ async fn preserves_callback_body_and_header_changes() {
("aCcEpT-EnCoDiNg".into(), "gzip, br".into()),
])
.collect();
ocr(prepared, headers, body.clone()).await.unwrap();
ocr(built, headers, body.clone()).await.unwrap();
let (headers, sent_body) = handle.join().unwrap();
let headers = headers.to_ascii_lowercase();
assert_eq!(sent_body, body);
@ -535,13 +532,13 @@ async fn preserves_callback_body_and_header_changes() {
#[tokio::test]
async fn settled_headers_are_the_only_headers_sent() {
let (base, handle) = server(200, "", "{}", Duration::ZERO);
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
api_base: Some(base),
..request()
})
.unwrap();
let body = body(&prepared);
ocr(prepared, Vec::new(), body).await.unwrap();
let body = body(&built);
ocr(built, Vec::new(), body).await.unwrap();
let (headers, _) = handle.join().unwrap();
assert!(!headers.to_ascii_lowercase().contains("authorization:"));
}
@ -552,13 +549,13 @@ async fn rejects_unsupported_inputs_before_io() {
listener.set_nonblocking(true).unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
for case in ["file", "local", "compression"] {
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
api_base: Some(base.clone()),
..request()
})
.unwrap();
let mut body = body(&prepared);
let mut headers = prepared.headers.clone();
let mut body = body(&built);
let mut headers = built.headers.clone();
match case {
"file" => body["document"] = json!({"type": "file", "file": "private"}),
"local" => body["document"]["document_url"] = json!("file:///private.pdf"),
@ -566,7 +563,7 @@ async fn rejects_unsupported_inputs_before_io() {
_ => unreachable!(),
}
assert!(matches!(
ocr(prepared, headers, body).await,
ocr(built, headers, body).await,
Err(Error::Unsupported(_))
));
assert_eq!(
@ -629,15 +626,15 @@ async fn handles_errors_compression_and_timeout_without_exposing_payloads() {
),
] {
let (base, handle) = server(status, headers, response_body, delay);
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
api_base: Some(base),
timeout_seconds: if delay.is_zero() { 2.0 } else { 0.05 },
..request()
})
.unwrap();
let body = body(&prepared);
let headers = prepared.headers.clone();
assert_eq!(ocr(prepared, headers, body).await.unwrap_err(), expected);
let body = body(&built);
let headers = built.headers.clone();
assert_eq!(ocr(built, headers, body).await.unwrap_err(), expected);
handle.join().unwrap();
}
}
@ -645,14 +642,14 @@ async fn handles_errors_compression_and_timeout_without_exposing_payloads() {
#[tokio::test]
async fn normalizes_missing_fields_and_accepts_identity_response() {
let (base, handle) = server(200, "Content-Encoding: Identity\r\n", "{}", Duration::ZERO);
let prepared = prepare(OcrRequest {
let built = build_pre_call_request(OcrRequest {
api_base: Some(base),
..request()
})
.unwrap();
let body = body(&prepared);
let headers = prepared.headers.clone();
let response = ocr(prepared, headers, body).await.unwrap().into_json();
let body = body(&built);
let headers = built.headers.clone();
let response = ocr(built, headers, body).await.unwrap().into_json();
handle.join().unwrap();
assert_eq!(
response,

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

@ -18,7 +18,6 @@ struct OperationBinding {
fn operation_binding(
operation: Operation,
asynchronous: bool,
supports_pre_call: bool,
route: &str,
) -> PyResult<OperationBinding> {
let binding = match operation {
@ -34,15 +33,10 @@ fn operation_binding(
method: "build_request",
awaiting: false,
},
Operation::PreCall if supports_pre_call => OperationBinding {
Operation::PreCall => 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,
@ -96,11 +90,10 @@ 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)?;
let binding = operation_binding(operation, asynchronous, route)?;
Ok((
binding.awaiting,
host.getattr(py, binding.method)?.call0(py)?,
@ -118,60 +111,39 @@ mod tests {
Python::initialize();
Python::attach(|_| {
let cases = [
(Operation::Setup, false, false, "setup", false),
(
Operation::DeploymentPre,
false,
false,
"deployment_pre",
true,
),
(
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),
(Operation::Setup, false, "setup", false),
(Operation::DeploymentPre, false, "deployment_pre", true),
(Operation::BuildRequest, false, "build_request", false),
(Operation::PreCall, false, "pre_call", false),
(Operation::Send, false, "send_sync", false),
(Operation::Send, true, "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::SyncSuccess, false, "sync_success", false),
(Operation::AsyncSuccess, 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),
(Operation::SyncFailure, false, "sync_failure", false),
(Operation::AsyncFailure, false, "async_failure", true),
(Operation::Restore, false, "restore", false),
];
for (operation, asynchronous, pre_call, method, awaiting) in cases {
for (operation, asynchronous, method, awaiting) in cases {
assert_eq!(
operation_binding(operation, asynchronous, pre_call, "test").unwrap(),
operation_binding(operation, asynchronous, "test").unwrap(),
OperationBinding { method, awaiting },
);
}
@ -182,19 +154,9 @@ mod tests {
fn invalid_operations_raise_route_specific_errors() {
Python::initialize();
Python::attach(|_| {
let pre_call = operation_binding(Operation::PreCall, false, false, "messages")
.expect_err("unsupported pre-call should fail");
assert_eq!(
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");
let complete =
operation_binding(Operation::Complete(Outcome::Success), false, "messages")
.expect_err("complete lifecycle should fail");
assert_eq!(
complete.to_string(),
"RuntimeError: messages lifecycle is complete"

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

@ -2,9 +2,10 @@ use litellm_core::Error;
use litellm_core::chat_completions::lifecycle::{
Admission, ChatCompletionsRoute, Observations, Operation, Options, machine,
};
use litellm_core::chat_completions::request::build_pre_call_request;
use litellm_core::chat_completions::types::ChatCompletionsRequest;
use litellm_core::chat_completions::{
chat_completions_decline_reason, execute_settled_with_terminal, prepare_callback_request,
chat_completions_decline_reason, execute_settled_with_terminal,
};
use litellm_core::lifecycle::{
CallLifecycleContext, ErrorDisposition, ExecutedCall, Lifecycle, Outcome, TerminalRecord,
@ -25,6 +26,8 @@ use crate::retained::RequestRoots;
#[pyclass]
struct ChatCompletionsState {
roots: Option<RequestRoots>,
logging: Option<Py<PyAny>>,
pre_call: Option<Py<PyDict>>,
pending: Option<PendingChatRequest>,
context: Option<CallLifecycleContext>,
terminal: Option<TerminalRecord>,
@ -36,7 +39,8 @@ impl ChatCompletionsState {
if let Some(roots) = &self.roots {
roots.traverse(&visit)?;
}
Ok(())
visit.call(&self.logging)?;
visit.call(&self.pre_call)
}
fn __clear__(slf: &Bound<'_, Self>) {
@ -44,6 +48,8 @@ impl ChatCompletionsState {
let mut state = slf.borrow_mut();
(
state.roots.take(),
state.logging.take(),
state.pre_call.take(),
state.pending.take(),
state.context.take(),
state.terminal.take(),
@ -158,7 +164,7 @@ fn invoke(
let machine = machine.borrow(py);
(machine.machine.operation(), machine.asynchronous)
};
crate::driver::invoke(py, operation, asynchronous, false, "chat completions", host)
crate::driver::invoke(py, operation, asynchronous, "chat completions", host)
}
enum PendingChatRequest {
@ -167,12 +173,12 @@ enum PendingChatRequest {
}
#[pyfunction]
fn prepare(
fn build_request(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
) -> PyResult<Py<ChatCompletionsState>> {
use litellm_core::chat_completions::types::ChatCallbackRequest;
use litellm_core::chat_completions::types::ChatPreCallRequest;
let bag = arguments.bind(py);
let admission = admission(bag)?;
@ -191,10 +197,10 @@ fn prepare(
scalar(bag, "litellm_call_id")?.unwrap_or_default(),
);
let extra_headers = optional_map(bag, "extra_headers")?;
let prepared = run_sync_value(
let built = run_sync_value(
py,
async move {
prepare_callback_request(ChatCompletionsRequest {
build_pre_call_request(ChatCompletionsRequest {
model: &admission.model,
messages: admission.messages,
optional_params: admission.optional_params,
@ -208,8 +214,8 @@ fn prepare(
},
core_error_to_pyerr,
)?;
let (body, pending, header_values) = match prepared {
ChatCallbackRequest::Live {
let (body, pending, header_values) = match built {
ChatPreCallRequest::Live {
endpoint,
generated,
parameter_fields,
@ -226,7 +232,7 @@ fn prepare(
}
(body.into_any(), PendingChatRequest::Live(endpoint), headers)
}
ChatCallbackRequest::Serialized {
ChatPreCallRequest::Serialized {
snapshot,
logging_body,
headers,
@ -266,9 +272,6 @@ fn prepare(
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 {
@ -277,6 +280,8 @@ fn prepare(
body.unbind(),
headers.unbind(),
)),
logging: Some(logging),
pre_call: Some(kwargs.unbind()),
pending: Some(pending),
context: Some(context),
terminal: None,
@ -284,6 +289,28 @@ fn prepare(
)
}
#[pyfunction]
fn pre_call(py: Python<'_>, state: Py<ChatCompletionsState>) -> PyResult<()> {
let (logging, arguments) = {
let state = state.borrow(py);
let logging = state
.logging
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("chat completions logging state was cleared"))?
.clone_ref(py);
let arguments = state
.pre_call
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("chat completions pre-call state was cleared"))?
.clone_ref(py);
(logging, arguments)
};
logging
.bind(py)
.call_method(pyo3::intern!(py, "pre_call"), (), Some(arguments.bind(py)))?;
Ok(())
}
struct OwnedRequest {
request: litellm_core::chat_completions::types::SettledChatRequest,
context: CallLifecycleContext,
@ -436,7 +463,8 @@ fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
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("build_request", wrap_pyfunction!(build_request, &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(
@ -484,7 +512,10 @@ mod tests {
Python::attach(|py| {
let module = PyModule::new(py, "chat_test").unwrap();
module
.add_function(wrap_pyfunction!(prepare, &module).unwrap())
.add_function(wrap_pyfunction!(build_request, &module).unwrap())
.unwrap();
module
.add_function(wrap_pyfunction!(pre_call, &module).unwrap())
.unwrap();
module
.add_function(wrap_pyfunction!(snapshot, &module).unwrap())
@ -525,7 +556,8 @@ arguments = dict(model='claude-opus-5', messages=messages,
extra_headers=headers, api_key='test',
custom_llm_provider='anthropic', opaque=opaque,
litellm_logging_obj=logger)
state = native.prepare(arguments, logger)
state = native.build_request(arguments, logger)
native.pre_call(state)
wire_body, wire_headers = native.snapshot(state)
assert wire_body['messages'][0]['content'][0]['text'] == 'body edit'
assert wire_body['stop_sequences'] == ['first', 'second']

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

@ -1,10 +1,11 @@
use litellm_core::lifecycle::FailureStage;
use litellm_core::lifecycle::{ErrorDisposition, Lifecycle, Outcome};
use litellm_core::messages::execute_provider_messages_request;
use litellm_core::messages::lifecycle::{MessagesRoute, Observations, Operation, Options, machine};
use litellm_core::messages::request::build_endpoint;
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::*;
@ -21,7 +22,9 @@ use crate::retained::RequestRoots;
#[pyclass]
struct MessagesState {
roots: Option<RequestRoots>,
prepared: Option<(MessagesEndpoint, MessagesBodySnapshot)>,
logging: Option<Py<PyAny>>,
pre_call: Option<Py<PyDict>>,
pending: Option<(MessagesEndpoint, MessagesBodySnapshot)>,
}
#[pymethods]
@ -30,13 +33,19 @@ impl MessagesState {
if let Some(roots) = &self.roots {
roots.traverse(&visit)?;
}
Ok(())
visit.call(&self.logging)?;
visit.call(&self.pre_call)
}
fn __clear__(slf: &Bound<'_, Self>) {
let roots = {
let mut state = slf.borrow_mut();
(state.roots.take(), state.prepared.take())
(
state.roots.take(),
state.logging.take(),
state.pre_call.take(),
state.pending.take(),
)
};
drop(roots);
}
@ -139,11 +148,11 @@ fn invoke(
let machine = machine.borrow(py);
(machine.machine.operation(), machine.asynchronous)
};
crate::driver::invoke(py, operation, asynchronous, false, "messages", host)
crate::driver::invoke(py, operation, asynchronous, "messages", host)
}
#[pyfunction]
fn prepare(
fn build_request(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
@ -151,7 +160,7 @@ fn prepare(
let bag = arguments.bind(py);
let options = decode_options(py, bag)?;
let endpoint = py
.detach(|| prepare_endpoint(options))
.detach(|| build_endpoint(options))
.map_err(core_error_to_pyerr)?;
let body = bag
.get_item("body")?
@ -176,9 +185,6 @@ fn prepare(
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))?;
Py::new(
py,
MessagesState {
@ -187,22 +193,46 @@ fn prepare(
body.unbind().into_any(),
headers.unbind().into_any(),
)),
prepared: Some((endpoint, snapshot)),
logging: Some(logging),
pre_call: Some(kwargs.unbind()),
pending: Some((endpoint, snapshot)),
},
)
}
#[pyfunction]
fn pre_call(py: Python<'_>, state: Py<MessagesState>) -> PyResult<()> {
let (logging, arguments) = {
let state = state.borrow(py);
let logging = state
.logging
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("messages logging state was cleared"))?
.clone_ref(py);
let arguments = state
.pre_call
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("messages pre-call state was cleared"))?
.clone_ref(py);
(logging, arguments)
};
logging
.bind(py)
.call_method(pyo3::intern!(py, "pre_call"), (), Some(arguments.bind(py)))?;
Ok(())
}
fn take_request(py: Python<'_>, state: &Py<MessagesState>) -> PyResult<ProviderMessagesRequest> {
let ((endpoint, snapshot), headers) = {
let mut state = state.borrow_mut(py);
let prepared = state.prepared.take().ok_or_else(|| {
let pending = state.pending.take().ok_or_else(|| {
PyRuntimeError::new_err("messages request was already sent or cleared")
})?;
let roots = state
.roots
.as_ref()
.ok_or_else(|| PyRuntimeError::new_err("messages roots were cleared"))?;
(prepared, roots.headers(py))
(pending, roots.headers(py))
};
let headers = headers
.cast::<PyDict>()?
@ -235,7 +265,7 @@ fn send(py: Python<'_>, state: Py<MessagesState>) -> PyResult<Bound<'_, PyAny>>
litellm_python_interop::run_async_py(py, async move {
let _state = state;
let response = run_async_value(
execute_prepared_messages_provider_call(request),
execute_provider_messages_request(request),
messages_provider_error_to_pyerr,
)
.await?;
@ -248,7 +278,7 @@ fn send_sync(py: Python<'_>, state: Py<MessagesState>) -> PyResult<Py<PyAny>> {
let request = take_request(py, &state)?;
let response = run_sync_value(
py,
execute_prepared_messages_provider_call(request),
execute_provider_messages_request(request),
messages_provider_error_to_pyerr,
)?;
Ok(Pythonized(response).into_pyobject(py)?.unbind().into_any())
@ -295,7 +325,8 @@ fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
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("build_request", wrap_pyfunction!(build_request, &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(
@ -331,11 +362,19 @@ mod tests {
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();
module
.add_function(wrap_pyfunction!(build_request, &module).unwrap())
.unwrap();
module
.add_function(wrap_pyfunction!(pre_call, &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"
py.run(
c"
import gc
import weakref
@ -361,7 +400,8 @@ 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)
state = native.build_request(arguments, logger)
native.pre_call(state)
wire_body, wire_headers = native.snapshot(state)
assert wire_body['messages'][0]['content'] == 'original'
assert body['messages'][0]['content'] == 'changed'
@ -382,7 +422,11 @@ assert alive() is not None
del state
gc.collect()
assert alive() is None
", Some(&globals), Some(&globals)).unwrap();
",
Some(&globals),
Some(&globals),
)
.unwrap();
});
}

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

@ -9,7 +9,7 @@ use litellm_core::lifecycle::{
};
use litellm_core::ocr::DefaultOcrServices;
use litellm_core::ocr::types::{
OcrAdmissionRequest, OcrDocumentProjection, OcrDraft, OcrEndpoint, SettledOcrRequest,
OcrAdmissionRequest, OcrDocumentProjection, OcrEndpoint, OcrPreCallRequest, SettledOcrRequest,
};
use litellm_core::routing_utils::provider::get_custom_llm_provider;
use litellm_python_interop::{Pythonized, from_py, to_py};
@ -309,11 +309,11 @@ fn invoke(
let machine = machine.borrow(py);
(machine.machine.operation(), machine.asynchronous)
};
crate::driver::invoke(py, operation, asynchronous, true, "OCR", host)
crate::driver::invoke(py, operation, asynchronous, "OCR", host)
}
#[pyfunction]
fn prepare(
fn build_request(
py: Python<'_>,
arguments: Py<PyDict>,
logging: Py<PyAny>,
@ -323,8 +323,8 @@ fn prepare(
let request = decode_request(py, bag)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let draft = py
.detach(|| litellm_core::ocr::prepare::prepare(request))
let pre_call_request = py
.detach(|| litellm_core::ocr::request::build_pre_call_request(request))
.map_err(|error| {
request_error_to_pyerr(py, error, &model, custom_llm_provider.as_deref())
})?;
@ -332,13 +332,13 @@ fn prepare(
.get_item("document")?
.ok_or_else(|| PyValueError::new_err("OCR requires document"))?
.cast_into::<PyDict>()?;
let OcrDraft {
let OcrPreCallRequest {
endpoint,
headers: draft_headers,
body: draft_body,
document_projection,
parameter_fields,
} = draft;
} = pre_call_request;
let body = PyDict::new(py);
for (name, value) in &draft_body {
match (name.as_str(), document_projection) {
@ -604,7 +604,7 @@ fn bindings(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> {
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("build_request", wrap_pyfunction!(build_request, &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)?)?;
@ -889,7 +889,7 @@ asyncio.run(exercise())
Python::attach(|py| {
let module = PyModule::new(py, "ocr_test").unwrap();
module
.add_function(wrap_pyfunction!(prepare, &module).unwrap())
.add_function(wrap_pyfunction!(build_request, &module).unwrap())
.unwrap();
module
.add_function(wrap_pyfunction!(pre_call, &module).unwrap())
@ -929,7 +929,7 @@ async def exercise():
port = server.sockets[0].getsockname()[1]
logger = Logger()
alive = weakref.ref(logger)
state = native.prepare(dict(
state = native.build_request(dict(
model='mistral/mistral-ocr-latest', api_key='test-key', timeout=5.0,
api_base=f'http://127.0.0.1:{port}', litellm_logging_obj=logger,
document={'type': 'document_url', 'document_url': 'https://example.test/doc.pdf'},
@ -983,7 +983,7 @@ asyncio.run(exercise())
Python::attach(|py| {
let module = PyModule::new(py, "ocr_test").unwrap();
module
.add_function(wrap_pyfunction!(prepare, &module).unwrap())
.add_function(wrap_pyfunction!(build_request, &module).unwrap())
.unwrap();
module
.add_function(wrap_pyfunction!(pre_call, &module).unwrap())
@ -1033,7 +1033,7 @@ logger = Logger()
arguments = dict(model='mistral/mistral-ocr-latest', document=document,
api_key='test-key', pages=pages, metadata=metadata,
opaque=opaque, litellm_logging_obj=logger, timeout=Timeout())
state = native.prepare(arguments, logger, False)
state = native.build_request(arguments, logger, False)
native.pre_call(state)
assert logger.calls == ['update', 'pre']
roots = gc.get_referents(state)

10
litellm/rust_bridge/chat_completions.py
View file

@ -616,9 +616,10 @@ class _ChatCompletionsBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[
[dict[str, object], bool, bool], NativeLifecycle
] # mutable-ok: native bridge retains and updates Python argument objects
prepare: Callable[
build_request: Callable[
[dict[str, object], object], object
] # mutable-ok: native bridge retains and updates Python argument objects
pre_call: Callable[[object], None]
send: Callable[[object], Awaitable[Mapping[str, object]]]
send_sync: Callable[[object], Mapping[str, object]]
terminal_record: Callable[[object], Mapping[str, object]]
@ -664,10 +665,13 @@ class _ChatCompletionsHost:
self.current = await deployment_pre(self.current, "acompletion")
self.current[LOGGING_OBJECT_KEY] = self.logger
def prepare(self) -> None:
def build_request(self) -> None:
if self.logger is None:
raise RuntimeError("chat completions logging was not initialized")
self.state = self.bindings.prepare(self.current, self.logger)
self.state = self.bindings.build_request(self.current, self.logger)
def pre_call(self) -> None:
self.bindings.pre_call(self.state)
def send_sync(self) -> None:
model_response: Final = self.arguments["model_response"]

10
litellm/rust_bridge/messages.py
View file

@ -303,9 +303,10 @@ class _MessagesLifecycle(NativeLifecycle, Protocol):
class _MessagesBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[[bool, bool], _MessagesLifecycle]
prepare: Callable[
build_request: Callable[
[dict[str, object], object], object
] # mutable-ok: native bridge retains and updates Python argument objects
pre_call: Callable[[object], None]
send: Callable[[object], Awaitable[AnthropicMessagesResponse]]
send_sync: Callable[[object], AnthropicMessagesResponse]
@ -349,10 +350,13 @@ class _MessagesHost:
self.current = await deployment_pre(self.current, "anthropic_messages")
self.current[LOGGING_OBJECT_KEY] = self.logger
def prepare(self) -> None:
def build_request(self) -> None:
if self.logger is None:
raise RuntimeError("messages logging was not initialized")
self.state = self.bindings.prepare(self.current, self.logger)
self.state = self.bindings.build_request(self.current, self.logger)
def pre_call(self) -> None:
self.bindings.pre_call(self.state)
def send_sync(self) -> None:
self.response = self.bindings.send_sync(self.state)

6
litellm/rust_bridge/ocr.py
View file

@ -85,7 +85,7 @@ class _OcrBindings(NativeLifecycleBindings, Protocol):
Lifecycle: Callable[
[dict[str, object], object | None, bool, bool], _OcrLifecycle
] # mutable-ok: native bridge retains and updates Python argument objects
prepare: Callable[
build_request: Callable[
[dict[str, object], object, bool], object
] # mutable-ok: native bridge retains and updates Python argument objects
pre_call: Callable[[object], None]
@ -140,10 +140,10 @@ class _OcrHost:
self.current["litellm_call_id"] = call_id
self.current["litellm_trace_id"] = trace_id
def prepare(self) -> None:
def build_request(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)
self.state = self.bindings.build_request(self.current, self.logger, self.asynchronous)
def pre_call(self) -> None:
self.bindings.pre_call(self.state)

4
tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py
View file

@ -24,7 +24,7 @@ _MISTRAL_OCR_TESTS: Final = "providers::mistral::ocr::transformation::tests"
_VERTEX_OCR_TESTS: Final = "providers::vertex_ai::ocr::transformation::tests"
_REDUCTO_OCR_TESTS: Final = "providers::reducto::ocr::tests"
_GATEWAY_OCR_TESTS: Final = "ocr::tests"
_GATEWAY_PREPARE_OCR_TESTS: Final = "ocr::prepare::tests"
_GATEWAY_REQUEST_OCR_TESTS: Final = "ocr::request::tests"
def _rust_test(target: RustTarget, module: str, test: str) -> RustTestIdentity:
@ -204,7 +204,7 @@ _REDUCTO_GATEWAY_MAPPING: Final = TestMapping(
_GATEWAY_PORT_MAPPINGS: Final = _test_mappings(
_GATEWAY_TARGET,
_GATEWAY_PREPARE_OCR_TESTS,
_GATEWAY_REQUEST_OCR_TESTS,
(
(
"tests/test_litellm/ocr/test_ocr_native_format.py::test_native_format_rejected_for_provider_without_support_as_bad_request",

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

@ -232,7 +232,7 @@ def restore_ocr_context(logger: object) -> None:
def drive_ocr_sync(arguments: dict[str, object], bindings: object) -> object:
logger: Final = initialize_ocr_logging(arguments, False)
try:
state: Final = bindings.prepare(arguments, logger, False)
state: Final = bindings.build_request(arguments, logger, False)
except RuntimeError as error:
raise NotImplementedError(str(error)) from error
bindings.pre_call(state)
@ -242,7 +242,7 @@ def drive_ocr_sync(arguments: dict[str, object], bindings: object) -> object:
async def drive_ocr_async(arguments: dict[str, object], bindings: object) -> object:
logger: Final = initialize_ocr_logging(arguments, True)
try:
state: Final = bindings.prepare(arguments, logger, True)
state: Final = bindings.build_request(arguments, logger, True)
except RuntimeError as error:
raise NotImplementedError(str(error)) from error
bindings.pre_call(state)
@ -255,22 +255,26 @@ class MessagesLogging:
def drive_messages_sync(arguments: dict[str, object], bindings: object) -> object:
state: Final = bindings.prepare(arguments, MessagesLogging())
state: Final = bindings.build_request(arguments, MessagesLogging())
bindings.pre_call(state)
return bindings.send_sync(state)
async def drive_messages_async(arguments: dict[str, object], bindings: object) -> object:
state: Final = bindings.prepare(arguments, MessagesLogging())
state: Final = bindings.build_request(arguments, MessagesLogging())
bindings.pre_call(state)
return await bindings.send(state)
def drive_chat_sync(arguments: dict[str, object], bindings: object) -> object:
state: Final = bindings.prepare(arguments, MessagesLogging())
state: Final = bindings.build_request(arguments, MessagesLogging())
bindings.pre_call(state)
return bindings.send_sync(state)
async def drive_chat_async(arguments: dict[str, object], bindings: object) -> object:
state: Final = bindings.prepare(arguments, MessagesLogging())
state: Final = bindings.build_request(arguments, MessagesLogging())
bindings.pre_call(state)
return await bindings.send(state)