bedrock: Fix subsequent bedrock tool calls fail (#33174)

Vladimir Kuznichenkov created 5 months ago

Closes #30714

Bedrock converse api expect to see tool options if at least one tool was
used in conversation in the past messages.

Right now if `LanguageModelToolChoice::None` isn't supported edit agent
[remove][1] tools from request. That point breaks Converse API of
Bedrock. As was proposed in [the issue][2] we won't drop tool choose but
instead will deny any of them if model will respond with a tool choose.

[1]:
https://github.com/x-qdo/zed/blob/fceba6c79540677c2504d2c22191963b6170591a/crates/assistant_tools/src/edit_agent.rs#L703
[2]:
https://github.com/zed-industries/zed/issues/30714#issuecomment-2886422716

Release Notes:

- Fixed bedrock tool calls in edit mode

Change summary

crates/language_models/src/provider/bedrock.rs | 37 ++++++++++++++++---
1 file changed, 31 insertions(+), 6 deletions(-)

Detailed changes

crates/language_models/src/provider/bedrock.rs

@@ -503,7 +503,8 @@ impl LanguageModel for BedrockModel {
             LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
                 self.model.supports_tool_use()
             }
-            LanguageModelToolChoice::None => false,
+            // Add support for None - we'll filter tool calls at response
+            LanguageModelToolChoice::None => self.model.supports_tool_use(),
         }
     }
 
@@ -549,6 +550,8 @@ impl LanguageModel for BedrockModel {
             }
         };
 
+        let deny_tool_calls = request.tool_choice == Some(LanguageModelToolChoice::None);
+
         let request = match into_bedrock(
             request,
             model_id,
@@ -565,11 +568,15 @@ impl LanguageModel for BedrockModel {
         let request = self.stream_completion(request, cx);
         let future = self.request_limiter.stream(async move {
             let response = request.map_err(|err| anyhow!(err))?.await;
-            Ok(map_to_language_model_completion_events(
-                response,
-                owned_handle,
-            ))
+            let events = map_to_language_model_completion_events(response, owned_handle);
+
+            if deny_tool_calls {
+                Ok(deny_tool_use_events(events).boxed())
+            } else {
+                Ok(events.boxed())
+            }
         });
+
         async move { Ok(future.await?.boxed()) }.boxed()
     }
 
@@ -578,6 +585,23 @@ impl LanguageModel for BedrockModel {
     }
 }
 
+fn deny_tool_use_events(
+    events: impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
+) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
+    events.map(|event| {
+        match event {
+            Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
+                // Convert tool use to an error message if model decided to call it
+                Ok(LanguageModelCompletionEvent::Text(format!(
+                    "\n\n[Error: Tool calls are disabled in this context. Attempted to call '{}']",
+                    tool_use.name
+                )))
+            }
+            other => other,
+        }
+    })
+}
+
 pub fn into_bedrock(
     request: LanguageModelRequest,
     model: String,
@@ -714,7 +738,8 @@ pub fn into_bedrock(
             BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
         }
         Some(LanguageModelToolChoice::None) => {
-            anyhow::bail!("LanguageModelToolChoice::None is not supported");
+            // For None, we still use Auto but will filter out tool calls in the response
+            BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
         }
     };
     let tool_config: BedrockToolConfig = BedrockToolConfig::builder()