chore: cleaned up redundant unwraps to make error handling more declarative

Author: iambenkay

llm-bedrock/src/client.rs

@@ -52,39 +52,39 @@ impl Bedrock {
         tool_results: Option<Vec<(llm::ToolCall, llm::ToolResult)>>,
     ) -> llm::ChatEvent {
         let bedrock_input = BedrockInput::from(messages, config, tool_results);
-        if let Err(err) = bedrock_input {
-            return llm::ChatEvent::Error(err);
-        }
-
-        let input = bedrock_input.unwrap();
 
         let runtime = get_async_runtime();
 
-        trace!("Sending request to AWS Bedrock: {input:?}");
-        runtime.block_on(async {
-            let model_id = input.model_id.clone();
-            let response = self
-                .init_converse(input)
-                .send()
-                .await
-                .map_err(|e| from_converse_sdk_error(model_id, e));
-
-            if let Err(err) = response {
-                return llm::ChatEvent::Error(err);
+        match bedrock_input {
+            Err(err) => llm::ChatEvent::Error(err),
+            Ok(input) => {
+                trace!("Sending request to AWS Bedrock: {input:?}");
+                runtime.block_on(async {
+                    let model_id = input.model_id.clone();
+                    let response = self
+                        .init_converse(input)
+                        .send()
+                        .await
+                        .map_err(|e| from_converse_sdk_error(model_id, e));
+
+                    match response {
+                        Err(err) => llm::ChatEvent::Error(err),
+                        Ok(response) => {
+                            let event = match response.stop_reason() {
+                                bedrock::types::StopReason::ToolUse => {
+                                    conversions::converse_output_to_tool_calls(response)
+                                        .map(llm::ChatEvent::ToolRequest)
+                                }
+                                _ => conversions::converse_output_to_complete_response(response)
+                                    .map(llm::ChatEvent::Message),
+                            };
+
+                            event.unwrap_or_else(llm::ChatEvent::Error)
+                        }
+                    }
+                })
             }
-            let response = response.unwrap();
-
-            let event = match response.stop_reason() {
-                bedrock::types::StopReason::ToolUse => {
-                    conversions::converse_output_to_tool_calls(response)
-                        .map(llm::ChatEvent::ToolRequest)
-                }
-                _ => conversions::converse_output_to_complete_response(response)
-                    .map(llm::ChatEvent::Message),
-            };
-
-            event.unwrap_or_else(llm::ChatEvent::Error)
-        })
+        }
     }
 
     pub fn converse_stream(
@@ -93,27 +93,28 @@ impl Bedrock {
         config: llm::Config,
     ) -> BedrockChatStream {
         let bedrock_input = BedrockInput::from(messages, config, None);
-        if let Err(err) = bedrock_input {
-            return BedrockChatStream::failed(err);
-        }
-        let input = bedrock_input.unwrap();
 
-        let runtime = get_async_runtime();
-        trace!("Sending request to AWS Bedrock: {input:?}");
-        runtime.block_on(async {
-            let model_id = input.model_id.clone();
-            let response = self
-                .init_converse_stream(input)
-                .send()
-                .await
-                .map_err(|e| from_converse_stream_sdk_error(model_id, e));
-
-            trace!("Creating AWS Bedrock event stream");
-            match response {
-                Ok(response) => BedrockChatStream::new(response.stream),
-                Err(error) => BedrockChatStream::failed(error),
+        match bedrock_input {
+            Err(err) => BedrockChatStream::failed(err),
+            Ok(input) => {
+                let runtime = get_async_runtime();
+                trace!("Sending request to AWS Bedrock: {input:?}");
+                runtime.block_on(async {
+                    let model_id = input.model_id.clone();
+                    let response = self
+                        .init_converse_stream(input)
+                        .send()
+                        .await
+                        .map_err(|e| from_converse_stream_sdk_error(model_id, e));
+
+                    trace!("Creating AWS Bedrock event stream");
+                    match response {
+                        Ok(response) => BedrockChatStream::new(response.stream),
+                        Err(error) => BedrockChatStream::failed(error),
+                    }
+                })
             }
-        })
+        }
     }
 
     fn init_converse(&self, input: conversions::BedrockInput) -> ConverseFluentBuilder {

llm-bedrock/src/conversions.rs

@@ -205,17 +205,18 @@ fn get_image_content_block_from_url(url: &str) -> Result<bedrock::types::Content
 
     let kind = infer::get(&bytes);
 
-    if kind.is_none() {
-        return Err(custom_error(
-            llm::ErrorCode::InvalidRequest,
-            format!(
-                "Could not infer the mime type of the image downloaded from url: {}",
-                url
-            ),
-        ));
-    }
-
-    let mime = str_to_bedrock_mime_type(kind.unwrap().mime_type())?;
+    let mime = match kind {
+        Some(kind) => str_to_bedrock_mime_type(kind.mime_type())?,
+        None => {
+            return Err(custom_error(
+                llm::ErrorCode::InvalidRequest,
+                format!(
+                    "Could not infer the mime type of the image downloaded from url: {}",
+                    url
+                ),
+            ));
+        }
+    };
 
     Ok(bedrock::types::ContentBlock::Image(
         ImageBlock::builder()
@@ -280,24 +281,22 @@ pub fn converse_output_to_tool_calls(
             .to_owned(),
     ))?;
 
-    if !output.is_message() {
-        return Err(custom_error(
+    match output.as_message() {
+        Err(_) => Err(custom_error(
             llm::ErrorCode::InternalError,
             "An error occurred while converting to tool calls: expected output to be a Message"
                 .to_owned(),
-        ));
-    }
-
-    let message = output.as_message().unwrap().clone();
-    let mut tool_calls: Vec<llm::ToolCall> = vec![];
-
-    for block in message.content {
-        if let bedrock::types::ContentBlock::ToolUse(tool) = block {
-            tool_calls.push(bedrock_tool_use_to_llm_tool_call(tool)?);
+        )),
+        Ok(message) => {
+            let mut tool_calls: Vec<llm::ToolCall> = vec![];
+            for block in message.content.clone() {
+                if let bedrock::types::ContentBlock::ToolUse(tool) = block {
+                    tool_calls.push(bedrock_tool_use_to_llm_tool_call(tool)?);
+                }
+            }
+            Ok(tool_calls)
         }
     }
-
-    Ok(tool_calls)
 }
 
 pub fn converse_output_to_complete_response(
@@ -309,44 +308,46 @@ pub fn converse_output_to_complete_response(
             .to_owned(),
     ))?;
 
-    if !output.is_message() {
-        return Err(custom_error(
+    match output.as_message() {
+        Err(_) => Err(custom_error(
             llm::ErrorCode::InternalError,
-            "An error occurred while converting to complete response: expected output to be a Message".to_owned(),
-        ));
-    }
-
-    let message = output.as_message().unwrap().clone();
-
-    let mut content_parts: Vec<llm::ContentPart> = vec![];
-    let mut tool_calls: Vec<llm::ToolCall> = vec![];
-    for block in message.content {
-        match block {
-            bedrock::types::ContentBlock::Text(text) => {
-                content_parts.push(llm::ContentPart::Text(text.to_owned()));
-            }
-            bedrock::types::ContentBlock::Image(image) => {
-                content_parts.push(bedrock_image_to_llm_content_part(image));
-            }
-            bedrock::types::ContentBlock::ToolUse(tool) => {
-                tool_calls.push(bedrock_tool_use_to_llm_tool_call(tool)?);
+            "An error occurred while converting to complete response: expected output to be a Message"
+               .to_owned(),
+        )),
+        Ok(message) => {
+            let mut content_parts: Vec<llm::ContentPart> = vec![];
+            let mut tool_calls: Vec<llm::ToolCall> = vec![];
+            for block in message.content.clone() {
+                match block {
+                    bedrock::types::ContentBlock::Text(text) => {
+                        content_parts.push(llm::ContentPart::Text(text.to_owned()));
+                    }
+                    bedrock::types::ContentBlock::Image(image) => {
+                        content_parts.push(bedrock_image_to_llm_content_part(image));
+                    }
+                    bedrock::types::ContentBlock::ToolUse(tool) => {
+                        tool_calls.push(bedrock_tool_use_to_llm_tool_call(tool)?);
+                    }
+                    _ => {}
+                }
             }
-            _ => {}
+        
+            let metadata = converse_output_to_response_metadata(&response);
+        
+            Ok(llm::CompleteResponse {
+                // bedrock does not return an id as part of the response struct.
+                // there may be one present in `additional_model_response_fields`
+                // but the schema varies depending on the model being invoked. Leaving it empty for now
+                // until we have a better solution for this.
+                id: "".to_owned(),
+                content: content_parts,
+                tool_calls,
+                metadata,
+            })
+
         }
     }
 
-    let metadata = converse_output_to_response_metadata(&response);
-
-    Ok(llm::CompleteResponse {
-        // bedrock does not return an id as part of the response struct.
-        // there may be one present in `additional_model_response_fields`
-        // but the schema varies depending on the model being invoked. Leaving it empty for now
-        // until we have a better solution for this.
-        id: "".to_owned(),
-        content: content_parts,
-        tool_calls,
-        metadata,
-    })
 }
 
 fn bedrock_tool_use_to_llm_tool_call(tool: ToolUseBlock) -> Result<llm::ToolCall, llm::Error> {
@@ -445,13 +446,12 @@ pub fn converse_stream_output_to_stream_event(
 
 fn process_content_block_start_event(block: ContentBlockStartEvent) -> Option<llm::StreamEvent> {
     if let Some(start_info) = block.start {
-        if start_info.is_tool_use() {
-            let tool_use = start_info.as_tool_use().unwrap().clone();
+        if let Ok(tool_use) = start_info.as_tool_use() {
             return Some(llm::StreamEvent::Delta(llm::StreamDelta {
                 content: None,
                 tool_calls: Some(vec![llm::ToolCall {
-                    id: tool_use.tool_use_id,
-                    name: tool_use.name,
+                    id: tool_use.tool_use_id.clone(),
+                    name: tool_use.name.clone(),
                     arguments_json: "".to_owned(),
                 }]),
             }));
@@ -462,20 +462,18 @@ fn process_content_block_start_event(block: ContentBlockStartEvent) -> Option<ll
 
 fn process_content_block_delta_event(block: ContentBlockDeltaEvent) -> Option<llm::StreamEvent> {
     if let Some(block_info) = block.delta {
-        if block_info.is_tool_use() {
-            let tool_use = block_info.as_tool_use().unwrap().clone();
+        if let Ok(tool_use) = block_info.as_tool_use() {
             return Some(llm::StreamEvent::Delta(llm::StreamDelta {
                 content: None,
                 tool_calls: Some(vec![llm::ToolCall {
                     id: "".to_owned(),
                     name: "".to_owned(),
-                    arguments_json: tool_use.input,
+                    arguments_json: tool_use.input.clone(),
                 }]),
             }));
-        } else if block_info.is_text() {
-            let text = block_info.as_text().unwrap().clone();
+        } else if let Ok(text) = block_info.as_text() {
             return Some(llm::StreamEvent::Delta(llm::StreamDelta {
-                content: Some(vec![llm::ContentPart::Text(text)]),
+                content: Some(vec![llm::ContentPart::Text(text.clone())]),
                 tool_calls: None,
             }));
         }

llm-bedrock/src/lib.rs

@@ -21,11 +21,11 @@ impl Guest for BedrockComponent {
         LOGGING_STATE.with_borrow_mut(|state| state.init());
 
         let bedrock = get_bedrock_client();
-        if let Err(err) = bedrock {
-            return ChatEvent::Error(err);
-        }
 
-        bedrock.unwrap().converse(messages, config, None)
+        match bedrock {
+            Ok(client) => client.converse(messages, config, None),
+            Err(err) => ChatEvent::Error(err),
+        }
     }
 
     fn continue_(
@@ -36,13 +36,11 @@ impl Guest for BedrockComponent {
         LOGGING_STATE.with_borrow_mut(|state| state.init());
 
         let bedrock = get_bedrock_client();
-        if let Err(err) = bedrock {
-            return ChatEvent::Error(err);
-        }
 
-        bedrock
-            .unwrap()
-            .converse(messages, config, Some(tool_results))
+        match bedrock {
+            Ok(client) => client.converse(messages, config, Some(tool_results)),
+            Err(err) => ChatEvent::Error(err),
+        }
     }
 
     fn stream(messages: Vec<Message>, config: Config) -> ChatStream {
@@ -58,11 +56,11 @@ impl ExtendedGuest for BedrockComponent {
         LOGGING_STATE.with_borrow_mut(|state| state.init());
 
         let bedrock = get_bedrock_client();
-        if let Err(err) = bedrock {
-            return BedrockChatStream::failed(err);
-        }
 
-        bedrock.unwrap().converse_stream(messages, config)
+        match bedrock {
+            Ok(client) => client.converse_stream(messages, config),
+            Err(err) => BedrockChatStream::failed(err),
+        }
     }
 
     fn subscribe(_stream: &Self::ChatStream) -> golem_rust::wasm_rpc::Pollable {

llm-bedrock/src/stream.rs

@@ -57,17 +57,20 @@ impl BedrockChatStream {
 
             runtime.block_on(async move {
                 let token = stream.recv().await;
+                log::trace!("Bedrock stream event: {token:?}");
 
                 match token {
                     Ok(Some(output)) => {
                         log::trace!("Processing bedrock stream event: {output:?}");
                         converse_stream_output_to_stream_event(output)
                     }
                     Ok(None) => {
+                        log::trace!("running set_finished on stream due to None event received");
                         self.set_finished();
                         None
                     }
                     Err(error) => {
+                        log::trace!("running set_finished on stream due to error: {error:?}");
                         self.set_finished();
                         Some(llm::StreamEvent::Error(custom_error(
                             llm::ErrorCode::InternalError,