1use std::collections::HashMap;
7use std::time::Duration;
8
9use async_trait::async_trait;
10use aws_sdk_bedrockruntime::Client;
11use aws_sdk_bedrockruntime::config::Region;
12use aws_sdk_bedrockruntime::types::{
13 CachePointBlock, CachePointType, CacheTtl as WireCacheTtl, ContentBlock, ConversationRole,
14 DocumentBlock, DocumentFormat, DocumentSource, ImageBlock, ImageFormat, ImageSource,
15 InferenceConfiguration, JsonSchemaDefinition, Message, OutputConfig, OutputFormat,
16 OutputFormatStructure, OutputFormatType, ReasoningContentBlock, ReasoningTextBlock,
17 SpecificToolChoice, StopReason, SystemContentBlock, Tool, ToolChoice, ToolConfiguration,
18 ToolInputSchema, ToolResultBlock, ToolResultContentBlock, ToolResultStatus, ToolSpecification,
19 ToolUseBlock,
20};
21use aws_smithy_types::{Blob, Document, Number};
22use serde_json::{Value, json};
23
24use super::{
25 AnthropicBudget, CacheTtl, Completion, ModelError, ModelId, ModelProvider,
26 ProviderContinuation, Request, SchemaMode, ToolDeclaration, ToolExchange, Usage,
27};
28
29const RESPOND_TOOL: &str = "__agentplane_respond";
30
31fn decoded_media(
32 block: &Value,
33 expected_type: &str,
34 model: &ModelId,
35) -> Result<(String, Vec<u8>), ModelError> {
36 let actual = block
37 .get("type")
38 .and_then(Value::as_str)
39 .unwrap_or_default();
40 if actual != expected_type {
41 return Err(Bedrock::refused(
42 model,
43 format!("unsupported Bedrock content block '{actual}'"),
44 ));
45 }
46 let media_type = block
47 .get("media_type")
48 .and_then(Value::as_str)
49 .ok_or_else(|| Bedrock::refused(model, "Bedrock media block has no media_type"))?;
50 let data = block
51 .get("data")
52 .and_then(Value::as_str)
53 .ok_or_else(|| Bedrock::refused(model, "Bedrock media block has no inline base64 data"))?;
54 let bytes = crate::core::b64::decode(data)
55 .ok_or_else(|| Bedrock::refused(model, "Bedrock media data is not canonical base64"))?;
56 Ok((media_type.to_owned(), bytes))
57}
58
59pub(crate) fn content_from_prompt_json(
60 block: &Value,
61 model: &ModelId,
62) -> Result<ContentBlock, ModelError> {
63 if let Some(text) = block.get("text").and_then(Value::as_str) {
64 return Ok(ContentBlock::Text(text.to_owned()));
65 }
66 match block.get("type").and_then(Value::as_str) {
67 Some("image") => {
68 let (media_type, bytes) = decoded_media(block, "image", model)?;
69 let format = match media_type.as_str() {
70 "image/gif" => ImageFormat::Gif,
71 "image/jpeg" => ImageFormat::Jpeg,
72 "image/png" => ImageFormat::Png,
73 "image/webp" => ImageFormat::Webp,
74 other => {
75 return Err(Bedrock::refused(
76 model,
77 format!("Bedrock does not support image media type '{other}'"),
78 ));
79 }
80 };
81 ImageBlock::builder()
82 .format(format)
83 .source(ImageSource::Bytes(Blob::new(bytes)))
84 .build()
85 .map(ContentBlock::Image)
86 .map_err(|error| Bedrock::refused(model, error.to_string()))
87 }
88 Some("document") => {
89 let (media_type, bytes) = decoded_media(block, "document", model)?;
90 let format = match media_type.as_str() {
91 "text/csv" => DocumentFormat::Csv,
92 "application/msword" => DocumentFormat::Doc,
93 "application/vnd.openxmlformats-officedocument.wordprocessingml.document" => {
94 DocumentFormat::Docx
95 }
96 "text/html" => DocumentFormat::Html,
97 "text/markdown" => DocumentFormat::Md,
98 "application/pdf" => DocumentFormat::Pdf,
99 "text/plain" => DocumentFormat::Txt,
100 "application/vnd.ms-excel" => DocumentFormat::Xls,
101 "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" => {
102 DocumentFormat::Xlsx
103 }
104 other => {
105 return Err(Bedrock::refused(
106 model,
107 format!("Bedrock does not support document media type '{other}'"),
108 ));
109 }
110 };
111 DocumentBlock::builder()
112 .format(format)
113 .name("document")
116 .source(DocumentSource::Bytes(Blob::new(bytes)))
117 .build()
118 .map(ContentBlock::Document)
119 .map_err(|error| Bedrock::refused(model, error.to_string()))
120 }
121 Some(other) => Err(Bedrock::refused(
122 model,
123 format!("unsupported Bedrock content block '{other}'"),
124 )),
125 None => Err(Bedrock::refused(
126 model,
127 "Bedrock content block has neither text nor a supported type",
128 )),
129 }
130}
131
132#[derive(Clone)]
153pub struct Bedrock {
154 client: Client,
155 region: String,
156 timeout: Duration,
157 schema_mode: SchemaMode,
158 stream: bool,
159 guardrail: Option<Guardrail>,
169 reasoning: Option<ReasoningDialect>,
174}
175
176#[derive(Debug, Clone, Copy, PartialEq, Eq)]
196#[non_exhaustive]
197pub enum ReasoningDialect {
198 Nova,
221}
222
223#[derive(Debug, Clone, PartialEq, Eq)]
225pub struct Guardrail {
226 pub identifier: String,
227 pub version: String,
231 pub trace: bool,
238}
239
240impl Guardrail {
241 #[must_use]
242 pub fn new(identifier: impl Into<String>, version: impl Into<String>) -> Self {
243 Self {
244 identifier: identifier.into(),
245 version: version.into(),
246 trace: false,
247 }
248 }
249
250 #[must_use]
253 pub const fn with_trace(mut self) -> Self {
254 self.trace = true;
255 self
256 }
257
258 fn stream_config(&self) -> aws_sdk_bedrockruntime::types::GuardrailStreamConfiguration {
267 use aws_sdk_bedrockruntime::types::{
268 GuardrailStreamConfiguration, GuardrailStreamProcessingMode, GuardrailTrace,
269 };
270 GuardrailStreamConfiguration::builder()
271 .guardrail_identifier(&self.identifier)
272 .guardrail_version(&self.version)
273 .trace(if self.trace {
274 GuardrailTrace::Enabled
275 } else {
276 GuardrailTrace::Disabled
277 })
278 .stream_processing_mode(GuardrailStreamProcessingMode::Sync)
279 .build()
280 }
281
282 fn config(&self) -> aws_sdk_bedrockruntime::types::GuardrailConfiguration {
283 use aws_sdk_bedrockruntime::types::{GuardrailConfiguration, GuardrailTrace};
284 GuardrailConfiguration::builder()
285 .guardrail_identifier(&self.identifier)
286 .guardrail_version(&self.version)
287 .trace(if self.trace {
288 GuardrailTrace::Enabled
289 } else {
290 GuardrailTrace::Disabled
291 })
292 .build()
293 }
294}
295
296impl std::fmt::Debug for Bedrock {
297 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
298 f.debug_struct("Bedrock")
299 .field("region", &self.region)
300 .field("timeout", &self.timeout)
301 .field("schema_mode", &self.schema_mode)
302 .field("stream", &self.stream)
303 .field("guardrail", &self.guardrail)
304 .finish_non_exhaustive()
305 }
306}
307
308impl Bedrock {
309 pub async fn from_env(region: impl Into<String>) -> Result<Self, String> {
325 let region = region.into();
326 if region.trim().is_empty() {
327 return Err("an AWS region is required for Bedrock".to_owned());
328 }
329 let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
330 .region(Region::new(region))
331 .load()
332 .await;
333 Self::from_client(Client::new(&config))
334 }
335
336 pub fn from_client(client: Client) -> Result<Self, String> {
354 let region = client
355 .config()
356 .region()
357 .map(|region| region.as_ref().trim().to_owned())
358 .filter(|region| !region.is_empty())
359 .ok_or_else(|| {
360 "the Bedrock client carries no region, so there is nothing true to put in the \
361 request profile — build it from a config with `.region(..)` set"
362 .to_owned()
363 })?;
364 Ok(Self {
365 client,
366 region,
367 timeout: Duration::from_mins(5),
368 schema_mode: SchemaMode::Native,
369 stream: true,
370 guardrail: None,
371 reasoning: None,
372 })
373 }
374
375 #[must_use]
385 pub fn guardrail(mut self, guardrail: Guardrail) -> Self {
386 self.guardrail = Some(guardrail);
387 self
388 }
389
390 #[must_use]
401 pub const fn reasoning(mut self, dialect: ReasoningDialect) -> Self {
402 self.reasoning = Some(dialect);
403 self
404 }
405
406 fn reasoning_config(
422 &self,
423 model: &ModelId,
424 effort: Option<super::ReasoningEffort>,
425 ) -> Result<Option<Document>, ModelError> {
426 use super::ReasoningEffort as E;
427
428 let Some(effort) = effort else {
429 return Ok(None);
430 };
431 let Some(dialect) = self.reasoning else {
432 return Err(Self::refused(
433 model,
434 "Bedrock Converse has no provider-neutral reasoning-effort mapping, because one \
435 envelope covers families that spell reasoning differently or not at all. \
436 Declare which one this driver speaks — `Bedrock::from_client(..).reasoning(\
437 ReasoningDialect::Nova)` — rather than having the driver guess and silently \
438 change the effort the manifest declared",
439 ));
440 };
441 match dialect {
442 ReasoningDialect::Nova => {
443 let level = match effort {
450 E::None => {
451 return Ok(Some(document_from_json(
452 &json!({ "reasoningConfig": { "type": "disabled" } }),
453 )));
454 }
455 E::Low => "low",
456 E::Medium => "medium",
457 E::High => "high",
458 E::Minimal | E::XHigh | E::Max => {
459 return Err(Self::refused(
460 model,
461 format!(
462 "Amazon Nova extended thinking has no counterpart for reasoning \
463 effort '{}' — it accepts low, medium and high. Declare one of \
464 those, or none to disable reasoning explicitly",
465 effort.as_str()
466 ),
467 ));
468 }
469 };
470 Ok(Some(document_from_json(&json!({
471 "reasoningConfig": { "type": "enabled", "maxReasoningEffort": level },
472 }))))
473 }
474 }
475 }
476
477 #[must_use]
479 pub const fn timeout(mut self, timeout: Duration) -> Self {
480 self.timeout = timeout;
481 self
482 }
483
484 #[must_use]
486 pub const fn structured_via(mut self, mode: SchemaMode) -> Self {
487 self.schema_mode = mode;
488 self
489 }
490
491 #[must_use]
493 pub const fn buffered(mut self) -> Self {
494 self.stream = false;
495 self
496 }
497
498 fn refused(model: &ModelId, detail: impl Into<String>) -> ModelError {
499 ModelError::Refused {
500 account: false,
501 model: model.clone(),
502 code: None,
503 detail: detail.into(),
504 }
505 }
506
507 fn messages(
508 model: &ModelId,
509 prompt: &Value,
510 exchanges: &[ToolExchange],
511 told: &[String],
512 continuation: Option<&ProviderContinuation>,
513 ) -> Result<Vec<Message>, ModelError> {
514 let mut messages = match crate::model::prompt::turns(prompt, &["messages", "input"]) {
515 crate::model::prompt::Turns::Native(items) => items
516 .iter()
517 .map(|item| Self::message_from_value(model, item))
518 .collect::<Result<Vec<_>, _>>()?,
519 crate::model::prompt::Turns::Parts(parts) => {
520 let mut content = Vec::with_capacity(parts.len());
521 for part in &parts {
522 content.push(ContentBlock::Text(part.text.clone().into_owned()));
523 if let Some(ttl) = part.cache {
524 content.push(ContentBlock::CachePoint(cache_point(ttl, model)?));
525 }
526 }
527 vec![
528 Message::builder()
529 .role(ConversationRole::User)
530 .set_content(Some(content))
531 .build()
532 .map_err(|error| Self::refused(model, error.to_string()))?,
533 ]
534 }
535 };
536
537 if exchanges.is_empty() && told.is_empty() {
538 return Ok(messages);
539 }
540 let prompt_len = messages.len();
541
542 if let Some(state) = continuation {
543 if state.provider != "bedrock" {
544 return Err(Self::refused(
545 model,
546 format!(
547 "a '{}' continuation cannot be sent to Bedrock",
548 state.provider
549 ),
550 ));
551 }
552 let transcript = state
553 .state
554 .as_array()
555 .ok_or_else(|| Self::refused(model, "the Bedrock continuation is not an array"))?
556 .iter()
557 .map(|message| continuation_message_from_json(message, model))
558 .collect::<Result<Vec<_>, _>>()?;
559 messages.extend(transcript);
560 } else if !exchanges.is_empty() {
561 let assistant = exchanges
562 .iter()
563 .map(|exchange| {
564 Ok(ContentBlock::ToolUse(
565 ToolUseBlock::builder()
566 .tool_use_id(exchange.call.id.clone())
567 .name(exchange.call.name.clone())
568 .input(document_from_json(&exchange.call.arguments))
569 .build()
570 .map_err(|error| Self::refused(model, error.to_string()))?,
571 ))
572 })
573 .collect::<Result<Vec<_>, ModelError>>()?;
574 messages.push(
575 Message::builder()
576 .role(ConversationRole::Assistant)
577 .set_content(Some(assistant))
578 .build()
579 .map_err(|error| Self::refused(model, error.to_string()))?,
580 );
581 }
582
583 messages.push(continuation_message_from_json(
584 &user_turn(exchanges, told),
585 model,
586 )?);
587 join_at_the_prompt(messages, prompt_len, model)
588 }
589
590 fn message_from_value(model: &ModelId, value: &Value) -> Result<Message, ModelError> {
591 let object = value
592 .as_object()
593 .ok_or_else(|| Self::refused(model, "a Bedrock message must be an object"))?;
594 let role = match object.get("role").and_then(Value::as_str) {
595 Some("user") => ConversationRole::User,
596 Some("assistant") => ConversationRole::Assistant,
597 Some(other) => {
598 return Err(Self::refused(
599 model,
600 format!("unsupported Bedrock message role '{other}'"),
601 ));
602 }
603 None => return Err(Self::refused(model, "a Bedrock message has no role")),
604 };
605 let content = object
606 .get("content")
607 .ok_or_else(|| Self::refused(model, "a Bedrock message has no content"))?;
608 match content {
609 Value::String(text) => Self::text_message(role, text, model),
610 Value::Array(blocks) => {
611 let blocks = blocks
612 .iter()
613 .map(|block| content_from_prompt_json(block, model))
614 .collect::<Result<Vec<_>, ModelError>>()?;
615 Message::builder()
616 .role(role)
617 .set_content(Some(blocks))
618 .build()
619 .map_err(|error| Self::refused(model, error.to_string()))
620 }
621 _ => Err(Self::refused(
622 model,
623 "Bedrock message content must be text or text blocks",
624 )),
625 }
626 }
627
628 fn text_message(
629 role: ConversationRole,
630 text: &str,
631 model: &ModelId,
632 ) -> Result<Message, ModelError> {
633 Message::builder()
634 .role(role)
635 .content(ContentBlock::Text(text.to_owned()))
636 .build()
637 .map_err(|error| Self::refused(model, error.to_string()))
638 }
639
640 fn tool_config(
641 model: &ModelId,
642 schema: Option<&Value>,
643 tools: &[ToolDeclaration],
644 mode: SchemaMode,
645 ) -> Result<Option<ToolConfiguration>, ModelError> {
646 if schema.is_some() && mode == SchemaMode::ForcedTool && !tools.is_empty() {
647 return Err(Self::refused(
648 model,
649 "Bedrock Converse obtains structured output by forcing a synthetic tool, so a \
650 response schema cannot be combined with callable tools",
651 ));
652 }
653 let mut builder = ToolConfiguration::builder();
654 let mut budget = AnthropicBudget::default();
658 if let Some(schema) = schema.filter(|_| mode == SchemaMode::ForcedTool) {
659 builder = builder
660 .tools(tool(
661 RESPOND_TOOL,
662 "Return the answer in the required shape.",
663 schema,
664 budget.tool(schema),
665 model,
666 )?)
667 .tool_choice(ToolChoice::Tool(
668 SpecificToolChoice::builder()
669 .name(RESPOND_TOOL)
670 .build()
671 .map_err(|error| Self::refused(model, error.to_string()))?,
672 ));
673 } else {
674 for declaration in tools {
675 builder = builder.tools(tool(
676 &declaration.name,
677 &declaration.description,
678 &declaration.parameters,
679 budget.tool(&declaration.parameters),
680 model,
681 )?);
682 if let Some(ttl) = declaration.cache {
683 builder = builder.tools(Tool::CachePoint(cache_point(ttl, model)?));
684 }
685 }
686 }
687 if (schema.is_none() || mode == SchemaMode::Native) && tools.is_empty() {
688 Ok(None)
689 } else {
690 builder
691 .build()
692 .map(Some)
693 .map_err(|error| Self::refused(model, error.to_string()))
694 }
695 }
696
697 fn output_config(
698 model: &ModelId,
699 schema: Option<&Value>,
700 mode: SchemaMode,
701 ) -> Result<Option<OutputConfig>, ModelError> {
702 let Some(schema) = schema.filter(|_| mode == SchemaMode::Native) else {
703 return Ok(None);
704 };
705 let definition = JsonSchemaDefinition::builder()
706 .schema(schema.to_string())
707 .name("agentplane_response")
708 .build()
709 .map_err(|error| Self::refused(model, error.to_string()))?;
710 let format = OutputFormat::builder()
711 .r#type(OutputFormatType::JsonSchema)
712 .structure(OutputFormatStructure::JsonSchema(definition))
713 .build()
714 .map_err(|error| Self::refused(model, error.to_string()))?;
715 Ok(Some(OutputConfig::builder().text_format(format).build()))
716 }
717
718 #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
719 async fn complete_streamed(
720 &self,
721 model: &ModelId,
722 prompt: &Value,
723 messages: Vec<Message>,
724 max_tokens: i32,
725 tools: Option<ToolConfiguration>,
726 output_config: Option<OutputConfig>,
727 reasoning_config: Option<Document>,
728 schema: Option<&Value>,
729 observer: Option<(&dyn super::ModelStreamObserver, &crate::core::Label)>,
730 stop: Option<&super::Stop>,
731 ) -> Result<Completion, ModelError> {
732 let deadline = tokio::time::Instant::now() + self.timeout;
733 let mut operation = self
734 .client
735 .converse_stream()
736 .model_id(&model.model)
737 .set_messages(Some(messages))
738 .inference_config(
739 InferenceConfiguration::builder()
740 .max_tokens(max_tokens)
741 .build(),
742 )
743 .set_tool_config(tools)
744 .set_output_config(output_config)
745 .set_guardrail_config(self.guardrail.as_ref().map(Guardrail::stream_config))
749 .set_additional_model_request_fields(reasoning_config);
750 let system = system_blocks(prompt, model)?;
751 if !system.is_empty() {
752 operation = operation.set_system(Some(system));
753 }
754 let output = tokio::time::timeout_at(deadline, operation.send())
755 .await
756 .map_err(|_| ModelError::Unavailable {
757 model: model.clone(),
758 detail: "Bedrock did not start the stream before the request deadline".to_owned(),
759 })?
760 .map_err(|error| {
761 classify_sdk(model, &error, |service, detail| {
762 classify_service(
763 model,
764 detail,
765 aws_sdk_bedrockruntime::error::ProvideErrorMetadata::code(service),
766 service.is_throttling_exception() || service.is_model_not_ready_exception(),
767 service.is_access_denied_exception()
768 || service.is_resource_not_found_exception()
769 || service.is_validation_exception(),
770 )
771 })
772 })?;
773
774 let mut stream = output.stream;
775 let mut accumulator = super::bedrock_stream::Accumulator::new();
776 loop {
777 let next = tokio::time::timeout_at(deadline, stream.recv());
778 let Some(next) = super::unless_stopped(next, stop).await else {
780 return Ok(Completion::stopped(accumulator.usage().unwrap_or_default()));
781 };
782 let event = next
783 .map_err(|_| severed_stream(model, &accumulator, "the request deadline elapsed"))?
784 .map_err(|error| {
785 let detail = error.to_string();
786 classify_stream_event(model, &accumulator, error.as_service_error(), &detail)
787 })?;
788 let Some(event) = event else {
789 break;
790 };
791 if let aws_sdk_bedrockruntime::types::ConverseStreamOutput::ContentBlockDelta(delta) =
792 &event
793 && let Some(aws_sdk_bedrockruntime::types::ContentBlockDelta::Text(text)) =
794 delta.delta()
795 && let Some((observer, label)) = observer
796 {
797 observer.event(crate::core::Tainted::with_label(
798 super::ModelStreamEvent::TextDelta(text.clone()),
799 label.clone(),
800 ));
801 }
802 accumulator.event(event);
803 }
804
805 if accumulator.stop_reason()
806 == Some(&aws_sdk_bedrockruntime::types::StopReason::GuardrailIntervened)
807 {
808 return Err(ModelError::Unusable {
814 model: model.clone(),
815 usage: accumulator.usage().unwrap_or_default(),
819 detail: "a Bedrock guardrail intervened on this call".to_owned(),
820 });
821 }
822 let Some(stop_reason) = accumulator.stop_reason().cloned() else {
823 return Err(severed_stream(
824 model,
825 &accumulator,
826 "the stream ended before messageStop",
827 ));
828 };
829 let Some(usage) = accumulator.usage() else {
830 return Err(severed_stream(
831 model,
832 &accumulator,
833 "the stream ended without usage metadata",
834 ));
835 };
836 if let Some((observer, label)) = observer {
837 observer.event(crate::core::Tainted::with_label(
838 super::ModelStreamEvent::Usage(usage),
839 label.clone(),
840 ));
841 }
842 let blocks = accumulator
843 .finish()
844 .map_err(|detail| ModelError::Unusable {
845 model: model.clone(),
846 usage,
847 detail,
848 })?;
849 let message = Message::builder()
850 .role(ConversationRole::Assistant)
851 .set_content(Some(blocks))
852 .build()
853 .map_err(|error| ModelError::Unusable {
854 model: model.clone(),
855 usage,
856 detail: error.to_string(),
857 })?;
858 let token_usage = aws_sdk_bedrockruntime::types::TokenUsage::builder()
859 .input_tokens(i32::try_from(usage.input_tokens).unwrap_or(i32::MAX))
860 .output_tokens(i32::try_from(usage.output_tokens).unwrap_or(i32::MAX))
861 .total_tokens(
862 i32::try_from(usage.input_tokens.saturating_add(usage.output_tokens))
863 .unwrap_or(i32::MAX),
864 )
865 .cache_write_input_tokens(i32::try_from(usage.cache_write_tokens).unwrap_or(i32::MAX))
866 .cache_read_input_tokens(i32::try_from(usage.cache_read_tokens).unwrap_or(i32::MAX))
867 .build()
868 .map_err(|error| ModelError::Unusable {
869 model: model.clone(),
870 usage,
871 detail: error.to_string(),
872 })?;
873 let assembled = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
874 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
875 message,
876 ))
877 .stop_reason(stop_reason)
878 .usage(token_usage)
879 .build()
880 .map_err(|error| ModelError::Unusable {
881 model: model.clone(),
882 usage,
883 detail: error.to_string(),
884 })?;
885 Self::interpret(model, schema, self.schema_mode, &assembled)
886 }
887
888 #[allow(clippy::too_many_lines)]
889 fn interpret(
890 model: &ModelId,
891 schema: Option<&Value>,
892 mode: SchemaMode,
893 output: &aws_sdk_bedrockruntime::operation::converse::ConverseOutput,
894 ) -> Result<Completion, ModelError> {
895 let usage = output.usage().map_or_else(Usage::default, usage_from);
900 let truncated = match output.stop_reason() {
911 StopReason::EndTurn | StopReason::ToolUse | StopReason::StopSequence => false,
912 StopReason::MaxTokens | StopReason::ModelContextWindowExceeded => true,
913 other => {
914 return Err(ModelError::Unusable {
915 model: model.clone(),
916 usage,
917 detail: format!(
918 "Bedrock stopped with '{}', which is not an answer: no fragment of \
919 what the model emitted is usable output",
920 other.as_str()
921 ),
922 });
923 }
924 };
925 let message = output
926 .output()
927 .and_then(|value| value.as_message().ok())
928 .ok_or_else(|| ModelError::Unusable {
929 model: model.clone(),
930 usage,
931 detail: "Bedrock returned no message output".to_owned(),
932 })?;
933
934 let mut text = String::new();
935 let mut calls = Vec::new();
936 let mut state = Vec::new();
937 let mut forced = None;
938 for block in message.content() {
939 state.push(content_to_json(block, model, usage)?);
940 match block {
941 ContentBlock::Text(value) => text.push_str(value),
942 ContentBlock::ToolUse(value) if value.name() == RESPOND_TOOL => {
943 forced = Some(json_from_document(value.input()));
944 }
945 ContentBlock::ToolUse(value) => calls.push(super::ToolCall {
946 id: value.tool_use_id().to_owned(),
947 name: value.name().to_owned(),
948 arguments: json_from_document(value.input()),
949 }),
950 ContentBlock::ReasoningContent(_) => {}
951 _ => {
952 return Err(ModelError::Unusable {
953 model: model.clone(),
954 usage,
955 detail: "Bedrock returned a content block this driver cannot preserve"
956 .to_owned(),
957 });
958 }
959 }
960 }
961
962 let emulating = schema.is_some() && mode == SchemaMode::ForcedTool;
963 let structured = if emulating {
964 let value = forced.ok_or_else(|| ModelError::Unusable {
965 model: model.clone(),
966 usage,
967 detail: "Bedrock did not honor the forced structured-output tool".to_owned(),
968 })?;
969 if let Some(schema) = schema {
970 super::validate_schema(schema, &value).map_err(|detail| ModelError::Unusable {
971 model: model.clone(),
972 usage,
973 detail,
974 })?;
975 }
976 text = value.to_string();
977 Some(value)
978 } else if schema.is_some() && calls.is_empty() {
979 let value = serde_json::from_str(&text).map_err(|error| ModelError::Unusable {
980 model: model.clone(),
981 usage,
982 detail: format!("Bedrock native structured output was not JSON: {error}"),
983 })?;
984 super::validate_schema(schema.expect("checked above"), &value).map_err(|detail| {
985 ModelError::Unusable {
986 model: model.clone(),
987 usage,
988 detail,
989 }
990 })?;
991 Some(value)
992 } else {
993 None
994 };
995 if text.is_empty() && calls.is_empty() && structured.is_none() {
996 return Err(ModelError::Unusable {
997 model: model.clone(),
998 usage,
999 detail: "Bedrock returned neither text nor a tool call".to_owned(),
1000 });
1001 }
1002 let stop_reason = output.stop_reason().as_str().to_owned();
1003 let continuation = (!emulating).then(|| {
1006 ProviderContinuation::new(
1007 "bedrock",
1008 json!([{ "role": "assistant", "content": state }]),
1009 )
1010 });
1011 Ok(Completion {
1012 stopped: false,
1013 text,
1014 tool_calls: calls,
1015 model: None,
1019 usage,
1020 stop_reason: Some(stop_reason),
1021 truncated,
1022 structured,
1023 continuation,
1024 reasoning: None,
1025 commentary: None,
1026 compacted: None,
1027 })
1028 }
1029}
1030
1031#[async_trait]
1032impl ModelProvider for Bedrock {
1033 fn gen_ai_provider(&self) -> Option<&str> {
1034 Some("aws.bedrock")
1035 }
1036 fn request_profile(&self, _model: &ModelId) -> Value {
1037 json!({
1038 "driver": "aws-bedrock-converse/v1",
1039 "region": self.region,
1040 "stream": self.stream,
1041 "schema_mode": match self.schema_mode {
1042 SchemaMode::Native => "native",
1043 SchemaMode::ForcedTool => "forced-tool",
1044 },
1045 "guardrail": self.guardrail.as_ref().map(|g| json!({
1052 "id": g.identifier,
1053 "version": g.version,
1054 })),
1055 "reasoning_dialect": self.reasoning.map(|d| match d {
1061 ReasoningDialect::Nova => "nova",
1062 }),
1063 })
1064 }
1065
1066 async fn complete(&self, request: Request<'_>) -> Result<Completion, ModelError> {
1067 let Request {
1068 model,
1069 prompt,
1070 max_output_tokens,
1071 reasoning_effort,
1072 reasoning_summary: _,
1075 schema,
1076 tools,
1077 exchanges,
1078 told,
1079 continuation,
1080 compaction,
1081 stream,
1082 stop,
1083 } = request;
1084 super::refuse_provider_side_media(prompt, model)?;
1085 super::refuse_in_thread_instructions(prompt, model)?;
1086 if compaction.is_some() {
1089 return Err(Self::refused(
1090 model,
1091 "Bedrock Converse has no server-side compaction",
1092 ));
1093 }
1094 let reasoning_config = self.reasoning_config(model, reasoning_effort)?;
1095 let max_tokens = i32::try_from(max_output_tokens)
1096 .map_err(|_| Self::refused(model, "max_output_tokens exceeds Bedrock's i32 limit"))?;
1097 let messages = Self::messages(model, prompt, exchanges, told, continuation)?;
1098 let tools = Self::tool_config(model, schema, tools, self.schema_mode)?;
1099 let output_config = Self::output_config(model, schema, self.schema_mode)?;
1100 if self.stream {
1101 let streamed = self
1102 .complete_streamed(
1103 model,
1104 prompt,
1105 messages,
1106 max_tokens,
1107 tools,
1108 output_config,
1109 reasoning_config,
1110 schema,
1111 stream,
1112 stop,
1113 )
1114 .await?;
1115 let mut completion = streamed;
1116 if !completion.stopped {
1117 accumulate_continuation(&mut completion, continuation, exchanges, told);
1118 }
1119 return Ok(completion);
1120 }
1121 let mut operation = self
1122 .client
1123 .converse()
1124 .model_id(&model.model)
1125 .set_messages(Some(messages))
1126 .inference_config(
1127 InferenceConfiguration::builder()
1128 .max_tokens(max_tokens)
1129 .build(),
1130 )
1131 .set_tool_config(tools)
1132 .set_output_config(output_config)
1133 .set_guardrail_config(self.guardrail.as_ref().map(Guardrail::config))
1137 .set_additional_model_request_fields(reasoning_config);
1138 let system = system_blocks(prompt, model)?;
1139 if !system.is_empty() {
1140 operation = operation.set_system(Some(system));
1141 }
1142
1143 let result = tokio::time::timeout(self.timeout, operation.send())
1144 .await
1145 .map_err(|_| ModelError::Unavailable {
1146 model: model.clone(),
1147 detail: format!("Bedrock did not complete within {:?}", self.timeout),
1148 })?;
1149 let output = result.map_err(|error| {
1150 classify_sdk(model, &error, |service, detail| {
1151 classify_service(
1152 model,
1153 detail,
1154 aws_sdk_bedrockruntime::error::ProvideErrorMetadata::code(service),
1155 service.is_throttling_exception() || service.is_model_not_ready_exception(),
1156 service.is_access_denied_exception()
1157 || service.is_resource_not_found_exception()
1158 || service.is_validation_exception(),
1159 )
1160 })
1161 })?;
1162 let mut completion = Self::interpret(model, schema, self.schema_mode, &output)?;
1163 if let Some((observer, label)) = stream {
1164 observer.event(crate::core::Tainted::with_label(
1165 super::ModelStreamEvent::Usage(completion.usage),
1166 label.clone(),
1167 ));
1168 }
1169 accumulate_continuation(&mut completion, continuation, exchanges, told);
1170 Ok(completion)
1171 }
1172}
1173
1174fn tool_result_json(exchange: &ToolExchange) -> Value {
1175 json!({
1176 "type": "tool_result",
1177 "tool_use_id": exchange.call.id,
1178 "output": exchange.output,
1179 "failed": exchange.failed,
1180 })
1181}
1182
1183fn join_at_the_prompt(
1187 mut messages: Vec<Message>,
1188 prompt_len: usize,
1189 model: &ModelId,
1190) -> Result<Vec<Message>, ModelError> {
1191 let user = |m: &Message| m.role() == &ConversationRole::User;
1192 if prompt_len == 0
1193 || prompt_len >= messages.len()
1194 || !user(&messages[prompt_len - 1])
1195 || !user(&messages[prompt_len])
1196 {
1197 return Ok(messages);
1198 }
1199 let next = messages.remove(prompt_len);
1200 let last = &messages[prompt_len - 1];
1201 let mut content = last.content().to_vec();
1202 content.extend(next.content().iter().cloned());
1203 messages[prompt_len - 1] = Message::builder()
1204 .role(ConversationRole::User)
1205 .set_content(Some(content))
1206 .build()
1207 .map_err(|error| Bedrock::refused(model, error.to_string()))?;
1208 Ok(messages)
1209}
1210
1211fn user_turn(exchanges: &[ToolExchange], told: &[String]) -> Value {
1214 json!({
1215 "role": "user",
1216 "content": exchanges
1217 .iter()
1218 .map(tool_result_json)
1219 .chain(told.iter().map(|text| json!({ "type": "text", "text": text })))
1220 .collect::<Vec<_>>(),
1221 })
1222}
1223
1224fn accumulate_continuation(
1225 completion: &mut Completion,
1226 prior: Option<&ProviderContinuation>,
1227 exchanges: &[ToolExchange],
1228 told: &[String],
1229) {
1230 let Some(current) = completion.continuation.as_mut() else {
1231 return;
1232 };
1233 let mut transcript = prior
1234 .and_then(|value| value.state.as_array())
1235 .cloned()
1236 .unwrap_or_default();
1237 if !exchanges.is_empty() || !told.is_empty() {
1238 transcript.push(user_turn(exchanges, told));
1239 }
1240 if let Some(messages) = current.state.as_array() {
1241 transcript.extend(messages.iter().cloned());
1242 }
1243 current.state = Value::Array(transcript);
1244}
1245
1246fn cache_point(ttl: CacheTtl, model: &ModelId) -> Result<CachePointBlock, ModelError> {
1249 CachePointBlock::builder()
1250 .r#type(CachePointType::Default)
1251 .ttl(match ttl {
1252 CacheTtl::FiveMinutes => WireCacheTtl::FiveMinutes,
1253 CacheTtl::OneHour => WireCacheTtl::OneHour,
1254 })
1255 .build()
1256 .map_err(|error| Bedrock::refused(model, error.to_string()))
1257}
1258
1259fn system_blocks(prompt: &Value, model: &ModelId) -> Result<Vec<SystemContentBlock>, ModelError> {
1262 let Some(system) = crate::model::prompt::system(prompt) else {
1263 return Ok(Vec::new());
1264 };
1265 let crate::model::prompt::System::Parts(parts) = system else {
1266 return Err(Bedrock::refused(
1267 model,
1268 "Bedrock's system instruction must be a string or a list of text parts",
1269 ));
1270 };
1271 let mut blocks = Vec::with_capacity(parts.len());
1272 for part in parts {
1273 blocks.push(SystemContentBlock::Text(part.text.into_owned()));
1274 if let Some(ttl) = part.cache {
1275 blocks.push(SystemContentBlock::CachePoint(cache_point(ttl, model)?));
1276 }
1277 }
1278 Ok(blocks)
1279}
1280
1281pub(crate) fn usage_from(usage: &aws_sdk_bedrockruntime::types::TokenUsage) -> Usage {
1284 let count = |n: i32| u64::try_from(n).unwrap_or_default();
1285 let one_hour = usage
1286 .cache_details()
1287 .iter()
1288 .filter(|detail| *detail.ttl() == WireCacheTtl::OneHour)
1289 .map(|detail| count(detail.input_tokens()))
1290 .sum::<u64>();
1291 Usage::with_cache_beside_input(
1292 count(usage.input_tokens()),
1293 count(usage.output_tokens()),
1294 count(usage.cache_write_input_tokens().unwrap_or_default()),
1295 count(usage.cache_read_input_tokens().unwrap_or_default()),
1296 )
1297 .with_cache_write_1h(one_hour)
1298}
1299
1300fn tool(
1301 name: &str,
1302 description: &str,
1303 schema: &Value,
1304 strict: bool,
1305 model: &ModelId,
1306) -> Result<Tool, ModelError> {
1307 let specification = ToolSpecification::builder()
1308 .name(name)
1309 .description(description)
1310 .input_schema(ToolInputSchema::Json(document_from_json(schema)))
1311 .strict(strict)
1312 .build()
1313 .map_err(|error| Bedrock::refused(model, error.to_string()))?;
1314 Ok(Tool::ToolSpec(specification))
1315}
1316
1317fn classify_sdk<E, R>(
1331 model: &ModelId,
1332 error: &aws_sdk_bedrockruntime::error::SdkError<E, R>,
1333 service: impl FnOnce(&E, String) -> ModelError,
1334) -> ModelError
1335where
1336 E: std::error::Error
1337 + aws_sdk_bedrockruntime::error::ProvideErrorMetadata
1338 + aws_sdk_bedrockruntime::operation::RequestId
1339 + 'static,
1340 R: std::fmt::Debug + 'static,
1341{
1342 use aws_sdk_bedrockruntime::error::{DisplayErrorContext, SdkError};
1343 if let SdkError::ServiceError(context) = error {
1344 let err = context.err();
1345 return service(err, service_detail(err));
1346 }
1347 let detail = DisplayErrorContext(error).to_string();
1348 if unresolvable_identity(error) {
1349 return ModelError::Refused {
1350 account: true,
1351 model: model.clone(),
1352 code: None,
1353 detail,
1354 };
1355 }
1356 match error {
1357 SdkError::ConstructionFailure(_) => ModelError::Unreachable {
1358 model: model.clone(),
1359 detail,
1360 },
1361 _ => ModelError::Unavailable {
1362 model: model.clone(),
1363 detail,
1364 },
1365 }
1366}
1367
1368fn service_detail<E>(error: &E) -> String
1370where
1371 E: std::error::Error
1372 + aws_sdk_bedrockruntime::error::ProvideErrorMetadata
1373 + aws_sdk_bedrockruntime::operation::RequestId,
1374{
1375 let code = error.code().unwrap_or("ServiceError");
1376 let message = error
1377 .message()
1378 .map_or_else(|| error.to_string(), str::to_owned);
1379 match error.request_id() {
1380 Some(id) => format!("{code}: {message} (request {id})"),
1381 None => format!("{code}: {message}"),
1382 }
1383}
1384
1385fn unresolvable_identity(error: &(dyn std::error::Error + 'static)) -> bool {
1401 use aws_credential_types::provider::error::{CredentialsError, TokenError};
1402 let mut verdict = None;
1403 let mut cause = Some(error);
1404 while let Some(inner) = cause {
1405 if let Some(credentials) = inner.downcast_ref::<CredentialsError>() {
1406 verdict = Some(matches!(
1407 credentials,
1408 CredentialsError::InvalidConfiguration(_) | CredentialsError::Unhandled(_)
1409 ));
1410 } else if verdict.is_some() && inner.is::<TokenError>() {
1411 return true;
1412 }
1413 cause = inner.source();
1414 }
1415 verdict.unwrap_or(false)
1416}
1417
1418const ACCOUNT_CODES: [&str; 4] = [
1423 "AccessDeniedException",
1424 "UnrecognizedClientException",
1425 "InvalidSignatureException",
1426 "ExpiredTokenException",
1427];
1428
1429fn classify_service(
1440 model: &ModelId,
1441 detail: String,
1442 code: Option<&str>,
1443 throttled: bool,
1444 refused: bool,
1445) -> ModelError {
1446 let account = code.is_some_and(|code| ACCOUNT_CODES.contains(&code));
1447 if throttled {
1448 ModelError::RateLimited {
1449 model: model.clone(),
1450 detail,
1451 retry_after: None,
1452 }
1453 } else if refused || account {
1454 ModelError::Refused {
1455 account,
1456 model: model.clone(),
1457 code: code.and_then(|code| super::bounded_code(code.to_owned())),
1458 detail,
1459 }
1460 } else {
1461 ModelError::Unavailable {
1462 model: model.clone(),
1463 detail,
1464 }
1465 }
1466}
1467
1468fn severed_stream(
1469 model: &ModelId,
1470 accumulator: &super::bedrock_stream::Accumulator,
1471 detail: &str,
1472) -> ModelError {
1473 if let Some(usage) = accumulator.usage() {
1474 return ModelError::Interrupted {
1475 model: model.clone(),
1476 usage,
1477 detail: detail.to_owned(),
1478 };
1479 }
1480 if accumulator.generated() {
1481 return ModelError::Unaccounted {
1482 model: model.clone(),
1483 detail: detail.to_owned(),
1484 };
1485 }
1486 ModelError::Unavailable {
1487 model: model.clone(),
1488 detail: detail.to_owned(),
1489 }
1490}
1491
1492fn classify_stream_event(
1493 model: &ModelId,
1494 accumulator: &super::bedrock_stream::Accumulator,
1495 error: Option<&aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError>,
1496 detail: &str,
1497) -> ModelError {
1498 if accumulator.started() || accumulator.generated() {
1499 return severed_stream(model, accumulator, detail);
1500 }
1501 if error.is_some_and(
1502 aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError::is_throttling_exception,
1503 ) {
1504 ModelError::RateLimited {
1505 model: model.clone(),
1506 detail: detail.to_owned(),
1507 retry_after: None,
1510 }
1511 } else if error.is_some_and(
1512 aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError::is_validation_exception,
1513 ) {
1514 ModelError::Refused {
1515 account: false,
1516 model: model.clone(),
1517 code: None,
1518 detail: detail.to_owned(),
1519 }
1520 } else {
1521 ModelError::Unavailable {
1522 model: model.clone(),
1523 detail: detail.to_owned(),
1524 }
1525 }
1526}
1527
1528pub(super) fn document_from_json(value: &Value) -> Document {
1529 match value {
1530 Value::Null => Document::Null,
1531 Value::Bool(value) => Document::Bool(*value),
1532 Value::String(value) => Document::String(value.clone()),
1533 Value::Array(values) => Document::Array(values.iter().map(document_from_json).collect()),
1534 Value::Object(values) => Document::Object(
1535 values
1536 .iter()
1537 .map(|(key, value)| (key.clone(), document_from_json(value)))
1538 .collect::<HashMap<_, _>>(),
1539 ),
1540 Value::Number(value) => {
1541 if let Some(value) = value.as_u64() {
1542 Document::Number(Number::PosInt(value))
1543 } else if let Some(value) = value.as_i64() {
1544 Document::Number(Number::NegInt(value))
1545 } else {
1546 Document::Number(Number::Float(
1547 value.as_f64().expect("JSON number is finite"),
1548 ))
1549 }
1550 }
1551 }
1552}
1553
1554pub(super) fn json_from_document(value: &Document) -> Value {
1555 match value {
1556 Document::Null => Value::Null,
1557 Document::Bool(value) => Value::Bool(*value),
1558 Document::String(value) => Value::String(value.clone()),
1559 Document::Array(values) => Value::Array(values.iter().map(json_from_document).collect()),
1560 Document::Object(values) => Value::Object(
1561 values
1562 .iter()
1563 .map(|(key, value)| (key.clone(), json_from_document(value)))
1564 .collect(),
1565 ),
1566 Document::Number(Number::PosInt(value)) => json!(value),
1567 Document::Number(Number::NegInt(value)) => json!(value),
1568 Document::Number(Number::Float(value)) => json!(value),
1569 }
1570}
1571
1572fn content_to_json(
1573 block: &ContentBlock,
1574 model: &ModelId,
1575 usage: Usage,
1576) -> Result<Value, ModelError> {
1577 match block {
1578 ContentBlock::Text(text) => Ok(json!({"type": "text", "text": text})),
1579 ContentBlock::ToolUse(tool) => Ok(json!({
1580 "type": "tool_use",
1581 "id": tool.tool_use_id(),
1582 "name": tool.name(),
1583 "input": json_from_document(tool.input()),
1584 })),
1585 ContentBlock::ReasoningContent(ReasoningContentBlock::ReasoningText(reasoning)) => {
1586 Ok(json!({
1587 "type": "reasoning",
1588 "text": reasoning.text(),
1589 "signature": reasoning.signature(),
1590 }))
1591 }
1592 ContentBlock::ReasoningContent(ReasoningContentBlock::RedactedContent(bytes)) => {
1593 Ok(json!({
1594 "type": "redacted_reasoning",
1595 "data": crate::core::b64::encode(bytes.as_ref()),
1596 }))
1597 }
1598 _ => Err(ModelError::Unusable {
1599 model: model.clone(),
1600 usage,
1601 detail: "Bedrock returned an unknown or unsupported continuation content block"
1602 .to_owned(),
1603 }),
1604 }
1605}
1606
1607fn content_from_json(value: &Value, model: &ModelId) -> Result<ContentBlock, ModelError> {
1608 let kind = value
1609 .get("type")
1610 .and_then(Value::as_str)
1611 .ok_or_else(|| Bedrock::refused(model, "a Bedrock continuation block has no type"))?;
1612 match kind {
1613 "text" => Ok(ContentBlock::Text(
1614 value
1615 .get("text")
1616 .and_then(Value::as_str)
1617 .ok_or_else(|| Bedrock::refused(model, "a text block has no text"))?
1618 .to_owned(),
1619 )),
1620 "tool_use" => Ok(ContentBlock::ToolUse(
1621 ToolUseBlock::builder()
1622 .tool_use_id(
1623 value
1624 .get("id")
1625 .and_then(Value::as_str)
1626 .ok_or_else(|| Bedrock::refused(model, "a tool block has no id"))?,
1627 )
1628 .name(
1629 value
1630 .get("name")
1631 .and_then(Value::as_str)
1632 .ok_or_else(|| Bedrock::refused(model, "a tool block has no name"))?,
1633 )
1634 .input(document_from_json(
1635 value.get("input").unwrap_or(&Value::Null),
1636 ))
1637 .build()
1638 .map_err(|error| Bedrock::refused(model, error.to_string()))?,
1639 )),
1640 "reasoning" => {
1641 let mut builder = ReasoningTextBlock::builder().text(
1642 value
1643 .get("text")
1644 .and_then(Value::as_str)
1645 .ok_or_else(|| Bedrock::refused(model, "a reasoning block has no text"))?,
1646 );
1647 if let Some(signature) = value.get("signature").and_then(Value::as_str) {
1648 builder = builder.signature(signature);
1649 }
1650 Ok(ContentBlock::ReasoningContent(
1651 ReasoningContentBlock::ReasoningText(
1652 builder
1653 .build()
1654 .map_err(|error| Bedrock::refused(model, error.to_string()))?,
1655 ),
1656 ))
1657 }
1658 "redacted_reasoning" => {
1659 let encoded = value
1660 .get("data")
1661 .and_then(Value::as_str)
1662 .ok_or_else(|| Bedrock::refused(model, "redacted reasoning has no data"))?;
1663 let bytes = crate::core::b64::decode(encoded).ok_or_else(|| {
1664 Bedrock::refused(model, "redacted reasoning data is not canonical base64")
1665 })?;
1666 Ok(ContentBlock::ReasoningContent(
1667 ReasoningContentBlock::RedactedContent(Blob::new(bytes)),
1668 ))
1669 }
1670 other => Err(Bedrock::refused(
1671 model,
1672 format!("unsupported Bedrock continuation block '{other}'"),
1673 )),
1674 }
1675}
1676
1677fn tool_result_from_json(value: &Value, model: &ModelId) -> Result<ContentBlock, ModelError> {
1678 let id = value
1679 .get("tool_use_id")
1680 .and_then(Value::as_str)
1681 .ok_or_else(|| Bedrock::refused(model, "a tool result has no tool_use_id"))?;
1682 let status = if value.get("failed").and_then(Value::as_bool) == Some(true) {
1683 ToolResultStatus::Error
1684 } else {
1685 ToolResultStatus::Success
1686 };
1687 Ok(ContentBlock::ToolResult(
1688 ToolResultBlock::builder()
1689 .tool_use_id(id)
1690 .content(tool_result_content(
1691 value.get("output").unwrap_or(&Value::Null),
1692 ))
1693 .status(status)
1694 .build()
1695 .map_err(|error| Bedrock::refused(model, error.to_string()))?,
1696 ))
1697}
1698
1699fn tool_result_content(output: &Value) -> ToolResultContentBlock {
1703 match output {
1704 Value::Object(_) => ToolResultContentBlock::Json(document_from_json(output)),
1705 Value::String(text) => ToolResultContentBlock::Text(text.clone()),
1706 other => ToolResultContentBlock::Text(other.to_string()),
1707 }
1708}
1709
1710fn continuation_message_from_json(value: &Value, model: &ModelId) -> Result<Message, ModelError> {
1711 let role = value
1712 .get("role")
1713 .and_then(Value::as_str)
1714 .ok_or_else(|| Bedrock::refused(model, "a continuation message has no role"))?;
1715 let blocks = value
1716 .get("content")
1717 .and_then(Value::as_array)
1718 .ok_or_else(|| Bedrock::refused(model, "a continuation message has no content array"))?;
1719 let (role, content) = match role {
1720 "assistant" => (
1721 ConversationRole::Assistant,
1722 blocks
1723 .iter()
1724 .map(|block| content_from_json(block, model))
1725 .collect::<Result<Vec<_>, _>>()?,
1726 ),
1727 "user" => (
1728 ConversationRole::User,
1729 blocks
1730 .iter()
1731 .map(|block| {
1732 if block.get("type").and_then(Value::as_str) == Some("text") {
1733 content_from_json(block, model)
1734 } else {
1735 tool_result_from_json(block, model)
1736 }
1737 })
1738 .collect::<Result<Vec<_>, _>>()?,
1739 ),
1740 other => {
1741 return Err(Bedrock::refused(
1742 model,
1743 format!("unsupported continuation role '{other}'"),
1744 ));
1745 }
1746 };
1747 Message::builder()
1748 .role(role)
1749 .set_content(Some(content))
1750 .build()
1751 .map_err(|error| Bedrock::refused(model, error.to_string()))
1752}
1753
1754#[cfg(test)]
1755mod tests {
1756 use super::*;
1757
1758 #[test]
1765 fn cache_tokens_are_folded_into_the_input_count() {
1766 let model = ModelId::new("bedrock", "test");
1767 let message = Message::builder()
1768 .role(ConversationRole::Assistant)
1769 .content(ContentBlock::Text("ok".to_owned()))
1770 .build()
1771 .expect("message");
1772 let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
1773 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
1774 message,
1775 ))
1776 .stop_reason(aws_sdk_bedrockruntime::types::StopReason::EndTurn)
1777 .usage(
1778 aws_sdk_bedrockruntime::types::TokenUsage::builder()
1779 .input_tokens(10)
1780 .output_tokens(4)
1781 .total_tokens(14)
1782 .cache_read_input_tokens(90)
1783 .cache_write_input_tokens(6)
1784 .build()
1785 .expect("usage"),
1786 )
1787 .build()
1788 .expect("converse output");
1789
1790 let completion =
1791 Bedrock::interpret(&model, None, SchemaMode::Native, &output).expect("an answer");
1792 assert_eq!(
1793 completion.usage.input_tokens, 106,
1794 "10 uncached + 90 cache-read + 6 cache-write"
1795 );
1796 assert_eq!(completion.usage.cache_read_tokens, 90);
1797 assert_eq!(completion.usage.cache_write_tokens, 6);
1798 }
1799
1800 #[test]
1807 fn a_buffered_guardrail_intervention_is_a_refusal_not_an_answer() {
1808 let model = ModelId::new("bedrock", "test");
1809 let message = Message::builder()
1810 .role(ConversationRole::Assistant)
1811 .content(ContentBlock::Text(
1812 "Sorry, I cannot answer that.".to_owned(),
1813 ))
1814 .build()
1815 .expect("message");
1816 let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
1817 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
1818 message,
1819 ))
1820 .stop_reason(aws_sdk_bedrockruntime::types::StopReason::GuardrailIntervened)
1821 .usage(
1822 aws_sdk_bedrockruntime::types::TokenUsage::builder()
1823 .input_tokens(7)
1824 .output_tokens(3)
1825 .total_tokens(10)
1826 .build()
1827 .expect("usage"),
1828 )
1829 .build()
1830 .expect("converse output");
1831
1832 let error = Bedrock::interpret(&model, None, SchemaMode::Native, &output)
1833 .expect_err("an intervention is a refusal, not an answer");
1834 match error {
1835 ModelError::Unusable { usage, .. } => {
1836 assert_eq!(usage.input_tokens, 7, "the assessment was billed");
1837 }
1838 other => panic!("expected a metered Unusable, got {other:?}"),
1839 }
1840 }
1841
1842 #[test]
1850 fn a_malformed_tool_use_stop_is_unusable_not_an_answer() {
1851 let model = ModelId::new("bedrock", "test");
1852 let message = Message::builder()
1853 .role(ConversationRole::Assistant)
1854 .content(ContentBlock::Text("{\"partial\": ".to_owned()))
1855 .build()
1856 .expect("message");
1857 let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
1858 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
1859 message,
1860 ))
1861 .stop_reason(aws_sdk_bedrockruntime::types::StopReason::MalformedToolUse)
1862 .usage(
1863 aws_sdk_bedrockruntime::types::TokenUsage::builder()
1864 .input_tokens(7)
1865 .output_tokens(3)
1866 .total_tokens(10)
1867 .build()
1868 .expect("usage"),
1869 )
1870 .build()
1871 .expect("converse output");
1872
1873 let error = Bedrock::interpret(&model, None, SchemaMode::Native, &output)
1874 .expect_err("a disowned turn is not an answer");
1875 match error {
1876 ModelError::Unusable { usage, .. } => {
1877 assert_eq!(usage.output_tokens, 3, "the generation was billed");
1878 }
1879 other => panic!("expected a metered Unusable, got {other:?}"),
1880 }
1881 }
1882
1883 #[test]
1886 fn a_filtered_or_unknown_stop_is_unusable_not_an_answer() {
1887 let model = ModelId::new("bedrock", "test");
1888 for reason in [
1889 aws_sdk_bedrockruntime::types::StopReason::ContentFiltered,
1890 aws_sdk_bedrockruntime::types::StopReason::from("a_reason_not_yet_invented"),
1891 ] {
1892 let message = Message::builder()
1893 .role(ConversationRole::Assistant)
1894 .content(ContentBlock::Text("looks whole".to_owned()))
1895 .build()
1896 .expect("message");
1897 let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
1898 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
1899 message,
1900 ))
1901 .stop_reason(reason.clone())
1902 .usage(
1903 aws_sdk_bedrockruntime::types::TokenUsage::builder()
1904 .input_tokens(7)
1905 .output_tokens(3)
1906 .total_tokens(10)
1907 .build()
1908 .expect("usage"),
1909 )
1910 .build()
1911 .expect("converse output");
1912 let error = Bedrock::interpret(&model, None, SchemaMode::Native, &output)
1913 .expect_err(reason.as_str());
1914 match error {
1915 ModelError::Unusable { usage, .. } => {
1916 assert_eq!(usage.output_tokens, 3, "the generation was billed");
1917 }
1918 other => panic!(
1919 "{}: expected a metered Unusable, got {other:?}",
1920 reason.as_str()
1921 ),
1922 }
1923 }
1924 }
1925
1926 #[test]
1927 fn smithy_documents_round_trip_json() {
1928 let value = json!({"s": "x", "n": -2, "u": 3, "f": 1.5, "a": [true, null]});
1929 assert_eq!(json_from_document(&document_from_json(&value)), value);
1930 }
1931
1932 #[test]
1933 fn governed_images_and_documents_become_inline_bedrock_blocks() {
1934 let model = ModelId::new("bedrock", "test");
1935 let image = content_from_prompt_json(
1936 &json!({
1937 "type": "image",
1938 "media_type": "image/png",
1939 "data": crate::core::b64::encode(b"png")
1940 }),
1941 &model,
1942 )
1943 .expect("image");
1944 assert_eq!(
1945 image.as_image().expect("image block").format(),
1946 &ImageFormat::Png
1947 );
1948
1949 let document = content_from_prompt_json(
1950 &json!({
1951 "type": "document",
1952 "media_type": "application/pdf",
1953 "data": crate::core::b64::encode(b"pdf")
1954 }),
1955 &model,
1956 )
1957 .expect("document");
1958 let document = document.as_document().expect("document block");
1959 assert_eq!(document.format(), &DocumentFormat::Pdf);
1960 assert_eq!(document.name(), "document");
1961 }
1962
1963 #[test]
1964 fn bedrock_media_never_accepts_a_remote_source() {
1965 let model = ModelId::new("bedrock", "test");
1966 assert!(
1967 content_from_prompt_json(
1968 &json!({
1969 "type": "image",
1970 "media_type": "image/png",
1971 "url": "https://example.test/image.png"
1972 }),
1973 &model,
1974 )
1975 .is_err()
1976 );
1977 }
1978
1979 #[test]
1980 fn signed_and_redacted_reasoning_round_trip() {
1981 let model = ModelId::new("bedrock", "test");
1982 for value in [
1983 json!({"type": "reasoning", "text": "private", "signature": "sig"}),
1984 json!({"type": "redacted_reasoning", "data": "AQID"}),
1985 ] {
1986 let block = content_from_json(&value, &model).expect("decode block");
1987 assert_eq!(
1988 content_to_json(&block, &model, Usage::default()).unwrap(),
1989 value
1990 );
1991 }
1992 }
1993
1994 #[test]
1997 fn a_guardrail_is_effect_identity_and_both_paths_send_the_same_one() {
1998 let config = aws_sdk_bedrockruntime::Config::builder()
1999 .region(Region::new("eu-west-1"))
2000 .behavior_version_latest()
2001 .http_client(crate::model::canned_http(200, None, ""))
2002 .build();
2003 let client = Client::from_conf(config);
2004 let plain = Bedrock::from_client(client.clone()).expect("a client with a region");
2005 let model = ModelId::new("bedrock", "m");
2006 let guarded = Bedrock::from_client(client)
2007 .expect("a client with a region")
2008 .guardrail(Guardrail::new("gr-7", "3"));
2009
2010 assert_eq!(plain.request_profile(&model)["guardrail"], Value::Null);
2011 assert_eq!(
2012 guarded.request_profile(&model)["guardrail"],
2013 json!({ "id": "gr-7", "version": "3" }),
2014 "a guardrail that does not enter the request profile can be turned \
2015 off between a run and its replay with nothing on the record"
2016 );
2017 assert_ne!(
2018 plain.request_profile(&model),
2019 guarded.request_profile(&model),
2020 "guarded and unguarded calls shared one effect identity"
2021 );
2022
2023 let g = Guardrail::new("gr-7", "3");
2027 assert_eq!(g.config().guardrail_identifier(), "gr-7");
2028 assert_eq!(g.stream_config().guardrail_identifier(), "gr-7");
2029 assert_eq!(g.config().guardrail_version(), "3");
2030 assert_eq!(g.stream_config().guardrail_version(), "3");
2031 assert_eq!(
2035 g.stream_config().stream_processing_mode().as_str(),
2036 "sync",
2037 "an asynchronous guardrail releases content before assessing it"
2038 );
2039
2040 assert_eq!(g.config().trace().as_str(), "disabled");
2045 assert_eq!(g.stream_config().trace().as_str(), "disabled");
2046 let traced = Guardrail::new("gr-7", "3").with_trace();
2047 assert_eq!(traced.config().trace().as_str(), "enabled");
2048 assert_eq!(
2049 traced.stream_config().trace().as_str(),
2050 "enabled",
2051 "with_trace must reach the streaming path too"
2052 );
2053 }
2054
2055 #[test]
2056 fn request_profile_commits_to_region_and_buffering() {
2057 let config = aws_sdk_bedrockruntime::Config::builder()
2058 .region(Region::new("eu-west-1"))
2059 .behavior_version_latest()
2060 .http_client(crate::model::canned_http(200, None, ""))
2066 .build();
2067 let driver =
2068 Bedrock::from_client(Client::from_conf(config)).expect("a client with a region");
2069 assert_eq!(
2070 driver.request_profile(&ModelId::new("bedrock", "m")),
2071 json!({
2072 "driver": "aws-bedrock-converse/v1",
2073 "region": "eu-west-1",
2074 "stream": true,
2075 "schema_mode": "native",
2076 "guardrail": Value::Null,
2080 "reasoning_dialect": Value::Null,
2081 })
2082 );
2083
2084 assert_eq!(
2085 driver.request_profile(&ModelId::new("bedrock", "m")),
2086 driver
2087 .clone()
2088 .timeout(Duration::from_secs(7))
2089 .request_profile(&ModelId::new("bedrock", "m")),
2090 "a timeout is how long this plane waits, not what it asks — it must not \
2091 enter the effect identity"
2092 );
2093
2094 assert_ne!(
2095 driver.request_profile(&ModelId::new("bedrock", "m")),
2096 driver
2097 .buffered()
2098 .request_profile(&ModelId::new("bedrock", "m")),
2099 "buffered and streamed Bedrock requests reused one effect identity"
2100 );
2101 }
2102
2103 #[test]
2118 fn the_attested_region_is_the_one_the_client_will_reach() {
2119 let stub = |region: Option<&str>| {
2120 let mut config = aws_sdk_bedrockruntime::Config::builder()
2121 .behavior_version_latest()
2122 .http_client(crate::model::canned_http(200, None, ""));
2123 if let Some(region) = region {
2124 config = config.region(Region::new(region.to_owned()));
2125 }
2126 Bedrock::from_client(Client::from_conf(config.build()))
2127 };
2128
2129 for region in ["us-east-1", "eu-west-1"] {
2130 let driver = stub(Some(region)).expect("a client carrying a region");
2131 assert_eq!(
2132 driver.request_profile(&ModelId::new("bedrock", "m"))["region"],
2133 json!(region),
2134 "the profile attested a region other than the one the client \
2135 resolves its endpoint from"
2136 );
2137 }
2138
2139 let err = stub(None).expect_err("a client with no region was accepted");
2143 assert!(
2144 err.contains("region"),
2145 "the refusal must name the missing region, got: {err}"
2146 );
2147 let err = stub(Some(" ")).expect_err("a blank region was accepted");
2148 assert!(
2149 err.contains("region"),
2150 "the refusal must name the blank region, got: {err}"
2151 );
2152 }
2153
2154 fn answering(
2157 status: u16,
2158 body: &'static str,
2159 credentials: aws_sdk_bedrockruntime::config::SharedCredentialsProvider,
2160 buffered: bool,
2161 ) -> Bedrock {
2162 let config = aws_sdk_bedrockruntime::Config::builder()
2163 .region(Region::new("eu-west-1"))
2164 .behavior_version_latest()
2165 .credentials_provider(credentials)
2166 .retry_config(aws_sdk_bedrockruntime::config::retry::RetryConfig::disabled())
2167 .http_client(crate::model::canned_http(
2168 status,
2169 Some("application/json"),
2170 body,
2171 ))
2172 .build();
2173 let driver = Bedrock::from_client(Client::from_conf(config)).expect("a region");
2174 if buffered { driver.buffered() } else { driver }
2175 }
2176
2177 async fn refusal_of(driver: &Bedrock) -> ModelError {
2178 let model = ModelId::new("bedrock", "anthropic.claude-x");
2179 let prompt = json!("hello");
2180 driver
2181 .complete(Request {
2182 stop: None,
2183 model: &model,
2184 prompt: &prompt,
2185 max_output_tokens: 64,
2186 reasoning_effort: None,
2187 reasoning_summary: false,
2188 schema: None,
2189 tools: &[],
2190 exchanges: &[],
2191 told: &[],
2192 continuation: None,
2193 compaction: None,
2194 stream: None,
2195 })
2196 .await
2197 .expect_err("every answer is an error")
2198 }
2199
2200 fn signing() -> aws_sdk_bedrockruntime::config::SharedCredentialsProvider {
2201 aws_sdk_bedrockruntime::config::SharedCredentialsProvider::new(
2202 aws_sdk_bedrockruntime::config::Credentials::for_tests(),
2203 )
2204 }
2205
2206 #[tokio::test]
2213 async fn a_refused_credential_refuses_the_account() {
2214 const EXPIRED: &str = r#"{"__type":"ExpiredTokenException","message":"The security token included in the request is expired"}"#;
2215 const UNKNOWN: &str = r#"{"__type":"UnrecognizedClientException","message":"The security token included in the request is invalid."}"#;
2216 const DENIED: &str = r#"{"__type":"AccessDeniedException","message":"You don't have access to the model with the specified model ID."}"#;
2217 const INVALID: &str = r#"{"__type":"ValidationException","message":"The provided model identifier is invalid."}"#;
2218 for buffered in [true, false] {
2219 for (status, body, code) in [
2220 (403, EXPIRED, "ExpiredTokenException"),
2221 (403, UNKNOWN, "UnrecognizedClientException"),
2222 (403, DENIED, "AccessDeniedException"),
2223 ] {
2224 let refused = refusal_of(&answering(status, body, signing(), buffered)).await;
2225 assert!(
2226 refused.refuses_the_account(),
2227 "{code} (buffered: {buffered}) is the account's: {refused:?}"
2228 );
2229 assert_eq!(refused.code().as_deref(), Some(code));
2230 let text = refused.to_string();
2231 assert!(text.contains(&format!("{code}: ")), "{text}");
2232 assert!(!text.contains("raw: Response"), "the SDK's dump: {text}");
2233 }
2234 let request = refusal_of(&answering(400, INVALID, signing(), buffered)).await;
2235 assert!(
2236 matches!(request, ModelError::Refused { account: false, .. }),
2237 "a validation failure is the request's: {request:?}"
2238 );
2239 }
2240 }
2241
2242 #[tokio::test]
2249 async fn a_credential_that_cannot_be_resolved_refuses_the_account() {
2250 use aws_credential_types::provider::error::CredentialsError;
2251 use aws_sdk_bedrockruntime::config::SharedCredentialsProvider;
2252 fn failing(error: fn() -> CredentialsError) -> SharedCredentialsProvider {
2253 SharedCredentialsProvider::new(
2254 aws_credential_types::credential_fn::provide_credentials_fn(move || async move {
2255 Err(error())
2256 }),
2257 )
2258 }
2259 for (error, account) in [
2260 (
2261 (|| CredentialsError::unhandled("the login token has expired"))
2262 as fn() -> CredentialsError,
2263 true,
2264 ),
2265 (
2266 || CredentialsError::invalid_configuration("no such profile"),
2267 true,
2268 ),
2269 (
2270 || CredentialsError::provider_error("STS answered 503"),
2271 false,
2272 ),
2273 (|| CredentialsError::not_loaded("nothing configured"), false),
2274 ] {
2275 let refused = refusal_of(&answering(200, "", failing(error), true)).await;
2276 if account {
2277 assert!(
2278 matches!(
2279 refused,
2280 ModelError::Refused {
2281 account: true,
2282 code: None,
2283 ..
2284 }
2285 ),
2286 "{refused:?}"
2287 );
2288 } else {
2289 assert!(
2290 matches!(refused, ModelError::Unavailable { .. }),
2291 "{refused:?}"
2292 );
2293 }
2294 }
2295 }
2296
2297 const PROFILE_CASE: &str = "AGENTPLANE_BEDROCK_PROFILE_CASE";
2299
2300 #[test]
2308 fn a_session_a_person_must_renew_refuses_the_account() {
2309 use std::fmt::Write as _;
2310 let home =
2311 std::env::temp_dir().join(format!("agentplane-aws-{}", crate::core::RunId::generate()));
2312 let cache = home.join(".aws/sso/cache");
2313 std::fs::create_dir_all(&cache).unwrap();
2314 let config = home.join(".aws/config");
2315 let mut profiles = String::new();
2316 for (profile, session) in [
2317 ("sso-expired", "ap"),
2318 ("sso-missing", "gone"),
2319 ("sso-live", "live"),
2320 ] {
2321 write!(
2322 profiles,
2323 "[profile {profile}]\nsso_session = {session}\nsso_account_id = 111122223333\n\
2324 sso_role_name = Runner\nregion = eu-west-1\n\n[sso-session {session}]\n\
2325 sso_start_url = https://example.awsapps.com/start\nsso_region = eu-west-1\n\n"
2326 )
2327 .unwrap();
2328 }
2329 profiles.push_str(
2330 "[profile login-missing]\nlogin_session = arn:aws:iam::111122223333:user/Runner\n\
2331 region = eu-west-1\n",
2332 );
2333 std::fs::write(&config, profiles).unwrap();
2334 for (digest, expires) in [
2336 (
2337 "ac78b022715c5b8357b4dca8045e8463b4de2124",
2338 "2020-01-01T00:00:00Z",
2339 ),
2340 (
2341 "98aadb37083eddd855c27feccb15dc8c5b127fd0",
2342 "2099-01-01T00:00:00Z",
2343 ),
2344 ] {
2345 std::fs::write(
2346 cache.join(format!("{digest}.json")),
2347 format!(
2348 r#"{{"accessToken":"token","expiresAt":"{expires}","region":"eu-west-1","startUrl":"https://example.awsapps.com/start"}}"#
2349 ),
2350 )
2351 .unwrap();
2352 }
2353 for (profile, verdict) in [
2354 ("sso-expired", "account"),
2355 ("sso-missing", "account"),
2356 ("login-missing", "account"),
2357 ("sso-live", "outage"),
2358 ] {
2359 let output = std::process::Command::new(std::env::current_exe().unwrap())
2360 .args([
2361 "model::bedrock::tests::a_profile_resolved_in_a_child",
2362 "--exact",
2363 "--ignored",
2364 "--nocapture",
2365 ])
2366 .env_clear()
2367 .env(PROFILE_CASE, verdict)
2368 .env("HOME", &home)
2369 .env("AWS_CONFIG_FILE", &config)
2370 .env("AWS_SHARED_CREDENTIALS_FILE", home.join("none"))
2371 .env("AWS_PROFILE", profile)
2372 .env("AWS_EC2_METADATA_DISABLED", "true")
2373 .env("AWS_ENDPOINT_URL", "http://127.0.0.1:9")
2375 .output()
2376 .expect("the child process");
2377 assert!(
2378 output.status.success(),
2379 "{profile} did not read as {verdict}:\n{}",
2380 String::from_utf8_lossy(&output.stdout)
2381 );
2382 }
2383 std::fs::remove_dir_all(&home).ok();
2384 }
2385
2386 #[tokio::test]
2387 #[ignore = "the child half of a_session_a_person_must_renew_refuses_the_account"]
2388 async fn a_profile_resolved_in_a_child() {
2389 let Ok(verdict) = std::env::var(PROFILE_CASE) else {
2390 return;
2391 };
2392 let driver = Bedrock::from_env("eu-west-1").await.expect("a driver");
2393 let refused = refusal_of(&driver.buffered()).await;
2394 let account = matches!(refused, ModelError::Refused { account: true, .. });
2395 assert_eq!(account, verdict == "account", "{refused:?}");
2396 }
2397
2398 fn stub_driver() -> Bedrock {
2399 let config = aws_sdk_bedrockruntime::Config::builder()
2400 .region(Region::new("eu-west-1"))
2401 .behavior_version_latest()
2402 .http_client(crate::model::canned_http(200, None, ""))
2403 .build();
2404 Bedrock::from_client(Client::from_conf(config)).expect("a client with a region")
2405 }
2406
2407 #[test]
2418 fn nova_reasoning_effort_is_rendered_the_way_aws_documents_it() {
2419 use super::super::ReasoningEffort as E;
2420
2421 let nova = stub_driver().reasoning(ReasoningDialect::Nova);
2422 let model = ModelId::new("bedrock", "us.amazon.nova-2-lite-v1:0");
2423
2424 for (effort, expected) in [(E::Low, "low"), (E::Medium, "medium"), (E::High, "high")] {
2425 let document = nova
2426 .reasoning_config(&model, Some(effort))
2427 .expect("a documented level maps")
2428 .expect("an effort produces a config");
2429 assert_eq!(
2430 json_from_document(&document),
2431 json!({ "reasoningConfig": { "type": "enabled", "maxReasoningEffort": expected } }),
2432 "the {} level did not render as AWS documents it",
2433 effort.as_str()
2434 );
2435 }
2436
2437 assert_eq!(
2442 json_from_document(
2443 &nova
2444 .reasoning_config(&model, Some(E::None))
2445 .expect("none maps")
2446 .expect("none is still a config")
2447 ),
2448 json!({ "reasoningConfig": { "type": "disabled" } })
2449 );
2450
2451 assert!(
2454 nova.reasoning_config(&model, None)
2455 .expect("no effort is not an error")
2456 .is_none()
2457 );
2458 }
2459
2460 #[test]
2472 fn nova_refuses_an_effort_it_has_no_counterpart_for() {
2473 use super::super::ReasoningEffort as E;
2474
2475 let nova = stub_driver().reasoning(ReasoningDialect::Nova);
2476 let model = ModelId::new("bedrock", "us.amazon.nova-2-lite-v1:0");
2477
2478 for effort in [E::Minimal, E::XHigh, E::Max] {
2479 let error = nova
2480 .reasoning_config(&model, Some(effort))
2481 .expect_err("an effort Nova cannot express must be refused");
2482 let text = error.to_string();
2483 assert!(
2484 text.contains(effort.as_str()) && text.contains("low, medium and high"),
2485 "the refusal must name the effort and the levels that exist: {text}"
2486 );
2487 }
2488 }
2489
2490 #[test]
2497 fn an_undeclared_dialect_refuses_rather_than_guessing() {
2498 let plain = stub_driver();
2499 let model = ModelId::new("bedrock", "m");
2500 let text = plain
2501 .reasoning_config(&model, Some(super::super::ReasoningEffort::Medium))
2502 .expect_err("no dialect must refuse")
2503 .to_string();
2504 assert!(
2505 text.contains("ReasoningDialect::Nova"),
2506 "a refusal that does not name the remedy leaves the reader stuck: {text}"
2507 );
2508 }
2509
2510 #[test]
2520 fn the_reasoning_dialect_is_effect_identity_and_reaches_both_paths() {
2521 let model = ModelId::new("bedrock", "us.amazon.nova-2-lite-v1:0");
2522 let plain = stub_driver();
2523 let nova = stub_driver().reasoning(ReasoningDialect::Nova);
2524
2525 assert_eq!(
2526 plain.request_profile(&model)["reasoning_dialect"],
2527 Value::Null
2528 );
2529 assert_eq!(nova.request_profile(&model)["reasoning_dialect"], "nova");
2530 assert_ne!(
2531 plain.request_profile(&model),
2532 nova.request_profile(&model),
2533 "a driver that renders reasoning differently shared one effect identity"
2534 );
2535
2536 assert_eq!(
2540 nova.clone().buffered().request_profile(&model)["reasoning_dialect"],
2541 nova.request_profile(&model)["reasoning_dialect"],
2542 );
2543 }
2544
2545 #[test]
2546 fn native_schema_can_coexist_with_tools_but_forced_fallback_cannot() {
2547 let model = ModelId::new("bedrock", "test");
2548 let schema = json!({"type": "object"});
2549 let tools = vec![ToolDeclaration::new(
2550 "lookup",
2551 "Look something up.",
2552 json!({"type": "object"}),
2553 )];
2554 assert!(
2555 Bedrock::tool_config(&model, Some(&schema), &tools, SchemaMode::Native)
2556 .expect("native schema and tools")
2557 .is_some()
2558 );
2559 assert!(
2560 Bedrock::tool_config(&model, Some(&schema), &tools, SchemaMode::ForcedTool).is_err(),
2561 "forced-tool structured output consumed the same tool channel as caller tools"
2562 );
2563 }
2564
2565 #[test]
2580 fn a_reasoning_tool_turn_with_no_text_is_an_answer_not_an_empty_one() {
2581 let model = ModelId::new("bedrock", "test");
2582 let message = Message::builder()
2583 .role(ConversationRole::Assistant)
2584 .content(ContentBlock::ReasoningContent(
2585 ReasoningContentBlock::ReasoningText(
2586 ReasoningTextBlock::builder()
2587 .text("private reasoning")
2588 .signature("sig-bytes")
2589 .build()
2590 .expect("reasoning block"),
2591 ),
2592 ))
2593 .content(ContentBlock::ToolUse(
2594 ToolUseBlock::builder()
2595 .tool_use_id("tool_1")
2596 .name("lookup")
2597 .input(document_from_json(&json!({ "id": "AC-1" })))
2598 .build()
2599 .expect("tool block"),
2600 ))
2601 .build()
2602 .expect("assistant message");
2603 let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
2604 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
2605 message,
2606 ))
2607 .stop_reason(aws_sdk_bedrockruntime::types::StopReason::ToolUse)
2608 .build()
2609 .expect("converse output");
2610
2611 let completion = Bedrock::interpret(&model, None, SchemaMode::Native, &output)
2612 .expect("a tool call with empty sibling text is a normal tool turn");
2613 assert_eq!(completion.tool_calls.len(), 1);
2614 assert_eq!(completion.tool_calls[0].name, "lookup");
2615 let state = completion
2616 .continuation
2617 .expect("a tool turn must carry its provider continuation")
2618 .state;
2619 assert_eq!(
2620 state,
2621 json!([{ "role": "assistant", "content": [
2622 { "type": "reasoning", "text": "private reasoning", "signature": "sig-bytes" },
2623 { "type": "tool_use", "id": "tool_1", "name": "lookup", "input": { "id": "AC-1" } },
2624 ]}]),
2625 "the continuation must carry the provider turn verbatim, reasoning \
2626 block and signature included"
2627 );
2628 }
2629
2630 #[test]
2636 fn an_answer_with_neither_text_nor_tool_calls_is_unusable() {
2637 let model = ModelId::new("bedrock", "test");
2638 let message = Message::builder()
2639 .role(ConversationRole::Assistant)
2640 .content(ContentBlock::ReasoningContent(
2641 ReasoningContentBlock::ReasoningText(
2642 ReasoningTextBlock::builder()
2643 .text("reasoned and then said nothing")
2644 .build()
2645 .expect("reasoning block"),
2646 ),
2647 ))
2648 .build()
2649 .expect("assistant message");
2650 let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
2651 .output(aws_sdk_bedrockruntime::types::ConverseOutput::Message(
2652 message,
2653 ))
2654 .stop_reason(aws_sdk_bedrockruntime::types::StopReason::EndTurn)
2655 .build()
2656 .expect("converse output");
2657 let error = Bedrock::interpret(&model, None, SchemaMode::Native, &output)
2658 .expect_err("reasoning with neither text nor a tool call is not usable");
2659 assert!(
2660 matches!(error, ModelError::Unusable { .. }),
2661 "wrong classification: {error}"
2662 );
2663 }
2664
2665 #[test]
2668 fn a_persons_turn_follows_the_history() {
2669 let model = ModelId::new("bedrock", "test");
2670 let continuation = ProviderContinuation::new(
2671 "bedrock",
2672 json!([{ "role": "assistant", "content": [{ "type": "text", "text": "42." }] }]),
2673 );
2674 let messages = Bedrock::messages(
2675 &model,
2676 &json!({"input": "balance?"}),
2677 &[],
2678 &["and in euros?".to_owned()],
2679 Some(&continuation),
2680 )
2681 .expect("a person's turn after an answer");
2682 assert_eq!(messages.len(), 3, "question, answer, the person's turn");
2683 assert_eq!(messages[2].role(), &ConversationRole::User);
2684 assert_eq!(
2685 messages[2].content(),
2686 &[ContentBlock::Text("and in euros?".to_owned())]
2687 );
2688 }
2689
2690 #[test]
2693 fn a_first_persons_turn_joins_the_prompts_message() {
2694 let model = ModelId::new("bedrock", "test");
2695 let messages = Bedrock::messages(
2696 &model,
2697 &json!({"input": "balance?"}),
2698 &[],
2699 &["in euros".to_owned()],
2700 None,
2701 )
2702 .expect("a first turn");
2703 assert_eq!(messages.len(), 1, "two user messages in a row");
2704 assert_eq!(messages[0].content().len(), 2);
2705 }
2706
2707 #[test]
2709 fn a_tool_result_is_json_only_when_it_is_an_object() {
2710 let model = ModelId::new("bedrock", "test");
2711 let call = super::super::ToolCall {
2712 id: "tool_1".to_owned(),
2713 name: "list_files".to_owned(),
2714 arguments: json!({}),
2715 };
2716 let content = |exchange: ToolExchange| {
2717 let messages =
2718 Bedrock::messages(&model, &json!({"input": "files?"}), &[exchange], &[], None)
2719 .expect("a tool turn");
2720 messages
2721 .iter()
2722 .flat_map(Message::content)
2723 .find_map(|block| block.as_tool_result().ok())
2724 .expect("a tool result")
2725 .content()[0]
2726 .clone()
2727 };
2728 let text = |text: &str| ToolResultContentBlock::Text(text.to_owned());
2729
2730 assert!(content(ToolExchange::ok(call.clone(), json!({"files": []}))).is_json());
2731 for (output, sent) in [
2732 (json!(["a.rs", "b.rs"]), r#"["a.rs","b.rs"]"#),
2733 (json!("a.rs"), "a.rs"),
2734 (json!(3), "3"),
2735 (json!(true), "true"),
2736 (Value::Null, "null"),
2737 ] {
2738 assert_eq!(content(ToolExchange::ok(call.clone(), output)), text(sent));
2739 }
2740 assert_eq!(
2741 content(ToolExchange::failed(call, "permission denied")),
2742 text("permission denied"),
2743 );
2744 }
2745
2746 #[test]
2747 fn cumulative_bedrock_transcript_decodes_every_prior_turn() {
2748 let model = ModelId::new("bedrock", "test");
2749 let continuation = ProviderContinuation::new(
2750 "bedrock",
2751 json!([
2752 {"role": "assistant", "content": [
2753 {"type": "reasoning", "text": "private", "signature": "sig"},
2754 {"type": "tool_use", "id": "tool_1", "name": "lookup", "input": {}}
2755 ]},
2756 {"role": "user", "content": [
2757 {"type": "tool_result", "tool_use_id": "tool_1", "output": {"v": 1}, "failed": false}
2758 ]}
2759 ]),
2760 );
2761 let messages = Bedrock::messages(
2762 &model,
2763 &json!({"input": "question"}),
2764 &[ToolExchange::ok(
2765 super::super::ToolCall {
2766 id: "tool_2".to_owned(),
2767 name: "lookup".to_owned(),
2768 arguments: json!({}),
2769 },
2770 json!({"v": 2}),
2771 )],
2772 &[],
2773 Some(&continuation),
2774 )
2775 .expect("decode cumulative transcript");
2776 assert_eq!(messages.len(), 4, "question plus three continuation turns");
2777 assert!(messages[1].content()[0].is_reasoning_content());
2778 assert!(messages[2].content()[0].is_tool_result());
2779 assert!(messages[3].content()[0].is_tool_result());
2780 }
2781}
2782
2783#[cfg(test)]
2784mod prompt_tests {
2785 use aws_sdk_bedrockruntime::types::{
2786 CacheDetail, CacheTtl as WireCacheTtl, ContentBlock, SystemContentBlock, TokenUsage, Tool,
2787 };
2788 use serde_json::json;
2789
2790 use super::{Bedrock, ModelId, SchemaMode, ToolDeclaration};
2791 use crate::core::Tainted;
2792 use crate::model::{CacheTtl, Prompt};
2793
2794 fn model() -> ModelId {
2795 ModelId::new("bedrock", "anthropic.claude-x")
2796 }
2797
2798 #[test]
2801 fn a_prompt_part_carries_its_breakpoint() {
2802 let prompt = Prompt::new()
2803 .system_cached(Tainted::trusted("policy".to_owned()), CacheTtl::OneHour)
2804 .input_cached(
2805 Tainted::trusted("context".to_owned()),
2806 CacheTtl::FiveMinutes,
2807 )
2808 .input(Tainted::trusted("task".to_owned()))
2809 .build();
2810 let system = super::system_blocks(prompt.peek(), &model()).expect("system");
2811 assert!(matches!(&system[0], SystemContentBlock::Text(t) if t == "policy"));
2812 assert!(
2813 matches!(&system[1], SystemContentBlock::CachePoint(p) if p.ttl() == Some(&WireCacheTtl::OneHour)),
2814 "{system:?}"
2815 );
2816 let messages =
2817 Bedrock::messages(&model(), prompt.peek(), &[], &[], None).expect("messages");
2818 let content = messages[0].content();
2819 assert!(matches!(&content[0], ContentBlock::Text(t) if t == "context"));
2820 assert!(
2821 matches!(&content[1], ContentBlock::CachePoint(p) if p.ttl() == Some(&WireCacheTtl::FiveMinutes)),
2822 "{content:?}"
2823 );
2824 assert!(matches!(&content[2], ContentBlock::Text(t) if t == "task"));
2825 }
2826
2827 #[test]
2830 fn a_tool_is_strict_only_where_the_subset_holds_it() {
2831 let tools = [
2832 ToolDeclaration::new(
2833 "closed",
2834 "d",
2835 json!({ "type": "object", "properties": {}, "required": [], "additionalProperties": false }),
2836 )
2837 .cache_breakpoint(CacheTtl::OneHour),
2838 ToolDeclaration::new("open", "d", json!({ "type": "object" })),
2839 ];
2840 let config = Bedrock::tool_config(&model(), None, &tools, SchemaMode::Native)
2841 .expect("tool config")
2842 .expect("tools were declared");
2843 let strict: Vec<Option<bool>> = config
2844 .tools()
2845 .iter()
2846 .filter_map(|tool| match tool {
2847 Tool::ToolSpec(spec) => Some(spec.strict()),
2848 _ => None,
2849 })
2850 .collect();
2851 assert_eq!(strict, [Some(true), Some(false)]);
2852 assert!(
2853 matches!(&config.tools()[1], Tool::CachePoint(p) if p.ttl() == Some(&WireCacheTtl::OneHour)),
2854 "{config:?}"
2855 );
2856 }
2857
2858 #[test]
2861 fn the_one_hour_writes_are_counted() {
2862 let usage = TokenUsage::builder()
2863 .input_tokens(10)
2864 .output_tokens(5)
2865 .total_tokens(15)
2866 .cache_write_input_tokens(300)
2867 .cache_details(
2868 CacheDetail::builder()
2869 .ttl(WireCacheTtl::OneHour)
2870 .input_tokens(200)
2871 .build()
2872 .expect("detail"),
2873 )
2874 .build()
2875 .expect("usage");
2876 let usage = super::usage_from(&usage);
2877 assert_eq!(usage.cache_write_tokens, 300);
2878 assert_eq!(usage.cache_write_1h_tokens, 200);
2879 }
2880}