diff --git a/.greptile/files.json b/.greptile/files.json index 170a60dbdeb..148881ac289 100644 --- a/.greptile/files.json +++ b/.greptile/files.json @@ -273,6 +273,13 @@ "litellm-rust/crates/llms/src/bedrock/chat/**" ] }, + { + "path": "litellm-rust/crates/llms/src/bedrock/claude_platform/messages/AGENTS.md", + "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/claude_platform/messages/", + "scope": [ + "litellm-rust/crates/llms/src/bedrock/claude_platform/messages/**" + ] + }, { "path": "litellm-rust/crates/llms/src/bedrock/messages/AGENTS.md", "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/messages/", @@ -280,6 +287,13 @@ "litellm-rust/crates/llms/src/bedrock/messages/**" ] }, + { + "path": "litellm-rust/crates/llms/src/bedrock_mantle/messages/AGENTS.md", + "description": "Conventions for code under litellm-rust/crates/llms/src/bedrock_mantle/messages/", + "scope": [ + "litellm-rust/crates/llms/src/bedrock_mantle/messages/**" + ] + }, { "path": "litellm-rust/crates/llms/src/cohere/ocr/AGENTS.md", "description": "Conventions for code under litellm-rust/crates/llms/src/cohere/ocr/", diff --git a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs index 4a42c7bb7b2..982831bcc5e 100644 --- a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs +++ b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs @@ -16,6 +16,8 @@ pub enum LlmProviders { AzureAi, #[strum(serialize = "bedrock")] Bedrock, + #[strum(serialize = "bedrock_mantle")] + BedrockMantle, #[strum(serialize = "cohere")] Cohere, #[strum(serialize = "deepseek")] diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index ede05248823..91b135a6a4a 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -11,6 +11,10 @@ use litellm_llms::{ minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG, vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG, }; +use litellm_llms::{ + bedrock::claude_platform::messages::transformation::BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG, + bedrock_mantle::messages::transformation::BEDROCK_MANTLE_MESSAGES_CONFIG, +}; use serde_json::{Map, Value}; use super::Error; @@ -42,7 +46,12 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option &TENCENT_MESSAGES_CONFIG, LlmProviders::Anthropic => &ANTHROPIC_MESSAGES_CONFIG, LlmProviders::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG, + LlmProviders::Bedrock if model.starts_with("claude_platform/") => { + &BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + } + LlmProviders::Bedrock if model.starts_with("mantle/") => &BEDROCK_MANTLE_MESSAGES_CONFIG, LlmProviders::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + LlmProviders::BedrockMantle => &BEDROCK_MANTLE_MESSAGES_CONFIG, LlmProviders::Deepseek => &DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, LlmProviders::Minimax => &MINIMAX_MESSAGES_CONFIG, LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => { diff --git a/litellm-rust/crates/inference-messages/src/prepare.rs b/litellm-rust/crates/inference-messages/src/prepare.rs index 50a079d5eb3..3ec756b65d0 100644 --- a/litellm-rust/crates/inference-messages/src/prepare.rs +++ b/litellm-rust/crates/inference-messages/src/prepare.rs @@ -76,6 +76,10 @@ fn prepare_provider_request( shaping, .. } = call; + let connection_params = litellm_router_types::LitellmParams { + api_base: api_base.clone().or(litellm_params.api_base.clone()), + ..litellm_params + }; let config = provider.config(); let env_lookup = |key: &str| secrets.get(key); @@ -95,25 +99,22 @@ fn prepare_provider_request( let forwarded = string_headers(Some( extra_headers.into_iter().flatten().chain(scoped).collect(), ))?; - let validated = config.validate_environment( + let forwarded = config.request_headers( + with_default_headers(forwarded, config.default_headers()), + &transformed, + ); + let environment = config.validate_environment( forwarded, api_key.as_deref(), &transformed.model, - &litellm_params, + &connection_params, &env_lookup, )?; - let environment = ValidatedEnvironment { - headers: config.request_headers( - with_default_headers(validated.headers, config.default_headers()), - &transformed, - ), - auth: validated.auth, - }; let url = config.get_complete_url( api_base.as_deref(), &transformed.model, - &litellm_params, + &connection_params, transformed.params.stream == Some(true), &env_lookup, )?; @@ -475,4 +476,44 @@ mod tests { )) ); } + #[rstest] + #[case::mantle( + "bedrock_mantle/claude-test", + "https://bedrock-mantle.us-west-2.api.aws", + "/anthropic/v1/messages", + "bedrock" + )] + #[case::platform( + "bedrock/claude_platform/claude-test", + "https://aws-external-anthropic.us-west-2.api.aws", + "/v1/messages", + "aws-external-anthropic" + )] + fn explicit_endpoint_region_is_used_by_the_signer( + shaping: MessagesShaping, + #[case] model: &str, + #[case] base: &str, + #[case] path: &str, + #[case] service: &str, + ) { + let prepared = prepare_with_secrets(MessagesCall { + body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})), + api_key: None, + api_base: Some(base.into()), + custom_llm_provider: None, + litellm_params: Default::default(), + extra_headers: None, + provider_specific_header: None, + timeout: None, + shaping, + }, &|name: &str| match name { + "AWS_REGION_NAME" => Some("us-east-1".into()), + "ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()), + _ => None, + }).unwrap(); + assert_eq!(prepared.url, format!("{base}{path}")); + assert!( + matches!(prepared.environment.auth, litellm_llms::base_llm::auth::AuthScheme::AwsSigV4 { region, service: actual, .. } if region == "us-west-2" && actual == service) + ); + } } diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs new file mode 100644 index 00000000000..5420477c34a --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -0,0 +1,323 @@ +use std::{sync::Arc, time::Duration}; + +use futures_util::TryStreamExt; +use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping}; +use litellm_inference_testing::{RecordingSecrets, no_secrets}; +use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse}; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; +use wiremock::ResponseTemplate; + +#[path = "../support/mod.rs"] +mod support; +use support::*; + +const MODEL: &str = "claude-test"; + +#[fixture] +fn call() -> MessagesCall { + MessagesCall { + body: serde_json::from_value(json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})).unwrap(), + api_key: None, + api_base: None, + custom_llm_provider: None, + litellm_params: Default::default(), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } +} + +fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall { + let body = serde_json::to_value(call.body).unwrap(); + MessagesCall { + body: serde_json::from_value(Value::Object( + body.as_object() + .unwrap() + .clone() + .into_iter() + .chain(fields.as_object().unwrap().clone()) + .collect(), + )) + .unwrap(), + ..call + } +} + +fn with_model(call: MessagesCall, model: &str) -> MessagesCall { + MessagesCall { + body: MessagesRequest { + model: model.into(), + ..call.body + }, + ..call + } +} + +fn message_response() -> ResponseTemplate { + json_response( + json!({"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}], "model": MODEL, "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 2}}), + ) +} + +async fn run_with( + secrets: Arc, + call: MessagesCall, +) -> Result { + messages_route(secrets).execute(call, &(), None).await +} + +async fn run(call: MessagesCall) -> Result { + run_with(Arc::new(RecordingSecrets::empty()), call).await +} + +async fn run_message(call: MessagesCall) -> MessagesResponse { + match run(call).await.unwrap() { + MessagesCallResponse::Complete(message) => *message, + MessagesCallResponse::Stream { .. } => panic!("expected a completed message"), + } +} + +#[rstest] +#[case::platform( + "bedrock/claude_platform/claude-test", + "/v1/messages", + "x-api-key", + "key" +)] +#[case::mantle( + "bedrock_mantle/anthropic.claude-test", + "/anthropic/v1/messages", + "authorization", + "Bearer key" +)] +#[case::mantle_alias( + "bedrock/mantle/anthropic.claude-test", + "/anthropic/v1/messages", + "authorization", + "Bearer key" +)] +#[tokio::test] +async fn prefixes_select_native_messages_and_strip_connection_fields( + call: MessagesCall, + #[case] model: &str, + #[case] path: &str, + #[case] header: &str, + #[case] credential: &str, +) { + let upstream = upstream([message_response()]).await; + let message = run_message(MessagesCall { + api_key: Some("key".into()), + api_base: Some(upstream.uri()), + custom_llm_provider: None, + ..with_fields( + call, + json!({"model": model, "workspace_id": "request-workspace"}), + ) + }) + .await; + assert_eq!(message.model, MODEL); + let request = only_request(&upstream).await; + assert_eq!(request.url.path(), path); + assert_eq!(request.header(header), Some(credential)); + assert_eq!(request.json()["model"], model.rsplit('/').next().unwrap()); + if path == "/v1/messages" { + assert_eq!( + request.header("anthropic-workspace-id"), + Some("request-workspace") + ); + assert_eq!(request.json().get("workspace_id"), None); + } +} + +#[rstest] +#[case::platform( + "bedrock/claude_platform/claude-test", + "ANTHROPIC_AWS_API_KEY", + "ANTHROPIC_AWS_BASE_URL", + "/v1/messages", + "x-api-key", + "managed" +)] +#[case::mantle( + "bedrock_mantle/claude-test", + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_BASE", + "/anthropic/v1/messages", + "authorization", + "Bearer managed" +)] +#[tokio::test] +async fn managed_keys_bases_and_workspaces_reach_the_upstream( + call: MessagesCall, + #[case] model: &str, + #[case] key_name: &str, + #[case] base_name: &str, + #[case] path: &str, + #[case] header: &str, + #[case] credential: &str, +) { + let upstream = upstream([message_response()]).await; + let base = upstream.uri(); + let secrets = Arc::new(RecordingSecrets::new([ + (key_name, "managed"), + (base_name, base.as_str()), + ("ANTHROPIC_AWS_WORKSPACE_ID", "managed-workspace"), + ])); + let output = run_with( + secrets, + MessagesCall { + custom_llm_provider: None, + ..with_model(call, model) + }, + ) + .await + .unwrap(); + assert!(matches!(output, MessagesCallResponse::Complete(_))); + let request = only_request(&upstream).await; + assert_eq!(request.url.path(), path); + assert_eq!(request.header(header), Some(credential)); + if path == "/v1/messages" { + assert_eq!( + request.header("anthropic-workspace-id"), + Some("managed-workspace") + ); + } +} + +#[rstest] +#[case::platform("bedrock/claude_platform/claude-test", "aws-external-anthropic")] +#[case::mantle("bedrock_mantle/claude-test", "bedrock")] +#[tokio::test] +async fn managed_aws_credentials_sign_the_exact_native_request( + call: MessagesCall, + #[case] model: &str, + #[case] service: &str, +) { + let upstream = upstream([message_response()]).await; + let secrets = Arc::new(RecordingSecrets::new([ + ("AWS_ACCESS_KEY_ID", "test-access"), + ("AWS_SECRET_ACCESS_KEY", "test-secret"), + ("AWS_SESSION_TOKEN", "test-session"), + ("AWS_REGION_NAME", "us-west-2"), + ("ANTHROPIC_AWS_WORKSPACE_ID", "managed-workspace"), + ])); + run_with( + secrets, + MessagesCall { + api_base: Some(upstream.uri()), + custom_llm_provider: None, + ..with_model(call, model) + }, + ) + .await + .unwrap(); + let request = only_request(&upstream).await; + let auth = request.header("authorization").unwrap(); + assert!(auth.starts_with("AWS4-HMAC-SHA256 Credential=test-access/")); + assert!(auth.contains(&format!("/us-west-2/{service}/aws4_request"))); + assert_eq!(request.header("x-amz-security-token"), Some("test-session")); + assert_eq!(request.header("x-api-key"), None); + assert_eq!(request.json()["model"], "claude-test"); +} + +#[rstest] +#[tokio::test] +async fn private_mantle_endpoint_is_used_and_signed(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let params = litellm_router_types::LitellmParams { + aws: litellm_auth::AwsParams { + aws_bedrock_runtime_endpoint: Some(upstream.uri()), + aws_access_key_id: Some("private-access".into()), + aws_secret_access_key: Some("private-secret".into()), + aws_region_name: Some("us-west-2".into()), + ..Default::default() + }, + ..Default::default() + }; + run(MessagesCall { + custom_llm_provider: None, + litellm_params: params, + ..with_model(call, "bedrock/mantle/claude-test") + }) + .await + .unwrap(); + let request = only_request(&upstream).await; + assert_eq!(request.url.path(), "/anthropic/v1/messages"); + assert!( + request + .header("authorization") + .unwrap() + .contains("/us-west-2/bedrock/aws4_request") + ); +} + +#[rstest] +#[case::platform("bedrock/claude_platform/claude-test")] +#[case::mantle("bedrock_mantle/claude-test")] +#[tokio::test] +async fn native_streams_relay_every_sse_byte(call: MessagesCall, #[case] model: &str) { + let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let upstream = + upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await; + let response = messages_route(no_secrets()) + .execute( + MessagesCall { + api_key: Some("key".into()), + api_base: Some(upstream.uri()), + custom_llm_provider: None, + ..with_fields( + call, + json!({"model": model, "stream": true, "workspace_id": "workspace"}), + ) + }, + &(), + None, + ) + .await + .unwrap(); + let MessagesCallResponse::Stream { chunks, .. } = response else { + panic!("expected SSE"); + }; + assert_eq!( + chunks.try_collect::>().await.unwrap().concat(), + sse.as_bytes() + ); + assert_eq!(only_request(&upstream).await.json()["stream"], true); +} + +#[rstest] +#[case::missing(false)] +#[case::configured(true)] +#[tokio::test] +async fn deployment_workspace_is_required_and_sent_as_a_header( + call: MessagesCall, + #[case] configured: bool, +) { + let upstream = upstream([message_response()]).await; + let params = if configured { + serde_json::from_value(json!({"model": "bedrock/claude_platform/claude-test", "workspace_id": "deployment-workspace"})).unwrap() + } else { + Default::default() + }; + let result = run(MessagesCall { + api_key: Some("key".into()), + api_base: Some(upstream.uri()), + litellm_params: params, + ..with_model(call, "bedrock/claude_platform/claude-test") + }) + .await; + if !configured { + assert!(matches!(result, Err(Error::Auth(_)))); + assert!(received(&upstream).await.is_empty()); + return; + } + assert!(matches!(result.unwrap(), MessagesCallResponse::Complete(_))); + let sent = only_request(&upstream).await; + assert_eq!( + sent.header("anthropic-workspace-id"), + Some("deployment-workspace") + ); + assert_eq!(sent.json().get("workspace_id"), None); +} diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/AGENTS.md b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/AGENTS.md new file mode 100644 index 00000000000..2357e8a9d13 --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/AGENTS.md @@ -0,0 +1,11 @@ +# rules + +- Reuse Anthropic Messages shaping, reasoning translation, and beta headers +- Require a workspace header and use an API key or AWS SigV4 +- Keep connection fields out of the upstream request body +- Drop `context_management` by default, with per-request `claude_platform_unsupported_params` overrides + +# references + +- Python counterpart `litellm/llms/bedrock/claude_platform/messages_transformation.py` +- https://platform.claude.com/docs/en/build-with-claude/claude-platform-on-aws diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/mod.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs new file mode 100644 index 00000000000..1b99eca975f --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -0,0 +1,395 @@ +use std::sync::LazyLock; + +use litellm_auth::{AwsParams, CredentialPlacement, ParamSpec, SecretValue}; +use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; +use litellm_core_utils::settings::resolve_non_empty; +use litellm_http::request::{header_value, with_header, without_headers}; +use litellm_llms_types::{ + formats::messages::MessagesRequest, + providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH}, +}; +use litellm_router_types::LitellmParams; +use serde_json::Value; + +use crate::{ + Error, ErrorDetail, + anthropic::messages::{ + handler::shape_anthropic_messages_request, + transformation::{transform_messages_request, update_headers_with_anthropic_beta}, + }, + base_llm::{ + auth::{AuthScheme, Headers, ValidatedEnvironment}, + messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + }, + bedrock::common_utils::validated_region, +}; + +const SERVICE: &str = "aws-external-anthropic"; +const BASE_ENVS: &[&str] = &["ANTHROPIC_AWS_BASE_URL", "ANTHROPIC_AWS_API_BASE"]; +const WORKSPACE_ENVS: &[&str] = &["ANTHROPIC_AWS_WORKSPACE_ID", "ANTHROPIC_WORKSPACE_ID"]; +const WORKSPACE_FIELDS: &[&str] = &[ + "workspace_id", + "aws_workspace_id", + "anthropic-workspace-id", + "anthropic_workspace_id", +]; +const WORKSPACE_HEADER: &str = "anthropic-workspace-id"; +static WORKSPACE: ParamSpec = ParamSpec { + setting: "workspace_id", + wire: WORKSPACE_FIELDS, + module_global: None, + env: WORKSPACE_ENVS, +}; +const UNSUPPORTED_OVERRIDE: &str = "claude_platform_unsupported_params"; + +pub struct BedrockClaudePlatformMessagesConfig; +pub const BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG: BedrockClaudePlatformMessagesConfig = + BedrockClaudePlatformMessagesConfig; + +impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { + fn shape_request( + &self, + request: MessagesRequest, + reasoning_auto_summary: bool, + ) -> Result { + shape_anthropic_messages_request(request, reasoning_auto_summary) + } + + fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + _model: &str, + params: &LitellmParams, + env: &dyn Fn(&str) -> Option, + ) -> Result { + let workspace = header_value(&headers, WORKSPACE_HEADER) + .filter(|value| !value.trim().is_empty()) + .or_else(|| { + WORKSPACE_FIELDS.iter().find_map(|field| { + params + .extra + .get(*field)? + .as_str() + .filter(|value| !value.trim().is_empty()) + }) + }); + let workspace = resolve_non_empty(workspace, env, WORKSPACE_ENVS) + .ok_or_else(|| Error::Auth(WORKSPACE.missing("Claude Platform on AWS")))?; + let headers = with_header(headers, WORKSPACE_HEADER, workspace); + let key = resolve_non_empty(api_key, env, &["ANTHROPIC_AWS_API_KEY"]); + let headers = without_headers(headers, &["authorization"]); + if header_value(&headers, "x-api-key").is_some_and(|value| !value.trim().is_empty()) { + return Ok(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + if let Some(key) = key { + return Ok(ValidatedEnvironment { + headers: without_headers(headers, &["x-api-key"]), + auth: AuthScheme::Credential { + placement: CredentialPlacement::Header("x-api-key"), + secret: SecretValue::new(key), + }, + }); + } + let base = resolve_non_empty(params.api_base.as_deref(), env, BASE_ENVS); + Ok(ValidatedEnvironment { + headers: without_headers(headers, &["x-api-key"]), + auth: AuthScheme::AwsSigV4 { + region: platform_region(base.as_deref(), params, env)?, + service: SERVICE, + credentials: Box::new(AwsCredentialSource::from_params(¶ms.aws, env)), + }, + }) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + _model: &str, + params: &LitellmParams, + _stream: bool, + env: &dyn Fn(&str) -> Option, + ) -> Result { + let base = resolve_non_empty( + api_base + .filter(|base| !base.trim().is_empty()) + .or(params.api_base.as_deref()), + env, + BASE_ENVS, + ); + let base = match base { + Some(base) => base, + None => format!( + "https://{SERVICE}.{}.api.aws", + platform_region(None, params, env)? + ), + }; + let base = base.trim_end_matches('/'); + Ok(if base.ends_with(MESSAGES_PATH) { + base.into() + } else { + format!("{base}{MESSAGES_PATH}") + }) + } + + fn transform_anthropic_messages_request( + &self, + request: MessagesRequest, + context: &MessagesTransformContext, + ) -> Result { + let body = serde_json::to_value(request).map_err(invalid_request)?; + let filtered = serde_json::from_value(filter_body(body, false)).map_err(invalid_request)?; + transform_messages_request(filtered, context) + } + + fn secret_names(&self) -> &'static [&'static str] { + static NAMES: LazyLock> = LazyLock::new(|| { + ["ANTHROPIC_AWS_API_KEY"] + .into_iter() + .chain(BASE_ENVS.iter().copied()) + .chain(WORKSPACE_ENVS.iter().copied()) + .chain(["AWS_DEFAULT_REGION"]) + .chain(AwsParams::secret_names().iter().copied()) + .collect() + }); + &NAMES + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + DEFAULT_HEADERS + } + + fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { + let headers = match WORKSPACE_FIELDS.iter().find_map(|field| { + request + .params + .extra + .get(*field)? + .as_str() + .filter(|value| !value.trim().is_empty()) + }) { + Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()), + None => headers, + }; + update_headers_with_anthropic_beta(headers, request) + } + + fn wire_body(&self, body: Value) -> Value { + filter_body(body, true) + } +} + +fn platform_region( + base: Option<&str>, + params: &LitellmParams, + env: &dyn Fn(&str) -> Option, +) -> Result { + let host_region = base.and_then(|base| { + let url = url::Url::parse(base).ok()?; + url.host_str()? + .strip_prefix("aws-external-anthropic.")? + .strip_suffix(".api.aws") + .map(str::to_string) + }); + let region = host_region + .or_else(|| resolve_aws_region(None, ¶ms.aws, env)) + .or_else(|| resolve_non_empty(None, env, &["AWS_DEFAULT_REGION"])) + .ok_or_else(|| Error::Auth(AwsParams::REGION.missing("Claude Platform on AWS")))?; + validated_region(region) +} + +fn filter_body(body: Value, strip_connection_fields: bool) -> Value { + let Value::Object(fields) = body else { + return body; + }; + let unsupported: Vec<&str> = fields + .get(UNSUPPORTED_OVERRIDE) + .and_then(Value::as_array) + .map(|values| values.iter().filter_map(Value::as_str).collect()) + .unwrap_or_else(|| vec!["context_management"]); + let retained = fields + .iter() + .filter(|(name, _)| { + !unsupported.contains(&name.as_str()) + && !(strip_connection_fields + && (WORKSPACE_FIELDS.contains(&name.as_str()) + || name.starts_with("aws_") + || name.as_str() == UNSUPPORTED_OVERRIDE)) + }) + .map(|(name, value)| { + let value = if name == "model" { + value + .as_str() + .map(|model| { + Value::String( + model + .strip_prefix("claude_platform/") + .unwrap_or(model) + .into(), + ) + }) + .unwrap_or_else(|| value.clone()) + } else { + value.clone() + }; + (name.clone(), value) + }) + .collect(); + Value::Object(retained) +} + +fn invalid_request(error: serde_json::Error) -> Error { + Error::InvalidRequest(ErrorDetail::invalid("Messages request", error)) +} + +#[cfg(test)] +mod tests { + use super::*; + use litellm_llms_types::providers::anthropic::AnthropicBeta; + use rstest::rstest; + use serde_json::json; + + #[rstest] + fn workspace_is_required_before_any_upstream_request() { + assert_eq!( + BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .validate_environment( + vec![], + Some("key"), + "claude-test", + &LitellmParams::default(), + &|_| None + ) + .unwrap_err(), + Error::Auth(WORKSPACE.missing("Claude Platform on AWS")) + ); + } + + #[rstest] + #[case::valid("forwarded", true)] + #[case::blank(" ", false)] + fn forwarded_keys_win_and_blank_keys_fall_back(#[case] key: &str, #[case] forwarded: bool) { + let environment = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .validate_environment( + vec![ + ("X-Api-Key".into(), key.into()), + ("authorization".into(), "stale".into()), + ], + Some("explicit"), + "claude-test", + &LitellmParams::default(), + &|name| (name == "ANTHROPIC_AWS_WORKSPACE_ID").then(|| "workspace".into()), + ) + .unwrap(); + assert_eq!( + header_value(&environment.headers, WORKSPACE_HEADER), + Some("workspace") + ); + assert_eq!(header_value(&environment.headers, "authorization"), None); + if forwarded { + assert_eq!(header_value(&environment.headers, "x-api-key"), Some(key)); + assert!(matches!(environment.auth, AuthScheme::Forwarded)); + } else { + assert_eq!(header_value(&environment.headers, "x-api-key"), None); + assert!( + matches!(environment.auth, AuthScheme::Credential { secret, placement: CredentialPlacement::Header("x-api-key") } if secret.expose() == "explicit") + ); + } + } + + #[rstest] + #[case::region("us-west-2", true)] + #[case::injection("us-east-1.attacker.test/steal?", false)] + fn generated_endpoint_regions_are_validated(#[case] region: &str, #[case] valid: bool) { + let result = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.get_complete_url( + None, + "claude-test", + &LitellmParams::default(), + false, + &|name| (name == "AWS_DEFAULT_REGION").then(|| region.into()), + ); + if valid { + assert_eq!( + result.unwrap(), + format!("https://{SERVICE}.{region}.api.aws/v1/messages") + ); + } else { + assert!(matches!(result, Err(Error::InvalidRequest(_)))); + } + } + + #[rstest] + #[case::explicit_regional_endpoint( + "https://aws-external-anthropic.us-west-2.api.aws", + "us-east-1", + "us-west-2" + )] + fn public_regional_host_precedes_configured_region( + #[case] api_base: &str, + #[case] configured_region: &str, + #[case] expected_region: &str, + ) { + let params = LitellmParams { + api_base: Some(api_base.into()), + aws: AwsParams { + aws_region_name: Some(configured_region.into()), + ..Default::default() + }, + ..Default::default() + }; + let env = |name: &str| (name == "ANTHROPIC_AWS_WORKSPACE_ID").then(|| "workspace".into()); + let url = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .get_complete_url(None, "claude-test", ¶ms, false, &env) + .unwrap(); + let environment = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .validate_environment(vec![], None, "claude-test", ¶ms, &env) + .unwrap(); + assert_eq!( + url, + "https://aws-external-anthropic.us-west-2.api.aws/v1/messages" + ); + assert!( + matches!(environment.auth, AuthScheme::AwsSigV4 { region, service: SERVICE, .. } if region == expected_region) + ); + } + + #[rstest] + #[case::default(json!({}), false)] + #[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)] + fn unsupported_fields_are_filtered_before_beta_selection( + #[case] extra: Value, + #[case] keep_context: bool, + ) { + let fields = json!({"model": "claude_platform/claude-test", "max_tokens": 16, "messages": [], "workspace_id": "workspace", "aws_region_name": "us-west-2", "context_management": {"edits": [{"type": "clear_thinking_20251015"}]}}); + let input = Value::Object( + fields + .as_object() + .unwrap() + .clone() + .into_iter() + .chain(extra.as_object().unwrap().clone()) + .collect(), + ); + let request = serde_json::from_value(input).unwrap(); + let transformed = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) + .unwrap(); + let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &transformed); + let expected_beta = AnthropicBeta::ContextManagement20250627.to_string(); + assert_eq!( + header_value(&headers, "anthropic-beta") + .is_some_and(|value| value.split(',').any(|beta| beta == expected_beta)), + keep_context + ); + let wire = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG + .wire_body(serde_json::to_value(transformed).unwrap()); + assert_eq!(wire["model"], "claude-test"); + assert_eq!(wire.get("context_management").is_some(), keep_context); + assert_eq!(wire.get("workspace_id"), None); + assert_eq!(wire.get("aws_region_name"), None); + assert_eq!(wire.get(UNSUPPORTED_OVERRIDE), None); + } +} diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/mod.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/mod.rs new file mode 100644 index 00000000000..ba63992f3cb --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/mod.rs @@ -0,0 +1 @@ +pub mod messages; diff --git a/litellm-rust/crates/llms/src/bedrock/common_utils.rs b/litellm-rust/crates/llms/src/bedrock/common_utils.rs new file mode 100644 index 00000000000..6c6fb8562f0 --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock/common_utils.rs @@ -0,0 +1,16 @@ +use crate::{Error, ErrorDetail}; + +pub(crate) fn validated_region(region: String) -> Result { + if !region.is_empty() + && region + .bytes() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-') + { + return Ok(region); + } + Err(Error::InvalidRequest(ErrorDetail::InvalidValue { + field: "aws_region_name", + expected: "lowercase letters, digits, and hyphens", + actual: region.into(), + })) +} diff --git a/litellm-rust/crates/llms/src/bedrock/mod.rs b/litellm-rust/crates/llms/src/bedrock/mod.rs index feed6e70e4d..e3160e547b4 100644 --- a/litellm-rust/crates/llms/src/bedrock/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/mod.rs @@ -1,3 +1,5 @@ pub mod audio_transcription; pub mod chat; +pub mod claude_platform; +pub(crate) mod common_utils; pub mod messages; diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/AGENTS.md b/litellm-rust/crates/llms/src/bedrock_mantle/messages/AGENTS.md new file mode 100644 index 00000000000..0f762efea16 --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/AGENTS.md @@ -0,0 +1,10 @@ +# rules + +- Reuse Anthropic Messages shaping and reasoning translation +- Use bearer credentials or AWS SigV4 for the Mantle Messages endpoint +- Filter and rename caller betas using the Mantle beta policy + +# references + +- Python counterpart `litellm/llms/bedrock/messages/mantle_transformation.py` +- https://docs.aws.amazon.com/bedrock/latest/userguide/inference-messages-api.html diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/mod.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs new file mode 100644 index 00000000000..798f6648250 --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -0,0 +1,528 @@ +use std::sync::LazyLock; + +use litellm_auth::{AwsParams, CredentialPlacement, SecretValue}; +use litellm_auth_aws::AwsCredentialSource; +use litellm_auth_aws::bedrock_model_id_and_region; +use litellm_core_utils::settings::resolve_non_empty; +use litellm_http::request::{with_header, without_headers}; +use litellm_llms_types::{ + formats::messages::MessagesRequest, + providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS}, +}; +use litellm_router_types::LitellmParams; +use serde_json::Value; + +use crate::{ + Error, + anthropic::{ + common_utils::{existing_betas, merge_beta_headers}, + messages::{ + handler::shape_anthropic_messages_request, + transformation::{provider_feature_betas, transform_messages_request}, + }, + }, + base_llm::{ + auth::AuthScheme, + messages::{ + context::MessagesTransformContext, + transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, + }, + }, + bedrock::common_utils::validated_region, +}; + +const API_BASE_ENV: &str = "BEDROCK_MANTLE_API_BASE"; +const REGION_ENVS: &[&str] = &[ + "BEDROCK_MANTLE_REGION", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", +]; +const KEY_ENVS: &[&str] = &["BEDROCK_MANTLE_API_KEY", "AWS_BEARER_TOKEN_BEDROCK"]; +const MESSAGES_PATH: &str = "/anthropic/v1/messages"; +const BASE_SUFFIXES: &[&str] = &[ + MESSAGES_PATH, + "/v1/messages", + "/messages", + "/anthropic/v1", + "/openai/v1", + "/v1", +]; + +pub struct BedrockMantleAnthropicMessagesConfig; +pub const BEDROCK_MANTLE_MESSAGES_CONFIG: BedrockMantleAnthropicMessagesConfig = + BedrockMantleAnthropicMessagesConfig; + +impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { + fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + model: &str, + params: &LitellmParams, + env: &dyn Fn(&str) -> Option, + ) -> Result { + let headers: Headers = headers + .into_iter() + .filter(|(name, _)| { + !name.eq_ignore_ascii_case("authorization") + && !name.eq_ignore_ascii_case("x-api-key") + }) + .collect(); + let headers = match params + .extra + .get("aws_bedrock_project_id") + .and_then(Value::as_str) + { + Some(project) => with_header(headers, "anthropic-workspace-id", project.to_string()), + None => headers, + }; + let auth = match resolve_non_empty(api_key, env, KEY_ENVS) { + Some(token) => AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(token), + }, + None => AuthScheme::AwsSigV4 { + region: mantle_region(params.api_base.as_deref(), model, params, env)?, + service: "bedrock", + credentials: Box::new(AwsCredentialSource::from_params(¶ms.aws, env)), + }, + }; + Ok(ValidatedEnvironment { headers, auth }) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + model: &str, + params: &LitellmParams, + _stream: bool, + env: &dyn Fn(&str) -> Option, + ) -> Result { + let region = mantle_region(api_base, model, params, env)?; + let configured = mantle_base(api_base, params, env) + .unwrap_or_else(|| format!("https://bedrock-mantle.{region}.api.aws")); + let trimmed = configured.trim_end_matches('/'); + let stripped = BASE_SUFFIXES + .iter() + .find_map(|suffix| trimmed.strip_suffix(suffix)) + .unwrap_or(trimmed); + let base = if mantle_host_region(stripped).is_some() { + format!("https://bedrock-mantle.{region}.api.aws") + } else { + stripped.to_string() + }; + Ok(format!("{base}{MESSAGES_PATH}")) + } + + fn transform_anthropic_messages_request( + &self, + request: MessagesRequest, + context: &MessagesTransformContext, + ) -> Result { + transform_messages_request(request, context) + } + + fn shape_request( + &self, + request: MessagesRequest, + reasoning_auto_summary: bool, + ) -> Result { + shape_anthropic_messages_request(request, reasoning_auto_summary) + } + + fn wire_body(&self, body: Value) -> Value { + match body { + Value::Object(fields) => Value::Object( + fields + .into_iter() + .filter_map(|(name, value)| match name.as_str() { + "anthropic_version" | "anthropic_beta" => None, + "model" => Some(( + name, + value + .as_str() + .map(|model| Value::String(mantle_model(model).0)) + .unwrap_or(value), + )), + _ => Some((name, value)), + }) + .collect(), + ), + other => other, + } + } + + fn secret_names(&self) -> &'static [&'static str] { + static NAMES: LazyLock> = LazyLock::new(|| { + [ + "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", + API_BASE_ENV, + ] + .into_iter() + .chain(REGION_ENVS.iter().copied()) + .chain( + AwsParams::secret_names() + .iter() + .copied() + .filter(|name| !REGION_ENVS.contains(name)), + ) + .collect() + }); + &NAMES + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + DEFAULT_HEADERS + } + + fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { + let body_betas: BetaSet = request + .params + .extra + .get("anthropic_beta") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_str) + .filter_map(|beta| beta.parse::().ok()) + .collect(); + let caller: BetaSet = existing_betas(&headers) + .union(body_betas) + .into_iter() + .filter_map(|beta| match beta { + AnthropicBeta::Other(_) => Some(beta), + _ => beta.on(BetaProvider::BedrockMantle), + }) + .collect(); + let safeguards: BetaSet = request + .params + .extra + .contains_key("safeguards") + .then_some(AnthropicBeta::DangerousToolUse20260903) + .into_iter() + .collect(); + merge_beta_headers( + without_headers(headers, &["anthropic-beta"]), + caller + .union(provider_feature_betas(request, BetaProvider::BedrockMantle)) + .union(safeguards), + ) + } +} + +fn mantle_model(model: &str) -> (String, Option) { + bedrock_model_id_and_region(model.strip_prefix("mantle/").unwrap_or(model)) +} + +fn mantle_host_region(base: &str) -> Option { + let url = url::Url::parse(base).ok()?; + if !matches!(url.scheme(), "https" | "http") { + return None; + } + url.host_str()? + .strip_prefix("bedrock-mantle.")? + .strip_suffix(".api.aws") + .filter(|region| !region.contains('.') && !region.is_empty()) + .map(str::to_string) +} + +fn mantle_base( + base: Option<&str>, + params: &LitellmParams, + env: &dyn Fn(&str) -> Option, +) -> Option { + let explicit = [ + base, + params.api_base.as_deref(), + params.aws.aws_bedrock_runtime_endpoint.as_deref(), + ] + .into_iter() + .flatten() + .find(|base| !base.trim().is_empty()); + resolve_non_empty(explicit, env, &[API_BASE_ENV]) +} + +fn mantle_region( + base: Option<&str>, + model: &str, + params: &LitellmParams, + env: &dyn Fn(&str) -> Option, +) -> Result { + let region = mantle_base(base, params, env) + .and_then(|base| mantle_host_region(&base)) + .or_else(|| resolve_non_empty(params.aws.aws_region_name.as_deref(), &|_| None, &[])) + .or_else(|| mantle_model(model).1) + .or_else(|| resolve_non_empty(None, env, REGION_ENVS)) + .unwrap_or_else(|| "us-east-1".into()); + validated_region(region) +} + +#[cfg(test)] +mod tests { + use super::*; + use litellm_http::request::header_value; + use rstest::rstest; + use serde_json::json; + + #[rstest] + #[case::root("https://private.test", "https://private.test/anthropic/v1/messages")] + #[case::openai( + "https://private.test/openai/v1/", + "https://private.test/anthropic/v1/messages" + )] + #[case::complete( + "https://private.test/anthropic/v1/messages/", + "https://private.test/anthropic/v1/messages" + )] + #[case::tenant( + "https://private.test/tenant/v1", + "https://private.test/tenant/anthropic/v1/messages" + )] + fn endpoints_are_normalized(#[case] base: &str, #[case] expected: &str) { + assert_eq!( + BEDROCK_MANTLE_MESSAGES_CONFIG + .get_complete_url( + Some(base), + "claude-test", + &LitellmParams::default(), + true, + &|_| None + ) + .unwrap(), + expected + ); + } + + #[rstest] + #[case::parameter("us-west-2", "claude-test")] + #[case::model("us-west-2", "us-west-2/claude-test")] + fn private_endpoints_keep_their_host_and_signing_region( + #[case] region: &str, + #[case] model: &str, + ) { + let params = LitellmParams { + aws: AwsParams { + aws_region_name: (model == "claude-test").then(|| region.into()), + aws_bedrock_runtime_endpoint: Some("https://private.test/anthropic/v1".into()), + ..Default::default() + }, + ..Default::default() + }; + assert_eq!( + BEDROCK_MANTLE_MESSAGES_CONFIG + .get_complete_url(None, model, ¶ms, false, &|_| None) + .unwrap(), + "https://private.test/anthropic/v1/messages" + ); + let environment = BEDROCK_MANTLE_MESSAGES_CONFIG + .validate_environment(vec![], None, model, ¶ms, &|_| None) + .unwrap(); + assert!( + matches!(environment.auth, AuthScheme::AwsSigV4 { region: actual, service: "bedrock", .. } if actual == region) + ); + } + + #[rstest] + #[case::explicit_regional_endpoint( + "https://bedrock-mantle.us-west-2.api.aws", + "us-east-1", + "us-west-2" + )] + fn public_regional_host_precedes_configured_region( + #[case] api_base: &str, + #[case] configured_region: &str, + #[case] expected_region: &str, + ) { + let params = LitellmParams { + api_base: Some(api_base.into()), + aws: AwsParams { + aws_region_name: Some(configured_region.into()), + ..Default::default() + }, + ..Default::default() + }; + let url = BEDROCK_MANTLE_MESSAGES_CONFIG + .get_complete_url(None, "claude-test", ¶ms, false, &|_| None) + .unwrap(); + let environment = BEDROCK_MANTLE_MESSAGES_CONFIG + .validate_environment(vec![], None, "claude-test", ¶ms, &|_| None) + .unwrap(); + assert_eq!( + url, + "https://bedrock-mantle.us-west-2.api.aws/anthropic/v1/messages" + ); + assert!( + matches!(environment.auth, AuthScheme::AwsSigV4 { region, service: "bedrock", .. } if region == expected_region) + ); + } + + #[rstest] + #[case::aws_default_region("AWS_DEFAULT_REGION", "eu-west-1")] + fn aws_default_region_selects_mantle_endpoint(#[case] env_name: &str, #[case] region: &str) { + let url = BEDROCK_MANTLE_MESSAGES_CONFIG + .get_complete_url( + None, + "claude-test", + &LitellmParams::default(), + false, + &|name| (name == env_name).then(|| region.into()), + ) + .unwrap(); + assert_eq!( + url, + "https://bedrock-mantle.eu-west-1.api.aws/anthropic/v1/messages" + ); + } + + #[rstest] + #[case::aws_default_region("AWS_DEFAULT_REGION")] + fn mantle_secret_names_include_region_fallbacks(#[case] env_name: &str) { + assert!( + BEDROCK_MANTLE_MESSAGES_CONFIG + .secret_names() + .contains(&env_name) + ); + } + + #[rstest] + #[case::explicit(Some("explicit"), "explicit")] + #[case::environment(None, "mantle")] + fn bearer_credentials_replace_forwarded_auth( + #[case] key: Option<&str>, + #[case] expected: &str, + ) { + let environment = BEDROCK_MANTLE_MESSAGES_CONFIG + .validate_environment( + vec![ + ("X-Api-Key".into(), "stale".into()), + ("Authorization".into(), "stale".into()), + ], + key, + "claude-test", + &LitellmParams::default(), + &|name| match name { + "BEDROCK_MANTLE_API_KEY" => Some("mantle".into()), + "AWS_BEARER_TOKEN_BEDROCK" => Some("aws".into()), + _ => None, + }, + ) + .unwrap(); + assert!(environment.headers.is_empty()); + assert!( + matches!(environment.auth, AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret } if secret.expose() == expected) + ); + } + + #[rstest] + #[case::url_injection("us-east-1.attacker.test/steal?")] + #[case::uppercase("US-EAST-1")] + fn invalid_regions_are_rejected_even_with_bearer_auth(#[case] region: &str) { + let params = LitellmParams { + aws: AwsParams { + aws_region_name: Some(region.into()), + ..Default::default() + }, + ..Default::default() + }; + assert!(matches!( + BEDROCK_MANTLE_MESSAGES_CONFIG.get_complete_url( + None, + "claude-test", + ¶ms, + false, + &|_| None + ), + Err(Error::InvalidRequest(_)) + )); + } + + #[rstest] + #[case::arn("arn:aws:bedrock:us-east-1.attacker.test:123:model/claude-test", None)] + #[case::environment("claude-test", Some("us-east-1/steal"))] + fn model_arn_and_environment_regions_are_validated( + #[case] model: &str, + #[case] region: Option<&str>, + ) { + let url = BEDROCK_MANTLE_MESSAGES_CONFIG.get_complete_url( + None, + model, + &LitellmParams::default(), + false, + &|name| { + (name == "BEDROCK_MANTLE_REGION") + .then(|| region.map(str::to_string)) + .flatten() + }, + ); + assert!(matches!(url, Err(Error::InvalidRequest(_)))); + } + + #[rstest] + fn caller_betas_are_filtered_renamed_and_merged_with_feature_betas() { + let request: MessagesRequest = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": [{"type": "dangerous_tool_use"}]})).unwrap(); + let headers = BEDROCK_MANTLE_MESSAGES_CONFIG.request_headers( + vec![( + "Anthropic-Beta".into(), + "advanced-tool-use-2025-11-20,code-execution-2025-08-25,custom-beta".into(), + )], + &request, + ); + let actual: std::collections::BTreeSet<&str> = header_value(&headers, "anthropic-beta") + .unwrap() + .split(',') + .collect(); + let tool_search = AnthropicBeta::ToolSearchTool20251019.to_string(); + let safeguards = AnthropicBeta::DangerousToolUse20260903.to_string(); + assert_eq!( + actual, + ["custom-beta", tool_search.as_str(), safeguards.as_str()] + .into_iter() + .collect() + ); + } + + #[rstest] + #[case::body_only( + json!(["custom-beta", "advanced-tool-use-2025-11-20"]), + vec![], + vec![ + AnthropicBeta::Other("custom-beta".into()), + AnthropicBeta::ToolSearchTool20251019, + ] + )] + #[case::header_and_body( + json!(["custom-b"]), + vec![("Anthropic-Beta".into(), "custom-a".into())], + vec![ + AnthropicBeta::Other("custom-a".into()), + AnthropicBeta::Other("custom-b".into()), + ] + )] + fn body_betas_are_filtered_renamed_and_merged_with_headers( + #[case] body_betas: Value, + #[case] headers: Headers, + #[case] expected: Vec, + ) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-test", + "max_tokens": 16, + "messages": [], + "anthropic_beta": body_betas + })) + .unwrap(); + let headers = BEDROCK_MANTLE_MESSAGES_CONFIG.request_headers(headers, &request); + let actual: BetaSet = header_value(&headers, "anthropic-beta") + .unwrap() + .parse() + .unwrap(); + let expected: BetaSet = expected.into_iter().collect(); + assert_eq!(actual, expected); + } + + #[rstest] + fn wire_body_keeps_streaming_and_removes_sdk_fields() { + assert_eq!(BEDROCK_MANTLE_MESSAGES_CONFIG.wire_body(json!({"model": "mantle/us-west-2/anthropic.claude-test", "stream": true, "anthropic_version": "sdk", "anthropic_beta": ["sdk"], "messages": []})), json!({"model": "anthropic.claude-test", "stream": true, "messages": []})); + } +} diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/mod.rs b/litellm-rust/crates/llms/src/bedrock_mantle/mod.rs new file mode 100644 index 00000000000..ba63992f3cb --- /dev/null +++ b/litellm-rust/crates/llms/src/bedrock_mantle/mod.rs @@ -0,0 +1 @@ +pub mod messages; diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 01e6f4328d3..1cf69ab5745 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -3,6 +3,7 @@ pub mod aws_textract; pub mod azure_ai; pub mod base_llm; pub mod bedrock; +pub mod bedrock_mantle; pub mod cohere; pub mod deepseek; mod error;