Skip to content

Commit 38cfe9b

Browse files
sxlijinsxlijin
authored
feat: bedrock video support (#2681)
make s3 URIs work
1 parent af2d872 commit 38cfe9b

File tree

46 files changed

+1912
-128
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

46 files changed

+1912
-128
lines changed

‎engine/baml-runtime/src/internal/llm_client/primitive/aws/aws_client.rs‎

Lines changed: 108 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ use secrecy::ExposeSecret;
4242
use serde::Deserialize;
4343
use serde_json::{json, Map};
4444
use shell_escape::escape;
45+
use url::Url;
4546
use uuid::Uuid;
4647
use web_time::{Instant, SystemTime};
4748

@@ -72,42 +73,46 @@ fn strip_mime_prefix(mime: &str) -> &str {
7273
}
7374

7475
fn media_to_content_block_json(media: &BamlMedia) -> Result<serde_json::Value> {
75-
match media.media_type {
76-
BamlMediaType::Image => match &media.content {
77-
BamlMediaContent::Base64(b64) => {
78-
let mut image_obj = serde_json::Map::new();
79-
if let Some(mime) = media.mime_type.as_deref() {
80-
image_obj.insert("format".into(), json!(strip_mime_prefix(mime)));
81-
}
82-
image_obj.insert("source".into(), json!({ "bytes": b64.base64 }));
83-
Ok(json!({ "image": serde_json::Value::Object(image_obj) }))
84-
}
85-
_ => anyhow::bail!("AWS Bedrock only supports base64 image inputs in modular requests"),
86-
},
87-
BamlMediaType::Pdf => match &media.content {
88-
BamlMediaContent::Base64(b64) => {
89-
let mut doc_obj = serde_json::Map::new();
90-
if let Some(mime) = media.mime_type.as_deref() {
91-
doc_obj.insert("format".into(), json!(strip_mime_prefix(mime)));
92-
}
93-
doc_obj.insert("name".into(), json!("document"));
94-
doc_obj.insert("source".into(), json!({ "bytes": b64.base64 }));
95-
Ok(json!({ "document": serde_json::Value::Object(doc_obj) }))
96-
}
97-
_ => anyhow::bail!("AWS Bedrock only supports base64 PDF inputs in modular requests"),
98-
},
99-
BamlMediaType::Video => match &media.content {
100-
BamlMediaContent::Base64(b64) => {
101-
let mut video_obj = serde_json::Map::new();
102-
if let Some(mime) = media.mime_type.as_deref() {
103-
video_obj.insert("format".into(), json!(strip_mime_prefix(mime)));
76+
let content_block = {
77+
let mut obj = Map::new();
78+
if let Some(mime) = media.mime_type.as_deref() {
79+
obj.insert("format".into(), json!(strip_mime_prefix(mime)));
80+
}
81+
let source = match &media.content {
82+
BamlMediaContent::File(media_file) => todo!(),
83+
BamlMediaContent::Url(url) => {
84+
let parsed = Url::parse(&url.url).with_context(|| {
85+
format!("Invalid S3 URI for AWS Bedrock video source: {url}")
86+
})?;
87+
88+
if parsed.scheme() != "s3" {
89+
anyhow::bail!("AWS Bedrock requires s3:// URIs, but got: {}", url.url);
10490
}
105-
video_obj.insert("source".into(), json!({ "bytes": b64.base64 }));
106-
Ok(json!({ "video": serde_json::Value::Object(video_obj) }))
91+
92+
// unimplemented!("make sure the test works")
93+
json!({
94+
"s3Location": {
95+
"uri": url.url,
96+
}
97+
})
10798
}
108-
_ => anyhow::bail!("AWS Bedrock only supports base64 video inputs in modular requests"),
109-
},
110-
BamlMediaType::Audio => anyhow::bail!("AWS Bedrock does not support audio media parts"),
99+
BamlMediaContent::Base64(base64) => json!({
100+
"bytes": base64.base64,
101+
}),
102+
};
103+
obj.insert("source".into(), source);
104+
obj
105+
};
106+
match media.media_type {
107+
// _ => anyhow::bail!("AWS Bedrock only supports base64 image inputs in modular requests"),
108+
BamlMediaType::Image => Ok(json!({ "image": content_block })),
109+
BamlMediaType::Pdf => {
110+
let mut content_block = content_block;
111+
content_block.insert("name".into(), json!("document"));
112+
Ok(json!({ "document": content_block }))
113+
}
114+
BamlMediaType::Video => Ok(json!({ "video": content_block })),
115+
BamlMediaType::Audio => Ok(json!({ "audio": content_block })),
111116
}
112117
}
113118

@@ -1079,36 +1084,46 @@ impl AwsClient {
10791084
media: &baml_types::BamlMedia,
10801085
) -> Result<bedrock::types::ContentBlock> {
10811086
match media.media_type {
1082-
BamlMediaType::Image => match &media.content {
1083-
BamlMediaContent::File(_) => {
1084-
anyhow::bail!(
1087+
BamlMediaType::Image => {
1088+
let format = bedrock::types::ImageFormat::from(
1089+
{
1090+
let mime_type = media.mime_type_as_ok()?;
1091+
match mime_type.strip_prefix("image/") {
1092+
Some(s) => s.to_string(),
1093+
None => mime_type,
1094+
}
1095+
}
1096+
.as_str(),
1097+
);
1098+
match &media.content {
1099+
BamlMediaContent::File(_) => {
1100+
anyhow::bail!(
10851101
"BAML internal error (AWSBedrock): file should have been resolved to base64"
10861102
)
1103+
}
1104+
BamlMediaContent::Url(url) => Ok(bedrock::types::ContentBlock::Image(
1105+
bedrock::types::ImageBlock::builder()
1106+
.set_format(Some(format))
1107+
.set_source(Some(bedrock::types::ImageSource::S3Location(
1108+
bedrock::types::S3Location::builder()
1109+
.set_uri(Some(url.url.clone()))
1110+
.build()
1111+
.context("Failed to build S3Location block")?,
1112+
)))
1113+
.build()
1114+
.context("Failed to build Image block")?,
1115+
)),
1116+
BamlMediaContent::Base64(b64_media) => Ok(bedrock::types::ContentBlock::Image(
1117+
bedrock::types::ImageBlock::builder()
1118+
.set_format(Some(format))
1119+
.set_source(Some(bedrock::types::ImageSource::Bytes(Blob::new(
1120+
aws_smithy_types::base64::decode(b64_media.base64.clone())?,
1121+
))))
1122+
.build()
1123+
.context("Failed to build image block")?,
1124+
)),
10871125
}
1088-
BamlMediaContent::Url(_) => {
1089-
anyhow::bail!(
1090-
"BAML internal error (AWSBedrock): media URL should have been resolved to base64"
1091-
)
1092-
}
1093-
BamlMediaContent::Base64(b64_media) => Ok(bedrock::types::ContentBlock::Image(
1094-
bedrock::types::ImageBlock::builder()
1095-
.set_format(Some(bedrock::types::ImageFormat::from(
1096-
{
1097-
let mime_type = media.mime_type_as_ok()?;
1098-
match mime_type.strip_prefix("image/") {
1099-
Some(s) => s.to_string(),
1100-
None => mime_type,
1101-
}
1102-
}
1103-
.as_str(),
1104-
)))
1105-
.set_source(Some(bedrock::types::ImageSource::Bytes(Blob::new(
1106-
aws_smithy_types::base64::decode(b64_media.base64.clone())?,
1107-
))))
1108-
.build()
1109-
.context("Failed to build image block")?,
1110-
)),
1111-
},
1126+
}
11121127
BamlMediaType::Pdf => {
11131128
match &media.content {
11141129
BamlMediaContent::File(_) => {
@@ -1148,46 +1163,44 @@ impl AwsClient {
11481163
}
11491164
}
11501165
BamlMediaType::Video => {
1166+
let format = bedrock::types::VideoFormat::from(
1167+
{
1168+
let mime_type = media.mime_type_as_ok()?;
1169+
match mime_type.strip_prefix("video/") {
1170+
Some(s) => s.to_string(),
1171+
None => mime_type,
1172+
}
1173+
}
1174+
.as_str(),
1175+
);
1176+
// AWS Bedrock supports video for Nova models with specific format
11511177
match &media.content {
11521178
BamlMediaContent::File(_) => {
11531179
anyhow::bail!(
11541180
"BAML internal error (AWSBedrock): video file should have been resolved to base64"
11551181
)
11561182
}
1157-
BamlMediaContent::Url(_) => {
1158-
anyhow::bail!(
1159-
"BAML internal error (AWSBedrock): video URL should have been resolved to base64"
1160-
)
1161-
}
1162-
BamlMediaContent::Base64(b64_media) => {
1163-
// AWS Bedrock supports video for Nova models with specific format
1164-
let mime_type = media.mime_type_as_ok()?;
1165-
let format = match mime_type.as_str() {
1166-
"video/mp4" => bedrock::types::VideoFormat::Mp4,
1167-
"video/mpeg" => bedrock::types::VideoFormat::Mpeg,
1168-
"video/mov" => bedrock::types::VideoFormat::Mov,
1169-
// "video/avi" => bedrock::types::VideoFormat::Avi,
1170-
"video/x-flv" => bedrock::types::VideoFormat::Flv,
1171-
"video/mkv" => bedrock::types::VideoFormat::Mkv,
1172-
"video/webm" => bedrock::types::VideoFormat::Webm,
1173-
_ => {
1174-
anyhow::bail!(
1175-
"AWS Bedrock video format not supported: {}. Supported formats: mp4, mpeg, mov, flv, mkv, webm",
1176-
mime_type
1177-
);
1178-
}
1179-
};
1180-
1181-
Ok(bedrock::types::ContentBlock::Video(
1182-
bedrock::types::VideoBlock::builder()
1183-
.set_format(Some(format))
1184-
.set_source(Some(bedrock::types::VideoSource::Bytes(Blob::new(
1185-
aws_smithy_types::base64::decode(b64_media.base64.clone())?,
1186-
))))
1187-
.build()
1188-
.context("Failed to build video block")?,
1189-
))
1190-
}
1183+
BamlMediaContent::Url(url) => Ok(bedrock::types::ContentBlock::Video(
1184+
bedrock::types::VideoBlock::builder()
1185+
.set_format(Some(format))
1186+
.set_source(Some(bedrock::types::VideoSource::S3Location(
1187+
bedrock::types::S3Location::builder()
1188+
.set_uri(Some(url.url.clone()))
1189+
.build()
1190+
.context("Failed to build S3Location block")?,
1191+
)))
1192+
.build()
1193+
.context("Failed to build Video document block")?,
1194+
)),
1195+
BamlMediaContent::Base64(b64_media) => Ok(bedrock::types::ContentBlock::Video(
1196+
bedrock::types::VideoBlock::builder()
1197+
.set_format(Some(format))
1198+
.set_source(Some(bedrock::types::VideoSource::Bytes(Blob::new(
1199+
aws_smithy_types::base64::decode(b64_media.base64.clone())?,
1200+
))))
1201+
.build()
1202+
.context("AWS Bedrock error: mime_type must be explicitly set on base64 videos")?,
1203+
)),
11911204
}
11921205
}
11931206
BamlMediaType::Audio => {

‎fern/03-reference/baml_client/video.mdx‎

Lines changed: 21 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎integ-tests/baml_src/clients.baml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -256,7 +256,8 @@ client<llm> AwsBedrock {
256256
}
257257
// max_tokens 100000
258258
// max_completion_tokens 100000
259-
model "us.anthropic.claude-3-5-haiku-20241022-v1:0"
259+
// model "us.anthropic.claude-3-5-haiku-20241022-v1:0"
260+
model "amazon.nova-lite-v1:0"
260261
// model "anthropic.claude-3-5-sonnet-20240620-v1:0"
261262
// model_id "anthropic.claude-3-haiku-20240307-v1:0"
262263
//model "arn:aws:bedrock:us-east-1:404337120808:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"

‎integ-tests/baml_src/test-files/providers/aws.baml‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,17 @@ function TestAws(input: string) -> string {
55
"#
66
}
77

8+
function TestAwsVideoDescribe(video_input: video) -> string {
9+
client AwsBedrock
10+
prompt #"
11+
{{ _.role("user") }}
12+
13+
Describe the following video in one sentence.
14+
15+
{{ video_input }}
16+
"#
17+
}
18+
819
/// my docs
920
class UniverseQuestion {
1021
question string

0 commit comments

Comments
 (0)