scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

chore(rust): prepare Bedrock messages interview starter

Browse source 
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-29 03:05:57 +00:00
parent ea7fc3747e
commit 78b8cf0693
8 changed files with 110 additions and 183 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

75
litellm-rust/INTERVIEW_TASK.md Normal file
View file

@ -0,0 +1,75 @@
# Bedrock Messages Interview Task
Complete the AWS Bedrock implementation behind the Anthropic-compatible
`POST /v1/messages` route. The gateway plumbing, provider registration, bearer
authentication setup, SigV4 support scaffolding, streaming response plumbing,
test harness, and live harness are already present. The remaining work is to
complete the shared Messages abstractions, Bedrock request and response
transforms, and Bedrock EventStream decoding.
## Given
The gateway is configured to use:
- `AWS_BEARER_TOKEN_BEDROCK` for the live bearer credential
- Region `us-west-2`
- Model `us.anthropic.claude-sonnet-4-5-20250929-v1:0`
Claude 3.5 and Claude 3.7 Bedrock model identifiers are end of life and should
not be used.
Bedrock streaming responses use
`application/vnd.amazon.eventstream`, not ordinary SSE. The EventStream
payload contains an object shaped like
`{"bytes": "<base64-encoded Anthropic event JSON>"}`. The
`aws-smithy-eventstream::MessageFrameDecoder` type is available for parsing
the binary framing.
## Running the gateway and live harness
From `litellm-rust/`, start the gateway with:
```bash
cargo run -p litellm-ai-gateway --features server
```
In another shell, with `AWS_BEARER_TOKEN_BEDROCK` available in the
environment, run:
```bash
./bedrock_messages_harness.sh
```
The harness exercises simple non-streaming output, streaming output, tool use,
and an invalid-model response.
## Tests and checks
Format and compile the workspace with:
```bash
cargo fmt --check
cargo build
```
Run the existing test suite with:
```bash
cargo test --workspace
```
The Bedrock-specific tests in the core transformation and messages route
modules describe the expected behavior. They are intentionally failing until
the implementation is completed.
## Required scenarios
Make all of these scenarios pass:
1. Non-streaming Bedrock InvokeModel requests use the correct encoded model
path, request body shape, authentication headers, and response mapping
2. Streaming Bedrock EventStream frames, including frames split across
transport chunks, become normalized Anthropic SSE
3. Tool-use requests and responses round trip correctly
4. Bedrock HTTP errors and EventStream exception frames map to the gateway's
expected error shapes

6
litellm-rust/crates/ai-gateway/src/messages/handler.rs
View file

