fix: bedrock provider to only include non user annotations and convert non supported docs to text (#1899)

This commit is contained in:
Zaki Ali
2025-03-28 14:20:08 -07:00
committed by GitHub
parent b22d173425
commit f72db13b4a

View File

@@ -62,6 +62,11 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result<bedrock::C
Ok(content) => Some( Ok(content) => Some(
content content
.iter() .iter()
// Filter out content items that have User in their audience
.filter(|c| {
c.audience()
.is_none_or(|audience| !audience.contains(&Role::User))
})
.map(|c| to_bedrock_tool_result_content_block(&tool_res.id, c)) .map(|c| to_bedrock_tool_result_content_block(&tool_res.id, c))
.collect::<Result<_>>()?, .collect::<Result<_>>()?,
), ),
@@ -89,9 +94,17 @@ pub fn to_bedrock_tool_result_content_block(
Ok(match content { Ok(match content {
Content::Text(text) => bedrock::ToolResultContentBlock::Text(text.text.to_string()), Content::Text(text) => bedrock::ToolResultContentBlock::Text(text.text.to_string()),
Content::Image(_) => bail!("Image content is not supported by Bedrock provider yet"), Content::Image(_) => bail!("Image content is not supported by Bedrock provider yet"),
Content::Resource(resource) => bedrock::ToolResultContentBlock::Document( Content::Resource(resource) => match &resource.resource {
to_bedrock_document(tool_use_id, &resource.resource)?, ResourceContents::TextResourceContents { text, .. } => {
), match to_bedrock_document(tool_use_id, &resource.resource)? {
Some(doc) => bedrock::ToolResultContentBlock::Document(doc),
None => bedrock::ToolResultContentBlock::Text(text.to_string()),
}
}
ResourceContents::BlobResourceContents { .. } => {
bail!("Blob resource content is not supported by Bedrock provider yet")
}
},
}) })
} }
@@ -149,7 +162,7 @@ pub fn to_bedrock_json(value: &Value) -> Document {
fn to_bedrock_document( fn to_bedrock_document(
tool_use_id: &str, tool_use_id: &str,
content: &ResourceContents, content: &ResourceContents,
) -> Result<bedrock::DocumentBlock> { ) -> Result<Option<bedrock::DocumentBlock>> {
let (uri, text) = match content { let (uri, text) = match content {
ResourceContents::TextResourceContents { uri, text, .. } => (uri, text), ResourceContents::TextResourceContents { uri, text, .. } => (uri, text),
ResourceContents::BlobResourceContents { .. } => { ResourceContents::BlobResourceContents { .. } => {
@@ -162,13 +175,13 @@ fn to_bedrock_document(
.and_then(|n| n.to_str()) .and_then(|n| n.to_str())
.unwrap_or(uri); .unwrap_or(uri);
// Return None if the file type is not supported
let (name, format) = match filename.split_once('.') { let (name, format) = match filename.split_once('.') {
Some((name, "txt")) => (name, bedrock::DocumentFormat::Txt), Some((name, "txt")) => (name, bedrock::DocumentFormat::Txt),
Some((name, "csv")) => (name, bedrock::DocumentFormat::Csv), Some((name, "csv")) => (name, bedrock::DocumentFormat::Csv),
Some((name, "md")) => (name, bedrock::DocumentFormat::Md), Some((name, "md")) => (name, bedrock::DocumentFormat::Md),
Some((name, "html")) => (name, bedrock::DocumentFormat::Html), Some((name, "html")) => (name, bedrock::DocumentFormat::Html),
Some((name, _)) => (name, bedrock::DocumentFormat::Txt), _ => return Ok(None), // Not a supported document type
_ => (filename, bedrock::DocumentFormat::Txt),
}; };
// Since we can't use the full path (due to character limit and also Bedrock does not accept `/` etc.), // Since we can't use the full path (due to character limit and also Bedrock does not accept `/` etc.),
@@ -176,12 +189,14 @@ fn to_bedrock_document(
// document names unique. // document names unique.
let name = format!("{tool_use_id}-{name}"); let name = format!("{tool_use_id}-{name}");
Ok(Some(
bedrock::DocumentBlock::builder() bedrock::DocumentBlock::builder()
.format(format) .format(format)
.name(name) .name(name)
.source(bedrock::DocumentSource::Bytes(text.as_bytes().into())) .source(bedrock::DocumentSource::Bytes(text.as_bytes().into()))
.build() .build()
.map_err(|err| anyhow!("Failed to construct Bedrock document: {}", err)) .map_err(|err| anyhow!("Failed to construct Bedrock document: {}", err))?,
))
} }
pub fn from_bedrock_message(message: &bedrock::Message) -> Result<Message> { pub fn from_bedrock_message(message: &bedrock::Message) -> Result<Message> {