diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index fe7e8bb4b80..de1a5f476ed 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -17,8 +17,10 @@ pub(super) async fn send( body: &Value, timeout: Option, ) -> Result { + let encoded = serde_json::to_vec(body) + .map_err(|err| Error::InvalidRequest(format!("failed to encode messages body: {err}")))?; let builder = headers.iter().fold( - http_client().post(url).json(body), + http_client().post(url).body(encoded), |builder, (key, value)| builder.header(key, value), ); let builder = match timeout { diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index fc1a9b63252..40aff185e81 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -6,7 +6,6 @@ use std::{ use bytes::Bytes; use litellm_auth::SecretValue; -use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{ event::{MachineEvent, RawResponse, RequestContext, WireRequest}, host::{Demand, Host}, @@ -22,7 +21,6 @@ use serde_json::{Map, Value}; use super::{ Error, - common_utils::messages_provider_config, handler::{decode_response, network, provider_error, send}, prepare::{prepare_provider_request, resolve_provider}, types::{MessagesRequest, MessagesShaping}, @@ -54,6 +52,11 @@ pub enum MessagesOutput { Streamed, } +/// The upstream response as the caller sees it at stream hand-off, before any chunk. +pub struct MessagesStreamHead { + pub headers: Vec<(String, String)>, +} + pub struct Messages; impl Protocol for Messages { @@ -62,7 +65,7 @@ impl Protocol for Messages { type Projection = MessagesCall; type Op = Infallible; type Chunk = Bytes; - type StreamHead = (); + type StreamHead = MessagesStreamHead; } impl From for Error { @@ -77,19 +80,6 @@ impl From for Error { pub type MessagesHost = HostChannel; pub type MessagesMachine = CallMachine; -/// Whether this route serves the request, decided before any callback runs so a host -/// can still run its own path. -pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool { - let provider = get_custom_llm_provider(model, custom_llm_provider) - .map(|resolved| resolved.custom_llm_provider) - .or(custom_llm_provider); - match provider { - Some(ANTHROPIC_MESSAGES_PROVIDER) => true, - Some(provider) => !stream && messages_provider_config(provider).is_some(), - None => false, - } -} - /// The in-process host for a request already in hand. It answers projection once and /// observes nothing. pub struct LocalMessagesHost { @@ -152,8 +142,11 @@ async fn execute( model: request.model.clone(), custom_llm_provider: request.provider.clone(), optional_params: Value::Object( - call.body - .iter() + request + .body + .as_object() + .into_iter() + .flatten() .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) .map(|(name, value)| (name.clone(), value.clone())) .collect(), @@ -193,7 +186,14 @@ async fn relay( host: &MessagesHost, mut response: reqwest::Response, ) -> Result { - if host.open(()).await? == Demand::Detached { + let head = MessagesStreamHead { + headers: response + .headers() + .iter() + .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) + .collect(), + }; + if host.open(head).await? == Demand::Detached { return Ok(MessagesOutput::Streamed); } while let Some(chunk) = response.chunk().await.map_err(network)? { diff --git a/litellm-rust/crates/core/tests/messages/host.rs b/litellm-rust/crates/core/tests/messages/host.rs new file mode 100644 index 00000000000..ca2aece5ebd --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/host.rs @@ -0,0 +1,210 @@ +use std::{convert::Infallible, sync::Mutex}; + +use litellm_core::messages::route::Messages; +use litellm_host::{ + event::{CallEvent, MachineEvent, RequestContext, WireRequest}, + host::Host, +}; +use litellm_llms::anthropic::common_utils::AnthropicModelCapabilities; +use rstest::rstest; + +use super::*; + +type Rewrite = Box Result + Send + Sync>; + +/// Projects like `LocalMessagesHost`, answers `before_send` through `rewrite`, and keeps +/// every event the driver emits. +struct RecordingHost { + call: LocalMessagesHost, + rewrite: Rewrite, + events: Mutex>, + optional_params: Mutex>, +} + +impl RecordingHost { + fn new(call: MessagesCall, rewrite: Rewrite) -> Self { + Self { + call: LocalMessagesHost::new(call), + rewrite, + events: Mutex::new(Vec::new()), + optional_params: Mutex::new(Vec::new()), + } + } + + fn passthrough(call: MessagesCall) -> Self { + Self::new(call, Box::new(Ok)) + } + + fn raw_responses(&self) -> Vec { + self.events + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + CallEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + Some(raw.body.clone()) + } + _ => None, + }) + .collect() + } +} + +impl Host for RecordingHost { + async fn project(&self) -> Result { + self.call.project().await + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} + } + + async fn before_send( + &self, + wire: WireRequest, + context: &RequestContext, + ) -> Result { + self.optional_params + .lock() + .unwrap() + .push(context.optional_params.clone()); + (self.rewrite)(wire) + } + + async fn emit(&self, event: &CallEvent) -> Result<(), Error> { + self.events.lock().unwrap().push(event.clone()); + Ok(()) + } +} + +async fn run_through(host: &RecordingHost) -> Result { + litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await +} + +fn authenticated(call: MessagesCall, api_base: String) -> MessagesCall { + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(api_base), + ..call + } +} + +#[rstest] +#[tokio::test] +async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::new( + authenticated(call, upstream.uri()), + Box::new(|wire| { + let mut body = wire.body; + body["system"] = json!("added by the host"); + Ok(WireRequest { + headers: wire + .headers + .into_iter() + .chain([("x-host".to_string(), "seen".to_string())]) + .collect(), + body, + ..wire + }) + }), + ); + + run_through(&host).await.expect("messages call succeeds"); + + let request = only_request(&upstream).await; + assert_eq!(request.json()["system"], "added by the host"); + assert_eq!(request.header("x-host"), Some("seen")); + assert_eq!(request.header("x-api-key"), Some("sk-ant")); +} + +#[rstest] +#[tokio::test] +async fn a_before_send_failure_never_sends(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let host = RecordingHost::new( + authenticated(call, upstream.uri()), + Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))), + ); + + let error = run_through(&host) + .await + .err() + .expect("the host failure fails the call"); + + assert_eq!(error, Error::InvalidRequest("vetoed by the host".into())); + assert!(received(&upstream).await.is_empty()); + assert!(host.raw_responses().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn the_raw_upstream_text_is_emitted_once_for_a_message(call: MessagesCall) { + let raw = message_body(); + let upstream = upstream([json_response(raw.clone())]).await; + let host = RecordingHost::passthrough(authenticated(call, upstream.uri())); + + let output = run_through(&host).await.expect("messages call succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let [emitted] = <[String; 1]>::try_from(host.raw_responses()) + .unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len())); + assert_eq!(serde_json::from_str::(&emitted).unwrap(), raw); +} + +#[rstest] +#[case::upstream_error(ResponseTemplate::new(500).set_body_string("boom"))] +#[case::stream(ResponseTemplate::new(200).set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream"))] +#[tokio::test] +async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( + call: MessagesCall, + #[case] response: ResponseTemplate, +) { + let upstream = upstream([response]).await; + let mut body = call.body.clone(); + body.insert("stream".into(), json!(true)); + let host = + RecordingHost::passthrough(authenticated(MessagesCall { body, ..call }, upstream.uri())); + + let _ = run_through(&host).await; + + assert_eq!(received(&upstream).await.len(), 1); + assert!(host.raw_responses().is_empty()); +} + +/// Python logs `optional_params` as what it is about to send, so a dropped param must +/// not resurface in callbacks. +#[rstest] +#[tokio::test] +async fn the_request_context_carries_the_shaped_params_without_model_or_messages( + call: MessagesCall, +) { + let upstream = upstream([message_response()]).await; + let body: Map = call + .body + .clone() + .into_iter() + .chain([("temperature".to_string(), json!(0.2))]) + .collect(); + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + body, + shaping: MessagesShaping { + capabilities: AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + }, + drop_params: true, + ..MessagesShaping::default() + }, + ..call + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_send runs once, saw {}", seen.len())); + assert_eq!(optional_params, json!({"max_tokens": 16})); +} diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs index 4e549bae309..21ee678ced3 100644 --- a/litellm-rust/crates/core/tests/messages/main.rs +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -14,6 +14,7 @@ use wiremock::ResponseTemplate; mod support; use support::*; +mod host; mod request; mod response; mod secrets; diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index 9353324d370..2927356b773 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -1,3 +1,7 @@ +use litellm_llms::anthropic::common_utils::{ + ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_OAUTH_BETA_HEADER, AnthropicModelCapabilities, + SupportedEffortTiers, beta, +}; use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; use rstest::rstest; @@ -132,8 +136,8 @@ async fn each_provider_posts_to_its_messages_endpoint( assert_eq!(request.method.as_str(), "POST"); assert_eq!(request.url.path(), path); assert_eq!(request.json()["model"], MODEL); - assert_eq!(request.header("anthropic-version"), Some("2023-06-01")); - assert_eq!(request.header("content-type"), Some("application/json")); + assert_eq!(request.header_values("anthropic-version"), ["2023-06-01"]); + assert_eq!(request.header_values("content-type"), ["application/json"]); } #[rstest] @@ -249,21 +253,423 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) assert_eq!(sent["top_k"], 3); } +fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall { + let body: Map = call.body.into_iter().chain(object(fields)).collect(); + MessagesCall { body, ..call } +} + +fn sent_betas(request: &wiremock::Request) -> Vec { + let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta")) + .unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}")); + header + .split(',') + .map(str::trim) + .map(str::to_string) + .collect() +} + #[rstest] -#[case::anthropic_streams(MODEL, Some("anthropic"), true, true)] -#[case::anthropic_prefix_streams("anthropic/claude-sonnet-4-5", None, true, true)] -#[case::azure_without_stream(MODEL, Some("azure_ai"), false, true)] -#[case::azure_stream(MODEL, Some("azure_ai"), true, false)] -#[case::other_provider(MODEL, Some("openai"), false, false)] -#[case::unresolvable_model("no-such-model", None, false, false)] -fn supports_matches_what_the_route_can_serve( - #[case] model: &str, - #[case] provider: Option<&str>, - #[case] stream: bool, - #[case] supported: bool, +#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] +#[case::fast_mode(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] +#[case::compaction(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] +#[case::context_management_edits( + json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}), + &[beta::CONTEXT_MANAGEMENT_2025_06_27] +)] +#[case::per_message_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &[beta::PER_TURN_CONTROL_2026_07_01] +)] +#[case::advisor_tool( + json!({"tools": [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": MODEL}]}), + &[beta::ADVISOR_TOOL_2026_03_01] +)] +#[case::several_features_at_once( + json!({"speed": "fast", "output_format": {"type": "json_schema"}}), + &[beta::STRUCTURED_OUTPUT, beta::FAST_MODE_2026_02_01] +)] +#[tokio::test] +async fn feature_betas_join_the_callers_betas_in_one_sorted_header( + call: MessagesCall, + #[case] fields: Value, + #[case] features: &[&str], ) { + let upstream = upstream([message_response()]).await; + let capabilities = AnthropicModelCapabilities { + supports_speed: true, + ..AnthropicModelCapabilities::default() + }; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..call + }, + fields, + )) + .await; + + let sent = sent_betas(&only_request(&upstream).await); + let mut expected: Vec = features + .iter() + .map(|feature| feature.to_string()) + .chain(["caller-beta-2025-01-01".to_string()]) + .collect(); + expected.sort(); + assert_eq!(sent, expected); +} + +#[rstest] +#[tokio::test] +async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + api_key: Some("sk-ant-oat01-token".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + let request = only_request(&upstream).await; assert_eq!( - litellm_core::messages::route::supports(model, provider, stream), - supported + request.header("anthropic-dangerous-direct-browser-access"), + Some("true") + ); + assert_eq!(sent_betas(&request), [ANTHROPIC_OAUTH_BETA_HEADER]); + assert_eq!(request.header("x-api-key"), None); +} + +#[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] +#[tokio::test] +async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([ + ("Anthropic-Version", "2024-01-01"), + ("Content-Type", "application/json; charset=utf-8"), + ]), + ..call + }) + .await; + + let request = only_request(&upstream).await; + assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]); + assert_eq!( + request.header_values("content-type"), + ["application/json; charset=utf-8"] ); } + +fn sampling_removed() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_sampling_params: false, + ..AnthropicModelCapabilities::default() + } +} + +#[rstest] +#[case::sampling_params(sampling_removed(), json!({"temperature": 0.2, "top_p": 0.9, "top_k": 5}), &["temperature", "top_p", "top_k"], "temperature=0.2")] +#[case::speed(AnthropicModelCapabilities::default(), json!({"speed": "fast"}), &["speed"], "speed='fast'")] +#[tokio::test] +async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_it( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, + #[case] dropped: &[&str], + #[case] rejected_as: &str, +) { + let upstream = upstream([message_response(), message_response()]).await; + let shaped = |drop_params: bool| { + with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: capabilities.clone(), + drop_params, + ..MessagesShaping::default() + }, + body: call.body.clone(), + custom_llm_provider: call.custom_llm_provider.clone(), + extra_headers: None, + provider_specific_header: None, + model: call.model.clone(), + timeout: call.timeout, + }, + fields.clone(), + ) + }; + + let error = run(shaped(false)) + .await + .err() + .expect("an unsupported param is rejected without drop_params"); + assert!( + matches!(&error, Error::InvalidRequest(message) if message.contains(rejected_as)), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); + + run_message(shaped(true)).await; + let sent = only_request(&upstream).await.json(); + for name in dropped { + assert_eq!(sent.get(*name), None, "{name} must be dropped"); + } + assert_eq!(sent["max_tokens"], 16); +} + +#[rstest] +#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))] +#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))] +#[tokio::test] +async fn reasoning_auto_summary_marks_active_thinking_on_the_wire( + call: MessagesCall, + #[case] thinking: Value, + #[case] expected: Value, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities: AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + ..AnthropicModelCapabilities::default() + }, + reasoning_auto_summary: true, + ..MessagesShaping::default() + }, + ..call + }, + json!({"thinking": thinking}), + )) + .await; + + assert_eq!(only_request(&upstream).await.json()["thinking"], expected); +} + +#[rstest] +#[case::reasoning_effort_on_an_adaptive_model( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { high: true, ..SupportedEffortTiers::default() }, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "high"}), + json!({"thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) +)] +#[case::reasoning_effort_on_a_legacy_model_caps_the_budget_below_max_tokens( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"reasoning_effort": "high"}), + json!({"thinking": {"type": "enabled", "budget_tokens": 2999}}) +)] +#[case::adaptive_payload_on_a_legacy_model_becomes_a_capped_budget( + AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, "temperature": 0}), + json!({"thinking": {"type": "enabled", "budget_tokens": 2999}}) +)] +#[case::adaptive_payload_on_a_model_without_reasoning_is_dropped( + AnthropicModelCapabilities::default(), + json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}), + json!({}) +)] +#[tokio::test] +async fn reasoning_is_translated_by_the_model_capabilities( + call: MessagesCall, + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, + #[case] expected: Value, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + shaping: MessagesShaping { + capabilities, + ..MessagesShaping::default() + }, + ..call + }, + [("max_tokens".to_string(), json!(3000))] + .into_iter() + .chain(object(fields)) + .collect(), + )) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("reasoning_effort"), None); + assert_eq!(sent.get("temperature"), None); + let reasoning: Map = ["thinking", "output_config"] + .into_iter() + .filter_map(|name| Some((name.to_string(), sent.get(name)?.clone()))) + .collect(); + assert_eq!(Value::Object(reasoning), expected); +} + +#[rstest] +#[case::empty_text_blocks( + json!([{"role": "assistant", "content": [{"type": "text", "text": " "}, {"type": "text", "text": "kept"}]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}]) +)] +#[case::provider_specific_fields( + json!([{"role": "assistant", "content": [{"type": "text", "text": "kept", "provider_specific_fields": {"x": 1}}]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}]) +)] +#[case::unencrypted_web_search_results_become_text( + json!([{"role": "assistant", "content": [{ + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "title": "T", "url": "https://e.x", "page_age": null}] + }]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "Web search results:\n\nTitle: T\nURL: https://e.x"}]}]) +)] +#[tokio::test] +async fn replayed_history_is_cleaned_before_sending( + call: MessagesCall, + #[case] history: Value, + #[case] expected: Value, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({"messages": history}), + )) + .await; + + assert_eq!(only_request(&upstream).await.json()["messages"], expected); +} + +#[rstest] +#[tokio::test] +async fn metadata_is_reduced_to_the_user_id(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({"metadata": {"user_id": "u-1", "trace_id": "internal", "tags": ["a"]}}), + )) + .await; + + assert_eq!( + only_request(&upstream).await.json()["metadata"], + json!({"user_id": "u-1"}) + ); +} + +#[rstest] +#[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))] +#[case::missing_max_tokens(json!({"max_tokens": null}))] +#[tokio::test] +async fn an_invalid_request_fails_before_sending(call: MessagesCall, #[case] fields: Value) { + let upstream = upstream([message_response()]).await; + + let error = run(with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + fields, + )) + .await + .err() + .expect("the request is rejected"); + + assert!(error.is_request(), "{error:?}"); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[tokio::test] +async fn azure_folds_system_role_messages_into_the_system_prompt(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + api_key: Some("sk-azure".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({ + "system": "top level", + "messages": [ + {"role": "system", "content": "from a message"}, + {"role": "user", "content": "hi"} + ] + }), + )) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!( + sent["system"], + json!([ + {"type": "text", "text": "top level"}, + {"type": "text", "text": "from a message"} + ]) + ); + assert_eq!(sent["messages"], json!([{"role": "user", "content": "hi"}])); +} + +#[rstest] +#[case::bare_model(MODEL, MODEL)] +#[case::one_prefix("anthropic/claude-sonnet-4-5", MODEL)] +#[case::doubled_prefix_loses_one_segment( + "anthropic/anthropic/claude-sonnet-4-5", + "anthropic/claude-sonnet-4-5" +)] +#[tokio::test] +async fn the_provider_prefix_is_stripped_exactly_once( + call: MessagesCall, + #[case] model: &str, + #[case] sent_model: &str, +) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + model: model.into(), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(only_request(&upstream).await.json()["model"], sent_model); +} diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs index 38a18c415ba..133b7d2b162 100644 --- a/litellm-rust/crates/core/tests/messages/response.rs +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -5,11 +5,14 @@ use rstest::rstest; use super::*; #[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] #[tokio::test] -async fn the_provider_message_is_returned(call: MessagesCall) { +async fn the_provider_message_is_returned(call: MessagesCall, #[case] provider: &str) { let upstream = upstream([message_response()]).await; let message = run_message(MessagesCall { + custom_llm_provider: Some(provider.into()), api_key: Some("sk".into()), api_base: Some(upstream.uri()), ..call @@ -21,6 +24,88 @@ async fn the_provider_message_is_returned(call: MessagesCall) { assert_eq!(message.stop_reason.as_deref(), Some("end_turn")); } +/// A refusal and fields the route does not model come back exactly as the provider sent +/// them, since the Python side returns the raw message and the router decides what to do. +#[rstest] +#[tokio::test] +async fn the_message_passes_through_losslessly(call: MessagesCall) { + let upstream_body = json!({ + "id": "msg_2", + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}}, + {"type": "text", "text": "no", "citations": [{"type": "web_search_result_location", "url": "https://e.x"}]} + ], + "stop_reason": "refusal", + "stop_sequence": null, + "stop_details": {"type": "safeguard", "safeguard_types": ["dangerous_tool_use"]}, + "container": {"id": "container_1", "expires_at": "2026-01-01T00:00:00Z"}, + "context_management": {"applied_edits": []}, + "usage": {"input_tokens": 1, "output_tokens": 2, "server_tool_use": {"web_search_requests": 1}}, + "unknown_future_field": {"nested": true} + }); + let upstream = upstream([json_response(upstream_body.clone())]).await; + + let message = run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(message.stop_reason.as_deref(), Some("refusal")); + assert_eq!(serde_json::to_value(&message).unwrap(), upstream_body); +} + +#[rstest] +#[tokio::test] +async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) { + let envelope = + json!({"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}); + let upstream = upstream([status_response(400, envelope.clone())]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("upstream error propagates"); + + let Error::Transport(TransportError::Http { status, body }) = error else { + panic!("{error:?}"); + }; + assert_eq!(status, 400); + assert_eq!(serde_json::from_str::(&body).unwrap(), envelope); +} + +#[rstest] +#[tokio::test] +async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) { + let long = "x".repeat(600); + let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("upstream error propagates"); + + assert_eq!( + error, + Error::Transport(TransportError::Http { + status: 500, + body: format!("{}... (truncated)", &long[..256]) + }) + ); +} + #[rstest] #[case::bad_request(400)] #[case::unauthorized(401)] diff --git a/litellm-rust/crates/core/tests/messages/secrets.rs b/litellm-rust/crates/core/tests/messages/secrets.rs index 419b6d6c753..55e510d00d3 100644 --- a/litellm-rust/crates/core/tests/messages/secrets.rs +++ b/litellm-rust/crates/core/tests/messages/secrets.rs @@ -1,23 +1,30 @@ -use litellm_llms::{ - anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, - azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, - base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, -}; use rstest::rstest; use super::*; #[rstest] -#[case::anthropic("anthropic", &ANTHROPIC_MESSAGES_CONFIG, "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "/v1/messages")] -#[case::azure_ai("azure_ai", &AZURE_ANTHROPIC_MESSAGES_CONFIG, "AZURE_API_KEY", "AZURE_API_BASE", "/anthropic/v1/messages")] +#[case::anthropic( + "anthropic", + "ANTHROPIC_API_KEY", + "ANTHROPIC_BASE_URL", + "/v1/messages", + &["ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"] +)] +#[case::azure_ai( + "azure_ai", + "AZURE_API_KEY", + "AZURE_API_BASE", + "/anthropic/v1/messages", + &["AZURE_API_KEY", "AZURE_API_BASE"] +)] #[tokio::test] async fn the_credential_and_base_come_from_the_secret_source( call: MessagesCall, #[case] provider: &str, - #[case] config: &dyn BaseAnthropicMessagesConfig, #[case] key_name: &str, #[case] base_name: &str, #[case] path: &str, + #[case] looked_up: &[&str], ) { let upstream = upstream([message_response()]).await; let base = upstream.uri(); @@ -40,7 +47,7 @@ async fn the_credential_and_base_come_from_the_secret_source( let request = only_request(&upstream).await; assert_eq!(request.url.path(), path); assert_eq!(request.header("x-api-key"), Some("sk-from-manager")); - assert_eq!(secrets.requested(), config.secret_names()); + assert_eq!(secrets.requested(), looked_up); } #[rstest] @@ -92,3 +99,102 @@ async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCa ); assert!(received(&upstream).await.is_empty()); } + +#[derive(Clone, Copy)] +enum Base { + Upstream, + Unreachable, + Blank, + Absent, +} + +fn base_value(base: Base, upstream: &str) -> Option { + match base { + Base::Upstream => Some(upstream.to_string()), + Base::Unreachable => Some(UNREACHABLE_BASE.to_string()), + Base::Blank => Some(" ".to_string()), + Base::Absent => None, + } +} + +#[rstest] +#[case::api_base_beats_base_url(Base::Upstream, Base::Unreachable)] +#[case::blank_api_base_falls_through_to_base_url(Base::Blank, Base::Upstream)] +#[case::base_url_alone(Base::Absent, Base::Upstream)] +#[tokio::test] +async fn the_anthropic_base_env_precedence_picks_the_upstream( + call: MessagesCall, + #[case] api_base: Base, + #[case] base_url: Base, +) { + let upstream = upstream([message_response()]).await; + let uri = upstream.uri(); + let values: Vec<(&str, &str)> = [ + ("ANTHROPIC_API_KEY", Some("sk-env".to_string())), + ("ANTHROPIC_API_BASE", base_value(api_base, &uri)), + ("ANTHROPIC_BASE_URL", base_value(base_url, &uri)), + ] + .iter() + .filter_map(|(name, value)| Some((*name, value.as_deref()?))) + .map(|(name, value)| (name, Box::leak(value.to_string().into_boxed_str()) as &str)) + .collect(); + + run_with(Arc::new(RecordingSecrets::new(values)), call) + .await + .expect("messages call reaches the upstream the precedence picks"); + + assert_eq!(only_request(&upstream).await.url.path(), "/v1/messages"); +} + +#[rstest] +#[case::auth_token_alone( + &[("ANTHROPIC_AUTH_TOKEN", "tok")], + ("authorization", "Bearer tok"), + "x-api-key" +)] +#[case::api_key_beats_the_auth_token( + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "tok")], + ("x-api-key", "sk-env"), + "authorization" +)] +#[tokio::test] +async fn the_auth_token_env_is_a_bearer_only_without_a_key( + call: MessagesCall, + #[case] values: &[(&str, &str)], + #[case] expected: (&str, &str), + #[case] absent: &str, +) { + let upstream = upstream([message_response()]).await; + + run_with( + Arc::new(RecordingSecrets::new(values.iter().copied())), + MessagesCall { + api_base: Some(upstream.uri()), + ..call + }, + ) + .await + .expect("messages call succeeds"); + + let request = only_request(&upstream).await; + let (name, value) = expected; + assert_eq!(request.header_values(name), [value]); + assert_eq!(request.header(absent), None); +} + +#[rstest] +#[tokio::test] +async fn azure_without_a_base_anywhere_fails_before_sending(call: MessagesCall) { + let error = run_with( + Arc::new(RecordingSecrets::new([("AZURE_API_KEY", "sk-azure")])), + MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + ..call + }, + ) + .await + .err() + .expect("azure needs a base"); + + assert_eq!(error, Error::Auth(litellm_auth::Error::MissingAzureApiBase)); +} diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs index ea23a9e8e38..c4be3127d66 100644 --- a/litellm-rust/crates/core/tests/messages/stream.rs +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -1,16 +1,25 @@ use std::{convert::Infallible, sync::Mutex}; use bytes::Bytes; -use litellm_core::messages::route::Messages; +use litellm_core::messages::route::{Messages, MessagesStreamHead}; use litellm_host::host::{Demand, Host}; use rstest::rstest; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; use super::*; +const UPSTREAM_HEADERS: [(&str, &str); 2] = [ + ("request-id", "req_upstream_123"), + ("anthropic-ratelimit-requests-remaining", "41"), +]; + const SSE_BODY: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; enum Seen { - Open, + Open(Vec<(String, String)>), Deliver(Bytes), } @@ -50,8 +59,8 @@ impl Host for RecordingStreamHost { match op {} } - async fn open(&self, (): ()) -> Result { - Ok(self.record(Seen::Open)) + async fn open(&self, head: MessagesStreamHead) -> Result { + Ok(self.record(Seen::Open(head.headers))) } async fn deliver(&self, chunk: Bytes) -> Result { @@ -71,7 +80,10 @@ fn streaming(call: MessagesCall, api_base: String) -> MessagesCall { } fn sse_response() -> ResponseTemplate { - ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream") + UPSTREAM_HEADERS.iter().fold( + ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream"), + |response, (name, value)| response.insert_header(*name, *value), + ) } async fn stream_through(host: &RecordingStreamHost) -> Result { @@ -80,7 +92,7 @@ async fn stream_through(host: &RecordingStreamHost) -> Result = headers + .iter() + .filter(|(name, _)| { + UPSTREAM_HEADERS + .iter() + .any(|(upstream, _)| upstream == name) + }) + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect(); + assert_eq!(surfaced, UPSTREAM_HEADERS); let delivered: Vec = chunks .iter() .flat_map(|step| match step { Seen::Deliver(chunk) => chunk.to_vec(), - Seen::Open => panic!("the stream opens exactly once"), + Seen::Open(_) => panic!("the stream opens exactly once"), }) .collect(); assert_eq!(delivered, SSE_BODY.as_bytes()); @@ -118,9 +140,18 @@ async fn a_detached_caller_receives_nothing_more(call: MessagesCall, #[case] det } #[rstest] +#[case::text_body(ResponseTemplate::new(429).set_body_string("slow down"), "slow down")] +#[case::json_envelope( + status_response(429, json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}})), + r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"# +)] #[tokio::test] -async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: MessagesCall) { - let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await; +async fn an_upstream_error_fails_the_call_without_opening_the_stream( + call: MessagesCall, + #[case] response: ResponseTemplate, + #[case] body: &str, +) { + let upstream = upstream([response]).await; let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); let error = stream_through(&host) @@ -128,16 +159,88 @@ async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: Messa .err() .expect("upstream error propagates"); - assert!( - matches!( - error, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "{error:?}" + assert_eq!( + error, + Error::Transport(litellm_http::transport::Error::Http { + status: 429, + body: body.into() + }) ); assert!(host.seen.into_inner().unwrap().is_empty()); } +/// The native route relays bytes as they are. Python's synthetic `api_error` for a stream +/// that never reaches `message_stop` lives in its SSE wrapper, above this route. +#[rstest] +#[tokio::test] +async fn a_stream_that_ends_without_message_stop_is_relayed_as_is(call: MessagesCall) { + const INCOMPLETE: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\n"; + let upstream = + upstream([ResponseTemplate::new(200).set_body_raw(INCOMPLETE, "text/event-stream")]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + + stream_through(&host).await.expect("streamed call succeeds"); + + let delivered: Vec = host + .seen + .into_inner() + .unwrap() + .iter() + .flat_map(|step| match step { + Seen::Deliver(chunk) => chunk.to_vec(), + Seen::Open(_) => Vec::new(), + }) + .collect(); + assert_eq!(delivered, INCOMPLETE.as_bytes()); +} + +/// Serves one SSE chunk and then holds the connection open without ever finishing. +async fn stalling_upstream() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n\ + 1f\r\nevent: message_start\ndata: {}\n\n\r\n", + ) + .await + .unwrap(); + std::future::pending::<()>().await; + }); + base +} + +#[rstest] +#[tokio::test] +async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) { + let base = stalling_upstream().await; + let host = RecordingStreamHost::new( + MessagesCall { + timeout: Some(Duration::from_millis(300)), + ..streaming(call, base) + }, + usize::MAX, + ); + + let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host)) + .await + .expect("the stalled stream gives up within the timeout") + .err() + .expect("a stalled body fails the call"); + + assert!(matches!(error, Error::Transport(_)), "{error:?}"); + let seen = host.seen.into_inner().unwrap(); + assert!( + matches!(seen.as_slice(), [Seen::Open(_), Seen::Deliver(chunk)] if chunk.as_ref() == b"event: message_start\ndata: {}\n\n"), + "the chunk before the stall reached the caller, saw {} ops", + seen.len() + ); +} + #[rstest] #[tokio::test] async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) { diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index 87481aa89b7..7f07475bc4c 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -134,6 +134,13 @@ pub trait ProtocolHost: Send + Sync { response: ::Response, ) -> PyResult>; + /// What the stream carries at hand-off, as the caller's stream receives it. + fn head( + &mut self, + py: Python<'_>, + head: ::StreamHead, + ) -> PyResult>; + /// One streamed chunk as the caller receives it. fn chunk( &mut self, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index aaa0752522b..372af2843bd 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -134,10 +134,10 @@ where } match driver.resume(None)? { ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Open => py + ExecutionStep::Open(head) => py .import("litellm.rust_bridge.lifecycle")? .getattr("SyncStream")? - .call1((Py::new(py, Execution::suspended(driver))?,)) + .call1((Py::new(py, Execution::suspended(driver))?, head)) .map(Bound::unbind), ExecutionStep::Await(_) | ExecutionStep::Yield(_) => { Err(PyRuntimeError::new_err("sync call suspended")) @@ -312,7 +312,7 @@ where Ok(_) => return Err(missing_state()), Err(error) => Err(error), }, - HostOp::Open(_, reply) => return self.opened(py, reply).map(Next::Return), + HostOp::Open(head, reply) => return self.opened(py, head, reply).map(Next::Return), HostOp::Deliver(chunk, reply) => { return self.delivered(py, chunk, reply).map(Next::Return); } @@ -340,12 +340,21 @@ where } } - fn opened(&mut self, py: Python<'_>, reply: Reply) -> PyResult { + fn opened( + &mut self, + py: Python<'_>, + head: as Protocol>::StreamHead, + reply: Reply, + ) -> PyResult { self.stage = Stage::Streaming; + let head = match self.host.head(py, head) { + Ok(head) => head, + Err(error) => return self.interrupt(py, error), + }; match self.adapter.opened(py) { Ok(()) => { self.pending = Some(Pending::Consumer(reply)); - Ok(ExecutionStep::Open) + Ok(ExecutionStep::Open(head)) } Err(error) => self.interrupt(py, error), } @@ -699,6 +708,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .map(|answer| reply.send(answer)) } + fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult> { + match head {} + } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { match chunk {} } @@ -945,6 +958,163 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + struct Streaming; + + impl Protocol for Streaming { + type Response = (); + type Error = Error; + type Projection = (); + type Op = std::convert::Infallible; + type Chunk = &'static str; + type StreamHead = Vec<(&'static str, &'static str)>; + } + + struct StreamingHost; + + impl ProtocolHost for StreamingHost { + type Protocol = Streaming; + type Failure = Classified; + + fn project( + &mut self, + _: Python<'_>, + _: &Bound<'_, PyDict>, + ) -> Result<(), InvokeError> { + Ok(()) + } + + fn invoke( + &mut self, + _: Python<'_>, + op: std::convert::Infallible, + ) -> Result<(), InvokeError> { + match op {} + } + + fn head( + &mut self, + py: Python<'_>, + head: Vec<(&'static str, &'static str)>, + ) -> PyResult> { + let headers = PyDict::new(py); + for (name, value) in head { + headers.set_item(name, value)?; + } + let hidden = PyDict::new(py); + hidden.set_item("additional_headers", headers)?; + Ok(hidden.into_any().unbind()) + } + + fn chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult> { + Ok(pyo3::types::PyString::new(py, chunk).into_any().unbind()) + } + + fn complete(&mut self, py: Python<'_>, (): ()) -> PyResult> { + Ok(py.None()) + } + + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + Ok(Classified(error.0)) + } + + fn host_error(error: &PyErr) -> Error { + Error(error.to_string()) + } + + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } + } + + fn streaming_machine() -> CallMachine { + CallMachine::new(|host| { + Box::pin(async move { + host.project().await?; + if host.open(vec![("request-id", "req_1")]).await? == Demand::Detached { + return Ok(()); + } + for chunk in ["first", "second"] { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(()) + }) + }) + } + + /// Drives a `Stream` (async) or `SyncStream` to completion from a sync test. + fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec { + if !asynchronous { + return stream + .try_iter() + .unwrap() + .map(|chunk| chunk.unwrap().extract().unwrap()) + .collect(); + } + std::iter::from_fn(|| { + let stop = stream + .call_method0("__anext__") + .unwrap() + .call_method1("send", (py.None(),)) + .unwrap_err(); + if stop.is_instance_of::(py) { + return None; + } + assert!(stop.is_instance_of::(py)); + Some(stop.value(py).getattr("value").unwrap().extract().unwrap()) + }) + .collect() + } + + #[test] + fn a_stream_carries_its_head_as_hidden_params_before_the_first_chunk() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + install_lifecycle_module(py); + for asynchronous in [false, true] { + let log = Log::default(); + let adapter = SyntheticAdapter { + log: Log(log.0.clone()), + script: AdapterScript::Plain, + }; + let handed = run_call( + py, + streaming_machine(), + StreamingHost, + Box::new(adapter), + PyDict::new(py).unbind(), + asynchronous, + ) + .unwrap(); + let stream = if asynchronous { + let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err(); + stop.value(py).getattr("value").unwrap() + } else { + handed.into_bound(py) + }; + let hidden: std::collections::HashMap< + String, + std::collections::HashMap, + > = stream.getattr("_hidden_params").unwrap().extract().unwrap(); + assert_eq!( + hidden["additional_headers"], + std::collections::HashMap::from([( + "request-id".to_string(), + "req_1".to_string() + )]) + ); + assert_eq!(log.entries(), ["started", "begin", "opened"]); + assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]); + } + }); + } + fn failing_machine() -> CallMachine { CallMachine::new(|host| { Box::pin(async move { @@ -1202,6 +1372,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) -> Result<(), InvokeError> { Err(missing_state().into()) } + fn head( + &mut self, + _: Python<'_>, + head: std::convert::Infallible, + ) -> PyResult> { + match head {} + } fn chunk( &mut self, _: Python<'_>, diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index 10abbadbda5..24adfd404d7 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -8,9 +8,9 @@ use pyo3::prelude::*; pub enum ExecutionStep { Return(Py), Await(Py), - /// The call streams: the caller gets a stream over this execution, which stays - /// suspended until the stream asks for a chunk. - Open, + /// The call streams: the caller gets a stream over this execution carrying this head, + /// and the execution stays suspended until the stream asks for a chunk. + Open(Py), Yield(Py), } @@ -75,7 +75,7 @@ impl Execution { let step = body.resume(result)?; let (tag, value, suspended) = match step { ExecutionStep::Await(value) => ("Await", value, true), - ExecutionStep::Open => ("Open", py.None(), true), + ExecutionStep::Open(head) => ("Open", head, true), ExecutionStep::Yield(value) => ("Yield", value, true), ExecutionStep::Return(value) => ("Complete", value, false), }; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 6de4e1320e1..a253f4f5670 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -3,7 +3,7 @@ use std::convert::Infallible; use bytes::Bytes; use litellm_core::messages::{ Error, - route::{Messages, MessagesCall, MessagesOutput}, + route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead}, types::MessagesShaping, }; use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py}; @@ -238,6 +238,13 @@ impl ProtocolHost for MessagesPythonHost { } } + fn head(&mut self, py: Python<'_>, head: MessagesStreamHead) -> PyResult> { + py.import(ROUTE_HOST_MODULE)? + .getattr("stream_hidden_params")? + .call1((to_py(py, &head.headers)?,)) + .map(Bound::unbind) + } + fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { Ok(PyBytes::new(py, &chunk).into_any().unbind()) } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 65040f31684..dae8623979a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -4,14 +4,12 @@ use host::MessagesPythonHost; use litellm_callbacks_legacy_python::{ LegacySurface, PassThroughStream, PublicCall, run_legacy_call, }; -use litellm_core::messages::route::{messages_machine, supports}; +use litellm_core::messages::route::messages_machine; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, }; -use crate::errors::RustBridgeDeclined; - const SURFACE: LegacySurface = LegacySurface { call_type: "anthropic_messages", input_description: "Messages", @@ -28,17 +26,6 @@ fn run_messages( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { - let model: String = request.getattr("model")?.extract()?; - let provider: Option = request.getattr("custom_llm_provider")?.extract()?; - let stream = request - .getattr("stream")? - .extract::>()? - .unwrap_or(false); - if !supports(&model, provider.as_deref(), stream) { - return Err(RustBridgeDeclined::new_err( - "the Rust Messages route does not serve this provider", - )); - } let secrets = crate::secrets::source(py)?; run_legacy_call( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 5a3806e61e3..dc01ced15a0 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -117,6 +117,10 @@ impl ProtocolHost for OcrPythonHost { .map(Bound::unbind) } + fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult> { + match head {} + } + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { match chunk {} } diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index a0c791a136c..13a030e7ebe 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -3,6 +3,8 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Itera from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.experimental_pass_through.messages import handler as main from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook @@ -71,10 +73,17 @@ def _public_request( ) +def _resolved_provider(request: LiteLLMMessagesRequest) -> str | None: + try: + return get_llm_provider(request.model, request.custom_llm_provider)[1] + except BadRequestError: + return request.custom_llm_provider + + def _context(request: LiteLLMMessagesRequest) -> RouteContext: return RouteContext( Route.MESSAGES, - provider=request.custom_llm_provider, + provider=_resolved_provider(request), model=request.model, delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED, ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index d9834adc7e8..6e455817194 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -109,6 +109,7 @@ RULES: Final[Rules] = ( RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index 2f243e8c212..7a6485a5f2c 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Awaitable, Iterator +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping from dataclasses import dataclass from typing import Final, Protocol @@ -17,7 +17,7 @@ class Complete: @dataclass(frozen=True, slots=True) class Open: - value: None + value: Mapping[str, object] | None @dataclass(frozen=True, slots=True) @@ -68,7 +68,7 @@ async def drive(execution: Execution) -> object: step: Final = await _settle(execution, execution.start()) if isinstance(step, Open): handed_off = True - return Stream(execution) + return Stream(execution, step.value) return step.value finally: if not handed_off: @@ -78,10 +78,10 @@ async def drive(execution: Execution) -> object: class Stream(AsyncIterator[object]): """A streamed native call: each read resumes the execution until its next chunk.""" - def __init__(self, execution: Execution) -> None: + def __init__(self, execution: Execution, hidden_params: Mapping[str, object] | None = None) -> None: self._execution: Final = execution self._done = False - self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place + self._hidden_params: dict[str, object] = dict(hidden_params or {}) # mutable-ok: header writers mutate it def __aiter__(self) -> Stream: return self @@ -115,10 +115,10 @@ class Stream(AsyncIterator[object]): class SyncStream(Iterator[object]): """The sync form of `Stream`; its execution never suspends on an awaitable.""" - def __init__(self, execution: Execution) -> None: + def __init__(self, execution: Execution, hidden_params: Mapping[str, object] | None = None) -> None: self._execution: Final = execution self._done = False - self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place + self._hidden_params: dict[str, object] = dict(hidden_params or {}) # mutable-ok: header writers mutate it def __iter__(self) -> SyncStream: return self diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index d49d7b75a6f..0a23989a59c 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from typing import Final, cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict +import httpx from pydantic import TypeAdapter, ValidationError import litellm @@ -53,6 +54,14 @@ def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: ) +def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]: + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + anthropic_messages_stream_hidden_params, + ) + + return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers))) + + def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]: return request.kwargs diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/test_litellm/rust_bridge/messages/test_route_host.py index f47333a45d9..c5a442e0709 100644 --- a/tests/test_litellm/rust_bridge/messages/test_route_host.py +++ b/tests/test_litellm/rust_bridge/messages/test_route_host.py @@ -110,3 +110,15 @@ def test_native_request_rejections_map_to_the_public_400() -> None: assert "does not support top_k=5" in mapped.message assert mapped.model == "claude-sonnet-5" assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError) + + +def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None: + hidden: Final = route_host.stream_hidden_params( + (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41")) + ) + + additional: Final = hidden["additional_headers"] + assert isinstance(additional, dict) + assert additional["llm_provider-request-id"] == "req_upstream_123" + assert additional["x-ratelimit-remaining-requests"] == "41" + assert "request-id" not in additional diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index 19043780eb6..dc66852d214 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -127,7 +127,7 @@ async def test_native_messages_stream_relays_provider_events_and_logs_success_on **arguments(messages_server, stream=True, callbacks=[recorder]) ) assert isinstance(stream, AsyncIterator) - assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} + assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true" first: Final = await anext(stream) await drain_logging() assert "async_log_success_event" not in recorder.names @@ -171,7 +171,7 @@ def test_native_sync_messages_stream_relays_provider_events_and_logs_success_onc stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) assert isinstance(stream, Iterator) - assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} + assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true" assert b"".join(stream) == sse_payload() assert_served_natively(messages_server) @@ -186,3 +186,56 @@ def test_native_sync_messages_returns_the_provider_message(messages_server: Reco assert_served_natively(messages_server) assert response["content"] == MESSAGES_RESPONSE["content"] assert len(recorder.wait_for("log_success_event")) == 1 + + +@pytest.mark.asyncio +async def test_native_messages_pre_call_sees_the_shaped_optional_params( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + await litellm.anthropic.messages.acreate( + **arguments(messages_server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True) + ) + + sent: Final = messages_server.requests[0].body + assert not {"temperature", "top_k"} & sent.keys() + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + optional_params: Final = pre_call["optional_params"] + assert isinstance(optional_params, dict) + assert not {"model", "messages", "temperature", "top_k"} & optional_params.keys() + assert optional_params["max_tokens"] == sent["max_tokens"] + + +@pytest.mark.asyncio +async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(messages_server: RecordingServer) -> None: + class Broken(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + raise RuntimeError("logger exploded") + + response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()])) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + + +@pytest.mark.asyncio +async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the_relayed_events( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + async for _ in stream: + pass + + success: Final = await recorder.wait_for_async("async_log_success_event") + usage: Final = success[0].response.usage + assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"] + assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] + assert success[0].response.choices[0].message.content == "Hello from native Messages"