Filter unsupported reasoning from Bedrock request history
This commit is contained in:
@@ -1,8 +1,9 @@
|
||||
use async_trait::async_trait;
|
||||
use aws_sdk_bedrockruntime::Client as AwsBedrockClient;
|
||||
use galaxy_agent_core::{
|
||||
AgentError, AgentErrorKind, AgentEventStream, AgentRuntime, RuntimeCapabilities,
|
||||
RuntimeDescriptor, RuntimeKind, TurnControl, TurnRequest,
|
||||
AgentError, AgentErrorKind, AgentEventStream, AgentRuntime, ContentPart, ConversationMessage,
|
||||
MessageContent, MessageRole, RuntimeCapabilities, RuntimeDescriptor, RuntimeKind, TurnControl,
|
||||
TurnRequest,
|
||||
};
|
||||
use rig_bedrock::client::Client as RigBedrockClient;
|
||||
use rig_bedrock::completion::CompletionModel;
|
||||
@@ -33,6 +34,7 @@ pub struct BedrockRuntime {
|
||||
client: RigBedrockClient,
|
||||
config: BedrockRigConfig,
|
||||
resolved_model: String,
|
||||
reasoning_model: String,
|
||||
descriptor: RuntimeDescriptor,
|
||||
}
|
||||
|
||||
@@ -53,6 +55,7 @@ impl BedrockRuntime {
|
||||
Ok(Self {
|
||||
client: RigBedrockClient::from(client),
|
||||
config,
|
||||
reasoning_model: resolved_model.clone(),
|
||||
resolved_model,
|
||||
descriptor,
|
||||
})
|
||||
@@ -62,6 +65,13 @@ impl BedrockRuntime {
|
||||
&self.resolved_model
|
||||
}
|
||||
|
||||
/// Use the discovered foundation model's capabilities when invocation targets an
|
||||
/// opaque application inference profile rather than a foundation-model ID.
|
||||
pub fn with_reasoning_model(mut self, model: impl Into<String>) -> Self {
|
||||
self.reasoning_model = model.into();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn completion_model(&self) -> CompletionModel {
|
||||
let model = self.client.completion_model(&self.resolved_model);
|
||||
if self.config.prompt_caching {
|
||||
@@ -84,8 +94,11 @@ impl AgentRuntime for BedrockRuntime {
|
||||
control: TurnControl,
|
||||
) -> Result<AgentEventStream, AgentError> {
|
||||
let max_output_tokens = request.max_output_tokens.or(self.config.max_output_tokens);
|
||||
let mut completion_request =
|
||||
build_bedrock_completion_request(request, self.config.max_output_tokens)?;
|
||||
let mut completion_request = build_bedrock_completion_request_for_model(
|
||||
request,
|
||||
self.config.max_output_tokens,
|
||||
&self.reasoning_model,
|
||||
)?;
|
||||
// Context markers and inference-profile expansion are Galaxy model
|
||||
// configuration, not identifiers that Rig should send unchanged.
|
||||
completion_request.model = Some(self.resolved_model.clone());
|
||||
@@ -103,9 +116,74 @@ pub fn build_bedrock_completion_request(
|
||||
request: TurnRequest,
|
||||
configured_max_output_tokens: Option<u64>,
|
||||
) -> Result<CompletionRequest, AgentError> {
|
||||
let model = request.model.as_str().to_owned();
|
||||
build_bedrock_completion_request_for_model(request, configured_max_output_tokens, &model)
|
||||
}
|
||||
|
||||
fn build_bedrock_completion_request_for_model(
|
||||
mut request: TurnRequest,
|
||||
configured_max_output_tokens: Option<u64>,
|
||||
model: &str,
|
||||
) -> Result<CompletionRequest, AgentError> {
|
||||
request.messages = prepare_bedrock_history(model, request.messages);
|
||||
build_completion_request(request, configured_max_output_tokens, true, true, None)
|
||||
}
|
||||
|
||||
/// Project durable history into Bedrock's input format without changing the saved
|
||||
/// transcript. Reasoning output is not universally accepted as reasoning input:
|
||||
/// DeepSeek, Nova and other providers return unsigned, display-only reasoning.
|
||||
/// Claude tool continuations need their complete signed blocks unchanged.
|
||||
pub fn prepare_bedrock_history(
|
||||
model_id: &str,
|
||||
messages: Vec<ConversationMessage>,
|
||||
) -> Vec<ConversationMessage> {
|
||||
let supports_signed_reasoning = accepts_signed_reasoning(model_id);
|
||||
messages
|
||||
.into_iter()
|
||||
.filter_map(|mut message| {
|
||||
if let MessageContent::MultiPart(parts) = &mut message.content {
|
||||
parts.retain(|part| match part {
|
||||
ContentPart::Reasoning { signature, .. } => {
|
||||
message.role == MessageRole::Assistant
|
||||
&& supports_signed_reasoning
|
||||
&& signature
|
||||
.as_ref()
|
||||
.is_some_and(|signature| !signature.trim().is_empty())
|
||||
}
|
||||
ContentPart::Text(_)
|
||||
| ContentPart::Image { .. }
|
||||
| ContentPart::ToolUse { .. }
|
||||
| ContentPart::ToolResult { .. } => true,
|
||||
});
|
||||
if parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
Some(message)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn accepts_signed_reasoning(model_id: &str) -> bool {
|
||||
let model_id = model_id.trim();
|
||||
// An opaque application profile may point to Claude. Preserve its signed
|
||||
// blocks when no foundation-model hint is available; unsigned output is
|
||||
// still omitted. Production discovery supplies that hint to the runtime.
|
||||
if model_id.starts_with("arn:") && model_id.contains(":application-inference-profile/") {
|
||||
return true;
|
||||
}
|
||||
let model = strip_context_marker(model_id.rsplit('/').next().unwrap_or(model_id));
|
||||
let model = INFERENCE_PROFILE_PREFIXES
|
||||
.iter()
|
||||
.find_map(|prefix| model.strip_prefix(prefix))
|
||||
.unwrap_or(model);
|
||||
model.starts_with("anthropic.claude-")
|
||||
&& !model.starts_with("anthropic.claude-instant-")
|
||||
&& !model.starts_with("anthropic.claude-v1")
|
||||
&& !model.starts_with("anthropic.claude-v2")
|
||||
&& (!model.starts_with("anthropic.claude-3-") || model.starts_with("anthropic.claude-3-7-"))
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_model_id(
|
||||
configured_model: &str,
|
||||
region: &str,
|
||||
|
||||
@@ -177,7 +177,7 @@ async fn cancellation_before_bedrock_stream_start_never_contacts_aws() {
|
||||
#[test]
|
||||
fn bedrock_request_preserves_system_image_reasoning_tool_and_token_semantics() {
|
||||
let mut request = TurnRequest::new(
|
||||
"anthropic.claude-test",
|
||||
"anthropic.claude-sonnet-4-6",
|
||||
vec![
|
||||
ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
@@ -278,3 +278,151 @@ fn bedrock_request_preserves_system_image_reasoning_tool_and_token_semantics() {
|
||||
Some(ToolResultContent::Text(text)) if text.text == "[ERROR] permission denied"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_follow_ups_omit_unsigned_reasoning_without_losing_tool_calls_or_results() {
|
||||
for model in [
|
||||
"deepseek.r1-v1:0",
|
||||
"us.amazon.nova-2-lite-v1:0",
|
||||
"openai.gpt-oss-120b-1:0",
|
||||
"qwen.qwen3-32b-v1:0",
|
||||
"anthropic.claude-sonnet-4-6",
|
||||
] {
|
||||
let request = TurnRequest::new(
|
||||
model,
|
||||
vec![
|
||||
ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::Text("Inspect the project".to_owned()),
|
||||
},
|
||||
ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::MultiPart(vec![
|
||||
ContentPart::Reasoning {
|
||||
text: "display-only reasoning".to_owned(),
|
||||
signature: None,
|
||||
},
|
||||
ContentPart::Text("Reading the manifest".to_owned()),
|
||||
ContentPart::ToolUse {
|
||||
tool_use_id: "manifest-call".to_owned(),
|
||||
name: "read_files".to_owned(),
|
||||
input: serde_json::json!({"files": ["Cargo.toml"]}),
|
||||
},
|
||||
]),
|
||||
},
|
||||
ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::ToolResult {
|
||||
tool_use_id: "manifest-call".to_owned(),
|
||||
content: "manifest contents".to_owned(),
|
||||
is_error: false,
|
||||
},
|
||||
},
|
||||
],
|
||||
);
|
||||
let converted = build_bedrock_completion_request(request, None).unwrap();
|
||||
let messages = converted.chat_history.iter().collect::<Vec<_>>();
|
||||
let [
|
||||
_,
|
||||
Message::Assistant { content, .. },
|
||||
Message::User { content: result },
|
||||
] = messages.as_slice()
|
||||
else {
|
||||
panic!("expected the tool exchange for {model}");
|
||||
};
|
||||
assert!(matches!(content.iter().collect::<Vec<_>>().as_slice(),
|
||||
[AssistantContent::Text(text), AssistantContent::ToolCall(call)]
|
||||
if text.text == "Reading the manifest" && call.id == "manifest-call"
|
||||
));
|
||||
assert!(
|
||||
matches!(result.iter().next(), Some(UserContent::ToolResult(result))
|
||||
if result.id == "manifest-call")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_model_switch_drops_foreign_signed_reasoning_and_empty_messages() {
|
||||
for model in [
|
||||
"amazon.nova-pro-v1:0",
|
||||
"us.deepseek.r1-v1:0",
|
||||
"openai.gpt-oss-120b-1:0",
|
||||
"anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"anthropic.claude-instant-v1",
|
||||
"anthropic.claude-v1",
|
||||
"anthropic.claude-v2:1",
|
||||
] {
|
||||
let messages = vec![
|
||||
ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::Text("Question".to_owned()),
|
||||
},
|
||||
ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::MultiPart(vec![ContentPart::Reasoning {
|
||||
text: "Earlier thinking".to_owned(),
|
||||
signature: Some("original-signature".to_owned()),
|
||||
}]),
|
||||
},
|
||||
ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::MultiPart(vec![
|
||||
ContentPart::Reasoning {
|
||||
text: "Earlier thinking".to_owned(),
|
||||
signature: Some("original-signature".to_owned()),
|
||||
},
|
||||
ContentPart::Text("Visible answer".to_owned()),
|
||||
]),
|
||||
},
|
||||
];
|
||||
// The runtime's selected model takes precedence over a historical request model.
|
||||
let request = TurnRequest::new("anthropic.claude-sonnet-4-6", messages);
|
||||
let converted = build_bedrock_completion_request_for_model(request, None, model).unwrap();
|
||||
assert_eq!(converted.chat_history.len(), 2);
|
||||
assert!(
|
||||
matches!(converted.chat_history.iter().nth(1), Some(Message::Assistant { content, .. })
|
||||
if matches!(content.iter().collect::<Vec<_>>().as_slice(), [AssistantContent::Text(text)] if text.text == "Visible answer"))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_preserves_signed_claude_blocks_exactly_across_identifier_forms() {
|
||||
let signed = ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::MultiPart(vec![
|
||||
ContentPart::Reasoning {
|
||||
text: "Unmodified thinking".to_owned(),
|
||||
signature: Some("signed-text".to_owned()),
|
||||
},
|
||||
ContentPart::Reasoning {
|
||||
text: String::new(),
|
||||
signature: Some("signature-only".to_owned()),
|
||||
},
|
||||
ContentPart::Text("Answer".to_owned()),
|
||||
]),
|
||||
};
|
||||
for model in [
|
||||
"anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"global.anthropic.claude-sonnet-4-6[1m]",
|
||||
"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-sonnet-4-6",
|
||||
"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.anthropic.claude-sonnet-4-6",
|
||||
"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque-profile",
|
||||
] {
|
||||
assert_eq!(
|
||||
prepare_bedrock_history(model, vec![signed.clone()]),
|
||||
vec![signed.clone()]
|
||||
);
|
||||
}
|
||||
for signature in [None, Some(String::new()), Some(" ".to_owned())] {
|
||||
let message = ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::MultiPart(vec![ContentPart::Reasoning {
|
||||
text: "Display only".to_owned(),
|
||||
signature,
|
||||
}]),
|
||||
};
|
||||
assert!(prepare_bedrock_history("anthropic.claude-sonnet-4-6", vec![message]).is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user