From ce37c2634ea7644fae6a77294566d14eb156995b Mon Sep 17 00:00:00 2001
From: "gcp-cherry-pick-bot[bot]"
 <98988430+gcp-cherry-pick-bot[bot]@users.noreply.github.com>
Date: Thu, 6 Mar 2025 16:49:58 -0500
Subject: [PATCH] language_models: Remove `.unwrap`s in Bedrock provider
 (cherry-pick #26238) (#26242)

Cherry-picked language_models: Remove `.unwrap`s in Bedrock provider
(#26238)

This PR removes a number of `.unwrap`s in the Bedrock provider.

We must not `.unwrap` in situations where it is not provably safe to do
so, which it was not in any of these cases.

Release Notes:

- Fixed some potential panics in the AWS Bedrock model provider.

Co-authored-by: Marshall Bowers <git@maxdeviant.com>
---
 .../language_models/src/provider/bedrock.rs   | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs
index fe4ae044aab8cfd500e169c79c70a64a52808ef2..0656e361e72853506083854d4f7ac54cf2bd8c1c 100644
--- a/crates/language_models/src/provider/bedrock.rs
+++ b/crates/language_models/src/provider/bedrock.rs
@@ -401,7 +401,7 @@ impl LanguageModel for BedrockModel {
 
         let request = self.stream_completion(request, cx);
         let future = self.request_limiter.stream(async move {
-            let response = request.map_err(|e| anyhow!(e)).unwrap().await;
+            let response = request.map_err(|err| anyhow!(err))?.await;
             Ok(map_to_language_model_completion_events(
                 response,
                 owned_handle,
@@ -425,26 +425,28 @@ impl LanguageModel for BedrockModel {
             self.model.max_output_tokens(),
         );
 
-        request.tool_choice = Some(BedrockToolChoice::Tool(
-            BedrockSpecificTool::builder()
-                .name(name.clone())
-                .build()
-                .unwrap(),
-        ));
+        request.tool_choice = BedrockSpecificTool::builder()
+            .name(name.clone())
+            .build()
+            .log_err()
+            .map(BedrockToolChoice::Tool);
 
-        request.tools = vec![BedrockTool::builder()
+        if let Some(tool) = BedrockTool::builder()
             .name(name.clone())
             .description(description.clone())
             .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema)))
             .build()
-            .unwrap()];
+            .log_err()
+        {
+            request.tools.push(tool);
+        }
 
         let handle = self.handler.clone();
 
         let request = self.stream_completion(request, _cx);
         self.request_limiter
             .run(async move {
-                let response = request.map_err(|e| anyhow!(e)).unwrap().await;
+                let response = request.map_err(|err| anyhow!(err))?.await;
                 Ok(extract_tool_args_from_events(name, response, handle)
                     .await?
                     .boxed())
@@ -758,7 +760,8 @@ pub fn map_to_language_model_completion_events(
                         None
                     })
                     .await
-                    .unwrap()
+                    .log_err()
+                    .flatten()
             }
         },
     )