diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 5102d063a8a..d61761b9348 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -6,6 +6,7 @@ use litellm_llms::{ azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::messages::transformation::BaseMessagesConfig, bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG, }; use serde_json::{Map, Value}; @@ -18,6 +19,7 @@ pub(crate) enum MessagesProvider { Anthropic, AzureAi, Bedrock, + VertexAi, } impl MessagesProvider { @@ -26,6 +28,7 @@ impl MessagesProvider { Self::Anthropic => LlmProviders::Anthropic, Self::AzureAi => LlmProviders::AzureAi, Self::Bedrock => LlmProviders::Bedrock, + Self::VertexAi => LlmProviders::VertexAi, } .into() } @@ -35,15 +38,21 @@ impl MessagesProvider { Self::Anthropic => &ANTHROPIC_MESSAGES_CONFIG, Self::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG, Self::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + Self::VertexAi => &VERTEX_ANTHROPIC_MESSAGES_CONFIG, } } } -pub(crate) fn messages_provider(provider: LlmProviders) -> Option { +/// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude +/// partner models on this route. +pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option { match provider { LlmProviders::Anthropic => Some(MessagesProvider::Anthropic), LlmProviders::AzureAi => Some(MessagesProvider::AzureAi), LlmProviders::Bedrock => Some(MessagesProvider::Bedrock), + LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => { + Some(MessagesProvider::VertexAi) + } _ => None, } } @@ -68,20 +77,26 @@ mod tests { #[case::anthropic("anthropic", MessagesProvider::Anthropic)] #[case::azure_ai("azure_ai", MessagesProvider::AzureAi)] #[case::bedrock("bedrock", MessagesProvider::Bedrock)] + #[case::vertex_ai("vertex_ai", MessagesProvider::VertexAi)] fn provider_round_trips_through_its_python_name( #[case] name: &str, #[case] provider: MessagesProvider, ) { assert_eq!( - messages_provider(name.parse::().unwrap()), + messages_provider(name.parse::().unwrap(), "claude-sonnet-4-5"), Some(provider) ); assert_eq!(provider.as_str(), name); } - #[test] - fn provider_without_a_messages_config_is_rejected() { - assert_eq!(messages_provider(LlmProviders::Openai), None); + #[rstest] + #[case::openai(LlmProviders::Openai, "gpt-5")] + #[case::vertex_gemini(LlmProviders::VertexAi, "gemini-2.5-pro")] + fn provider_without_a_messages_config_is_rejected( + #[case] provider: LlmProviders, + #[case] model: &str, + ) { + assert_eq!(messages_provider(provider, model), None); } #[test] diff --git a/litellm-rust/crates/inference-messages/src/prepare.rs b/litellm-rust/crates/inference-messages/src/prepare.rs index 0c189d0a309..50a079d5eb3 100644 --- a/litellm-rust/crates/inference-messages/src/prepare.rs +++ b/litellm-rust/crates/inference-messages/src/prepare.rs @@ -51,7 +51,7 @@ fn resolve_provider( custom_llm_provider: Option<&str>, ) -> Result { let resolved = resolve_llm_provider(model, custom_llm_provider, "messages")?; - let provider = messages_provider(resolved.provider) + let provider = messages_provider(resolved.provider, resolved.model) .ok_or_else(|| Error::InvalidProvider(<&str>::from(resolved.provider).to_string()))?; Ok(ResolvedProvider { model: resolved.model.to_string(), diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 523b25f81d6..b1a989713ec 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -1,8 +1,10 @@ +use litellm_auth::VertexParams; use litellm_llms::base_llm::messages::context::{MessagesModelCapabilities, SupportedEffortTiers}; use litellm_llms_types::{ headers::{ProviderSpecificHeader, ProviderSpecificHeaders}, providers::anthropic::{AnthropicBeta, BetaSet}, }; +use litellm_router_types::LitellmParams; use rstest::rstest; use super::*; @@ -599,6 +601,55 @@ async fn metadata_is_reduced_to_the_user_id(call: MessagesCall, #[case] provider ); } +#[rstest] +#[case::not_streaming(false, ":rawPredict")] +#[case::streaming(true, ":streamRawPredict?alt=sse")] +#[tokio::test] +async fn vertex_ai_addresses_the_model_in_the_url_and_not_in_the_body( + call: MessagesCall, + #[case] stream: bool, + #[case] suffix: &str, +) { + let streamed = ResponseTemplate::new(200).set_body_raw( + "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "text/event-stream", + ); + let upstream = upstream([if stream { streamed } else { message_response() }]).await; + + run(with_fields( + MessagesCall { + custom_llm_provider: Some("vertex_ai".into()), + litellm_params: LitellmParams { + vertex: VertexParams { + vertex_project: Some("proj".into()), + vertex_location: Some("us-east5".into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, + api_base: Some(upstream.uri()), + extra_headers: headers([("Authorization", "Bearer caller-token")]), + ..with_model(call, "claude-sonnet-4-5@20250929") + }, + json!({"stream": stream}), + )) + .await + .expect("messages call succeeds"); + + let request = only_request(&upstream).await; + assert_eq!( + request.url.path().to_string() + request.url.query().map_or("", |_| "?alt=sse"), + format!( + "/v1/projects/proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-5@20250929{suffix}" + ) + ); + assert_eq!(request.header("authorization"), Some("Bearer caller-token")); + assert_eq!(request.header("anthropic-version"), None); + let body = request.json(); + assert_eq!(body.get("model"), None); + assert_eq!(body["anthropic_version"], json!("vertex-2023-10-16")); +} + #[rstest] #[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))] #[case::missing_max_tokens(json!({"max_tokens": null}))] diff --git a/litellm-rust/crates/llms/src/vertex_ai/common_utils.rs b/litellm-rust/crates/llms/src/vertex_ai/common_utils.rs new file mode 100644 index 00000000000..45abf5cd9f2 --- /dev/null +++ b/litellm-rust/crates/llms/src/vertex_ai/common_utils.rs @@ -0,0 +1,71 @@ +use crate::Error; + +const GLOBAL_LOCATION: &str = "global"; +pub const DEFAULT_VERTEX_LOCATION: &str = "us-central1"; + +fn is_location_token(location: &str) -> bool { + !location.is_empty() + && location + .bytes() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-') + && location + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && location + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric) +} + +/// Python's `validate_vertex_location`: a location reaches a hostname, so only `global` or a +/// lowercase alphanumeric token with hyphens is allowed. +pub fn validate_vertex_location(location: &str) -> Result<&str, Error> { + if location == GLOBAL_LOCATION || is_location_token(location) { + return Ok(location); + } + Err(Error::Auth(litellm_auth::Error::InvalidConfiguration( + litellm_auth::ErrorDetail::InvalidType { + field: "vertex_location".into(), + expected: "global or a lowercase alphanumeric token with hyphens", + }, + ))) +} + +/// Python's `get_vertex_base_url`: the global control plane, a multi-region geography, or a +/// regional host. +pub fn get_vertex_base_url(location: &str) -> Result { + let location = validate_vertex_location(location)?; + if location == GLOBAL_LOCATION { + return Ok("https://aiplatform.googleapis.com".to_string()); + } + if !location.contains('-') { + return Ok(format!("https://aiplatform.{location}.rep.googleapis.com")); + } + Ok(format!("https://{location}-aiplatform.googleapis.com")) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::regional("us-east5", "https://us-east5-aiplatform.googleapis.com")] + #[case::geography("eu", "https://aiplatform.eu.rep.googleapis.com")] + #[case::global("global", "https://aiplatform.googleapis.com")] + fn base_url_follows_the_location_kind(#[case] location: &str, #[case] expected: &str) { + assert_eq!(get_vertex_base_url(location).unwrap(), expected); + } + + #[rstest] + #[case::host_injection("attacker.example/")] + #[case::fragment("evil.com#")] + #[case::uppercase("US-EAST5")] + #[case::leading_hyphen("-us")] + #[case::empty("")] + fn a_location_that_is_not_a_token_is_rejected(#[case] location: &str) { + assert!(get_vertex_base_url(location).is_err()); + } +} diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs new file mode 100644 index 00000000000..8bedea4d100 --- /dev/null +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -0,0 +1,682 @@ +use litellm_auth::VertexParams; +use litellm_auth_gcp::{ + VertexConfig, get_vertex_ai_location, get_vertex_ai_project, + get_vertex_ai_project_from_credentials, secret_names, +}; +use litellm_http::request::has_header; +use litellm_llms_types::{ + formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool, OutputConfig}, + providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet}, + recognized::Recognized, +}; +use litellm_router_types::LitellmParams; +use serde_json::{Value, json}; + +use crate::{ + Error, + anthropic::{ + common_utils::{merge_beta_headers, supports_effort_param}, + messages::{ + handler::shape_anthropic_messages_request, + transformation::{provider_feature_betas, transform_messages_request}, + }, + }, + base_llm::{ + auth::{AuthScheme, Headers, ValidatedEnvironment}, + messages::{ + context::MessagesTransformContext, + normalization::{normalize_system_role_messages, strip_cache_control_scope}, + transformation::BaseMessagesConfig, + }, + }, + vertex_ai::common_utils::{DEFAULT_VERTEX_LOCATION, get_vertex_base_url}, +}; + +pub const VERTEX_ANTHROPIC_VERSION: &str = "vertex-2023-10-16"; +const RAW_PREDICT: &str = "rawPredict"; +const STREAM_RAW_PREDICT: &str = "streamRawPredict"; +const WEB_SEARCH_TOOL_PREFIX: &str = "web_search"; + +/// Claude on Vertex AI, Python's `VertexAIPartnerModelsAnthropicMessagesConfig`: the Anthropic +/// payload addressed to a project and location, authenticated with a Google access token. +pub struct VertexAiPartnerModelsAnthropicMessagesConfig; + +pub const VERTEX_ANTHROPIC_MESSAGES_CONFIG: VertexAiPartnerModelsAnthropicMessagesConfig = + VertexAiPartnerModelsAnthropicMessagesConfig; + +impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { + fn shape_request( + &self, + request: MessagesRequest, + reasoning_auto_summary: bool, + ) -> Result { + shape_anthropic_messages_request(request, reasoning_auto_summary) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + model: &str, + litellm_params: &LitellmParams, + stream: bool, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let config = VertexConfig::from_params(&litellm_params.vertex); + let project = get_vertex_ai_project(&config, env_lookup) + .or_else(|| get_vertex_ai_project_from_credentials(&config, env_lookup)) + .ok_or_else(|| Error::Auth(VertexParams::PROJECT.missing("Vertex AI")))?; + let location = get_vertex_ai_location(&config, env_lookup) + .unwrap_or_else(|| DEFAULT_VERTEX_LOCATION.to_string()); + complete_vertex_anthropic_url(api_base, &project, &location, model, stream) + } + + fn transform_anthropic_messages_request( + &self, + request: MessagesRequest, + context: &MessagesTransformContext, + ) -> Result { + let request = transform_messages_request(request, context)?; + let request = strip_cache_control_scope(normalize_system_role_messages( + request, + context + .thinking + .capabilities + .supports_mid_conversation_system, + )); + let accepts_effort = supports_effort_param(&context.thinking.capabilities); + Ok(MessagesRequest { + params: MessagesOptionalParams { + output_config: request + .params + .output_config + .and_then(|config| sanitize_output_config(config, accepts_effort)), + extra: request + .params + .extra + .into_iter() + .chain([( + "anthropic_version".to_string(), + json!(VERTEX_ANTHROPIC_VERSION), + )]) + .collect(), + ..request.params + }, + ..request + }) + } + + fn secret_names(&self) -> &'static [&'static str] { + secret_names() + } + + /// A forwarded bearer is sent as is; otherwise a Google access token for the configured + /// credentials (`vertex_credentials`, `VERTEXAI_CREDENTIALS`, `GOOGLE_APPLICATION_CREDENTIALS`, + /// then application default credentials) is acquired when the request is sent. + fn validate_environment( + &self, + headers: Headers, + _api_key: Option<&str>, + _model: &str, + litellm_params: &LitellmParams, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + if has_header(&headers, "authorization") { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::GcpAccessToken { + config: Box::new(VertexConfig::from_params(&litellm_params.vertex)), + }, + }) + } + + fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { + merge_beta_headers(headers, vertex_feature_betas(request)) + } + + /// Vertex addresses the model in the URL and rejects it in the body. + fn wire_body(&self, mut body: Value) -> Value { + if let Value::Object(fields) = &mut body { + fields.remove("model"); + } + body + } +} + +/// Python's `create_vertex_url` plus `_check_custom_proxy` for the Claude partner: the +/// publisher path under the location host, or under a custom base that has no path of its +/// own, else the custom base with the predict verb appended. A custom base also gets +/// `alt=sse` when streaming. +pub fn complete_vertex_anthropic_url( + api_base: Option<&str>, + project: &str, + location: &str, + model: &str, + stream: bool, +) -> Result { + let verb = if stream { + STREAM_RAW_PREDICT + } else { + RAW_PREDICT + }; + let path = format!( + "/v1/projects/{project}/locations/{location}/publishers/anthropic/models/{model}:{verb}" + ); + let custom = api_base.map(str::trim).filter(|base| !base.is_empty()); + let Some(custom) = custom else { + return Ok(format!("{}{path}", get_vertex_base_url(location)?)); + }; + let custom = custom.trim_end_matches('/'); + let authority_and_path = custom.split_once("://").map_or(custom, |(_, rest)| rest); + let base_path = authority_and_path + .find('/') + .map_or("", |index| &authority_and_path[index..]); + let url = match base_path { + "" => format!("{custom}{path}"), + "/v1" | "/v1beta1" => format!("{custom}{}", path.trim_start_matches("/v1")), + _ => format!("{custom}:{verb}"), + }; + if !stream { + return Ok(url); + } + Ok(match url.contains('?') { + true => format!("{url}&alt=sse"), + false => format!("{url}?alt=sse"), + }) +} + +/// Python's `sanitize_vertex_anthropic_output_params`: `effort` only reaches models whose +/// capabilities accept it, a malformed `output_config` is dropped, and an emptied one is omitted. +fn sanitize_output_config( + config: Recognized, + accepts_effort: bool, +) -> Option> { + let Recognized::Known(config) = config else { + return None; + }; + let sanitized = OutputConfig { + effort: config.effort.filter(|_| accepts_effort), + ..config + }; + (!sanitized.is_empty()).then_some(Recognized::Known(sanitized)) +} + +fn has_web_search_tool(tools: Option<&[Recognized]>) -> bool { + tools.unwrap_or_default().iter().any(|tool| match tool { + Recognized::Unrecognized(value) => value + .get("type") + .and_then(Value::as_str) + .is_some_and(|kind| kind.starts_with(WEB_SEARCH_TOOL_PREFIX)), + Recognized::Known(_) => false, + }) +} + +/// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature +/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled +/// out. +fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet { + let params = &request.params; + provider_feature_betas(request, BetaProvider::VertexAi).union( + [ + has_web_search_tool(params.tools.as_deref()) + .then_some(AnthropicBeta::WebSearch20250305), + params + .extra + .get("safeguards") + .is_some_and(|value| !value.is_null()) + .then_some(AnthropicBeta::DangerousToolUse20260903), + ] + .into_iter() + .flatten() + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::base_llm::messages::context::{MessagesModelCapabilities, ThinkingContext}; + + fn no_env(_: &str) -> Option { + None + } + + fn params(project: Option<&str>, location: Option<&str>) -> LitellmParams { + LitellmParams { + vertex: VertexParams { + vertex_project: project.map(Into::into), + vertex_location: location.map(Into::into), + ..VertexParams::default() + }, + ..LitellmParams::default() + } + } + + fn url( + api_base: Option<&str>, + litellm_params: &LitellmParams, + stream: bool, + env: &dyn Fn(&str) -> Option, + ) -> Result { + VERTEX_ANTHROPIC_MESSAGES_CONFIG.get_complete_url( + api_base, + "claude-sonnet-4-5@20250929", + litellm_params, + stream, + env, + ) + } + + const PATH: &str = "/v1/projects/proj/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-5@20250929"; + + #[rstest] + #[case::regional( + None, + false, + "https://us-east5-aiplatform.googleapis.com", + ":rawPredict" + )] + #[case::streaming( + None, + true, + "https://us-east5-aiplatform.googleapis.com", + ":streamRawPredict" + )] + #[case::custom_host_gets_the_default_path( + Some("https://proxy.example/"), + false, + "https://proxy.example", + ":rawPredict" + )] + #[case::custom_host_streaming_adds_alt_sse( + Some("https://proxy.example"), + true, + "https://proxy.example", + ":streamRawPredict?alt=sse" + )] + #[case::custom_versioned_base_is_grafted( + Some("https://proxy.example/v1"), + false, + "https://proxy.example", + ":rawPredict" + )] + fn the_url_is_the_publisher_path_under_the_host( + #[case] api_base: Option<&str>, + #[case] stream: bool, + #[case] host: &str, + #[case] suffix: &str, + ) { + assert_eq!( + url( + api_base, + ¶ms(Some("proj"), Some("us-east5")), + stream, + &no_env + ) + .unwrap(), + format!("{host}{PATH}{suffix}") + ); + } + + #[rstest] + #[case::not_streaming(false, "https://proxy.example/custom/path:rawPredict")] + #[case::streaming(true, "https://proxy.example/custom/path:streamRawPredict?alt=sse")] + fn a_custom_base_with_its_own_path_only_gets_the_verb( + #[case] stream: bool, + #[case] expected: &str, + ) { + assert_eq!( + url( + Some("https://proxy.example/custom/path"), + ¶ms(Some("proj"), Some("us-east5")), + stream, + &no_env + ) + .unwrap(), + expected + ); + } + + #[rstest] + #[case::global( + "global", + "https://aiplatform.googleapis.com/v1/projects/proj/locations/global/" + )] + #[case::geography( + "eu", + "https://aiplatform.eu.rep.googleapis.com/v1/projects/proj/locations/eu/" + )] + fn the_host_follows_the_location_kind(#[case] location: &str, #[case] prefix: &str) { + assert!( + url(None, ¶ms(Some("proj"), Some(location)), false, &no_env) + .unwrap() + .starts_with(prefix) + ); + } + + #[rstest] + fn a_location_that_could_name_another_host_is_rejected() { + assert!(matches!( + url( + None, + ¶ms(Some("proj"), Some("attacker.example/")), + false, + &no_env + ), + Err(Error::Auth(_)) + )); + } + + #[rstest] + #[case::param_wins( + params(Some("from-param"), None), + &[("VERTEXAI_PROJECT", "from-env"), ("VERTEXAI_LOCATION", "europe-west4")], + "https://europe-west4-aiplatform.googleapis.com/v1/projects/from-param/locations/europe-west4/" + )] + #[case::legacy_spelling( + LitellmParams { + vertex: VertexParams { + vertex_ai_project: Some("legacy".into()), + vertex_ai_location: Some("us-east5".into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, + &[], + "https://us-east5-aiplatform.googleapis.com/v1/projects/legacy/locations/us-east5/" + )] + #[case::env_then_default_location( + LitellmParams::default(), + &[("VERTEXAI_PROJECT", "from-env")], + "https://us-central1-aiplatform.googleapis.com/v1/projects/from-env/locations/us-central1/" + )] + #[case::project_from_the_credentials( + LitellmParams { + vertex: VertexParams { + vertex_credentials: Some(r#"{"type": "service_account", "project_id": "from-creds"}"#.into()), + ..VertexParams::default() + }, + ..LitellmParams::default() + }, + &[("VERTEX_LOCATION", "us-east5")], + "https://us-east5-aiplatform.googleapis.com/v1/projects/from-creds/locations/us-east5/" + )] + fn project_and_location_follow_the_python_precedence( + #[case] litellm_params: LitellmParams, + #[case] env: &[(&str, &str)], + #[case] prefix: &str, + ) { + let lookup = |name: &str| { + env.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + assert!( + url(None, &litellm_params, false, &lookup) + .unwrap() + .starts_with(prefix) + ); + } + + #[rstest] + fn without_a_project_anywhere_the_call_fails_before_sending() { + let error = url(None, &LitellmParams::default(), false, &no_env).unwrap_err(); + + assert!(matches!( + error, + Error::Auth(litellm_auth::Error::MissingParam { + provider: "Vertex AI", + spec: &VertexParams::PROJECT, + }) + )); + assert_eq!( + error.to_string(), + "Missing Vertex AI project - pass vertex_project or vertex_ai_project, set litellm.vertex_project, or set VERTEXAI_PROJECT" + ); + } + + fn validated(forwarded: &[(&str, &str)]) -> ValidatedEnvironment { + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .validate_environment( + forwarded + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + Some("ignored-api-key"), + "claude-sonnet-4-5", + ¶ms(Some("proj"), None), + &no_env, + ) + .unwrap() + } + + #[rstest] + fn the_credential_is_a_google_access_token_for_the_configured_project() { + assert!(matches!( + validated(&[]).auth, + AuthScheme::GcpAccessToken { config } if config.project_id() == Some("proj") + )); + } + + #[rstest] + fn a_forwarded_bearer_is_sent_as_is() { + assert!(matches!( + validated(&[("Authorization", "Bearer caller")]).auth, + AuthScheme::Forwarded + )); + } + + #[rstest] + fn default_headers_carry_no_anthropic_version() { + assert_eq!( + VERTEX_ANTHROPIC_MESSAGES_CONFIG.default_headers(), + &[("content-type", "application/json")] + ); + } + + fn transformed(value: Value, context: &MessagesTransformContext) -> Value { + let request: MessagesRequest = serde_json::from_value(value).unwrap(); + serde_json::to_value( + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request(request, context) + .unwrap(), + ) + .unwrap() + } + + #[rstest] + fn the_body_carries_the_vertex_version_and_the_wire_drops_the_model() { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5@20250929", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + }), + &MessagesTransformContext::default(), + ); + assert_eq!(body["anthropic_version"], json!("vertex-2023-10-16")); + assert_eq!(body["model"], json!("claude-sonnet-4-5@20250929")); + let wire = VERTEX_ANTHROPIC_MESSAGES_CONFIG.wire_body(body); + assert_eq!(wire.get("model"), None); + assert_eq!(wire["anthropic_version"], json!("vertex-2023-10-16")); + } + + #[rstest] + fn leading_system_roles_are_hoisted_billing_blocks_drop_and_scope_is_stripped() { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "system": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}, + {"type": "text", "text": "be terse", "cache_control": {"type": "ephemeral", "scope": "global"}} + ], + "messages": [ + {"role": "system", "content": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2"}, + {"type": "text", "text": "leading system"} + ]}, + {"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}} + ]} + ] + }), + &MessagesTransformContext::default(), + ); + assert_eq!( + body["system"], + json!([ + {"type": "text", "text": "be terse", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "leading system"} + ]) + ); + assert_eq!( + body["messages"], + json!([{"role": "user", "content": [ + {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ]}]) + ); + } + + #[rstest] + #[case::converted_in_place_without_the_capability(false, "user")] + #[case::kept_in_place_with_the_capability(true, "system")] + fn a_later_system_turn_never_moves_into_the_prompt_prefix( + #[case] supports_mid_conversation_system: bool, + #[case] role: &str, + ) { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "system": "be terse", + "messages": [ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": "c"} + ] + }), + &MessagesTransformContext { + thinking: ThinkingContext { + capabilities: MessagesModelCapabilities { + supports_mid_conversation_system, + ..MessagesModelCapabilities::default() + }, + ..ThinkingContext::default() + }, + drop_params: false, + }, + ); + assert_eq!(body["system"], json!("be terse")); + assert_eq!(body["messages"][2]["role"], json!(role)); + assert_eq!(body["messages"].as_array().unwrap().len(), 4); + } + + fn with_output_config() -> MessagesTransformContext { + MessagesTransformContext { + thinking: ThinkingContext { + capabilities: MessagesModelCapabilities { + supports_output_config: true, + ..MessagesModelCapabilities::default() + }, + ..ThinkingContext::default() + }, + drop_params: false, + } + } + + #[rstest] + #[case::effort_kept_when_the_model_accepts_it( + with_output_config(), + json!({"effort": "high", "format": {"type": "json_schema"}}), + Some(json!({"effort": "high", "format": {"type": "json_schema"}})) + )] + #[case::effort_dropped_otherwise( + MessagesTransformContext::default(), + json!({"effort": "high", "format": {"type": "json_schema"}}), + Some(json!({"format": {"type": "json_schema"}})) + )] + #[case::emptied_config_is_omitted( + MessagesTransformContext::default(), + json!({"effort": "high"}), + None + )] + #[case::malformed_config_is_dropped( + with_output_config(), + json!("not an object"), + None + )] + fn output_config_is_sanitized_for_vertex( + #[case] context: MessagesTransformContext, + #[case] output_config: Value, + #[case] expected: Option, + ) { + let body = transformed( + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "output_config": output_config, + "messages": [{"role": "user", "content": "hi"}] + }), + &context, + ); + assert_eq!(body.get("output_config").cloned(), expected); + } + + fn betas(fields: Value) -> Vec { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + let request: MessagesRequest = serde_json::from_value({ + let mut base = serde_json::to_value(request).unwrap(); + base.as_object_mut() + .unwrap() + .extend(fields.as_object().unwrap().clone()); + base + }) + .unwrap(); + VERTEX_ANTHROPIC_MESSAGES_CONFIG + .request_headers(Vec::new(), &request) + .into_iter() + .filter(|(name, _)| name == "anthropic-beta") + .flat_map(|(_, value)| value.split(',').map(str::to_string).collect::>()) + .collect() + } + + #[rstest] + #[case::web_search_tool( + json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), + &["web-search-2025-03-05"] + )] + #[case::tool_search_uses_the_vertex_header( + json!({"tools": [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}]}), + &["tool-search-tool-2025-10-19"] + )] + #[case::safeguards(json!({"safeguards": {"mode": "strict"}}), &["dangerous-tool-use-2026-09-03"])] + #[case::context_management_edits( + json!({"context_management": {"edits": [ + {"type": "compact_20260112"}, + {"type": "clear_tool_uses_20250919"} + ]}}), + &["compact-2026-01-12", "context-management-2025-06-27"] + )] + #[case::advisor_tool_is_not_a_vertex_beta( + json!({"tools": [{"type": "advisor_20260301", "name": "advisor"}]}), + &[] + )] + #[case::fast_mode_is_not_a_vertex_beta(json!({"speed": "fast"}), &[])] + fn beta_headers_follow_the_vertex_policy(#[case] fields: Value, #[case] expected: &[&str]) { + assert_eq!(betas(fields), expected); + } +} diff --git a/litellm-rust/crates/llms/src/vertex_ai/mod.rs b/litellm-rust/crates/llms/src/vertex_ai/mod.rs index 3621ff6a2fd..138b893e4ca 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/mod.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/mod.rs @@ -1 +1,3 @@ +pub mod common_utils; +pub mod messages; pub mod ocr; diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index c7865132840..e895f0e72d1 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -72,6 +72,13 @@ def _resolved_provider(request: NativeCall) -> str | None: return optional_str(request.resolved.get("custom_llm_provider")) +def _python_serves(request: NativeCall) -> bool: + """Vertex AI serves only Claude through the Anthropic Messages config, as + ``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's + translation through Chat Completions.""" + return _resolved_provider(request) == "vertex_ai" and "claude" not in str(request.resolved["model"]).lower() + + def _context(request: NativeCall) -> RouteContext: return RouteContext( Route.MESSAGES, @@ -84,13 +91,14 @@ _DISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs), context=_context, - bypass=lambda request: request.kwargs.get("is_async") is True, + bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request), ) _ADISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs), context=_context, + bypass=_python_serves, ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 73e95c24de4..5024e58b4ca 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -70,7 +70,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.RUST_OPT_IN, providers=frozenset({"anthropic"})), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "vertex_ai"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), diff --git a/tests/unit/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py index faaefa1437e..23400bf6440 100644 --- a/tests/unit/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -196,6 +196,53 @@ def test_internal_async_marker_bypasses_native() -> None: assert captured == [(args, kwargs)] +@pytest.mark.parametrize( + ("model", "custom_llm_provider"), + [ + ("vertex_ai/gemini-2.5-pro", None), + ("gemini-2.5-pro", "vertex_ai"), + ], +) +def test_vertex_models_other_than_claude_stay_on_python(model: str, custom_llm_provider: str | None) -> None: + args: Final[tuple[object, ...]] = (16, MESSAGES, model) + kwargs: Final[Mapping[str, object]] = {"custom_llm_provider": custom_llm_provider} + expected: Final = response(model) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape + return expected + + def native(request: NativeCall) -> AnthropicMessagesResponse: + pytest.fail("Rust serves only Claude on Vertex AI") + + result: Final = _DISPATCH.run( + args, + kwargs, + python=python, + binding=messages_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + assert result is expected + + +def test_vertex_claude_reaches_native() -> None: + args: Final[tuple[object, ...]] = (16, MESSAGES, "vertex_ai/claude-sonnet-4-5@20250929") + expected: Final = response("vertex_ai/claude-sonnet-4-5@20250929") + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape + pytest.fail("Claude on Vertex AI is served natively") + + result: Final = _DISPATCH.run( + args, + {}, + python=python, + binding=messages_binding(lambda request: expected), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + assert result is expected + + @pytest.mark.parametrize( ("args", "kwargs"), ( diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 4242294a33a..ea3e61ecc3d 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -27,7 +27,9 @@ def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: @pytest.mark.parametrize("route", tuple(Route)) -@pytest.mark.parametrize("provider", (None, "bedrock", "mistral", "anthropic", "openai", "azure_ai", "unknown")) +@pytest.mark.parametrize( + "provider", (None, "bedrock", "mistral", "anthropic", "vertex_ai", "openai", "azure_ai", "unknown") +) @pytest.mark.parametrize("process", (None, False, True)) @pytest.mark.parametrize("environment", (None, "0", "1")) def test_shipped_decisions( @@ -45,7 +47,7 @@ def test_shipped_decisions( if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider == "anthropic": + elif route is Route.MESSAGES and provider in ("anthropic", "vertex_ai"): assert catalog.rollout(context) is Rollout.RUST_OPT_IN opted_in: Final = environment == "1" or (environment is None and process is True) assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON)