diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index cd969eef394..7c5609a6a11 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -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 { - 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. diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 7ee3e4c67d3..1611f014961 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -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; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/request.rs similarity index 88% rename from litellm-rust/crates/core/src/chat_completions/prepare.rs rename to litellm-rust/crates/core/src/chat_completions/request.rs index 6a88baa6147..0449beec83e 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/request.rs @@ -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 { 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 { +) -> Result { 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 ¶meter_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, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index f8594dee447..afbd698d5dc 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -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 { - 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}")); } } diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index 0581f9b8898..c24f5485188 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -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 }, diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 805fb90df77..3cb5ad9cf6b 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -69,7 +69,7 @@ pub struct SettledChatRequest { pub(super) headers: Vec<(String, String)>, } -pub enum ChatCallbackRequest { +pub enum ChatPreCallRequest { Live { endpoint: ChatEndpoint, generated: Map, diff --git a/litellm-rust/crates/core/src/lifecycle/ocr.rs b/litellm-rust/crates/core/src/lifecycle/ocr.rs index 312752b4da2..53394d7d263 100644 --- a/litellm-rust/crates/core/src/lifecycle/ocr.rs +++ b/litellm-rust/crates/core/src/lifecycle/ocr.rs @@ -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, 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 diff --git a/litellm-rust/crates/core/src/lifecycle/program.rs b/litellm-rust/crates/core/src/lifecycle/program.rs index 860a637acfc..a76fbf043f7 100644 --- a/litellm-rust/crates/core/src/lifecycle/program.rs +++ b/litellm-rust/crates/core/src/lifecycle/program.rs @@ -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 + ); + } } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 13db329b663..e9de5d2b08d 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -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 { - 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 { 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 { - 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(), diff --git a/litellm-rust/crates/core/src/messages/lifecycle.rs b/litellm-rust/crates/core/src/messages/lifecycle.rs index b03014c6dff..b8925b0ec56 100644 --- a/litellm-rust/crates/core/src/messages/lifecycle.rs +++ b/litellm-rust/crates/core/src/messages/lifecycle.rs @@ -83,6 +83,61 @@ impl Lifecycle { } } +#[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 + TerminalDispatcher + Clock { diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 308abc855b8..4e52bf62f16 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -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 { diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/request.rs similarity index 93% rename from litellm-rust/crates/core/src/messages/prepare.rs rename to litellm-rust/crates/core/src/messages/request.rs index 6496fd6bd13..598d743ad05 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/request.rs @@ -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 { - let endpoint = prepare_endpoint(MessagesOptions { +pub fn build_provider_request(request: MessagesRequest) -> Result { + 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 { +pub fn build_endpoint(request: MessagesOptions) -> Result { let provider_info = get_custom_llm_provider(&request.model, request.custom_llm_provider.as_deref()) .or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 05f2f419547..7776ab65661 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -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( 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), diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/request.rs similarity index 93% rename from litellm-rust/crates/core/src/ocr/prepare.rs rename to litellm-rust/crates/core/src/ocr/request.rs index bb1a1c6bf2c..78623705165 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/request.rs @@ -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 { +pub fn build_pre_call_request(request: OcrAdmissionRequest) -> Result { 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 { 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] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index e54fe2adf27..b463dccafcd 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -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, diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index e60cdef99e5..9b73852f15b 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -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 { - 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, diff --git a/litellm-rust/crates/python-bridge/src/driver.rs b/litellm-rust/crates/python-bridge/src/driver.rs index e21f872a585..ff97325371d 100644 --- a/litellm-rust/crates/python-bridge/src/driver.rs +++ b/litellm-rust/crates/python-bridge/src/driver.rs @@ -18,7 +18,6 @@ struct OperationBinding { fn operation_binding( operation: Operation, asynchronous: bool, - supports_pre_call: bool, route: &str, ) -> PyResult { 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, ) -> PyResult<(bool, Py)> { - 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" diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index e6e22117de2..1e0ae74903b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -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, + logging: Option>, + pre_call: Option>, pending: Option, context: Option, terminal: Option, @@ -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, logging: Py, ) -> PyResult> { - 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) -> 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::())?; 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'] diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index af7f017318d..056e2f4d1eb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -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, - prepared: Option<(MessagesEndpoint, MessagesBodySnapshot)>, + logging: Option>, + pre_call: Option>, + 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, logging: Py, @@ -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) -> 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) -> PyResult { 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::()? @@ -235,7 +265,7 @@ fn send(py: Python<'_>, state: Py) -> PyResult> 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) -> PyResult> { 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::())?; 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(); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index d8950b440cf..e7dc9cde42d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -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, logging: Py, @@ -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::()?; - 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::())?; 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) diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index e5710f503d1..da344a27675 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -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"] diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 76ba561461a..3e4e7257e07 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -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) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index f4dd5233880..b2dc51fada3 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -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) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py index 08d5c5a630c..60fefc8ccb0 100644 --- a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -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", diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 86e58b2e579..1e2b41bced6 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -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)