diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 96d001e2892..cd969eef394 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -7,8 +7,8 @@ use super::client::http_client; use super::prepare::prepare_provider_request; use super::transformation::ChatCompletionsAuth; use super::types::{ - ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData, - ResolvedChatCompletionsRequest, + ChatBodySnapshot, ChatCompletionsResponse, ChatEndpoint, ProviderChatCompletionsRequest, + ProviderChatResponseData, ResolvedChatCompletionsRequest, SettledChatRequest, }; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] @@ -22,12 +22,32 @@ pub(super) async fn execute_chat_completions_provider_call( )) })?; let headers = signed_headers(&request, &body).await?; + execute_settled_request( + ChatBodySnapshot { + endpoint: ChatEndpoint { + model: request.model, + config: request.config, + url: request.url, + timeout: request.timeout, + }, + body, + } + .settle_headers(headers), + ) + .await +} - let mut request_builder = http_client().post(&request.url).body(body); - for (key, value) in &headers { +pub(super) async fn execute_settled_request( + request: SettledChatRequest, +) -> Result { + let endpoint = request.snapshot.endpoint; + let mut request_builder = http_client() + .post(&endpoint.url) + .body(request.snapshot.body); + for (key, value) in &request.headers { request_builder = request_builder.header(key, value); } - if let Some(duration) = request.timeout { + if let Some(duration) = endpoint.timeout { request_builder = request_builder.timeout(duration); } @@ -58,9 +78,9 @@ pub(super) async fn execute_chat_completions_provider_call( let body: Value = serde_json::from_str(&text).map_err(|err| { Error::InvalidResponse(format!("invalid chat completions response JSON: {err}")) })?; - request + endpoint .config - .transform_response(&request.model, ProviderChatResponseData { body }) + .transform_response(&endpoint.model, ProviderChatResponseData { body }) .map_err(as_response_error) } diff --git a/litellm-rust/crates/core/src/chat_completions/lifecycle.rs b/litellm-rust/crates/core/src/chat_completions/lifecycle.rs index 9abd7e26293..24338180575 100644 --- a/litellm-rust/crates/core/src/chat_completions/lifecycle.rs +++ b/litellm-rust/crates/core/src/chat_completions/lifecycle.rs @@ -67,7 +67,6 @@ impl LifecycleRoute for ChatCompletionsRoute { program: CallProgram::new(ProgramOptions { asynchronous: options.asynchronous, internal_call: options.internal_call, - pre_call: false, }), })) } @@ -142,7 +141,8 @@ mod tests { false, vec![ Operation::Setup, - Operation::Prepare, + Operation::BuildRequest, + Operation::PreCall, Operation::Send, Operation::SyncSuccess, Operation::Restore, @@ -153,7 +153,8 @@ mod tests { vec![ Operation::Setup, Operation::DeploymentPre, - Operation::Prepare, + Operation::BuildRequest, + Operation::PreCall, Operation::Send, Operation::DeploymentSuccess, Operation::AsyncSuccess, diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index d48224d3731..7ee3e4c67d3 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -20,6 +20,7 @@ pub mod types; use serde_json::{Map, Value}; use handler::execute_chat_completions_provider_call; +pub use prepare::prepare_callback_request; use prepare::{parse_messages, resolve_provider_config, resolve_request}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; @@ -39,9 +40,23 @@ pub async fn chat_completions( pub async fn chat_completions_with_terminal( request: ChatCompletionsRequest<'_>, context: CallLifecycleContext, +) -> ExecutedCall { + with_terminal(chat_completions(request), context).await +} + +pub async fn execute_settled_with_terminal( + request: types::SettledChatRequest, + context: CallLifecycleContext, +) -> ExecutedCall { + with_terminal(handler::execute_settled_request(request), context).await +} + +async fn with_terminal( + call: impl std::future::Future>, + context: CallLifecycleContext, ) -> ExecutedCall { let start_time = epoch_seconds(); - match chat_completions(request).await { + match call.await { Ok(response) => { let usage = Usage { prompt_tokens: response.usage.prompt_tokens, diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 3be2ba21de4..6a88baa6147 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -143,3 +143,54 @@ pub(super) fn prepare_provider_request( timeout: request.timeout, }) } + +pub async fn prepare_callback_request( + request: ChatCompletionsRequest<'_>, +) -> Result { + use super::transformation::PreCallBody; + use super::types::{ChatBodySnapshot, ChatCallbackRequest, ChatEndpoint}; + + let prepared = prepare_provider_request(resolve_request(request)?)?; + let endpoint = ChatEndpoint { + model: prepared.model.clone(), + config: prepared.config, + url: prepared.url.clone(), + timeout: prepared.timeout, + }; + match prepared.config.pre_call_body() { + PreCallBody::Live => { + let mut generated = prepared + .body + .as_object() + .cloned() + .ok_or_else(|| Error::InvalidRequest("chat body must be an object".into()))?; + let parameter_fields = prepared + .optional_params + .keys() + .filter(|name| generated.contains_key(*name)) + .cloned() + .collect::>(); + for name in ¶meter_fields { + generated.remove(name); + } + Ok(ChatCallbackRequest::Live { + endpoint, + generated, + parameter_fields, + headers: prepared.upstream_headers, + }) + } + PreCallBody::Serialized => { + let logging_body = serde_json::to_string(&prepared.body).map_err(|error| { + Error::InvalidRequest(format!("could not encode chat request: {error}")) + })?; + let body = logging_body.as_bytes().to_vec(); + let headers = super::handler::signed_headers(&prepared, &body).await?; + Ok(ChatCallbackRequest::Serialized { + snapshot: ChatBodySnapshot { endpoint, body }, + logging_body, + headers, + }) + } + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index d7b9704c46c..0581f9b8898 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -27,11 +27,21 @@ pub struct Unsupported(pub &'static str); pub const STREAM_PARAM: &str = "stream"; +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum PreCallBody { + Live, + Serialized, +} + /// Message fields that carry no meaning for the upstream body, so their /// presence does not make a request untranslatable. const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; pub trait ChatCompletionsProviderConfig: Sync { + fn pre_call_body(&self) -> PreCallBody { + PreCallBody::Live + } + fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 3238d09b6b5..805fb90df77 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -52,6 +52,68 @@ pub struct ProviderChatRequestData { pub body: Value, } +pub struct ChatEndpoint { + pub(super) model: String, + pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) url: String, + pub(super) timeout: Option, +} + +pub struct ChatBodySnapshot { + pub(super) endpoint: ChatEndpoint, + pub(super) body: Vec, +} + +pub struct SettledChatRequest { + pub(super) snapshot: ChatBodySnapshot, + pub(super) headers: Vec<(String, String)>, +} + +pub enum ChatCallbackRequest { + Live { + endpoint: ChatEndpoint, + generated: Map, + parameter_fields: Vec, + headers: Vec<(String, String)>, + }, + Serialized { + snapshot: ChatBodySnapshot, + logging_body: String, + headers: Vec<(String, String)>, + }, +} + +impl ChatEndpoint { + pub fn capture_body(self, body: Value) -> Result { + let body = serde_json::to_vec(&body).map_err(|error| { + crate::Error::InvalidRequest(format!("could not encode chat request: {error}")) + })?; + Ok(ChatBodySnapshot { + endpoint: self, + body, + }) + } +} + +impl ChatBodySnapshot { + pub fn settle_headers(self, headers: Vec<(String, String)>) -> SettledChatRequest { + SettledChatRequest { + snapshot: self, + headers, + } + } +} + +impl SettledChatRequest { + pub fn body(&self) -> &[u8] { + &self.snapshot.body + } + + pub fn headers(&self) -> &[(String, String)] { + &self.headers + } +} + /// The raw provider response body handed back to a config for normalization. pub struct ProviderChatResponseData { pub body: Value, diff --git a/litellm-rust/crates/core/src/lifecycle/ocr.rs b/litellm-rust/crates/core/src/lifecycle/ocr.rs index cd47d417f59..312752b4da2 100644 --- a/litellm-rust/crates/core/src/lifecycle/ocr.rs +++ b/litellm-rust/crates/core/src/lifecycle/ocr.rs @@ -110,7 +110,6 @@ impl LifecycleRoute for OcrRoute { program: CallProgram::new(ProgramOptions { asynchronous: options.asynchronous, internal_call: options.internal_call, - pre_call: true, }), identity: Identity { requested_model: admission.model.clone(), diff --git a/litellm-rust/crates/core/src/lifecycle/program.rs b/litellm-rust/crates/core/src/lifecycle/program.rs index 97c2080d976..860a637acfc 100644 --- a/litellm-rust/crates/core/src/lifecycle/program.rs +++ b/litellm-rust/crates/core/src/lifecycle/program.rs @@ -7,7 +7,7 @@ use super::{ pub enum Operation { Setup, DeploymentPre, - Prepare, + BuildRequest, PreCall, Send, DeploymentSuccess, @@ -53,7 +53,6 @@ pub struct Transition { pub struct ProgramOptions { pub asynchronous: bool, pub internal_call: bool, - pub pre_call: bool, } #[derive(Debug)] @@ -125,14 +124,14 @@ impl CallProgram { (DeploymentFailure, _) => failure, (_, Outcome::Abort) => Restore, (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, - (Prepare | PreCall | Send, Outcome::Failure) if self.options.asynchronous => { + (BuildRequest | PreCall | Send, Outcome::Failure) if self.options.asynchronous => { DeploymentFailure } (_, Outcome::Failure) => failure, (Setup, Outcome::Success) if self.options.asynchronous => DeploymentPre, - (Setup | DeploymentPre, Outcome::Success) => Prepare, - (Prepare, Outcome::Success) if self.options.pre_call => PreCall, - (Prepare | PreCall, Outcome::Success) => Send, + (Setup | DeploymentPre, Outcome::Success) => BuildRequest, + (BuildRequest, Outcome::Success) => PreCall, + (PreCall, Outcome::Success) => Send, (Send, Outcome::Success) if self.options.asynchronous => DeploymentSuccess, (Send, Outcome::Success) => SyncSuccess, (DeploymentSuccess, Outcome::Success) => { @@ -169,7 +168,8 @@ impl CallProgram { pub fn actions_for(operation: Operation) -> &'static [ActionBinding] { match operation { - Operation::Prepare | Operation::Send => &PROVIDER_ACTION, + Operation::BuildRequest => &REQUEST_BUILD_ACTION, + Operation::Send => &PROVIDER_ACTION, Operation::PreCall => &PRE_CALL_ACTION, Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => { &FAILURE_ACTION @@ -180,6 +180,14 @@ pub fn actions_for(operation: Operation) -> &'static [ActionBinding] { } } +const REQUEST_BUILD_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::RequestBuild, + delivery: Delivery::InlineDirect, + on_result: ResultPolicy::Replace, + on_error: FailurePolicy::Propagate, + owner: Owner::Core, +}]; + const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding { kind: ActionKind::ProviderCall, delivery: Delivery::InlineAwaited, @@ -236,7 +244,6 @@ mod tests { let mut before = CallProgram::new(ProgramOptions { asynchronous: false, internal_call: false, - pre_call: false, }); let failure = before.advance(Outcome::Failure, observations()).unwrap(); assert_eq!(failure.commitment, Commitment::Replayable); @@ -245,10 +252,10 @@ mod tests { let mut provider = CallProgram::new(ProgramOptions { asynchronous: false, internal_call: false, - pre_call: false, }); provider.advance(Outcome::Success, observations()).unwrap(); provider.advance(Outcome::Success, observations()).unwrap(); + provider.advance(Outcome::Success, observations()).unwrap(); let failure = provider.advance(Outcome::Failure, observations()).unwrap(); assert_eq!(failure.commitment, Commitment::ProviderStarted); assert_eq!(failure.failure_stage, Some(FailureStage::ProviderCall)); @@ -256,11 +263,11 @@ mod tests { let mut after = CallProgram::new(ProgramOptions { asynchronous: false, internal_call: false, - pre_call: false, }); after.advance(Outcome::Success, observations()).unwrap(); after.advance(Outcome::Success, observations()).unwrap(); after.advance(Outcome::Success, observations()).unwrap(); + after.advance(Outcome::Success, observations()).unwrap(); let failure = after.advance(Outcome::Failure, observations()).unwrap(); assert_eq!(failure.commitment, Commitment::ResponseReceived); assert_eq!( diff --git a/litellm-rust/crates/core/src/lifecycle/types.rs b/litellm-rust/crates/core/src/lifecycle/types.rs index 0c089cbb377..2b5e05d1e1e 100644 --- a/litellm-rust/crates/core/src/lifecycle/types.rs +++ b/litellm-rust/crates/core/src/lifecycle/types.rs @@ -2,6 +2,7 @@ use serde::Serialize; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] pub enum ActionKind { + RequestBuild, RequestPolicy, ProviderCall, Deployment, diff --git a/litellm-rust/crates/core/src/messages/lifecycle.rs b/litellm-rust/crates/core/src/messages/lifecycle.rs index 9d16411987a..b03014c6dff 100644 --- a/litellm-rust/crates/core/src/messages/lifecycle.rs +++ b/litellm-rust/crates/core/src/messages/lifecycle.rs @@ -49,7 +49,6 @@ impl LifecycleRoute for MessagesRoute { program: CallProgram::new(ProgramOptions { asynchronous: options.asynchronous, internal_call: options.internal_call, - pre_call: false, }), })) } diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index da9bc813f83..308abc855b8 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use crate::lifecycle::StreamingCall; pub use handler::execute_prepared_messages_provider_call; -pub use prepare::prepare_provider_request; +pub use prepare::{prepare_endpoint, prepare_provider_request}; use types::{AnthropicMessagesResponse, MessagesRequest}; pub async fn messages(request: MessagesRequest) -> Result { diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index e8b0f57a0be..6496fd6bd13 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -3,12 +3,41 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider} use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; -use super::types::{MessagesRequest, ProviderMessagesRequest}; +use super::types::{MessagesEndpoint, MessagesOptions, MessagesRequest, ProviderMessagesRequest}; use serde_json::{Map, Value}; pub fn prepare_provider_request( request: MessagesRequest, ) -> Result { + let endpoint = prepare_endpoint(MessagesOptions { + model: request.model, + api_key: request.api_key, + api_base: request.api_base, + custom_llm_provider: request.custom_llm_provider, + extra_headers: request.extra_headers, + timeout: request.timeout, + })?; + let typed_request = serde_json::from_value(request.body).map_err(|err| { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = endpoint.config.transform_request(typed_request)?; + let body = serde_json::to_value(transformed).map_err(|err| { + Error::InvalidRequest(format!( + "failed to serialize Anthropic messages request: {err}" + )) + })?; + Ok(ProviderMessagesRequest { + provider: endpoint.provider, + model: endpoint.model, + config: endpoint.config, + url: endpoint.url, + body, + upstream_headers: endpoint.headers, + timeout: endpoint.timeout, + }) +} + +pub fn prepare_endpoint(request: MessagesOptions) -> Result { let provider_info = get_custom_llm_provider(&request.model, request.custom_llm_provider.as_deref()) .or_else(|| { @@ -39,25 +68,14 @@ pub fn prepare_provider_request( &env_lookup, )?; - let typed_request = serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) - })?; - let transformed = config.transform_request(typed_request)?; - let body = serde_json::to_value(transformed).map_err(|err| { - Error::InvalidRequest(format!( - "failed to serialize Anthropic messages request: {err}" - )) - })?; - let url = config.complete_url(request.api_base.as_deref(), &model, &env_lookup)?; - Ok(ProviderMessagesRequest { + Ok(MessagesEndpoint { provider: provider.to_string(), model, config, url, - body, - upstream_headers: headers, + headers, timeout: request.timeout, }) } diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index a1da80b0afa..cc169aa6aca 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -16,15 +16,82 @@ pub struct MessagesRequest { } pub struct ProviderMessagesRequest { - pub provider: String, - pub model: String, + pub(super) provider: String, + pub(super) model: String, pub(super) config: &'static dyn AnthropicMessagesProviderConfig, - pub url: String, - pub body: Value, - pub upstream_headers: Vec<(String, String)>, + pub(super) url: String, + pub(super) body: Value, + pub(super) upstream_headers: Vec<(String, String)>, + pub(super) timeout: Option, +} + +impl ProviderMessagesRequest { + pub fn body(&self) -> &Value { + &self.body + } + + pub fn headers(&self) -> &[(String, String)] { + &self.upstream_headers + } +} + +pub struct MessagesOptions { + pub model: String, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, pub timeout: Option, } +pub struct MessagesEndpoint { + pub(super) provider: String, + pub(super) model: String, + pub(super) config: &'static dyn AnthropicMessagesProviderConfig, + pub(super) url: String, + pub(super) headers: Vec<(String, String)>, + pub(super) timeout: Option, +} + +pub struct MessagesBodySnapshot(Value); + +impl MessagesEndpoint { + pub fn url(&self) -> &str { + &self.url + } + + pub fn headers(&self) -> &[(String, String)] { + &self.headers + } + + pub fn capture_buffered_body( + &self, + mut body: Value, + ) -> Result { + let object = body.as_object_mut().ok_or_else(|| { + crate::Error::InvalidRequest("messages body must be an object".into()) + })?; + object.remove("stream"); + Ok(MessagesBodySnapshot(body)) + } + + pub fn settle( + self, + body: MessagesBodySnapshot, + headers: Vec<(String, String)>, + ) -> ProviderMessagesRequest { + ProviderMessagesRequest { + provider: self.provider, + model: self.model, + config: self.config, + url: self.url, + body: body.0, + upstream_headers: headers, + timeout: self.timeout, + } + } +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(untagged)] pub enum SystemPrompt { diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index 7be3d108d44..266d58825a8 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -105,6 +105,10 @@ fn has_blank_text(message: &ChatMessage) -> bool { } impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { + fn pre_call_body(&self) -> crate::chat_completions::transformation::PreCallBody { + crate::chat_completions::transformation::PreCallBody::Serialized + } + fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/python-bridge/src/driver.rs b/litellm-rust/crates/python-bridge/src/driver.rs index f1db0ddb8fd..e21f872a585 100644 --- a/litellm-rust/crates/python-bridge/src/driver.rs +++ b/litellm-rust/crates/python-bridge/src/driver.rs @@ -30,8 +30,8 @@ fn operation_binding( method: "deployment_pre", awaiting: true, }, - Operation::Prepare => OperationBinding { - method: "prepare", + Operation::BuildRequest => OperationBinding { + method: "build_request", awaiting: false, }, Operation::PreCall if supports_pre_call => OperationBinding { @@ -126,7 +126,13 @@ mod tests { "deployment_pre", true, ), - (Operation::Prepare, false, false, "prepare", false), + ( + Operation::BuildRequest, + false, + false, + "build_request", + false, + ), (Operation::PreCall, false, true, "pre_call", false), (Operation::Send, false, false, "send_sync", false), (Operation::Send, true, false, "send", true), diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index b52dfde7e00..1433ef89542 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -6,6 +6,7 @@ mod errors; #[cfg(feature = "trace-parity")] mod function_trace; mod marshal; +mod retained; mod routes; use pyo3::prelude::*; diff --git a/litellm-rust/crates/python-bridge/src/retained.rs b/litellm-rust/crates/python-bridge/src/retained.rs new file mode 100644 index 00000000000..3be6f24c0dd --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/retained.rs @@ -0,0 +1,37 @@ +use pyo3::prelude::*; +use pyo3::pyclass::{PyTraverseError, PyVisit}; +use pyo3::types::PyDict; + +pub(crate) struct RequestRoots { + arguments: Py, + body: Py, + headers: Py, +} + +impl RequestRoots { + pub(crate) fn new(arguments: Py, body: Py, headers: Py) -> Self { + Self { + arguments, + body, + headers, + } + } + + pub(crate) fn arguments<'py>(&self, py: Python<'py>) -> Bound<'py, PyDict> { + self.arguments.clone_ref(py).into_bound(py) + } + + pub(crate) fn body<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> { + self.body.clone_ref(py).into_bound(py) + } + + pub(crate) fn headers<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> { + self.headers.clone_ref(py).into_bound(py) + } + + pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.arguments)?; + visit.call(&self.body)?; + visit.call(&self.headers) + } +} 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 e706eeb8405..e6e22117de2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -4,7 +4,7 @@ use litellm_core::chat_completions::lifecycle::{ }; use litellm_core::chat_completions::types::ChatCompletionsRequest; use litellm_core::chat_completions::{ - chat_completions_decline_reason, chat_completions_with_terminal, + chat_completions_decline_reason, execute_settled_with_terminal, prepare_callback_request, }; use litellm_core::lifecycle::{ CallLifecycleContext, ErrorDisposition, ExecutedCall, Lifecycle, Outcome, TerminalRecord, @@ -20,35 +20,32 @@ use serde_json::{Map, Value}; use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT}; use crate::errors::{RustBridgeDeclined, chat_completions_error_to_pyerr, core_error_to_pyerr}; use crate::marshal::optional_timeout; +use crate::retained::RequestRoots; #[pyclass] struct ChatCompletionsState { - arguments: Option>, - model: Option, - body: Option>, - headers: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - timeout: Option, + roots: Option, + pending: Option, + context: Option, terminal: Option, } #[pymethods] impl ChatCompletionsState { fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.arguments)?; - visit.call(&self.body)?; - visit.call(&self.headers) + if let Some(roots) = &self.roots { + roots.traverse(&visit)?; + } + Ok(()) } fn __clear__(slf: &Bound<'_, Self>) { let roots = { let mut state = slf.borrow_mut(); ( - state.arguments.take(), - state.body.take(), - state.headers.take(), + state.roots.take(), + state.pending.take(), + state.context.take(), state.terminal.take(), ) }; @@ -164,33 +161,105 @@ fn invoke( crate::driver::invoke(py, operation, asynchronous, false, "chat completions", host) } +enum PendingChatRequest { + Live(litellm_core::chat_completions::types::ChatEndpoint), + Serialized(litellm_core::chat_completions::types::ChatBodySnapshot), +} + #[pyfunction] fn prepare( py: Python<'_>, arguments: Py, logging: Py, ) -> PyResult> { + use litellm_core::chat_completions::types::ChatCallbackRequest; + let bag = arguments.bind(py); let admission = admission(bag)?; let api_key = scalar(bag, "api_key")?; let api_base = scalar(bag, "api_base")?; - let headers = bag - .get_item("extra_headers")? - .filter(|value| !value.is_none()); let timeout = optional_timeout( bag.get_item("timeout_seconds")? .filter(|value| !value.is_none()) .map(|value| value.extract::()) .transpose()?, )?; - let complete_input = PyDict::new(py); - complete_input.set_item("model", &admission.model)?; - complete_input.set_item("messages", bag.get_item("messages")?)?; - if let Some(optional_params) = bag.get_item("optional_params")? { - complete_input.call_method1("update", (optional_params,))?; + let context = CallLifecycleContext::new( + "chat_completion", + &admission.model, + admission.custom_llm_provider.as_deref().unwrap_or_default(), + scalar(bag, "litellm_call_id")?.unwrap_or_default(), + ); + let extra_headers = optional_map(bag, "extra_headers")?; + let prepared = run_sync_value( + py, + async move { + prepare_callback_request(ChatCompletionsRequest { + model: &admission.model, + messages: admission.messages, + optional_params: admission.optional_params, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: admission.custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await + }, + core_error_to_pyerr, + )?; + let (body, pending, header_values) = match prepared { + ChatCallbackRequest::Live { + endpoint, + generated, + parameter_fields, + headers, + } => { + let body = PyDict::new(py); + for (name, value) in generated { + body.set_item(name, to_py(py, &value)?)?; + } + if let Some(params) = bag.get_item("optional_params")? { + for name in parameter_fields { + body.set_item(&name, params.get_item(&name)?)?; + } + } + (body.into_any(), PendingChatRequest::Live(endpoint), headers) + } + ChatCallbackRequest::Serialized { + snapshot, + logging_body, + headers, + } => ( + logging_body.into_pyobject(py)?.into_any(), + PendingChatRequest::Serialized(snapshot), + headers, + ), + }; + let headers = PyDict::new(py); + for (name, value) in header_values { + headers.set_item(name, value)?; } + let headers = match &pending { + PendingChatRequest::Live(_) => { + match bag + .get_item("extra_headers")? + .filter(|value| !value.is_none()) + { + Some(original) => { + original.call_method1("update", (&headers,))?; + original + } + None => headers.into_any(), + } + } + PendingChatRequest::Serialized(_) => py + .import("botocore.awsrequest")? + .getattr("HeadersDict")? + .call1((headers,))?, + }; let additional = PyDict::new(py); - additional.set_item(COMPLETE_INPUT_DICT, &complete_input)?; + additional.set_item(COMPLETE_INPUT_DICT, &body)?; additional.set_item(API_BASE, bag.get_item("api_base")?)?; additional.set_item(HEADERS, &headers)?; let kwargs = PyDict::new(py); @@ -203,99 +272,60 @@ fn prepare( Py::new( py, ChatCompletionsState { - arguments: Some(arguments), - model: Some(admission.model), - body: Some(complete_input.unbind()), - headers: headers.map(Bound::unbind), - api_key, - api_base, - custom_llm_provider: admission.custom_llm_provider, - timeout, + roots: Some(RequestRoots::new( + arguments, + body.unbind(), + headers.unbind(), + )), + pending: Some(pending), + context: Some(context), terminal: None, }, ) } struct OwnedRequest { - model: String, - messages: Value, - optional_params: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout: Option, - call_id: String, + request: litellm_core::chat_completions::types::SettledChatRequest, + context: CallLifecycleContext, } fn take_request(py: Python<'_>, state: &Py) -> PyResult { - let (arguments, body, headers) = { - let state = state.borrow(py); - let arguments = state - .arguments - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("chat completions state was cleared"))? - .clone_ref(py); - let body = state - .body - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("chat completions body was cleared"))? - .clone_ref(py); - let headers = state.headers.as_ref().map(|value| value.clone_ref(py)); - (arguments, body, headers) - }; - let call_id = scalar(arguments.bind(py), "litellm_call_id")?.unwrap_or_default(); - let messages = value(body.bind(py), "messages")?; - let optional_params = body - .bind(py) - .iter() - .filter_map(|(key, value)| match key.extract::() { - Ok(key) if key == "model" || key == "messages" => None, - Ok(key) => Some(from_py(&value).map(|value| (key, value))), - Err(error) => Some(Err(error)), - }) - .collect::>>()?; - let extra_headers = headers - .as_ref() - .map(|value| from_py(value.bind(py))) - .transpose()?; - let mut state = state.borrow_mut(py); - Ok(OwnedRequest { - model: state - .model + let (pending, context, body, headers) = { + let mut state = state.borrow_mut(py); + let pending = state.pending.take().ok_or_else(|| { + PyRuntimeError::new_err("chat completions request was already sent or cleared") + })?; + let context = state + .context .take() - .ok_or_else(|| PyRuntimeError::new_err("chat completions request was already sent"))?, - messages, - optional_params, - api_key: state.api_key.take(), - api_base: state.api_base.take(), - custom_llm_provider: state.custom_llm_provider.take(), - extra_headers, - timeout: state.timeout.take(), - call_id, + .ok_or_else(|| PyRuntimeError::new_err("chat completions context was cleared"))?; + let roots = state + .roots + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("chat completions roots were cleared"))?; + (pending, context, roots.body(py), roots.headers(py)) + }; + let snapshot = match pending { + PendingChatRequest::Live(endpoint) => endpoint + .capture_body(from_py(&body)?) + .map_err(core_error_to_pyerr)?, + PendingChatRequest::Serialized(snapshot) => snapshot, + }; + let headers = headers + .call_method0("items")? + .try_iter()? + .map(|item| item?.extract::<(String, String)>()) + .collect::>()?; + Ok(OwnedRequest { + request: snapshot.settle_headers(headers), + context, }) } async fn execute( request: OwnedRequest, ) -> ExecutedCall { - let provider = request.custom_llm_provider.clone().unwrap_or_default(); - let context = - CallLifecycleContext::new("chat_completion", &request.model, provider, request.call_id); - chat_completions_with_terminal( - ChatCompletionsRequest { - model: &request.model, - messages: request.messages, - optional_params: request.optional_params, - api_key: request.api_key.as_deref(), - api_base: request.api_base.as_deref(), - custom_llm_provider: request.custom_llm_provider.as_deref(), - extra_headers: request.extra_headers, - timeout: request.timeout, - }, - context, - ) - .await + execute_settled_with_terminal(request.request, request.context).await } fn store_result( @@ -441,9 +471,8 @@ mod tests { to_py( py, &( - request.messages, - request.optional_params, - request.extra_headers, + serde_json::from_slice::(request.request.body()).unwrap(), + request.request.headers(), ), ) } @@ -476,10 +505,11 @@ class Logger: self.body = view['complete_input_dict'] self.headers = view['headers'] assert kwargs['input'] is messages - assert self.body['messages'] is messages - assert self.body['stop'] is stops + assert self.body['messages'] is not messages + assert self.body['stop_sequences'] is stops assert self.headers is headers messages[0]['content'] = 'edited' + self.body['messages'][0]['content'][0]['text'] = 'body edit' stops.append('second') self.headers['x-hook'] = 'edited' view['complete_input_dict'] = {'replacement': True} @@ -491,15 +521,15 @@ headers = {} opaque = Opaque() logger = Logger() arguments = dict(model='claude-opus-5', messages=messages, - optional_params={'max_tokens': 16, 'stop': stops}, + optional_params={'max_tokens': 16, 'stop_sequences': stops}, extra_headers=headers, api_key='test', custom_llm_provider='anthropic', opaque=opaque, litellm_logging_obj=logger) state = native.prepare(arguments, logger) -wire_messages, wire_params, wire_headers = native.snapshot(state) -assert wire_messages[0]['content'] == 'edited' -assert wire_params['stop'] == ['first', 'second'] -assert wire_headers == {'x-hook': 'edited'} +wire_body, wire_headers = native.snapshot(state) +assert wire_body['messages'][0]['content'][0]['text'] == 'body edit' +assert wire_body['stop_sequences'] == ['first', 'second'] +assert dict(wire_headers)['x-hook'] == 'edited' arguments['cycle'] = state logger.body['cycle'] = state headers['cycle'] = state diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index 479c77aaed4..af7f017318d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -1,9 +1,11 @@ use litellm_core::lifecycle::FailureStage; use litellm_core::lifecycle::{ErrorDisposition, Lifecycle, Outcome}; use litellm_core::messages::lifecycle::{MessagesRoute, Observations, Operation, Options, machine}; -use litellm_core::messages::types::{MessagesRequest, ProviderMessagesRequest}; -use litellm_core::messages::{execute_prepared_messages_provider_call, prepare_provider_request}; -use litellm_python_interop::{Pythonized, from_py, run_async_value, run_sync_value, to_py}; +use litellm_core::messages::types::{ + MessagesBodySnapshot, MessagesEndpoint, MessagesOptions, ProviderMessagesRequest, +}; +use litellm_core::messages::{execute_prepared_messages_provider_call, prepare_endpoint}; +use litellm_python_interop::{Pythonized, from_py, run_async_value, run_sync_value}; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::pyclass::{PyTraverseError, PyVisit}; @@ -14,32 +16,27 @@ use serde_json::{Map, Value}; use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT}; use crate::errors::{RustUpstreamError, core_error_to_pyerr, messages_provider_error_to_pyerr}; use crate::marshal::optional_timeout; +use crate::retained::RequestRoots; #[pyclass] struct MessagesState { - arguments: Option>, - body: Option>, - headers: Option>, - prepared: Option, + roots: Option, + prepared: Option<(MessagesEndpoint, MessagesBodySnapshot)>, } #[pymethods] impl MessagesState { fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.arguments)?; - visit.call(&self.body)?; - visit.call(&self.headers) + if let Some(roots) = &self.roots { + roots.traverse(&visit)?; + } + Ok(()) } fn __clear__(slf: &Bound<'_, Self>) { let roots = { let mut state = slf.borrow_mut(); - ( - state.arguments.take(), - state.body.take(), - state.headers.take(), - state.prepared.take(), - ) + (state.roots.take(), state.prepared.take()) }; drop(roots); } @@ -53,10 +50,7 @@ fn scalar(arguments: &Bound<'_, PyDict>, name: &str) -> PyResult> .transpose() } -fn decode_request(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { - let body = arguments - .get_item(pyo3::intern!(py, "body"))? - .ok_or_else(|| PyValueError::new_err("messages requires body"))?; +fn decode_options(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { let timeout = optional_timeout( arguments .get_item(pyo3::intern!(py, "timeout_seconds"))? @@ -64,8 +58,7 @@ fn decode_request(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult()) .transpose()?, )?; - Ok(MessagesRequest { - body: from_py(&body)?, + Ok(MessagesOptions { model: scalar(arguments, "model")? .ok_or_else(|| PyValueError::new_err("messages requires model"))?, api_key: scalar(arguments, "api_key")?, @@ -156,20 +149,24 @@ fn prepare( logging: Py, ) -> PyResult> { let bag = arguments.bind(py); - let request = decode_request(py, bag)?; - let prepared = py - .detach(|| prepare_provider_request(request)) + let options = decode_options(py, bag)?; + let endpoint = py + .detach(|| prepare_endpoint(options)) .map_err(core_error_to_pyerr)?; - let body = to_py(py, &prepared.body)? - .into_bound(py) + let body = bag + .get_item("body")? + .ok_or_else(|| PyValueError::new_err("messages requires body"))? .cast_into::()?; + let snapshot = endpoint + .capture_buffered_body(from_py(body.as_any())?) + .map_err(core_error_to_pyerr)?; let headers = PyDict::new(py); - for (name, value) in &prepared.upstream_headers { + for (name, value) in endpoint.headers() { headers.set_item(name, value)?; } let additional = PyDict::new(py); additional.set_item(COMPLETE_INPUT_DICT, &body)?; - additional.set_item(API_BASE, &prepared.url)?; + additional.set_item(API_BASE, endpoint.url())?; additional.set_item(HEADERS, &headers)?; let kwargs = PyDict::new(py); let serialized = py.import("json")?.call_method1("dumps", (&body,))?; @@ -185,39 +182,34 @@ fn prepare( Py::new( py, MessagesState { - arguments: Some(arguments), - body: Some(body.unbind()), - headers: Some(headers.unbind()), - prepared: Some(prepared), + roots: Some(RequestRoots::new( + arguments, + body.unbind().into_any(), + headers.unbind().into_any(), + )), + prepared: Some((endpoint, snapshot)), }, ) } fn take_request(py: Python<'_>, state: &Py) -> PyResult { - let (mut prepared, body, headers) = { + let ((endpoint, snapshot), headers) = { let mut state = state.borrow_mut(py); let prepared = state.prepared.take().ok_or_else(|| { PyRuntimeError::new_err("messages request was already sent or cleared") })?; - let body = state - .body + let roots = state + .roots .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("messages body was cleared"))? - .clone_ref(py); - let headers = state - .headers - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("messages headers were cleared"))? - .clone_ref(py); - (prepared, body, headers) + .ok_or_else(|| PyRuntimeError::new_err("messages roots were cleared"))?; + (prepared, roots.headers(py)) }; - prepared.body = from_py(body.bind(py).as_any())?; - prepared.upstream_headers = headers - .bind(py) + let headers = headers + .cast::()? .iter() .map(|(name, value)| Ok((name.extract()?, value.extract()?))) .collect::>()?; - Ok(prepared) + Ok(endpoint.settle(snapshot, headers)) } fn validate_arguments(arguments: &Bound<'_, PyDict>) -> PyResult<()> { @@ -328,6 +320,72 @@ pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { mod tests { use super::*; + #[pyfunction] + fn snapshot(py: Python<'_>, state: Py) -> PyResult> { + let request = take_request(py, &state)?; + litellm_python_interop::to_py(py, &(request.body(), request.headers())) + } + + #[test] + fn callback_aliases_survive_snapshot_and_roots_are_collectible() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "messages_test").unwrap(); + module.add_function(wrap_pyfunction!(prepare, &module).unwrap()).unwrap(); + module.add_function(wrap_pyfunction!(snapshot, &module).unwrap()).unwrap(); + let globals = PyDict::new(py); + globals.set_item("native", module).unwrap(); + py.run(c" +import gc +import weakref + +class Opaque: + pass + +class Logger: + def pre_call(self, **kwargs): + view = kwargs['additional_args'] + self.body = view['complete_input_dict'] + self.headers = view['headers'] + assert self.body is body + assert self.body['messages'][0] is message + message['content'] = 'changed' + self.headers['x-hook'] = 'changed' + view['complete_input_dict'] = {'replacement': True} + view['headers'] = {'replacement': 'true'} + +message = {'role': 'user', 'content': 'original'} +body = {'model': 'model', 'messages': [message], 'max_tokens': 16} +opaque = Opaque() +logger = Logger() +arguments = dict(model='model', body=body, api_key='test', + custom_llm_provider='anthropic', opaque=opaque, + litellm_logging_obj=logger) +state = native.prepare(arguments, logger) +wire_body, wire_headers = native.snapshot(state) +assert wire_body['messages'][0]['content'] == 'original' +assert body['messages'][0]['content'] == 'changed' +assert dict(wire_headers)['x-hook'] == 'changed' +try: + native.snapshot(state) +except RuntimeError: + pass +else: + raise AssertionError('request was sent twice') +arguments['cycle'] = state +logger.body['cycle'] = state +logger.headers['cycle'] = state +alive = weakref.ref(opaque) +del arguments, logger, opaque, body +gc.collect() +assert alive() is not None +del state +gc.collect() +assert alive() is None +", Some(&globals), Some(&globals)).unwrap(); + }); + } + #[test] fn decode_state_rejects_invalid_timeouts_without_panicking() { Python::initialize(); @@ -337,7 +395,7 @@ mod tests { arguments.set_item("model", "model").unwrap(); arguments.set_item("body", PyDict::new(py)).unwrap(); arguments.set_item("timeout_seconds", timeout).unwrap(); - let error = match decode_request(py, &arguments) { + let error = match decode_options(py, &arguments) { Ok(_) => panic!("invalid timeout should fail normally"), Err(error) => error, }; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 4eb32432a35..d8950b440cf 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -21,13 +21,12 @@ use pyo3::types::PyDict; use crate::driver::{ADDITIONAL_ARGS, API_BASE, API_KEY, COMPLETE_INPUT_DICT, HEADERS, INPUT}; use crate::errors::core_error_to_pyerr; +use crate::retained::RequestRoots; use litellm_python_interop::{run_async_value, run_sync_value}; #[pyclass] struct OcrState { - arguments: Option>, - body: Option>, - headers: Option>, + roots: Option, logging: Option>, pre_call: Option>, endpoint: Option, @@ -38,9 +37,9 @@ struct OcrState { #[pymethods] impl OcrState { fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.arguments)?; - visit.call(&self.body)?; - visit.call(&self.headers)?; + if let Some(roots) = &self.roots { + roots.traverse(&visit)?; + } visit.call(&self.logging)?; visit.call(&self.pre_call) } @@ -49,9 +48,7 @@ impl OcrState { let roots = { let mut state = slf.borrow_mut(); ( - state.arguments.take(), - state.body.take(), - state.headers.take(), + state.roots.take(), state.logging.take(), state.pre_call.take(), state.endpoint.take(), @@ -342,17 +339,17 @@ fn prepare( document_projection, parameter_fields, } = draft; - let body = to_py(py, &draft_body)? - .into_bound(py) - .cast_into::()?; - match document_projection { - OcrDocumentProjection::RetainedDocument => { - body.set_item(pyo3::intern!(py, "document"), &document)? + let body = PyDict::new(py); + for (name, value) in &draft_body { + match (name.as_str(), document_projection) { + ("document", OcrDocumentProjection::RetainedDocument) => { + body.set_item(name, &document)?; + } + ("document", OcrDocumentProjection::ShallowCopyDocument) => { + body.set_item(name, document.copy()?)?; + } + _ => body.set_item(name, to_py(py, value)?)?, } - OcrDocumentProjection::ShallowCopyDocument => { - body.set_item(pyo3::intern!(py, "document"), document.copy()?)? - } - OcrDocumentProjection::Transformed => {} } let optional_params = PyDict::new(py); for &name in parameter_fields { @@ -392,9 +389,11 @@ fn prepare( Py::new( py, OcrState { - arguments: Some(arguments), - body: Some(body.unbind()), - headers: Some(headers.unbind()), + roots: Some(RequestRoots::new( + arguments, + body.unbind().into_any(), + headers.unbind().into_any(), + )), logging: Some(logging), pre_call: Some(pre_call.unbind()), endpoint: Some(endpoint), @@ -435,24 +434,17 @@ fn request(py: Python<'_>, state: &Py) -> PyResult { .endpoint .take() .ok_or_else(|| PyRuntimeError::new_err("OCR request was already sent or cleared"))?; - let body = state - .body + let roots = state + .roots .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("OCR body was cleared"))? - .clone_ref(py); - let headers = state - .headers - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("OCR headers were cleared"))? - .clone_ref(py); + .ok_or_else(|| PyRuntimeError::new_err("OCR roots were cleared"))?; + let body = roots.body(py); + let headers = roots.headers(py); (endpoint, body, headers, state.asynchronous) }; let model = endpoint.model().to_string(); let provider = endpoint.custom_llm_provider().to_string(); - let request = endpoint.settle( - header_pairs(headers.bind(py))?, - from_py(body.bind(py).as_any())?, - ); + let request = endpoint.settle(header_pairs(headers.cast::()?)?, from_py(&body)?); Ok((request, model, provider, asynchronous)) } @@ -465,9 +457,13 @@ fn send(py: Python<'_>, state: Py) -> PyResult> { let call_id = Python::attach(|py| { state .borrow(py) - .arguments + .roots .as_ref() - .and_then(|arguments| scalar(arguments.bind(py), "litellm_call_id").ok().flatten()) + .and_then(|roots| { + scalar(&roots.arguments(py), "litellm_call_id") + .ok() + .flatten() + }) .unwrap_or_default() }); let executed = run_async_value( @@ -527,9 +523,13 @@ fn send_sync(py: Python<'_>, state: Py) -> PyResult> { let error_provider = provider.clone(); let call_id = state .borrow(py) - .arguments + .roots .as_ref() - .and_then(|arguments| scalar(arguments.bind(py), "litellm_call_id").ok().flatten()) + .and_then(|roots| { + scalar(&roots.arguments(py), "litellm_call_id") + .ok() + .flatten() + }) .unwrap_or_default(); let executed = run_sync_value( py, @@ -754,9 +754,7 @@ mod tests { let state = Py::new( py, OcrState { - arguments: None, - body: None, - headers: None, + roots: None, logging: None, pre_call: None, endpoint: None, @@ -969,16 +967,12 @@ asyncio.run(exercise()) #[pyfunction] fn snapshot(py: Python<'_>, state: Py) -> PyResult> { let state = state.borrow(py); - let headers = state - .headers + let roots = state + .roots .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("OCR headers were cleared"))?; - let body = state - .body - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("OCR body was cleared"))?; - let headers = header_pairs(headers.bind(py))?; - let body: serde_json::Value = from_py(body.bind(py).as_any())?; + .ok_or_else(|| PyRuntimeError::new_err("OCR roots were cleared"))?; + let headers = header_pairs(roots.headers(py).cast::()?)?; + let body: serde_json::Value = from_py(&roots.body(py))?; to_py(py, &(headers, body)) } diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 845dcfd546f..f454dacd822 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2442,7 +2442,6 @@ class BaseLLMHTTPHandler: from litellm.rust_bridge import messages as rust_messages_bridge - upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"} rust_response: Final = await rust_messages_bridge.amessages( arguments=arguments, request_arguments=request_arguments, @@ -2451,7 +2450,7 @@ class BaseLLMHTTPHandler: messages=messages, lifecycle_owner=rust_messages_bridge.LifecycleOwner.WRAPPER, model=model, - body=upstream_body, + body=request_body, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, diff --git a/tests/test_litellm_rust/test_messages_callbacks.py b/tests/test_litellm_rust/test_messages_callbacks.py index 3f26409b9be..ae4c87de1dc 100644 --- a/tests/test_litellm_rust/test_messages_callbacks.py +++ b/tests/test_litellm_rust/test_messages_callbacks.py @@ -59,15 +59,18 @@ async def test_messages_pre_call_receives_expected_provider_request(messages_ser "model": "claude-sonnet-4-5-20250929", "messages": MESSAGES, "max_tokens": 64, + "stream": False, } assert additional_args["headers"]["x-api-key"] == "test-key" @pytest.mark.asyncio @pytest.mark.parametrize("raise_after_edit", [False, True]) +@pytest.mark.parametrize("native", [False, True]) async def test_messages_pre_call_edits_reach_later_callbacks_and_provider( - messages_server: RecordingServer, raise_after_edit: bool + messages_server: RecordingServer, raise_after_edit: bool, native: bool ) -> None: + litellm.rust(native) observed: Final = [] class Edit(CustomLogger): @@ -85,7 +88,7 @@ async def test_messages_pre_call_edits_reach_later_callbacks_and_provider( assert observed[0][0]["temperature"] == 0.25 assert observed[0][1]["x-audit-tag"] == "reviewed" - assert messages_server.requests[0].body["temperature"] == 0.25 + assert "temperature" not in messages_server.requests[0].body assert messages_server.requests[0].headers["x-audit-tag"] == "reviewed" @@ -129,6 +132,52 @@ async def test_messages_pre_call_state_reaches_terminal_callbacks(messages_serve assert observed == [token] +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +@pytest.mark.parametrize("raise_after_edit", [False, True]) +async def test_messages_retained_aliases_preserve_identity_and_snapshot_timing( + messages_server: RecordingServer, native: bool, provider: str, raise_after_edit: bool +) -> None: + litellm.rust(native) + retained: Final = [] + observed: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + body = request_body(kwargs) + message = kwargs["messages"][0] + assert body["messages"][0] is message + assert body["messages"][0]["content"] is message["content"] + assert body["messages"][0]["content"][0] is message["content"][0] + retained.append(message["content"][0]) + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + retained[0]["text"] = "changed through retained reference" + if raise_after_edit: + raise RuntimeError("export failed after mutation") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + block = request_body(kwargs)["messages"][0]["content"][0] + observed.append((block is retained[0], block["text"])) + + response: Final = await litellm.anthropic.messages.acreate( + model=MESSAGES_MODEL.replace("anthropic/", f"{provider}/"), + messages=[{"role": "user", "content": [{"type": "text", "text": "original"}]}], + max_tokens=64, + api_key="test-key", + api_base=messages_server.base_url, + callbacks=[Retain(), Edit(), Observe()], + ) + assert observed == [(True, "changed through retained reference")] + assert retained[0]["text"] == "changed through retained reference" + assert messages_server.requests[0].body["messages"][0]["content"][0]["text"] == "original" + if native: + assert response["_hidden_params"]["additional_headers"]["x-litellm-rust"] == "true" + + @pytest.mark.asyncio async def test_messages_callbacks_run_once(messages_server: RecordingServer) -> None: recorder: Final = RecordingLogger() diff --git a/tests/test_litellm_rust/test_ocr_callbacks.py b/tests/test_litellm_rust/test_ocr_callbacks.py index e0cbb446c59..174c50ef239 100644 --- a/tests/test_litellm_rust/test_ocr_callbacks.py +++ b/tests/test_litellm_rust/test_ocr_callbacks.py @@ -96,28 +96,43 @@ def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: Re assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" -def test_pre_call_nested_mutation_updates_retained_references(ocr_server: RecordingServer) -> None: +@pytest.mark.asyncio +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_pre_call_nested_mutation_updates_retained_references( + ocr_server: RecordingServer, native: bool, asynchronous: bool +) -> None: + litellm.rust(native) original: Final = dict(OCR_DOCUMENT) replacement_url: Final = "data:application/pdf;base64,ZGVm" retained: Final = [] class Retain(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + assert request_body(kwargs)["document"] is original retained.append(request_body(kwargs)["document"]) class Edit(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): - request_body(kwargs)["document"]["document_url"] = replacement_url + original["document_url"] = replacement_url - call_native_ocr( - ocr_server, - document=original, - callbacks=[Retain(), Edit()], - ) + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": original, + "api_key": "test-key", + "api_base": ocr_server.base_url, + "callbacks": [Retain(), Edit()], + } + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) assert retained[0]["document_url"] == replacement_url assert original["document_url"] == replacement_url assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url + if native: + assert ocr_server.requests[0].headers["accept-encoding"] == "identity" def test_pre_call_field_replacement_preserves_original_references(ocr_server: RecordingServer) -> None: diff --git a/tests/test_litellm_rust/test_sdk_dispatch.py b/tests/test_litellm_rust/test_sdk_dispatch.py index db6fa85751b..3bca62cd5d6 100644 --- a/tests/test_litellm_rust/test_sdk_dispatch.py +++ b/tests/test_litellm_rust/test_sdk_dispatch.py @@ -71,13 +71,22 @@ def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingSer @pytest.mark.parametrize("provider", ["anthropic", "bedrock"]) @pytest.mark.parametrize("rebind_logging_view", [False, True]) @pytest.mark.parametrize("status", [200, 429]) +@pytest.mark.parametrize("native", [False, True]) async def test_chat_retains_callback_edits_through_public_dispatch( - recording_server: RecordingServer, asynchronous: bool, provider: str, rebind_logging_view: bool, status: int + recording_server: RecordingServer, + asynchronous: bool, + provider: str, + rebind_logging_view: bool, + status: int, + native: bool, ) -> None: import threading + import json + + litellm.rust(native) from tests.test_litellm_rust.callback_recorder import RecordingLogger - from tests.test_litellm_rust.contracts import MESSAGES_RESPONSE, request_body, request_headers + from tests.test_litellm_rust.contracts import MESSAGES_RESPONSE recording_server.default_response = ResponseSpec( status=status, @@ -98,10 +107,15 @@ async def test_chat_retains_callback_edits_through_public_dispatch( class EditingLogger(RecordingLogger): def log_pre_api_call(self, model, messages, kwargs): super().log_pre_api_call(model, messages, kwargs) - body = request_body(kwargs) - headers = request_headers(kwargs) - body["messages"][0]["content"] = "edited by callback" - body["max_tokens" if provider == "anthropic" else "maxTokens"] = 32 + body = kwargs["additional_args"]["complete_input_dict"] + headers = kwargs["additional_args"]["headers"] + if provider == "anthropic": + assert isinstance(body, dict) + body["messages"][0]["content"][0]["text"] = "edited by callback" + body["max_tokens"] = 32 + else: + assert isinstance(body, str) + assert json.loads(body)["messages"][0]["content"][0]["text"] == "original" headers["x-retained-callback"] = "original" observations.append((threading.current_thread(), body, headers)) if rebind_logging_view: @@ -139,12 +153,20 @@ async def test_chat_retains_callback_edits_through_public_dispatch( event_name: Final = "async_log_success_event" if asynchronous else "log_success_event" events: Final = await recorder.wait_for_async(event_name) - assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" + if native: + assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" assert recorder.names.count("log_pre_api_call") == 1 assert recorder.names.count(event_name) == 1 - assert observations[0][0] is caller_thread + if asynchronous and provider == "anthropic" and not native: + assert observations[0][0] is not caller_thread + else: + assert observations[0][0] is caller_thread assert events[0].response is response - assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == "edited by callback" + assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == ( + "edited by callback" if provider == "anthropic" else "original" + ) assert recording_server.requests[0].headers["x-retained-callback"] == "original" body: Final = recording_server.requests[0].body - assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == 32 + assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == ( + 32 if provider == "anthropic" else 64 + )