use std::collections::HashMap;
use anyhow::Context as _;
use async_trait::async_trait;
use aws_sdk_bedrockruntime::{
operation::converse::ConverseOutput,
types::{
AudioBlock, AudioFormat, AudioSource, AutoToolChoice, ContentBlock, ContentBlockDelta,
ContentBlockStart, ConversationRole, ConverseOutput as ConverseResult,
ConverseStreamOutput, DocumentBlock, DocumentFormat, DocumentSource, ImageBlock,
ImageFormat, ImageSource, InferenceConfiguration, Message, ReasoningContentBlock,
ReasoningContentBlockDelta, ReasoningTextBlock, S3Location, StopReason, SystemContentBlock,
Tool, ToolChoice, ToolConfiguration, ToolInputSchema, ToolResultBlock,
ToolResultContentBlock, ToolResultStatus, ToolSpecification, ToolUseBlock, VideoBlock,
VideoFormat, VideoSource,
},
};
use aws_smithy_json::{
deserialize::{json_token_iter, token::expect_document},
serialize::JsonValueWriter,
};
use aws_smithy_types::{Blob, Document};
use base64::Engine as _;
use futures_util::stream;
#[cfg(feature = "langfuse")]
use serde_json::{Value, json};
use swiftide_core::{
ChatCompletion, ChatCompletionStream,
chat_completion::{
ChatCompletionRequest, ChatCompletionResponse, ChatMessage, ChatMessageContentPart,
ChatMessageContentSource, ReasoningItem, ToolCall, ToolOutput, ToolSpec,
errors::LanguageModelError,
},
};
use tracing_futures::Instrument;
use super::tool_schema::AwsBedrockToolSchema;
use super::{AwsBedrock, Options};
type ConverseInputParts = (
Vec<Message>,
Option<Vec<SystemContentBlock>>,
Option<InferenceConfiguration>,
Option<ToolConfiguration>,
);
type ExtractedMessage = (Option<String>, Option<Vec<ToolCall>>, Vec<ReasoningItem>);
#[async_trait]
impl ChatCompletion for AwsBedrock {
#[cfg_attr(not(feature = "langfuse"), tracing::instrument(skip_all, err))]
#[cfg_attr(
feature = "langfuse",
tracing::instrument(skip_all, err, fields(langfuse.type = "GENERATION"))
)]
async fn complete(
&self,
request: &ChatCompletionRequest<'_>,
) -> Result<ChatCompletionResponse, LanguageModelError> {
let model = self.prompt_model()?;
#[cfg(feature = "langfuse")]
let fallback_tracking_request = Some(json!({
"model": model,
"messages": request.messages(),
"tools_spec": request.tools_spec(),
}));
#[cfg(not(feature = "langfuse"))]
let fallback_tracking_request: Option<serde_json::Value> = None;
let (messages, system, inference_config, tool_config) =
match build_converse_input(request, &self.default_options) {
Ok(parts) => parts,
Err(error) => {
Self::track_failure(
model,
fallback_tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return Err(error);
}
};
let additional_model_request_fields =
match super::additional_model_request_fields_from_options(model, &self.default_options)
{
Ok(fields) => fields,
Err(error) => {
#[cfg(feature = "langfuse")]
let tracking_request = Some(tracking_request_from_converse_parts(
model,
&messages,
system.as_deref(),
inference_config.as_ref(),
tool_config.as_ref(),
None,
self.default_options
.additional_model_response_field_paths
.as_deref(),
));
#[cfg(not(feature = "langfuse"))]
let tracking_request: Option<serde_json::Value> = None;
Self::track_failure(
model,
tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return Err(error);
}
};
#[cfg(feature = "langfuse")]
let tracking_request = Some(tracking_request_from_converse_parts(
model,
&messages,
system.as_deref(),
inference_config.as_ref(),
tool_config.as_ref(),
additional_model_request_fields.as_ref(),
self.default_options
.additional_model_response_field_paths
.as_deref(),
));
#[cfg(not(feature = "langfuse"))]
let tracking_request: Option<serde_json::Value> = None;
tracing::debug!(
model = model,
inference_config = ?inference_config,
has_tool_config = tool_config.is_some(),
"[ChatCompletion] Request to bedrock converse"
);
let response = match self
.client
.converse(
model,
messages,
system,
inference_config,
tool_config,
None,
additional_model_request_fields,
self.default_options
.additional_model_response_field_paths
.clone(),
)
.await
{
Ok(response) => response,
Err(error) => {
Self::track_failure(
model,
tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return Err(error);
}
};
tracing::debug!(response = ?response, "[ChatCompletion] Response from bedrock converse");
let completion = match response_to_chat_completion(&response) {
Ok(completion) => completion,
Err(error) => {
Self::track_failure(
model,
tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return Err(error);
}
};
if let Some(error) = super::context_length_exceeded_if_empty(
completion.message.is_some(),
completion.tool_calls.is_some(),
completion
.reasoning
.as_ref()
.is_some_and(|reasoning| !reasoning.is_empty()),
Some(response.stop_reason()),
) {
Self::track_failure(model, tracking_request.as_ref(), Some(&completion), &error);
return Err(error);
}
self.track_completion(
model,
completion.usage.as_ref(),
tracking_request.as_ref(),
Some(&completion),
)
.await?;
Ok(completion)
}
#[cfg_attr(not(feature = "langfuse"), tracing::instrument(skip_all))]
#[cfg_attr(
feature = "langfuse",
tracing::instrument(skip_all, fields(langfuse.type = "GENERATION"))
)]
async fn complete_stream(&self, request: &ChatCompletionRequest<'_>) -> ChatCompletionStream {
let model = match self.prompt_model() {
Ok(model) => model.to_string(),
Err(error) => return error.into(),
};
#[cfg(feature = "langfuse")]
let fallback_tracking_request = Some(json!({
"model": model,
"messages": request.messages(),
"tools_spec": request.tools_spec(),
}));
#[cfg(not(feature = "langfuse"))]
let fallback_tracking_request: Option<serde_json::Value> = None;
let (messages, system, inference_config, tool_config) =
match build_converse_input(request, &self.default_options) {
Ok(parts) => parts,
Err(error) => {
Self::track_failure(
&model,
fallback_tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return error.into();
}
};
let additional_model_request_fields =
match super::additional_model_request_fields_from_options(&model, &self.default_options)
{
Ok(fields) => fields,
Err(error) => {
#[cfg(feature = "langfuse")]
let tracking_request = Some(tracking_request_from_converse_parts(
&model,
&messages,
system.as_deref(),
inference_config.as_ref(),
tool_config.as_ref(),
None,
self.default_options
.additional_model_response_field_paths
.as_deref(),
));
#[cfg(not(feature = "langfuse"))]
let tracking_request: Option<serde_json::Value> = None;
Self::track_failure(
&model,
tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return error.into();
}
};
#[cfg(feature = "langfuse")]
let tracking_request = Some(tracking_request_from_converse_parts(
&model,
&messages,
system.as_deref(),
inference_config.as_ref(),
tool_config.as_ref(),
additional_model_request_fields.as_ref(),
self.default_options
.additional_model_response_field_paths
.as_deref(),
));
#[cfg(not(feature = "langfuse"))]
let tracking_request: Option<serde_json::Value> = None;
let stream_output = match self
.client
.converse_stream(
&model,
messages,
system,
inference_config,
tool_config,
additional_model_request_fields,
self.default_options
.additional_model_response_field_paths
.clone(),
)
.await
{
Ok(stream_output) => stream_output,
Err(error) => {
Self::track_failure(
&model,
tracking_request.as_ref(),
None::<&serde_json::Value>,
&error,
);
return error.into();
}
};
let self_for_stream = self.clone();
let event_stream = stream_output.stream;
let stream = stream::unfold(
(
event_stream,
ChatCompletionResponse::default(),
None::<StopReason>,
false,
model,
tracking_request,
),
move |(
mut event_stream,
mut response,
mut stop_reason,
finished,
model,
tracking_request,
)| {
let self_for_stream = self_for_stream.clone();
async move {
if finished {
return None;
}
match event_stream.recv().await {
Ok(Some(event)) => {
apply_stream_event(&event, &mut response, &mut stop_reason);
Some((
Ok(response.clone()),
(
event_stream,
response,
stop_reason,
false,
model,
tracking_request,
),
))
}
Ok(None) => {
if let Some(error) = super::context_length_exceeded_if_empty(
response.message.is_some(),
response.tool_calls.is_some(),
response
.reasoning
.as_ref()
.is_some_and(|reasoning| !reasoning.is_empty()),
stop_reason.as_ref(),
) {
Self::track_failure(
&model,
tracking_request.as_ref(),
Some(&response),
&error,
);
return Some((
Err(error),
(
event_stream,
response,
stop_reason,
true,
model,
tracking_request,
),
));
}
if let Err(error) = self_for_stream
.track_completion(
&model,
response.usage.as_ref(),
tracking_request.as_ref(),
Some(&response),
)
.await
{
return Some((
Err(error),
(
event_stream,
response,
stop_reason,
true,
model,
tracking_request,
),
));
}
Some((
Ok(response.clone()),
(
event_stream,
response,
stop_reason,
true,
model,
tracking_request,
),
))
}
Err(error) => {
let error =
super::converse_stream_output_error_to_language_model_error(error);
Self::track_failure(
&model,
tracking_request.as_ref(),
Some(&response),
&error,
);
Some((
Err(error),
(
event_stream,
response,
stop_reason,
true,
model,
tracking_request,
),
))
}
}
}
},
);
let span = if cfg!(feature = "langfuse") {
tracing::info_span!("stream", langfuse.type = "GENERATION")
} else {
tracing::info_span!("stream")
};
Box::pin(Instrument::instrument(stream, span))
}
}
fn build_converse_input(
request: &ChatCompletionRequest<'_>,
options: &Options,
) -> Result<ConverseInputParts, LanguageModelError> {
let source_messages = request.messages();
let mut messages = Vec::with_capacity(source_messages.len());
let mut system = Vec::new();
let mut source_messages = source_messages.iter().peekable();
while let Some(message) = source_messages.next() {
match message {
ChatMessage::System(text) => {
system.push(SystemContentBlock::Text(text.clone()));
}
ChatMessage::Summary(text) | ChatMessage::User(text) => {
messages.push(user_message_from_text(text.clone())?);
}
ChatMessage::UserWithParts(parts) => messages.push(user_message_from_parts(parts)?),
ChatMessage::Assistant(content, maybe_tool_calls) => {
let mut blocks = Vec::with_capacity(
usize::from(content.as_ref().is_some_and(|text| !text.is_empty()))
+ maybe_tool_calls.as_ref().map_or(0, Vec::len),
);
if let Some(content) = content.as_ref()
&& !content.is_empty()
{
blocks.push(ContentBlock::Text(content.clone()));
}
if let Some(tool_calls) = maybe_tool_calls.as_ref() {
for tool_call in tool_calls {
let input =
tool_call_args_to_document(tool_call.args()).with_context(|| {
format!("Invalid JSON args for tool call {}", tool_call.name())
})?;
let tool_use = ToolUseBlock::builder()
.tool_use_id(tool_call.id())
.name(tool_call.name())
.input(input)
.build()
.map_err(LanguageModelError::permanent)?;
blocks.push(ContentBlock::ToolUse(tool_use));
}
}
if !blocks.is_empty() {
messages.push(message_from_blocks(ConversationRole::Assistant, blocks)?);
}
}
ChatMessage::ToolOutput(tool_call, output) => {
let mut blocks = vec![tool_output_to_result_block(tool_call, output)?];
while let Some(ChatMessage::ToolOutput(tool_call, output)) = source_messages.peek()
{
blocks.push(tool_output_to_result_block(tool_call, output)?);
source_messages.next();
}
messages.push(message_from_blocks(ConversationRole::User, blocks)?);
}
ChatMessage::Reasoning(item) => {
if let Some(reasoning_message) = assistant_reasoning_message_from_item(item)? {
messages.push(reasoning_message);
}
}
}
}
if messages.is_empty() {
return Err(LanguageModelError::permanent(
"Bedrock Converse requires at least one non-system message",
));
}
Ok((
messages,
(!system.is_empty()).then_some(system),
super::inference_config_from_options(options),
tool_config_from_specs(request.tools_spec().iter(), options.tool_strict_enabled())?,
))
}
#[cfg(feature = "langfuse")]
fn tracking_request_from_converse_parts(
model: &str,
messages: &[Message],
system: Option<&[SystemContentBlock]>,
inference_config: Option<&InferenceConfiguration>,
tool_config: Option<&ToolConfiguration>,
additional_model_request_fields: Option<&Document>,
additional_model_response_field_paths: Option<&[String]>,
) -> Value {
let mut payload = serde_json::Map::new();
payload.insert("model".into(), Value::String(model.to_string()));
payload.insert(
"messages".into(),
Value::Array(messages.iter().map(message_to_json).collect()),
);
if let Some(system) = system {
payload.insert(
"system".into(),
Value::Array(system.iter().map(system_content_block_to_json).collect()),
);
}
if let Some(inference_config) = inference_config {
payload.insert(
"inferenceConfig".into(),
inference_config_to_json(inference_config),
);
}
if let Some(tool_config) = tool_config {
payload.insert("toolConfig".into(), tool_config_to_json(tool_config));
}
if let Some(fields) = additional_model_request_fields {
payload.insert(
"additionalModelRequestFields".into(),
document_to_json_value(fields),
);
}
if let Some(paths) = additional_model_response_field_paths {
payload.insert(
"additionalModelResponseFieldPaths".into(),
Value::Array(paths.iter().cloned().map(Value::String).collect()),
);
}
Value::Object(payload)
}
#[cfg(feature = "langfuse")]
fn inference_config_to_json(config: &InferenceConfiguration) -> Value {
let mut value = serde_json::Map::new();
if let Some(max_tokens) = config.max_tokens() {
value.insert("maxTokens".into(), Value::from(max_tokens));
}
let stop_sequences = config.stop_sequences();
if !stop_sequences.is_empty() {
value.insert(
"stopSequences".into(),
Value::Array(stop_sequences.iter().cloned().map(Value::String).collect()),
);
}
if let Some(temperature) = config.temperature() {
value.insert("temperature".into(), Value::from(temperature));
}
if let Some(top_p) = config.top_p() {
value.insert("topP".into(), Value::from(top_p));
}
Value::Object(value)
}
#[cfg(feature = "langfuse")]
fn tool_config_to_json(config: &ToolConfiguration) -> Value {
let mut value = serde_json::Map::new();
let tools = config.tools();
if !tools.is_empty() {
value.insert(
"tools".into(),
Value::Array(tools.iter().map(tool_to_json).collect()),
);
}
if let Some(choice) = config.tool_choice() {
value.insert("toolChoice".into(), tool_choice_to_json(choice));
}
Value::Object(value)
}
#[cfg(feature = "langfuse")]
fn tool_choice_to_json(choice: &ToolChoice) -> Value {
match choice {
ToolChoice::Auto(_) => json!({ "auto": {} }),
ToolChoice::Any(_) => json!({ "any": {} }),
ToolChoice::Tool(tool) => json!({ "tool": { "name": tool.name() } }),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn tool_to_json(tool: &Tool) -> Value {
match tool {
Tool::ToolSpec(spec) => {
let mut value = serde_json::Map::new();
value.insert("name".into(), Value::String(spec.name().to_string()));
if let Some(description) = spec.description() {
value.insert("description".into(), Value::String(description.to_string()));
}
if let Some(input_schema) = spec.input_schema() {
value.insert(
"inputSchema".into(),
tool_input_schema_to_json(input_schema),
);
}
if let Some(strict) = spec.strict() {
value.insert("strict".into(), Value::Bool(strict));
}
json!({ "toolSpec": Value::Object(value) })
}
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn tool_input_schema_to_json(schema: &ToolInputSchema) -> Value {
match schema {
ToolInputSchema::Json(document) => json!({ "json": document_to_json_value(document) }),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn system_content_block_to_json(block: &SystemContentBlock) -> Value {
match block {
SystemContentBlock::Text(text) => json!({ "text": text }),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn message_to_json(message: &Message) -> Value {
json!({
"role": message.role().as_str(),
"content": message.content().iter().map(content_block_to_json).collect::<Vec<_>>(),
})
}
#[cfg(feature = "langfuse")]
fn content_block_to_json(block: &ContentBlock) -> Value {
match block {
ContentBlock::Text(text) => json!({ "text": text }),
ContentBlock::ToolUse(tool_use) => json!({
"toolUse": {
"toolUseId": tool_use.tool_use_id(),
"name": tool_use.name(),
"input": document_to_json_value(tool_use.input()),
}
}),
ContentBlock::ToolResult(tool_result) => json!({
"toolResult": {
"toolUseId": tool_result.tool_use_id(),
"content": tool_result.content().iter().map(tool_result_content_block_to_json).collect::<Vec<_>>(),
"status": tool_result.status().map(|status| format!("{status:?}")),
}
}),
ContentBlock::ReasoningContent(reasoning) => {
json!({ "reasoningContent": reasoning_content_block_to_json(reasoning) })
}
ContentBlock::Image(image) => json!({
"image": {
"format": format!("{:?}", image.format()),
"source": image.source().map_or(Value::Null, image_source_to_json),
}
}),
ContentBlock::Document(document) => json!({
"document": {
"format": format!("{:?}", document.format()),
"name": document.name(),
"source": document.source().map_or(Value::Null, document_source_to_json),
}
}),
ContentBlock::Audio(audio) => json!({
"audio": {
"format": format!("{:?}", audio.format()),
"source": audio.source().map_or(Value::Null, audio_source_to_json),
}
}),
ContentBlock::Video(video) => json!({
"video": {
"format": format!("{:?}", video.format()),
"source": video.source().map_or(Value::Null, video_source_to_json),
}
}),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn tool_result_content_block_to_json(block: &ToolResultContentBlock) -> Value {
match block {
ToolResultContentBlock::Text(text) => json!({ "text": text }),
ToolResultContentBlock::Json(document) => {
json!({ "json": document_to_json_value(document) })
}
ToolResultContentBlock::Image(image) => json!({
"image": {
"format": format!("{:?}", image.format()),
"source": image.source().map_or(Value::Null, image_source_to_json),
}
}),
ToolResultContentBlock::Document(document) => json!({
"document": {
"format": format!("{:?}", document.format()),
"name": document.name(),
"source": document.source().map_or(Value::Null, document_source_to_json),
}
}),
ToolResultContentBlock::Video(video) => json!({
"video": {
"format": format!("{:?}", video.format()),
"source": video.source().map_or(Value::Null, video_source_to_json),
}
}),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn reasoning_content_block_to_json(block: &ReasoningContentBlock) -> Value {
match block {
ReasoningContentBlock::ReasoningText(text) => json!({
"reasoningText": {
"text": text.text(),
"signature": text.signature(),
}
}),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn image_source_to_json(source: &ImageSource) -> Value {
match source {
ImageSource::Bytes(blob) => {
json!({ "bytes": base64::engine::general_purpose::STANDARD.encode(blob.as_ref()) })
}
ImageSource::S3Location(location) => s3_location_to_json(location),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn document_source_to_json(source: &DocumentSource) -> Value {
match source {
DocumentSource::Bytes(blob) => {
json!({ "bytes": base64::engine::general_purpose::STANDARD.encode(blob.as_ref()) })
}
DocumentSource::S3Location(location) => s3_location_to_json(location),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn audio_source_to_json(source: &AudioSource) -> Value {
match source {
AudioSource::Bytes(blob) => {
json!({ "bytes": base64::engine::general_purpose::STANDARD.encode(blob.as_ref()) })
}
AudioSource::S3Location(location) => s3_location_to_json(location),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn video_source_to_json(source: &VideoSource) -> Value {
match source {
VideoSource::Bytes(blob) => {
json!({ "bytes": base64::engine::general_purpose::STANDARD.encode(blob.as_ref()) })
}
VideoSource::S3Location(location) => s3_location_to_json(location),
other => Value::String(format!("{other:?}")),
}
}
#[cfg(feature = "langfuse")]
fn s3_location_to_json(location: &S3Location) -> Value {
json!({
"s3Location": {
"uri": location.uri(),
"bucketOwner": location.bucket_owner(),
}
})
}
#[cfg(feature = "langfuse")]
fn document_to_json_value(document: &Document) -> Value {
serde_json::from_str(&document_to_json_string(document))
.unwrap_or_else(|error| Value::String(format!("<<failed to convert document: {error}>>")))
}
fn user_message_from_text(text: String) -> Result<Message, LanguageModelError> {
message_from_blocks(ConversationRole::User, vec![ContentBlock::Text(text)])
}
fn user_message_from_parts(
parts: &[ChatMessageContentPart],
) -> Result<Message, LanguageModelError> {
let mut blocks = Vec::with_capacity(parts.len());
let mut has_text = false;
let mut has_document = false;
for part in parts {
match part {
ChatMessageContentPart::Text { text } => {
if !text.is_empty() {
blocks.push(ContentBlock::Text(text.clone()));
has_text = true;
}
}
ChatMessageContentPart::Image { source, format } => {
blocks.push(ContentBlock::Image(image_block_from_part(
source,
format.as_deref(),
)?));
}
ChatMessageContentPart::Document {
source,
format,
name,
} => {
blocks.push(ContentBlock::Document(document_block_from_part(
source,
format.as_deref(),
name.as_deref(),
)?));
has_document = true;
}
ChatMessageContentPart::Audio { source, format } => {
blocks.push(ContentBlock::Audio(audio_block_from_part(
source,
format.as_deref(),
)?));
}
ChatMessageContentPart::Video { source, format } => {
blocks.push(ContentBlock::Video(video_block_from_part(
source,
format.as_deref(),
)?));
}
}
}
if blocks.is_empty() {
return Err(LanguageModelError::permanent(
"UserWithParts requires at least one content part",
));
}
if has_document && !has_text {
return Err(LanguageModelError::permanent(
"Bedrock document parts require at least one text part in the same message",
));
}
message_from_blocks(ConversationRole::User, blocks)
}
fn image_block_from_part(
source: &ChatMessageContentSource,
format: Option<&str>,
) -> Result<ImageBlock, LanguageModelError> {
let format = image_format_from_source(format, source)?;
let source = image_source_from_content_source(source)?;
ImageBlock::builder()
.format(format)
.source(source)
.build()
.map_err(LanguageModelError::permanent)
}
fn document_block_from_part(
source: &ChatMessageContentSource,
format: Option<&str>,
name: Option<&str>,
) -> Result<DocumentBlock, LanguageModelError> {
let format = document_format_from_source(format, source)?;
let source = document_source_from_content_source(source)?;
let name = name.unwrap_or("document");
DocumentBlock::builder()
.format(format)
.name(name)
.source(source)
.build()
.map_err(LanguageModelError::permanent)
}
fn audio_block_from_part(
source: &ChatMessageContentSource,
format: Option<&str>,
) -> Result<AudioBlock, LanguageModelError> {
let format = audio_format_from_source(format, source)?;
let source = audio_source_from_content_source(source)?;
AudioBlock::builder()
.format(format)
.source(source)
.build()
.map_err(LanguageModelError::permanent)
}
fn video_block_from_part(
source: &ChatMessageContentSource,
format: Option<&str>,
) -> Result<VideoBlock, LanguageModelError> {
let format = video_format_from_source(format, source)?;
let source = video_source_from_content_source(source)?;
VideoBlock::builder()
.format(format)
.source(source)
.build()
.map_err(LanguageModelError::permanent)
}
fn image_source_from_content_source(
source: &ChatMessageContentSource,
) -> Result<ImageSource, LanguageModelError> {
source_from_content_source(source, "image", ImageSource::Bytes, ImageSource::S3Location)
}
fn document_source_from_content_source(
source: &ChatMessageContentSource,
) -> Result<DocumentSource, LanguageModelError> {
source_from_content_source(
source,
"document",
DocumentSource::Bytes,
DocumentSource::S3Location,
)
}
fn audio_source_from_content_source(
source: &ChatMessageContentSource,
) -> Result<AudioSource, LanguageModelError> {
source_from_content_source(source, "audio", AudioSource::Bytes, AudioSource::S3Location)
}
fn video_source_from_content_source(
source: &ChatMessageContentSource,
) -> Result<VideoSource, LanguageModelError> {
source_from_content_source(source, "video", VideoSource::Bytes, VideoSource::S3Location)
}
fn source_from_content_source<T>(
source: &ChatMessageContentSource,
label: &str,
from_bytes: impl Fn(Blob) -> T,
from_s3: impl Fn(S3Location) -> T,
) -> Result<T, LanguageModelError> {
match source {
ChatMessageContentSource::Bytes { data, .. } => Ok(from_bytes(Blob::new(data.clone()))),
ChatMessageContentSource::S3 { uri, bucket_owner } => {
Ok(from_s3(s3_location(uri, bucket_owner.as_deref())?))
}
ChatMessageContentSource::Url { url } => {
if is_s3_url(url) {
Ok(from_s3(s3_location(url, None)?))
} else if let Some((_, encoded)) = parse_data_url(url) {
Ok(from_bytes(Blob::new(decode_data_url_bytes(encoded)?)))
} else {
Err(LanguageModelError::permanent(format!(
"Bedrock {label} source URL must be data: or s3://"
)))
}
}
ChatMessageContentSource::FileId { .. } => Err(LanguageModelError::permanent(format!(
"Bedrock does not support file_id {label} sources"
))),
}
}
fn image_format_from_source(
format: Option<&str>,
source: &ChatMessageContentSource,
) -> Result<ImageFormat, LanguageModelError> {
resolve_format(
format,
source,
infer_image_format_from_source,
|value| ImageFormat::try_parse(value).ok(),
"image",
)
}
fn document_format_from_source(
format: Option<&str>,
source: &ChatMessageContentSource,
) -> Result<DocumentFormat, LanguageModelError> {
resolve_format(
format,
source,
infer_document_format_from_source,
|value| DocumentFormat::try_parse(value).ok(),
"document",
)
}
fn audio_format_from_source(
format: Option<&str>,
source: &ChatMessageContentSource,
) -> Result<AudioFormat, LanguageModelError> {
resolve_format(
format,
source,
infer_audio_format_from_source,
|value| AudioFormat::try_parse(value).ok(),
"audio",
)
}
fn video_format_from_source(
format: Option<&str>,
source: &ChatMessageContentSource,
) -> Result<VideoFormat, LanguageModelError> {
resolve_format(
format,
source,
infer_video_format_from_source,
|value| VideoFormat::try_parse(value).ok(),
"video",
)
}
fn resolve_format<T>(
explicit_format: Option<&str>,
source: &ChatMessageContentSource,
infer: impl Fn(&ChatMessageContentSource) -> Option<&'static str>,
parse: impl Fn(&str) -> Option<T>,
label: &str,
) -> Result<T, LanguageModelError> {
let value = explicit_format.or_else(|| infer(source)).ok_or_else(|| {
LanguageModelError::permanent(format!("Bedrock {label} format is required"))
})?;
parse(value).ok_or_else(|| {
LanguageModelError::permanent(format!("Unsupported Bedrock {label} format: {value}"))
})
}
fn infer_image_format_from_source(source: &ChatMessageContentSource) -> Option<&'static str> {
infer_format_from_source(
source,
IMAGE_MEDIA_TYPE_FORMATS,
IMAGE_EXTENSION_FORMATS,
None,
)
}
fn infer_document_format_from_source(source: &ChatMessageContentSource) -> Option<&'static str> {
infer_format_from_source(
source,
DOCUMENT_MEDIA_TYPE_FORMATS,
DOCUMENT_EXTENSION_FORMATS,
Some("txt"),
)
}
fn infer_audio_format_from_source(source: &ChatMessageContentSource) -> Option<&'static str> {
infer_format_from_source(
source,
AUDIO_MEDIA_TYPE_FORMATS,
AUDIO_EXTENSION_FORMATS,
None,
)
}
fn infer_video_format_from_source(source: &ChatMessageContentSource) -> Option<&'static str> {
infer_format_from_source(
source,
VIDEO_MEDIA_TYPE_FORMATS,
VIDEO_EXTENSION_FORMATS,
None,
)
}
fn infer_format_from_source(
source: &ChatMessageContentSource,
media_type_mappings: &[(&'static str, &'static str)],
extension_mappings: &[(&'static str, &'static str)],
fallback: Option<&'static str>,
) -> Option<&'static str> {
match source {
ChatMessageContentSource::Bytes { media_type, .. } => media_type
.as_deref()
.and_then(|media_type| mapped_format(media_type, media_type_mappings))
.or(fallback),
ChatMessageContentSource::Url { url } => if let Some((media_type, _)) = parse_data_url(url)
{
mapped_format(media_type, media_type_mappings)
} else {
extension_from_url(url)
.and_then(|extension| mapped_format(extension, extension_mappings))
}
.or(fallback),
ChatMessageContentSource::S3 { uri, .. } => extension_from_url(uri)
.and_then(|extension| mapped_format(extension, extension_mappings))
.or(fallback),
ChatMessageContentSource::FileId { .. } => fallback,
}
}
fn s3_location(uri: &str, bucket_owner: Option<&str>) -> Result<S3Location, LanguageModelError> {
let mut builder = S3Location::builder().uri(uri);
if let Some(bucket_owner) = bucket_owner {
builder = builder.bucket_owner(bucket_owner);
}
builder.build().map_err(LanguageModelError::permanent)
}
fn is_s3_url(url: &str) -> bool {
url.starts_with("s3://")
}
fn parse_data_url(url: &str) -> Option<(&str, &str)> {
let rest = url.strip_prefix("data:")?;
let (header, data) = rest.split_once(',')?;
let media_type = header.strip_suffix(";base64")?;
Some((media_type, data))
}
fn decode_data_url_bytes(encoded: &str) -> Result<Vec<u8>, LanguageModelError> {
base64::engine::general_purpose::STANDARD
.decode(encoded)
.map_err(LanguageModelError::permanent)
}
fn extension_from_url(url: &str) -> Option<&str> {
let without_query = url.split(['?', '#']).next()?;
let filename = without_query.rsplit('/').next()?;
let (_, extension) = filename.rsplit_once('.')?;
Some(extension)
}
fn mapped_format(value: &str, mappings: &[(&'static str, &'static str)]) -> Option<&'static str> {
mappings
.iter()
.find_map(|(input, output)| input.eq_ignore_ascii_case(value).then_some(*output))
}
fn message_from_blocks(
role: ConversationRole,
blocks: Vec<ContentBlock>,
) -> Result<Message, LanguageModelError> {
Message::builder()
.role(role)
.set_content(Some(blocks))
.build()
.map_err(LanguageModelError::permanent)
}
fn tool_output_to_result_block(
tool_call: &ToolCall,
output: &ToolOutput,
) -> Result<ContentBlock, LanguageModelError> {
let status = match output {
ToolOutput::Fail(_) => Some(ToolResultStatus::Error),
_ => Some(ToolResultStatus::Success),
};
let tool_result = ToolResultBlock::builder()
.tool_use_id(tool_call.id())
.content(tool_output_to_content_block(output)?)
.set_status(status)
.build()
.map_err(LanguageModelError::permanent)?;
Ok(ContentBlock::ToolResult(tool_result))
}
fn tool_output_to_content_block(
output: &ToolOutput,
) -> Result<ToolResultContentBlock, LanguageModelError> {
match output {
ToolOutput::Text(text) | ToolOutput::Fail(text) => {
Ok(ToolResultContentBlock::Text(text.clone()))
}
ToolOutput::FeedbackRequired(Some(value))
| ToolOutput::Stop(Some(value))
| ToolOutput::AgentFailed(Some(value)) => {
Ok(ToolResultContentBlock::Json(json_value_to_document(value)?))
}
_ => Ok(ToolResultContentBlock::Text(output.to_string())),
}
}
fn tool_call_args_to_document(args: Option<&str>) -> Result<Document, LanguageModelError> {
match args.map(str::trim) {
Some(args) if !args.is_empty() => parse_document_json_bytes(args.as_bytes())
.with_context(|| format!("Failed to parse tool args as JSON: {args}"))
.map_err(LanguageModelError::permanent),
_ => Ok(Document::Object(HashMap::new())),
}
}
fn tool_config_from_specs<'a>(
tool_specs: impl IntoIterator<Item = &'a ToolSpec>,
strict: bool,
) -> Result<Option<ToolConfiguration>, LanguageModelError> {
let tools = tool_specs
.into_iter()
.map(|spec| tool_spec_to_bedrock(spec, strict))
.collect::<Result<Vec<_>, _>>()?;
if tools.is_empty() {
return Ok(None);
}
let tool_config = ToolConfiguration::builder()
.set_tools(Some(tools))
.tool_choice(ToolChoice::Auto(AutoToolChoice::builder().build()))
.build()
.map_err(LanguageModelError::permanent)?;
Ok(Some(tool_config))
}
fn tool_spec_to_bedrock(spec: &ToolSpec, strict: bool) -> Result<Tool, LanguageModelError> {
let schema_value = AwsBedrockToolSchema::try_from(spec)
.map(AwsBedrockToolSchema::into_value)
.map_err(LanguageModelError::permanent)?;
let input_schema = ToolInputSchema::Json(json_value_to_document(&schema_value)?);
let mut builder = ToolSpecification::builder()
.name(spec.name.clone())
.input_schema(input_schema)
.strict(strict);
if !spec.description.is_empty() {
builder = builder.description(spec.description.clone());
}
let tool_spec = builder.build().map_err(LanguageModelError::permanent)?;
Ok(Tool::ToolSpec(tool_spec))
}
pub(super) fn response_to_chat_completion(
response: &ConverseOutput,
) -> Result<ChatCompletionResponse, LanguageModelError> {
let (message, tool_calls, reasoning) =
if let Some(ConverseResult::Message(message)) = response.output() {
extract_message_and_tool_calls(message)?
} else {
(None, None, Vec::new())
};
let mut builder = ChatCompletionResponse::builder()
.maybe_message(message)
.maybe_tool_calls(tool_calls)
.to_owned();
if !reasoning.is_empty() {
builder.reasoning(reasoning);
}
if let Some(usage) = response.usage() {
builder.usage(super::usage_from_bedrock(usage));
}
builder.build().map_err(LanguageModelError::from)
}
fn extract_message_and_tool_calls(
message: &Message,
) -> Result<ExtractedMessage, LanguageModelError> {
let mut text = String::new();
let mut has_text = false;
let mut tool_calls = Vec::with_capacity(message.content().len());
let mut reasoning = Vec::new();
for (content_block_index, block) in message.content().iter().enumerate() {
match block {
ContentBlock::Text(block_text) => {
text.push_str(block_text);
has_text = true;
}
ContentBlock::ToolUse(tool_use) => {
let args = document_to_json_string(tool_use.input());
let tool_call = ToolCall::builder()
.id(tool_use.tool_use_id())
.name(tool_use.name())
.args(args)
.build()
.map_err(LanguageModelError::permanent)?;
tool_calls.push(tool_call);
}
ContentBlock::ReasoningContent(ReasoningContentBlock::ReasoningText(
reasoning_text,
)) => {
reasoning.push(reasoning_item_from_reasoning_text(
content_block_index,
reasoning_text.text(),
reasoning_text.signature(),
));
}
_ => {}
}
}
let message = has_text.then_some(text);
let tool_calls = (!tool_calls.is_empty()).then_some(tool_calls);
Ok((message, tool_calls, reasoning))
}
fn document_to_json_string(document: &Document) -> String {
let mut output = String::new();
JsonValueWriter::new(&mut output).document(document);
output
}
fn apply_stream_event(
event: &ConverseStreamOutput,
response: &mut ChatCompletionResponse,
stop_reason: &mut Option<StopReason>,
) {
match event {
ConverseStreamOutput::ContentBlockStart(event) => {
if let (Some(ContentBlockStart::ToolUse(tool_use)), Ok(index)) =
(event.start(), usize::try_from(event.content_block_index()))
{
response.append_tool_call_delta(
index,
Some(tool_use.tool_use_id()),
Some(tool_use.name()),
None,
);
}
}
ConverseStreamOutput::ContentBlockDelta(event) => {
let Ok(index) = usize::try_from(event.content_block_index()) else {
return;
};
let Some(delta) = event.delta() else {
return;
};
match delta {
ContentBlockDelta::Text(text) => {
response.append_message_delta(Some(text));
}
ContentBlockDelta::ToolUse(delta) => {
response.append_tool_call_delta(index, None, None, Some(delta.input()));
}
ContentBlockDelta::ReasoningContent(delta) => {
apply_reasoning_delta(response, index, delta);
}
_ => {}
}
}
ConverseStreamOutput::MessageStop(event) => {
*stop_reason = Some(event.stop_reason().clone());
}
ConverseStreamOutput::Metadata(event) => {
if let Some(usage) = event.usage() {
response.usage = Some(super::usage_from_bedrock(usage));
}
}
_ => {}
}
}
fn assistant_reasoning_message_from_item(
item: &ReasoningItem,
) -> Result<Option<Message>, LanguageModelError> {
let text = item
.content
.as_ref()
.and_then(|content| content.first())
.map(String::as_str)
.filter(|text| !text.is_empty());
let signature = item
.encrypted_content
.as_deref()
.filter(|value| !value.is_empty());
let (Some(text), Some(signature)) = (text, signature) else {
return Ok(None);
};
let reasoning_text_block = ReasoningTextBlock::builder()
.text(text)
.signature(signature)
.build()
.map_err(LanguageModelError::permanent)?;
message_from_blocks(
ConversationRole::Assistant,
vec![ContentBlock::ReasoningContent(
ReasoningContentBlock::ReasoningText(reasoning_text_block),
)],
)
.map(Some)
}
fn reasoning_item_from_reasoning_text(
content_block_index: usize,
text: &str,
signature: Option<&str>,
) -> ReasoningItem {
ReasoningItem {
id: Some(format!("bedrock_reasoning_{content_block_index}")),
summary: Vec::new(),
content: Some(vec![text.to_string()]),
encrypted_content: signature.map(ToString::to_string),
status: None,
}
}
fn apply_reasoning_delta(
response: &mut ChatCompletionResponse,
content_block_index: usize,
delta: &ReasoningContentBlockDelta,
) {
let reasoning_item = ensure_reasoning_item(response, content_block_index);
match delta {
ReasoningContentBlockDelta::Text(text) => {
let content = reasoning_item
.content
.get_or_insert_with(|| vec![String::new()]);
if content.is_empty() {
content.push(String::new());
}
content[0].push_str(text);
}
ReasoningContentBlockDelta::Signature(signature) => {
reasoning_item.encrypted_content = Some(signature.clone());
}
_ => {}
}
}
fn ensure_reasoning_item(
response: &mut ChatCompletionResponse,
content_block_index: usize,
) -> &mut ReasoningItem {
let reasoning = response.reasoning.get_or_insert_with(Vec::new);
let reasoning_id = format!("bedrock_reasoning_{content_block_index}");
if let Some(position) = reasoning
.iter()
.position(|item| item.id.as_deref() == Some(reasoning_id.as_str()))
{
return reasoning
.get_mut(position)
.expect("position from iter().position must exist");
}
reasoning.push(ReasoningItem {
id: Some(reasoning_id),
summary: Vec::new(),
content: None,
encrypted_content: None,
status: None,
});
reasoning
.last_mut()
.expect("pushed reasoning item must exist")
}
fn json_value_to_document(value: &serde_json::Value) -> Result<Document, LanguageModelError> {
let bytes = serde_json::to_vec(value).map_err(LanguageModelError::permanent)?;
parse_document_json_bytes(&bytes).map_err(LanguageModelError::permanent)
}
fn parse_document_json_bytes(input: &[u8]) -> anyhow::Result<Document> {
let mut tokens = json_token_iter(input).peekable();
let document = expect_document(&mut tokens)?;
if tokens.next().transpose()?.is_some() {
anyhow::bail!("JSON input must contain exactly one value");
}
Ok(document)
}
const IMAGE_MEDIA_TYPE_FORMATS: &[(&str, &str)] = &[
("image/gif", "gif"),
("image/jpeg", "jpeg"),
("image/jpg", "jpeg"),
("image/png", "png"),
("image/webp", "webp"),
];
const IMAGE_EXTENSION_FORMATS: &[(&str, &str)] = &[
("gif", "gif"),
("jpeg", "jpeg"),
("jpg", "jpeg"),
("png", "png"),
("webp", "webp"),
];
const DOCUMENT_MEDIA_TYPE_FORMATS: &[(&str, &str)] = &[
("text/csv", "csv"),
("application/msword", "doc"),
(
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"docx",
),
("text/html", "html"),
("text/markdown", "md"),
("text/x-markdown", "md"),
("application/pdf", "pdf"),
("text/plain", "txt"),
("application/vnd.ms-excel", "xls"),
(
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"xlsx",
),
];
const DOCUMENT_EXTENSION_FORMATS: &[(&str, &str)] = &[
("csv", "csv"),
("doc", "doc"),
("docx", "docx"),
("html", "html"),
("htm", "html"),
("md", "md"),
("markdown", "md"),
("pdf", "pdf"),
("txt", "txt"),
("xls", "xls"),
("xlsx", "xlsx"),
];
const AUDIO_MEDIA_TYPE_FORMATS: &[(&str, &str)] = &[
("audio/aac", "aac"),
("audio/flac", "flac"),
("audio/m4a", "m4a"),
("audio/mka", "mka"),
("audio/x-matroska", "mkv"),
("audio/mpeg", "mp3"),
("audio/mp3", "mp3"),
("audio/mp4", "mp4"),
("audio/ogg", "ogg"),
("audio/opus", "opus"),
("audio/wav", "wav"),
("audio/x-wav", "wav"),
("audio/wave", "wav"),
("audio/webm", "webm"),
("audio/x-aac", "x-aac"),
];
const AUDIO_EXTENSION_FORMATS: &[(&str, &str)] = &[
("aac", "aac"),
("flac", "flac"),
("m4a", "m4a"),
("mka", "mka"),
("mkv", "mkv"),
("mp3", "mp3"),
("mp4", "mp4"),
("mpeg", "mpeg"),
("mpga", "mpga"),
("ogg", "ogg"),
("opus", "opus"),
("pcm", "pcm"),
("wav", "wav"),
("webm", "webm"),
("x-aac", "x-aac"),
];
const VIDEO_MEDIA_TYPE_FORMATS: &[(&str, &str)] = &[
("video/x-flv", "flv"),
("video/x-matroska", "mkv"),
("video/quicktime", "mov"),
("video/mp4", "mp4"),
("video/mpeg", "mpeg"),
("video/3gpp", "three_gp"),
("video/webm", "webm"),
("video/x-ms-wmv", "wmv"),
];
const VIDEO_EXTENSION_FORMATS: &[(&str, &str)] = &[
("flv", "flv"),
("mkv", "mkv"),
("mov", "mov"),
("mp4", "mp4"),
("mpeg", "mpeg"),
("mpg", "mpg"),
("3gp", "three_gp"),
("webm", "webm"),
("wmv", "wmv"),
];
#[cfg(test)]
mod tests {
use aws_sdk_bedrockruntime::Client;
use aws_sdk_bedrockruntime::{
operation::converse::ConverseOutput,
types::{
ContentBlockDeltaEvent, ContentBlockStart, ContentBlockStartEvent,
ConverseOutput as ConverseResult, Message, MessageStopEvent, ReasoningContentBlock,
ReasoningContentBlockDelta, ReasoningTextBlock, StopReason, TokenUsage,
ToolUseBlockDelta, ToolUseBlockStart,
},
};
use futures_util::StreamExt as _;
use schemars::{JsonSchema, schema_for};
use serde_json::{Value, json};
use swiftide_core::chat_completion::{
ChatMessage, ChatMessageContentPart, ChatMessageContentSource, ReasoningItem, ToolSpec,
};
use wiremock::{
Mock, MockServer, Request, Respond, ResponseTemplate,
matchers::{method, path},
};
use super::*;
#[cfg(feature = "langfuse")]
use crate::aws_bedrock_v2::test_utils::run_with_langfuse_event_capture;
use crate::aws_bedrock_v2::{
AwsBedrock, MockBedrockConverse, ReasoningEffort,
test_utils::{TEST_MODEL_ID, bedrock_client_for_mock_server, converse_stream_event},
};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, JsonSchema)]
struct WeatherArgs {
#[schemars(description = "City or region to retrieve weather for")]
location: String,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct NestedCommentArgs {
request: NestedCommentRequest,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct NestedCommentRequest {
#[serde(default, skip_serializing_if = "Option::is_none")]
body: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
page_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
block_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
discussion_id: Option<String>,
}
fn response_with_text_and_tool_call() -> ConverseOutput {
let mut args = HashMap::new();
args.insert(
"location".to_string(),
Document::String("Amsterdam".to_string()),
);
ConverseOutput::builder()
.output(ConverseResult::Message(
Message::builder()
.role(ConversationRole::Assistant)
.content(ContentBlock::Text("Working on it".to_string()))
.content(ContentBlock::ToolUse(
ToolUseBlock::builder()
.tool_use_id("call_1")
.name("get_weather")
.input(Document::Object(args))
.build()
.unwrap(),
))
.build()
.unwrap(),
))
.usage(
TokenUsage::builder()
.input_tokens(10)
.output_tokens(8)
.total_tokens(18)
.build()
.unwrap(),
)
.stop_reason(StopReason::ToolUse)
.build()
.unwrap()
}
#[test_log::test(tokio::test)]
async fn test_complete_maps_text_and_tool_calls() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock
.expect_converse()
.once()
.withf(
|model_id,
messages,
system,
inference_config,
tool_config,
output_config,
_additional_model_request_fields,
_additional_model_response_field_paths| {
model_id == "anthropic.claude-3-5-sonnet-20241022-v2:0"
&& messages.len() == 1
&& system.is_none()
&& inference_config.is_none()
&& tool_config.is_none()
&& output_config.is_none()
},
)
.returning(|_, _, _, _, _, _, _, _| Ok(response_with_text_and_tool_call()));
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("Check weather".into())])
.build()
.unwrap();
let response = bedrock.complete(&request).await.unwrap();
assert_eq!(response.message.as_deref(), Some("Working on it"));
let tool_call = response
.tool_calls
.as_ref()
.and_then(|calls| calls.first())
.expect("tool call");
assert_eq!(tool_call.id(), "call_1");
assert_eq!(tool_call.name(), "get_weather");
assert_eq!(
serde_json::from_str::<serde_json::Value>(tool_call.args().unwrap()).unwrap(),
serde_json::json!({"location":"Amsterdam"})
);
assert_eq!(response.usage.unwrap().total_tokens, 18);
}
#[cfg(feature = "langfuse")]
#[test]
fn test_complete_tracks_langfuse_failure_metadata_on_converse_error() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock
.expect_converse()
.once()
.returning(|_, _, _, _, _, _, _, _| {
Err(LanguageModelError::permanent("bedrock request failed"))
});
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![
ChatMessage::System("You are a careful helper".into()),
ChatMessage::User("Trace this failure".into()),
])
.tool_specs(vec![
ToolSpec::builder()
.name("get_weather")
.description("Look up the weather")
.parameters_schema(schema_for!(WeatherArgs))
.build()
.unwrap(),
])
.build()
.unwrap();
let (result, events) =
run_with_langfuse_event_capture(|| async { bedrock.complete(&request).await });
let error = result.expect_err("request should fail");
assert!(error.to_string().contains("bedrock request failed"));
let failure_event = events
.iter()
.find(|event| event.contains_key("langfuse.status_message"))
.expect("langfuse failure event");
assert_eq!(
failure_event
.get("langfuse.model")
.map(std::string::String::as_str),
Some("anthropic.claude-3-5-sonnet-20241022-v2:0")
);
let input = failure_event
.get("langfuse.input")
.expect("langfuse input should be tracked");
assert!(
input.contains("Trace this failure"),
"expected tracked input to include the user prompt: {input}"
);
assert!(
input.contains("\"system\""),
"expected tracked input to include built bedrock system blocks: {input}"
);
assert!(
input.contains("\"toolConfig\""),
"expected tracked input to include built bedrock tool config: {input}"
);
assert!(
!input.contains("tools_spec"),
"expected tracked input to use built bedrock request shape, not fallback tracking payload: {input}"
);
assert!(
failure_event
.get("langfuse.status_message")
.is_some_and(|message| message.contains("bedrock request failed"))
);
}
#[test_log::test(tokio::test)]
async fn test_complete_passes_additional_model_fields() {
let mut bedrock_mock = MockBedrockConverse::new();
let mut thinking = HashMap::new();
thinking.insert("type".to_string(), Document::String("enabled".to_string()));
thinking.insert("budget_tokens".to_string(), Document::from(512_u64));
let mut request_fields = HashMap::new();
request_fields.insert("thinking".to_string(), Document::Object(thinking));
let request_fields = Document::Object(request_fields);
bedrock_mock
.expect_converse()
.once()
.withf(
|model_id,
_,
_,
_,
_,
_,
additional_model_request_fields,
additional_model_response_field_paths| {
model_id == "anthropic.claude-3-5-sonnet-20241022-v2:0"
&& additional_model_request_fields
.as_ref()
.is_some_and(|fields| {
fields
.as_object()
.and_then(|map| map.get("thinking"))
.and_then(Document::as_object)
.and_then(|thinking| thinking.get("type"))
.and_then(Document::as_string)
== Some("enabled")
})
&& additional_model_response_field_paths
.as_ref()
.is_some_and(|paths| paths == &vec!["/thinking".to_string()])
},
)
.returning(|_, _, _, _, _, _, _, _| Ok(response_with_text_and_tool_call()));
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.default_options(Options {
additional_model_request_fields: Some(request_fields),
additional_model_response_field_paths: Some(vec!["/thinking".to_string()]),
..Default::default()
})
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("Hello".into())])
.build()
.unwrap();
let _ = bedrock.complete(&request).await.unwrap();
}
#[test_log::test(tokio::test)]
async fn test_complete_passes_reasoning_effort_for_claude_opus_4_5() {
let mut bedrock_mock = MockBedrockConverse::new();
let mut thinking = HashMap::new();
thinking.insert("type".to_string(), Document::String("enabled".to_string()));
thinking.insert("budget_tokens".to_string(), Document::from(512_u64));
let mut request_fields = HashMap::new();
request_fields.insert("thinking".to_string(), Document::Object(thinking));
let request_fields = Document::Object(request_fields);
bedrock_mock
.expect_converse()
.once()
.withf(
|model_id,
_,
_,
_,
_,
_,
additional_model_request_fields,
_additional_model_response_field_paths| {
model_id == "anthropic.claude-opus-4-5-20251101-v1:0"
&& additional_model_request_fields
.as_ref()
.is_some_and(|fields| {
let Some(fields) = fields.as_object() else {
return false;
};
let effort_matches = fields
.get("output_config")
.and_then(Document::as_object)
.and_then(|output_config| output_config.get("effort"))
.and_then(Document::as_string)
== Some("medium");
let thinking_matches = fields
.get("thinking")
.and_then(Document::as_object)
.and_then(|thinking| thinking.get("type"))
.and_then(Document::as_string)
== Some("enabled");
let beta_matches = fields
.get("anthropic_beta")
.and_then(Document::as_array)
.is_some_and(|betas| {
betas.iter().any(|beta| {
beta.as_string() == Some("effort-2025-11-24")
})
});
effort_matches && thinking_matches && beta_matches
})
},
)
.returning(|_, _, _, _, _, _, _, _| Ok(response_with_text_and_tool_call()));
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-opus-4-5-20251101-v1:0")
.default_options(Options {
reasoning_effort: Some(ReasoningEffort::Medium),
additional_model_request_fields: Some(request_fields),
..Default::default()
})
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("Hello".into())])
.build()
.unwrap();
let _ = bedrock.complete(&request).await.unwrap();
}
#[test_log::test(tokio::test)]
#[allow(deprecated)]
async fn test_complete_respects_tool_strict_option() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock
.expect_converse()
.once()
.withf(
|model_id,
_,
_,
_,
tool_config,
output_config,
_additional_model_request_fields,
_additional_model_response_field_paths| {
model_id == "anthropic.claude-3-5-sonnet-20241022-v2:0"
&& output_config.is_none()
&& tool_config
.as_ref()
.and_then(|config| config.tools().first())
.is_some_and(|tool| match tool {
Tool::ToolSpec(spec) => spec.strict() == Some(false),
_ => false,
})
},
)
.returning(|_, _, _, _, _, _, _, _| Ok(response_with_text_and_tool_call()));
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.default_options(Options {
tool_strict: Some(false),
..Default::default()
})
.build()
.unwrap();
let tool_spec = ToolSpec::builder()
.name("get_weather")
.description("Get weather")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("Check weather".into())])
.tools_spec([tool_spec])
.build()
.unwrap();
let _ = bedrock.complete(&request).await.unwrap();
}
#[test_log::test(tokio::test)]
async fn test_complete_stream_requires_model() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock.expect_converse_stream().never();
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user("Hello")])
.build()
.unwrap();
let mut stream = bedrock.complete_stream(&request).await;
let first = stream.next().await.expect("stream should yield one item");
assert!(matches!(first, Err(LanguageModelError::PermanentError(_))));
assert!(stream.next().await.is_none());
}
#[cfg(feature = "langfuse")]
#[test]
fn test_complete_stream_tracks_langfuse_failure_metadata_on_stream_error() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock
.expect_converse_stream()
.once()
.returning(|_, _, _, _, _, _, _| {
Err(LanguageModelError::transient("bedrock stream failed"))
});
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user("Stream this failure")])
.build()
.unwrap();
let (first_item, events) = run_with_langfuse_event_capture(|| async {
let mut stream = bedrock.complete_stream(&request).await;
stream.next().await.expect("stream should yield an error")
});
let error = first_item.expect_err("stream should fail");
assert!(error.to_string().contains("bedrock stream failed"));
let failure_event = events
.iter()
.find(|event| event.contains_key("langfuse.status_message"))
.expect("langfuse failure event");
assert!(
failure_event
.get("langfuse.input")
.is_some_and(|input| input.contains("Stream this failure"))
);
assert!(
failure_event
.get("langfuse.status_message")
.is_some_and(|message| message.contains("bedrock stream failed"))
);
}
#[test_log::test(tokio::test)]
async fn test_complete_stream_rejects_system_only_messages() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock.expect_converse_stream().never();
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_system("You are a helper")])
.build()
.unwrap();
let mut stream = bedrock.complete_stream(&request).await;
let first = stream.next().await.expect("stream should yield one item");
assert!(matches!(first, Err(LanguageModelError::PermanentError(_))));
assert!(stream.next().await.is_none());
}
#[test_log::test(tokio::test)]
async fn test_complete_stream_returns_upstream_stream_error() {
let mut bedrock_mock = MockBedrockConverse::new();
bedrock_mock
.expect_converse_stream()
.once()
.withf(
|model_id,
messages,
system,
inference_config,
tool_config,
_additional_model_request_fields,
_additional_model_response_field_paths| {
model_id == "anthropic.claude-3-5-sonnet-20241022-v2:0"
&& messages.len() == 1
&& matches!(messages[0].role(), ConversationRole::User)
&& matches!(messages[0].content().first(), Some(ContentBlock::Text(text)) if text == "Hello")
&& system.is_none()
&& inference_config.is_none()
&& tool_config.is_none()
},
)
.returning(|_, _, _, _, _, _, _| {
Err(LanguageModelError::transient(anyhow::anyhow!(
"stream init failed"
)))
});
let bedrock = AwsBedrock::builder()
.test_client(bedrock_mock)
.default_prompt_model("anthropic.claude-3-5-sonnet-20241022-v2:0")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user("Hello")])
.build()
.unwrap();
let mut stream = bedrock.complete_stream(&request).await;
let first = stream.next().await.expect("stream should yield one item");
assert!(matches!(first, Err(LanguageModelError::TransientError(_))));
assert!(stream.next().await.is_none());
}
#[test_log::test(tokio::test)]
async fn test_complete_green_path_with_wiremock() {
struct ValidateConverseRequest;
impl Respond for ValidateConverseRequest {
fn respond(&self, request: &Request) -> ResponseTemplate {
let payload: Value = serde_json::from_slice(&request.body).expect("request json");
assert_eq!(payload["messages"][0]["role"], "user");
assert_eq!(payload["messages"][0]["content"][0]["text"], "Hello");
ResponseTemplate::new(200).set_body_json(json!({
"output": {
"message": {
"role": "assistant",
"content": [
{"text": "Hello from bedrock"}
]
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 2,
"outputTokens": 5,
"totalTokens": 7
}
}))
}
}
let mock_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path(format!("/model/{TEST_MODEL_ID}/converse")))
.respond_with(ValidateConverseRequest)
.mount(&mock_server)
.await;
let client: Client = bedrock_client_for_mock_server(&mock_server.uri());
let bedrock = AwsBedrock::builder()
.client(client)
.default_prompt_model(TEST_MODEL_ID)
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user("Hello")])
.build()
.unwrap();
let response = bedrock.complete(&request).await.unwrap();
assert_eq!(response.message(), Some("Hello from bedrock"));
assert_eq!(
response.usage.as_ref().map(|usage| usage.total_tokens),
Some(7)
);
}
#[test_log::test(tokio::test)]
async fn test_complete_stream_green_path_with_wiremock() {
struct ValidateConverseStreamRequest {
stream_body: Vec<u8>,
}
impl Respond for ValidateConverseStreamRequest {
fn respond(&self, request: &Request) -> ResponseTemplate {
let payload: Value = serde_json::from_slice(&request.body).expect("request json");
assert_eq!(payload["messages"][0]["role"], "user");
assert_eq!(payload["messages"][0]["content"][0]["text"], "Hello");
ResponseTemplate::new(200).set_body_raw(
self.stream_body.clone(),
"application/vnd.amazon.eventstream",
)
}
}
let mock_server = MockServer::start().await;
let stream_body = [
converse_stream_event(
"contentBlockDelta",
&json!({
"contentBlockIndex": 0,
"delta": {"text": "Hello stream"}
}),
),
converse_stream_event(
"metadata",
&json!({
"usage": {
"inputTokens": 4,
"outputTokens": 5,
"totalTokens": 9
}
}),
),
converse_stream_event(
"messageStop",
&json!({
"stopReason": "end_turn"
}),
),
]
.concat();
Mock::given(method("POST"))
.and(path(format!("/model/{TEST_MODEL_ID}/converse-stream")))
.respond_with(ValidateConverseStreamRequest { stream_body })
.mount(&mock_server)
.await;
let client: Client = bedrock_client_for_mock_server(&mock_server.uri());
let bedrock = AwsBedrock::builder()
.client(client)
.default_prompt_model(TEST_MODEL_ID)
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user("Hello")])
.build()
.unwrap();
let responses = bedrock
.complete_stream(&request)
.await
.collect::<Vec<_>>()
.await;
let last = responses
.last()
.expect("stream should yield")
.as_ref()
.expect("last response ok");
assert_eq!(last.message(), Some("Hello stream"));
assert_eq!(last.usage.as_ref().map(|usage| usage.total_tokens), Some(9));
}
#[test]
fn test_tool_config_from_specs_builds_schema() {
let tool_spec = ToolSpec::builder()
.name("get_weather")
.description("Get weather by location")
.parameters_schema(schema_for!(WeatherArgs))
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hi".into())])
.tool_specs([tool_spec])
.build()
.unwrap();
let tool_config = tool_config_from_specs(request.tools_spec().iter(), true)
.unwrap()
.expect("tool config");
assert_eq!(tool_config.tools().len(), 1);
let Tool::ToolSpec(spec) = &tool_config.tools()[0] else {
panic!("expected tool spec");
};
assert_eq!(spec.name(), "get_weather");
assert_eq!(spec.description(), Some("Get weather by location"));
assert_eq!(spec.strict(), Some(true));
assert!(matches!(
spec.input_schema(),
Some(ToolInputSchema::Json(Document::Object(schema)))
if schema.get("type") == Some(&Document::String("object".to_string()))
&& schema.get("additionalProperties") == Some(&Document::Bool(false))
));
let Some(ToolInputSchema::Json(Document::Object(schema))) = spec.input_schema() else {
panic!("expected JSON object schema");
};
let Some(Document::Object(properties)) = schema.get("properties") else {
panic!("expected schema properties");
};
let Some(Document::Object(location)) = properties.get("location") else {
panic!("expected location property schema");
};
assert_eq!(
location.get("description"),
Some(&Document::String(
"City or region to retrieve weather for".to_string()
))
);
}
#[test]
fn test_tool_config_from_specs_can_disable_strict() {
let tool_spec = ToolSpec::builder()
.name("get_weather")
.description("Get weather")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hi".into())])
.tool_specs([tool_spec])
.build()
.unwrap();
let tool_config = tool_config_from_specs(request.tools_spec().iter(), false)
.unwrap()
.expect("tool config");
let Tool::ToolSpec(spec) = &tool_config.tools()[0] else {
panic!("expected tool spec");
};
assert_eq!(spec.strict(), Some(false));
assert!(matches!(
spec.input_schema(),
Some(ToolInputSchema::Json(Document::Object(schema)))
if schema.get("type") == Some(&Document::String("object".to_string()))
&& schema.get("additionalProperties") == Some(&Document::Bool(false))
));
}
#[test]
fn test_tool_config_from_specs_does_not_apply_openai_required_workaround() {
let tool_spec = ToolSpec::builder()
.name("create_comment")
.description("Create a comment")
.parameters_schema(schema_for!(NestedCommentArgs))
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hi".into())])
.tool_specs([tool_spec])
.build()
.unwrap();
let tool_config = tool_config_from_specs(request.tools_spec().iter(), true)
.unwrap()
.expect("tool config");
let Tool::ToolSpec(spec) = &tool_config.tools()[0] else {
panic!("expected tool spec");
};
let Some(ToolInputSchema::Json(Document::Object(schema))) = spec.input_schema() else {
panic!("expected JSON object schema");
};
assert_eq!(
schema.get("type"),
Some(&Document::String("object".to_string()))
);
assert_eq!(
schema.get("additionalProperties"),
Some(&Document::Bool(false))
);
assert_eq!(
schema.get("required"),
Some(&Document::Array(vec![Document::String(
"request".to_string()
)]))
);
let Some(Document::Object(properties)) = schema.get("properties") else {
panic!("expected properties map");
};
let Some(Document::String(nested_ref)) = properties
.get("request")
.and_then(Document::as_object)
.and_then(|request| request.get("$ref"))
else {
panic!("expected nested request $ref");
};
let nested_name = nested_ref
.rsplit('/')
.next()
.expect("nested request ref name");
let Some(Document::Object(defs)) = schema.get("$defs") else {
panic!("expected defs map");
};
let Some(Document::Object(nested_schema)) = defs.get(nested_name) else {
panic!("expected nested request schema");
};
assert!(!nested_schema.contains_key("required"));
}
#[test]
fn test_tool_config_from_specs_orders_tools_deterministically() {
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::User("hi".into())])
.tool_specs([
ToolSpec::builder()
.name("z_tool")
.description("later")
.build()
.unwrap(),
ToolSpec::builder()
.name("a_tool")
.description("earlier")
.build()
.unwrap(),
])
.build()
.unwrap();
let tool_config = tool_config_from_specs(request.tools_spec().iter(), true)
.unwrap()
.expect("tool config");
let tool_names = tool_config
.tools()
.iter()
.map(|tool| match tool {
Tool::ToolSpec(spec) => spec.name(),
_ => panic!("expected tool spec"),
})
.collect::<Vec<_>>();
assert_eq!(tool_names, vec!["a_tool", "z_tool"]);
}
#[test]
fn test_response_to_chat_completion_maps_reasoning_content() {
let response = ConverseOutput::builder()
.output(ConverseResult::Message(
Message::builder()
.role(ConversationRole::Assistant)
.content(ContentBlock::ReasoningContent(
ReasoningContentBlock::ReasoningText(
ReasoningTextBlock::builder()
.text("I should call a weather tool")
.signature("sig_123")
.build()
.unwrap(),
),
))
.content(ContentBlock::Text("Working on it".to_string()))
.build()
.unwrap(),
))
.stop_reason(StopReason::EndTurn)
.build()
.unwrap();
let completion = response_to_chat_completion(&response).unwrap();
assert_eq!(completion.message.as_deref(), Some("Working on it"));
let reasoning = completion.reasoning.expect("reasoning items");
assert_eq!(reasoning.len(), 1);
assert_eq!(reasoning[0].id.as_deref(), Some("bedrock_reasoning_0"));
assert_eq!(
reasoning[0].content.as_ref().and_then(|c| c.first()),
Some(&"I should call a weather tool".to_string())
);
assert_eq!(reasoning[0].encrypted_content.as_deref(), Some("sig_123"));
}
#[test]
fn test_build_converse_input_replays_reasoning_items() {
let request = ChatCompletionRequest::builder()
.messages(vec![
ChatMessage::Reasoning(ReasoningItem {
id: Some("r1".to_string()),
summary: Vec::new(),
content: Some(vec!["I should call a weather tool".to_string()]),
encrypted_content: Some("sig_123".to_string()),
status: None,
}),
ChatMessage::new_user("Use tool"),
])
.build()
.unwrap();
let (messages, _system, _inference, _tool_config) =
build_converse_input(&request, &Options::default()).unwrap();
assert_eq!(messages.len(), 2);
assert!(matches!(messages[0].role(), ConversationRole::Assistant));
let reasoning = messages[0]
.content()
.first()
.and_then(|content| content.as_reasoning_content().ok())
.and_then(|content| content.as_reasoning_text().ok())
.expect("reasoning content");
assert_eq!(reasoning.text(), "I should call a weather tool");
assert_eq!(reasoning.signature(), Some("sig_123"));
}
#[test]
fn test_build_converse_input_groups_adjacent_tool_outputs() {
let first_tool = ToolCall::builder()
.id("tool_1")
.name("shell_command")
.args("{\"cmd\":\"pwd\"}")
.build()
.unwrap();
let second_tool = ToolCall::builder()
.id("tool_2")
.name("git")
.args("{\"command\":\"status\"}")
.build()
.unwrap();
let request = ChatCompletionRequest::builder()
.messages(vec![
ChatMessage::Assistant(None, Some(vec![first_tool.clone(), second_tool.clone()])),
ChatMessage::new_tool_output(
first_tool,
ToolOutput::Text("pwd output".to_string()),
),
ChatMessage::new_tool_output(
second_tool,
ToolOutput::Text("git output".to_string()),
),
])
.build()
.unwrap();
let (messages, _system, _inference, _tool_config) =
build_converse_input(&request, &Options::default()).unwrap();
assert_eq!(messages.len(), 2);
assert!(matches!(messages[0].role(), ConversationRole::Assistant));
assert!(matches!(messages[1].role(), ConversationRole::User));
assert_eq!(messages[1].content().len(), 2);
let first_result = messages[1]
.content()
.first()
.and_then(|block| block.as_tool_result().ok())
.expect("first tool result");
let second_result = messages[1]
.content()
.get(1)
.and_then(|block| block.as_tool_result().ok())
.expect("second tool result");
assert_eq!(first_result.tool_use_id(), "tool_1");
assert_eq!(second_result.tool_use_id(), "tool_2");
}
#[test]
fn test_build_converse_input_maps_image_part() {
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user_with_parts(vec![
ChatMessageContentPart::text("Describe this image"),
ChatMessageContentPart::image("data:image/png;base64,AA=="),
])])
.build()
.unwrap();
let (messages, _system, _inference, _tool_config) =
build_converse_input(&request, &Options::default()).unwrap();
assert_eq!(messages.len(), 1);
assert!(matches!(messages[0].role(), ConversationRole::User));
assert_eq!(messages[0].content().len(), 2);
let image = messages[0]
.content()
.get(1)
.and_then(|content| content.as_image().ok())
.expect("image block");
assert!(matches!(image.format(), ImageFormat::Png));
assert!(
image
.source()
.is_some_and(aws_sdk_bedrockruntime::types::ImageSource::is_bytes)
);
}
#[test]
fn test_build_converse_input_maps_audio_part() {
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user_with_parts(vec![
ChatMessageContentPart::text("Transcribe this"),
ChatMessageContentPart::audio(ChatMessageContentSource::bytes(
vec![1_u8, 2_u8, 3_u8],
Some("audio/mpeg".to_string()),
)),
])])
.build()
.unwrap();
let (messages, _system, _inference, _tool_config) =
build_converse_input(&request, &Options::default()).unwrap();
let audio = messages[0]
.content()
.get(1)
.and_then(|content| content.as_audio().ok())
.expect("audio block");
assert!(matches!(audio.format(), AudioFormat::Mp3));
assert!(
audio
.source()
.is_some_and(aws_sdk_bedrockruntime::types::AudioSource::is_bytes)
);
}
#[test]
fn test_build_converse_input_maps_video_part() {
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user_with_parts(vec![
ChatMessageContentPart::text("Describe this clip"),
ChatMessageContentPart::video("s3://bucket/video.mp4"),
])])
.build()
.unwrap();
let (messages, _system, _inference, _tool_config) =
build_converse_input(&request, &Options::default()).unwrap();
let video = messages[0]
.content()
.get(1)
.and_then(|content| content.as_video().ok())
.expect("video block");
assert!(matches!(video.format(), VideoFormat::Mp4));
assert!(
video
.source()
.is_some_and(aws_sdk_bedrockruntime::types::VideoSource::is_s3_location)
);
}
#[test]
fn test_build_converse_input_rejects_audio_http_url() {
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user_with_parts(vec![
ChatMessageContentPart::text("Transcribe this"),
ChatMessageContentPart::audio("https://example.com/audio.mp3"),
])])
.build()
.unwrap();
let error = build_converse_input(&request, &Options::default()).unwrap_err();
assert!(format!("{error}").contains("audio source URL must be data: or s3://"));
}
#[test]
fn test_build_converse_input_rejects_document_without_text() {
let request = ChatCompletionRequest::builder()
.messages(vec![ChatMessage::new_user_with_parts(vec![
ChatMessageContentPart::document(ChatMessageContentSource::bytes(
vec![1_u8, 2_u8],
Some("text/plain".to_string()),
)),
])])
.build()
.unwrap();
let error = build_converse_input(&request, &Options::default()).unwrap_err();
assert!(format!("{error}").contains("require at least one text part"));
}
#[test]
#[allow(clippy::too_many_lines)]
fn test_apply_stream_event_accumulates_deltas() {
let mut response = ChatCompletionResponse::default();
let mut stop_reason = None;
apply_stream_event(
&ConverseStreamOutput::ContentBlockStart(
ContentBlockStartEvent::builder()
.content_block_index(0)
.start(ContentBlockStart::ToolUse(
ToolUseBlockStart::builder()
.tool_use_id("call_1")
.name("get_weather")
.build()
.unwrap(),
))
.build()
.unwrap(),
),
&mut response,
&mut stop_reason,
);
apply_stream_event(
&ConverseStreamOutput::ContentBlockDelta(
ContentBlockDeltaEvent::builder()
.content_block_index(0)
.delta(ContentBlockDelta::ToolUse(
ToolUseBlockDelta::builder()
.input("{\"location\":\"Amsterdam\"}")
.build()
.unwrap(),
))
.build()
.unwrap(),
),
&mut response,
&mut stop_reason,
);
apply_stream_event(
&ConverseStreamOutput::ContentBlockDelta(
ContentBlockDeltaEvent::builder()
.content_block_index(1)
.delta(ContentBlockDelta::Text("Tool call created".to_string()))
.build()
.unwrap(),
),
&mut response,
&mut stop_reason,
);
apply_stream_event(
&ConverseStreamOutput::ContentBlockDelta(
ContentBlockDeltaEvent::builder()
.content_block_index(2)
.delta(ContentBlockDelta::ReasoningContent(
ReasoningContentBlockDelta::Text("Thinking...".to_string()),
))
.build()
.unwrap(),
),
&mut response,
&mut stop_reason,
);
apply_stream_event(
&ConverseStreamOutput::ContentBlockDelta(
ContentBlockDeltaEvent::builder()
.content_block_index(2)
.delta(ContentBlockDelta::ReasoningContent(
ReasoningContentBlockDelta::Signature("sig_123".to_string()),
))
.build()
.unwrap(),
),
&mut response,
&mut stop_reason,
);
apply_stream_event(
&ConverseStreamOutput::Metadata(
aws_sdk_bedrockruntime::types::ConverseStreamMetadataEvent::builder()
.usage(
TokenUsage::builder()
.input_tokens(5)
.output_tokens(3)
.total_tokens(8)
.build()
.unwrap(),
)
.build(),
),
&mut response,
&mut stop_reason,
);
apply_stream_event(
&ConverseStreamOutput::MessageStop(
MessageStopEvent::builder()
.stop_reason(StopReason::ToolUse)
.build()
.unwrap(),
),
&mut response,
&mut stop_reason,
);
assert_eq!(response.message.as_deref(), Some("Tool call created"));
let tool_call = response
.tool_calls
.as_ref()
.and_then(|calls| calls.first())
.expect("tool call");
assert_eq!(tool_call.id(), "call_1");
assert_eq!(tool_call.name(), "get_weather");
assert_eq!(tool_call.args(), Some("{\"location\":\"Amsterdam\"}"));
let reasoning = response.reasoning.expect("reasoning item");
assert_eq!(reasoning.len(), 1);
assert_eq!(reasoning[0].id.as_deref(), Some("bedrock_reasoning_2"));
assert_eq!(
reasoning[0].content.as_ref().and_then(|c| c.first()),
Some(&"Thinking...".to_string())
);
assert_eq!(reasoning[0].encrypted_content.as_deref(), Some("sig_123"));
assert_eq!(response.usage.unwrap().total_tokens, 8);
assert!(matches!(stop_reason, Some(StopReason::ToolUse)));
}
}