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:
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user