@ -3,7 +3,6 @@ use std::time::SystemTime;
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::messages::transformation::MessagesAuthStrategy;
use litellm_core::providers::bedrock::aws_base::{
AwsAuthConfig, resolve_credentials, sign_bedrock_post,
};
@ -24,10 +23,7 @@ async fn signed_request(
request: &ProviderMessagesRequest,
body: &[u8],
) -> CoreResult<Vec<(String, String)>> {
if !matches!(
request.config.auth_strategy(),
MessagesAuthStrategy::AwsSigV4
) {
if request.provider != BEDROCK_MESSAGES_PROVIDER {
return Ok(request.upstream_headers.clone());
}
if let Some(token) = &request.bearer_token {

15
litellm-rust/crates/ai-gateway/src/messages/prepare.rs
View file

@ -6,6 +6,7 @@ use serde_json::Value;
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
use super::types::{MessagesRequest, ProviderMessagesRequest};
use crate::constants::BEDROCK_MESSAGES_PROVIDER;
pub(super) fn prepare_messages_call(
request: MessagesRequest<'_>,
@ -33,8 +34,13 @@ pub(super) fn prepare_messages_call(
let mut headers = string_headers(request.extra_headers)?;
let auth_strategy = config.auth_strategy();
let bearer_token = if matches!(auth_strategy, MessagesAuthStrategy::AwsSigV4) {
let is_bedrock = provider == BEDROCK_MESSAGES_PROVIDER;
let auth_strategy = if is_bedrock {
MessagesAuthStrategy::Header("authorization")
} else {
config.auth_strategy()
};
let bearer_token = if is_bedrock {
request
.api_key
.map(str::to_string)
@ -44,7 +50,6 @@ pub(super) fn prepare_messages_call(
None
};
let auth_header = match auth_strategy {
MessagesAuthStrategy::AwsSigV4 => None,
MessagesAuthStrategy::Bearer
if has_header(&headers, "authorization")
|| (config.accepts_bearer_auth() && has_bearer_auth(&headers)) =>
@ -76,8 +81,8 @@ pub(super) fn prepare_messages_call(
}
}
let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
let url = config.complete_url(request.api_base, &model, stream, &env_lookup)?;
let _stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
let url = config.complete_url(request.api_base, &model, &env_lookup)?;
let signing_region = config.signing_region(request.api_base, &env_lookup);
let typed_request = serde_json::from_value(request.body).map_err(|err| {
CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}"))

59
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -2,7 +2,7 @@
mod service;
use aws_smithy_eventstream::frame::{DecodedFrame, MessageFrameDecoder};
use aws_smithy_eventstream::frame::MessageFrameDecoder;
use axum::Router;
use axum::body::Body;
use axum::extract::{Json, State};
@ -88,6 +88,7 @@ fn stream_response(
})
}
#[allow(dead_code)]
struct EventStreamState {
upstream: BoxStream<'static, Result<Bytes, reqwest::Error>>,
buffer: bytes::BytesMut,
@ -105,63 +106,14 @@ fn bedrock_sse_stream(
decoder: MessageFrameDecoder::new(),
terminated: false,
},
|mut state| async move {
if state.terminated {
return None;
}
loop {
match state.decoder.decode_frame(&mut state.buffer) {
Ok(DecodedFrame::Complete(message)) => {
let bytes = message
.headers()
.iter()
.find(|header| header.name().as_str() == ":message-type")
.and_then(|header| header.value().as_string().ok())
.map_or_else(
|| sse_data(message.payload()),
|message_type| {
if message_type.as_str() == "exception"
|| message_type.as_str() == "error"
{
sse_error(message.payload())
} else {
sse_data(message.payload())
}
},
);
return Some((Ok(Bytes::from(bytes)), state));
}
Ok(DecodedFrame::Incomplete) => {}
Err(error) => {
state.terminated = true;
return Some((
Ok(Bytes::from(sse_error(error.to_string().as_bytes()))),
state,
));
}
}
match state.upstream.next().await {
Some(Ok(chunk)) => state.buffer.extend_from_slice(&chunk),
Some(Err(error)) => {
return Some((Err(std::io::Error::other(error.to_string())), state));
}
None if state.buffer.is_empty() => return None,
None => {
state.terminated = true;
return Some((
Ok(Bytes::from(sse_error(
b"incomplete Bedrock event stream frame",
))),
state,
));
}
}
}
|_state| async move {
todo!("decode Bedrock EventStream frames and emit normalized Anthropic SSE")
},
)
.boxed()
}
#[allow(dead_code)]
fn sse_data(payload: &[u8]) -> String {
let value = serde_json::from_slice::<serde_json::Value>(payload)
.ok()
@ -178,6 +130,7 @@ fn sse_data(payload: &[u8]) -> String {
format!("event: {event}\ndata: {value}\n\n")
}
#[allow(dead_code)]
fn sse_error(payload: &[u8]) -> String {
let message = String::from_utf8_lossy(payload);
format!(

3
litellm-rust/crates/core/src/messages/transformation.rs
View file

@ -6,7 +6,6 @@ use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
pub enum MessagesAuthStrategy {
Bearer,
Header(&'static str),
AwsSigV4,
}
impl MessagesAuthStrategy {
@ -14,7 +13,6 @@ impl MessagesAuthStrategy {
match self {
Self::Bearer => "authorization",
Self::Header(header_name) => header_name,
Self::AwsSigV4 => "",
}
}
}
@ -24,7 +22,6 @@ pub trait AnthropicMessagesProviderConfig: Sync {
&self,
api_base: Option<&str>,
model: &str,
stream: bool,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;

1
litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs
View file

@ -51,7 +51,6 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
&self,
api_base: Option<&str>,
_model: &str,
_stream: bool,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
Ok(complete_anthropic_url(api_base, env_lookup))

1
litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs
View file

@ -146,7 +146,6 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
&self,
api_base: Option<&str>,
_model: &str,
_stream: bool,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
complete_azure_anthropic_url(api_base, env_lookup)

133
litellm-rust/crates/core/src/providers/bedrock/messages/transformation.rs
View file

@ -1,100 +1,18 @@
use crate::error::{CoreError, CoreResult};
use crate::error::CoreResult;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::messages::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
use crate::providers::anthropic::messages::transformation::non_empty;
use crate::providers::bedrock::constants::{
AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, DEFAULT_BEDROCK_REGION,
};
use serde_json::Value;
const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION";
const API_BASE_SCHEME: &str = "https://";
const MODEL_PATH_PREFIX: &str = "/model/";
const INVOKE_PATH: &str = "/invoke";
const STREAM_PATH: &str = "/invoke-with-response-stream";
const ANTHROPIC_VERSION_FIELD: &str = "anthropic_version";
const ANTHROPIC_VERSION: &str = "bedrock-2023-05-31";
const UNSUPPORTED_FIELDS: &[&str] = &[
"metadata",
"service_tier",
"container",
"mcp_servers",
"context_management",
"output_format",
"output_config",
"speed",
"inference_geo",
];
pub struct BedrockMessagesConfig;
pub const BEDROCK_MESSAGES_CONFIG: BedrockMessagesConfig = BedrockMessagesConfig;
fn resolve_region(api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option<String>) -> String {
api_base
.and_then(bedrock_region_from_api_base)
.or_else(|| env_lookup(AWS_REGION_NAME))
.or_else(|| env_lookup(AWS_REGION))
.or_else(|| env_lookup(AWS_DEFAULT_REGION))
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
fn bedrock_region_from_api_base(api_base: &str) -> Option<String> {
let host = api_base
.trim()
.trim_start_matches("https://")
.trim_start_matches("http://")
.split('/')
.next()?
.split(':')
.next()?;
let region = host
.strip_prefix("bedrock-runtime.")?
.strip_suffix(".amazonaws.com")?;
(!region.is_empty()).then(|| region.to_string())
}
fn encode_path_segment(value: &str) -> String {
value.bytes().fold(String::new(), |mut encoded, byte| {
if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') {
encoded.push(byte as char);
} else {
encoded.push('%');
encoded.push_str(&format!("{byte:02X}"));
}
encoded
})
}
fn endpoint_base(api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option<String>) -> String {
non_empty(api_base)
.map(str::to_string)
.unwrap_or_else(|| {
BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", &resolve_region(None, env_lookup))
})
.trim_end_matches('/')
.to_string()
}
pub fn complete_bedrock_url(
api_base: Option<&str>,
model: &str,
stream: bool,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let model = non_empty(Some(model))
.ok_or_else(|| CoreError::InvalidRequest("Bedrock model cannot be empty".to_string()))?;
let suffix = if stream { STREAM_PATH } else { INVOKE_PATH };
let base = endpoint_base(api_base, env_lookup);
let base = if base.starts_with(API_BASE_SCHEME) || base.starts_with("http://") {
base
} else {
format!("{API_BASE_SCHEME}{base}")
};
Ok(format!(
"{base}{MODEL_PATH_PREFIX}{}{suffix}",
encode_path_segment(model)
))
let _ = (api_base, model, stream, env_lookup);
todo!("implement Bedrock InvokeModel URL construction and model path encoding")
}
impl AnthropicMessagesProviderConfig for BedrockMessagesConfig {
@ -102,10 +20,10 @@ impl AnthropicMessagesProviderConfig for BedrockMessagesConfig {
&self,
api_base: Option<&str>,
model: &str,
stream: bool,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
complete_bedrock_url(api_base, model, stream, env_lookup)
let _ = (api_base, model, env_lookup);
todo!("extend the shared URL contract for streaming, then implement Bedrock URLs")
}
fn signing_region(
@ -113,7 +31,8 @@ impl AnthropicMessagesProviderConfig for BedrockMessagesConfig {
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
Some(resolve_region(api_base, env_lookup))
let _ = (api_base, env_lookup);
todo!("implement Bedrock region precedence and endpoint parsing")
}
fn resolve_api_key(
@ -125,50 +44,34 @@ impl AnthropicMessagesProviderConfig for BedrockMessagesConfig {
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
MessagesAuthStrategy::AwsSigV4
todo!("extend MessagesAuthStrategy for Bedrock bearer and SigV4 authentication")
}
fn transform_request(
&self,
mut request: AnthropicMessagesRequest,
request: AnthropicMessagesRequest,
) -> CoreResult<AnthropicMessagesRequest> {
request.model.clear();
request.stream = None;
request.metadata = None;
request.service_tier = None;
request.container = None;
request.mcp_servers = None;
request.context_management = None;
request.output_format = None;
request.output_config = None;
request.speed = None;
request.inference_geo = None;
request
.extra
.retain(|key, _| !UNSUPPORTED_FIELDS.contains(&key.as_str()));
request.extra.insert(
ANTHROPIC_VERSION_FIELD.to_string(),
Value::String(ANTHROPIC_VERSION.to_string()),
);
Ok(request)
let _ = request;
todo!("implement Bedrock request body filtering and anthropic_version injection")
}
fn transform_response(
&self,
model: &str,
mut response: AnthropicMessagesResponse,
response: AnthropicMessagesResponse,
) -> CoreResult<AnthropicMessagesResponse> {
if response.model.trim().is_empty() {
response.model = model.to_string();
}
Ok(response)
let _ = (model, response);
todo!("restamp the requested model when Bedrock omits it")
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use crate::providers::bedrock::constants::{AWS_REGION, AWS_REGION_NAME};
use serde_json::{Value, json};
const ANTHROPIC_VERSION: &str = "bedrock-2023-05-31";
fn request(value: Value) -> AnthropicMessagesRequest {
serde_json::from_value(value).expect("valid request")