fix(bedrock): send inference config (max_tokens, temperature) on Converse (#9889)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>
This commit is contained in:
Lucas Kim
2026-07-01 07:59:40 +09:00
committed by GitHub
parent 12c3e234f5
commit e8891b7bb0
3 changed files with 195 additions and 14 deletions
@@ -618,7 +618,8 @@ pub fn thinking_budget_tokens(model_config: &ModelConfig) -> i32 {
// Anthropic counts thinking tokens against max_tokens, so the budget must leave
// room for a response. Clamp it to preserve at least this many answer tokens, and
// drop thinking only when even a minimal budget wouldn't fit under the cap.
const MIN_ANSWER_TOKENS: i32 = 1024;
// Shared with the Bedrock formatter, which applies the same clamp.
pub const MIN_ANSWER_TOKENS: i32 = 1024;
fn apply_thinking_config(
payload: &mut Value,
+8 -4
View File
@@ -27,8 +27,8 @@ use serde_json::Value;
use smithy_transport_reqwest::ReqwestHttpClient;
use super::formats::bedrock::{
bedrock_anthropic_thinking_fields, from_bedrock_message, from_bedrock_usage,
to_bedrock_message_with_caching, to_bedrock_tool_config,
bedrock_anthropic_thinking_fields, bedrock_inference_config, from_bedrock_message,
from_bedrock_usage, to_bedrock_message_with_caching, to_bedrock_tool_config,
};
pub(crate) const BEDROCK_PROVIDER_NAME: &str = "aws_bedrock";
@@ -75,6 +75,7 @@ struct ConverseRequestParts {
messages: Vec<bedrock::Message>,
tool_config: Option<bedrock::ToolConfiguration>,
thinking_fields: Option<aws_smithy_types::Document>,
inference_config: bedrock::InferenceConfiguration,
}
impl BedrockProvider {
@@ -331,6 +332,7 @@ impl BedrockProvider {
messages: bedrock_messages,
tool_config,
thinking_fields: bedrock_anthropic_thinking_fields(model),
inference_config: bedrock_inference_config(model),
})
}
@@ -349,7 +351,8 @@ impl BedrockProvider {
.converse()
.set_system(Some(parts.system_blocks))
.model_id(&model.model_name)
.set_messages(Some(parts.messages));
.set_messages(Some(parts.messages))
.inference_config(parts.inference_config);
if let Some(fields) = parts.thinking_fields {
request = request.additional_model_request_fields(fields);
@@ -444,7 +447,8 @@ impl BedrockProvider {
.converse_stream()
.set_system(Some(parts.system_blocks))
.model_id(&model.model_name)
.set_messages(Some(parts.messages));
.set_messages(Some(parts.messages))
.inference_config(parts.inference_config);
if let Some(fields) = parts.thinking_fields {
request = request.additional_model_request_fields(fields);
+185 -9
View File
@@ -15,9 +15,11 @@ use rmcp::model::{
use serde_json::Value;
use crate::conversation::message::{Message, MessageContent};
use crate::providers::bedrock::BEDROCK_PROVIDER_NAME;
use crate::providers::canonical::maybe_get_canonical_model;
use crate::providers::formats::anthropic::{
adaptive_output_effort, thinking_budget_tokens, thinking_type_for_provider, ThinkingType,
ANTHROPIC_PROVIDER_NAME,
adaptive_output_effort, model_supports_temperature, thinking_budget_tokens,
thinking_type_for_provider, ThinkingType, ANTHROPIC_PROVIDER_NAME, MIN_ANSWER_TOKENS,
};
use goose_providers::conversation::token_usage::Usage;
use goose_providers::model::ModelConfig;
@@ -33,13 +35,27 @@ pub fn bedrock_anthropic_thinking_fields(model_config: &ModelConfig) -> Option<D
"type".to_string(),
Document::String("adaptive".to_string()),
)])),
ThinkingType::Enabled => Document::Object(HashMap::from([
("type".to_string(), Document::String("enabled".to_string())),
(
"budget_tokens".to_string(),
Document::Number(Number::PosInt(thinking_budget_tokens(model_config) as u64)),
),
])),
ThinkingType::Enabled => {
// Thinking tokens count against `maxTokens`, which `bedrock_inference_config`
// now sends when explicitly configured. Mirror the Anthropic formatter: clamp
// the budget to leave room for an answer, and drop thinking entirely when even
// a minimal budget wouldn't fit under the cap. When max_tokens is unset, Bedrock
// applies its per-model default so there is nothing to clamp against.
let mut budget_tokens = thinking_budget_tokens(model_config);
if let Some(max_tokens) = model_config.max_tokens {
budget_tokens = budget_tokens.min(max_tokens.saturating_sub(MIN_ANSWER_TOKENS));
if budget_tokens < MIN_ANSWER_TOKENS {
return None;
}
}
Document::Object(HashMap::from([
("type".to_string(), Document::String("enabled".to_string())),
(
"budget_tokens".to_string(),
Document::Number(Number::PosInt(budget_tokens as u64)),
),
]))
}
ThinkingType::Disabled => return None,
};
@@ -81,6 +97,59 @@ fn strip_bedrock_version_suffix(model_name: &str) -> String {
.into_owned()
}
/// Build the Bedrock `InferenceConfiguration` (`maxTokens`, `temperature`) for
/// a request from the active [`ModelConfig`].
///
/// Without this the `Converse`/`ConverseStream` APIs fall back to per-model
/// server defaults, so a configured `max_tokens`/`temperature` is silently
/// dropped. Each field is sent only when the user has configured it, so that
/// unset values continue to use Bedrock's per-model server defaults rather than
/// being pinned to a generic fallback:
/// - `max_tokens` is sent only when explicitly set (`model_config.max_tokens`).
/// Using [`ModelConfig::max_output_tokens`] here would forward its `4096`
/// fallback for every model whose id is not in the canonical catalog (e.g.
/// cross-region ids like `us.anthropic.claude-...`), capping models whose
/// real output limit is far higher.
/// - `temperature` is sent only when set and the model supports it. Support is
/// resolved against the Anthropic canonical registry for `anthropic.*` model
/// ids (the same mapping used for thinking) and the Bedrock canonical registry
/// for other known Bedrock ids, so models that reject a custom temperature keep
/// the server default.
pub fn bedrock_inference_config(model_config: &ModelConfig) -> bedrock::InferenceConfiguration {
let mut builder = bedrock::InferenceConfiguration::builder();
if let Some(max_tokens) = model_config.max_tokens {
builder = builder.max_tokens(max_tokens);
}
if let Some(temperature) = model_config.temperature {
if bedrock_model_supports_temperature(model_config) {
builder = builder.temperature(temperature);
}
}
builder.build()
}
/// Whether `temperature` may be sent for this Bedrock model. For `anthropic.*`
/// ids we resolve against the Anthropic canonical registry (mapping the model
/// name the same way [`bedrock_anthropic_thinking_type`] does); for other known
/// Bedrock ids we consult the Bedrock canonical registry and otherwise keep the
/// permissive fallback used by [`model_supports_temperature`].
fn bedrock_model_supports_temperature(model_config: &ModelConfig) -> bool {
if let Some((_, anthropic_model)) = model_config.model_name.rsplit_once("anthropic.") {
let anthropic_config = ModelConfig {
model_name: strip_bedrock_version_suffix(anthropic_model),
..model_config.clone()
};
model_supports_temperature(ANTHROPIC_PROVIDER_NAME, &anthropic_config)
} else {
maybe_get_canonical_model(BEDROCK_PROVIDER_NAME, &model_config.model_name)
.and_then(|model| model.temperature)
.unwrap_or(true)
}
}
pub fn to_bedrock_message_with_caching(
message: &Message,
enable_caching: bool,
@@ -563,6 +632,44 @@ mod tests {
);
}
#[test]
fn test_bedrock_anthropic_thinking_fields_clamped_to_max_tokens() {
// budget (4000) exceeds the room left under an explicit max_tokens, so it
// is clamped to max_tokens - MIN_ANSWER_TOKENS, matching the Anthropic
// formatter. Without max_tokens set there is nothing to clamp against.
let mut params = HashMap::new();
params.insert("thinking_effort".to_string(), json!("low"));
let mut config = ModelConfig::new("us.anthropic.claude-3-7-sonnet-20250219-v1:0");
config.request_params = Some(params);
config.reasoning = Some(true);
config.max_tokens = Some(3000);
let fields = bedrock_anthropic_thinking_fields(&config).expect("thinking fields");
assert_eq!(
from_bedrock_json(&fields).unwrap(),
json!({
"thinking": {
"type": "enabled",
"budget_tokens": 3000 - 1024
}
})
);
}
#[test]
fn test_bedrock_anthropic_thinking_fields_dropped_when_no_room() {
// When even a minimal budget wouldn't leave MIN_ANSWER_TOKENS under the
// cap, thinking is dropped rather than emitting an unsatisfiable request.
let mut params = HashMap::new();
params.insert("thinking_effort".to_string(), json!("low"));
let mut config = ModelConfig::new("us.anthropic.claude-3-7-sonnet-20250219-v1:0");
config.request_params = Some(params);
config.reasoning = Some(true);
config.max_tokens = Some(1500);
assert!(bedrock_anthropic_thinking_fields(&config).is_none());
}
#[test]
fn test_bedrock_anthropic_thinking_fields_disabled() {
let mut config = ModelConfig::new("us.anthropic.claude-3-7-sonnet-20250219-v1:0");
@@ -1228,4 +1335,73 @@ mod tests {
Ok(())
}
#[test]
fn test_bedrock_inference_config_sets_max_tokens_and_temperature() {
let mut config = ModelConfig::new("us.anthropic.claude-sonnet-4-5-20250929-v1:0");
config.max_tokens = Some(8192);
config.temperature = Some(0.5);
let inference_config = bedrock_inference_config(&config);
assert_eq!(inference_config.max_tokens(), Some(8192));
assert_eq!(inference_config.temperature(), Some(0.5));
}
#[test]
fn test_bedrock_inference_config_omits_max_tokens_without_config() {
let mut config = ModelConfig::new("us.anthropic.claude-sonnet-4-5-20250929-v1:0");
config.max_tokens = None;
config.temperature = None;
let inference_config = bedrock_inference_config(&config);
// When max_tokens is not explicitly configured we leave it unset so
// Bedrock applies its per-model server default. Forwarding
// ModelConfig::max_output_tokens() here would pin every model without a
// canonical-catalog entry (e.g. cross-region ids) to the generic 4096
// fallback, capping models whose real output limit is much higher.
assert_eq!(inference_config.max_tokens(), None);
assert_eq!(inference_config.temperature(), None);
}
#[test]
fn test_bedrock_inference_config_sends_explicit_max_tokens() {
let mut config = ModelConfig::new("us.anthropic.claude-sonnet-4-5-20250929-v1:0");
config.max_tokens = Some(4096);
let inference_config = bedrock_inference_config(&config);
// An explicitly configured value is always forwarded.
assert_eq!(inference_config.max_tokens(), Some(4096));
}
#[test]
fn test_bedrock_inference_config_omits_temperature_for_unsupported_model() {
// The Anthropic canonical registry maps this id and reports whether a
// custom temperature may be sent; when it cannot, temperature is left
// unset so the server default is used.
let mut config = ModelConfig::new("us.anthropic.claude-sonnet-4-5-20250929-v1:0");
config.temperature = Some(0.5);
let supported = bedrock_model_supports_temperature(&config);
let inference_config = bedrock_inference_config(&config);
if supported {
assert_eq!(inference_config.temperature(), Some(0.5));
} else {
assert_eq!(inference_config.temperature(), None);
}
}
#[test]
fn test_bedrock_inference_config_omits_temperature_for_bedrock_registry_unsupported_model() {
let mut config = ModelConfig::new("openai.gpt-5.4");
config.temperature = Some(0.5);
let inference_config = bedrock_inference_config(&config);
assert!(!bedrock_model_supports_temperature(&config));
assert_eq!(inference_config.temperature(), None);
}
}