From fa40353fc57054639c65f13232905d545a9624f1 Mon Sep 17 00:00:00 2001
From: Shardul Vaidya <31039336+5herlocked@users.noreply.github.com>
Date: Tue, 29 Apr 2025 12:18:32 -0400
Subject: [PATCH] bedrock: Preserve thinking blocks for Bedrock (#29602)

Fixes a regression from #29055, resolves #29290

Release Notes:

- agent: Fixed a regression that rendered Claude 3.7 Thinking unusable
on Bedrock.
---
 crates/bedrock/src/bedrock.rs                 |  2 ++
 .../language_models/src/provider/bedrock.rs   | 27 ++++++++++++++++---
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/crates/bedrock/src/bedrock.rs b/crates/bedrock/src/bedrock.rs
index 5b34388a9a..278430f7f2 100644
--- a/crates/bedrock/src/bedrock.rs
+++ b/crates/bedrock/src/bedrock.rs
@@ -11,12 +11,14 @@ pub use aws_sdk_bedrockruntime::types::{
     Tool as BedrockTool, ToolChoice as BedrockToolChoice, ToolConfiguration as BedrockToolConfig,
     ToolInputSchema as BedrockToolInputSchema, ToolSpecification as BedrockToolSpec,
 };
+pub use aws_smithy_types::Blob as BedrockBlob;
 use aws_smithy_types::{Document, Number as AwsNumber};
 pub use bedrock::operation::converse_stream::ConverseStreamInput as BedrockStreamingRequest;
 pub use bedrock::types::{
     ContentBlock as BedrockRequestContent, ConversationRole as BedrockRole,
     ConverseOutput as BedrockResponse, ConverseStreamOutput as BedrockStreamingResponse,
     ImageBlock as BedrockImageBlock, Message as BedrockMessage,
+    ReasoningContentBlock as BedrockThinkingBlock, ReasoningTextBlock as BedrockThinkingTextBlock,
     ResponseStream as BedrockResponseStream, ToolResultBlock as BedrockToolResultBlock,
     ToolResultContentBlock as BedrockToolResultContentBlock,
     ToolResultStatus as BedrockToolResultStatus, ToolUseBlock as BedrockToolUseBlock,
diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs
index b9793ac16e..70c92b39e8 100644
--- a/crates/language_models/src/provider/bedrock.rs
+++ b/crates/language_models/src/provider/bedrock.rs
@@ -15,10 +15,11 @@ use bedrock::bedrock_client::types::{
     StopReason,
 };
 use bedrock::{
-    BedrockAutoToolChoice, BedrockError, BedrockInnerContent, BedrockMessage, BedrockModelMode,
-    BedrockStreamingResponse, BedrockTool, BedrockToolChoice, BedrockToolConfig,
-    BedrockToolInputSchema, BedrockToolResultBlock, BedrockToolResultContentBlock,
-    BedrockToolResultStatus, BedrockToolSpec, BedrockToolUseBlock, Model, value_to_aws_document,
+    BedrockAutoToolChoice, BedrockBlob, BedrockError, BedrockInnerContent, BedrockMessage,
+    BedrockModelMode, BedrockStreamingResponse, BedrockThinkingBlock, BedrockThinkingTextBlock,
+    BedrockTool, BedrockToolChoice, BedrockToolConfig, BedrockToolInputSchema,
+    BedrockToolResultBlock, BedrockToolResultContentBlock, BedrockToolResultStatus,
+    BedrockToolSpec, BedrockToolUseBlock, Model, value_to_aws_document,
 };
 use collections::{BTreeMap, HashMap};
 use credentials_provider::CredentialsProvider;
@@ -626,6 +627,24 @@ pub fn into_bedrock(
                                 None
                             }
                         }
+                        MessageContent::Thinking { text, signature } => {
+                            let thinking = BedrockThinkingTextBlock::builder()
+                                .text(text)
+                                .set_signature(signature)
+                                .build()
+                                .context("failed to build reasoning block")
+                                .log_err()?;
+
+                            Some(BedrockInnerContent::ReasoningContent(
+                                BedrockThinkingBlock::ReasoningText(thinking),
+                            ))
+                        }
+                        MessageContent::RedactedThinking(blob) => {
+                            let redacted =
+                                BedrockThinkingBlock::RedactedContent(BedrockBlob::new(blob));
+
+                            Some(BedrockInnerContent::ReasoningContent(redacted))
+                        }
                         MessageContent::ToolUse(tool_use) => BedrockToolUseBlock::builder()
                             .name(tool_use.name.to_string())
                             .tool_use_id(tool_use.id.to_string())