Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

language_models: Remove .unwraps in Bedrock provider (#26238)

Browse source 
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.
This commit is contained in:
Marshall Bowers 2025-03-06 16:02:37 -05:00 • committed by GitHub
parent 7e964290bf
commit 97c0a0a86e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 14 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

21
crates/language_models/src/provider/bedrock.rs
View file

@ -400,7 +400,7 @@ impl LanguageModel for BedrockModel {
let request = self.stream_completion(request, cx); let request = self.stream_completion(request, cx);
let future = self.request_limiter.stream(async move { 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( Ok(map_to_language_model_completion_events(
response, response,
owned_handle, owned_handle,
@ -424,26 +424,28 @@ impl LanguageModel for BedrockModel {
self.model.max_output_tokens(), self.model.max_output_tokens(),
); );
request.tool_choice = Some(BedrockToolChoice::Tool( request.tool_choice = BedrockSpecificTool::builder()
BedrockSpecificTool::builder()
.name(name.clone()) .name(name.clone())
.build() .build()
.unwrap(), .log_err()
)); .map(BedrockToolChoice::Tool);
request.tools = vec![BedrockTool::builder() if let Some(tool) = BedrockTool::builder()
.name(name.clone()) .name(name.clone())
.description(description.clone()) .description(description.clone())
.input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema))) .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema)))
.build() .build()
.unwrap()]; .log_err()
{
request.tools.push(tool);
}
let handle = self.handler.clone(); let handle = self.handler.clone();
let request = self.stream_completion(request, _cx); let request = self.stream_completion(request, _cx);
self.request_limiter self.request_limiter
.run(async move { .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) Ok(extract_tool_args_from_events(name, response, handle)
.await? .await?
.boxed()) .boxed())
@ -757,7 +759,8 @@ pub fn map_to_language_model_completion_events(
None None
}) })
.await .await
.unwrap() .log_err()
.flatten()
} }
}, },
) )