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

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]:
fceba6c795/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
This commit is contained in:
Vladimir Kuznichenkov 2025-06-25 10:37:07 +03:00 committed by GitHub
parent 96409965e4
commit 098896146e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -503,7 +503,8 @@ impl LanguageModel for BedrockModel {
LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => { LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
self.model.supports_tool_use() 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( let request = match into_bedrock(
request, request,
model_id, model_id,
@ -565,11 +568,15 @@ 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(|err| anyhow!(err))?.await; let response = request.map_err(|err| anyhow!(err))?.await;
Ok(map_to_language_model_completion_events( let events = map_to_language_model_completion_events(response, owned_handle);
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() 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( pub fn into_bedrock(
request: LanguageModelRequest, request: LanguageModelRequest,
model: String, model: String,
@ -714,7 +738,8 @@ pub fn into_bedrock(
BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build()) BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
} }
Some(LanguageModelToolChoice::None) => { 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() let tool_config: BedrockToolConfig = BedrockToolConfig::builder()