rig_bedrock/types/
converse_output.rs

1//! Types that replace the AWS Bedrock Runtime SDK's `ConverseOutput` type.
2//! This is required so that we can impl Serialize and Deserialize.
3//!
4//! Rig's normalized [`CompletionResponse`](rig_core::completion::CompletionResponse)
5//! reads only part of the Converse response, but this type is what
6//! `raw_completion` hands back — the escape hatch whose whole purpose is that
7//! nothing the provider sent has been thrown away. Model-specific extras
8//! (`additional_model_response_fields`) are carried as plain
9//! [`serde_json::Value`].
10//!
11//! The guardrail trace, performance configuration and service tier keep the
12//! SDK's own types rather than gaining hand-written mirrors: they are deeply
13//! nested (a guardrail assessment alone is a dozen types), and a mirror that
14//! drifts from the SDK would reintroduce exactly the silent loss this exists
15//! to prevent. Those three are `#[serde(skip)]` because the SDK types are not
16//! `Serialize`, so a serialized `InternalConverseOutput` — a cassette fixture,
17//! a persisted response — omits them while an in-process caller reads them in
18//! full.
19use std::fmt;
20
21use aws_sdk_bedrockruntime::types as aws_bedrock;
22use serde::{Deserialize, Serialize};
23
24use super::errors::TypeConversionError;
25use super::json::AwsDocument;
26
27/// Our own implementation of the AWS Bedrock runtime "converse" operation output.
28/// The reason why we need to implement this is that we need to impl Deserialize/Serialize on top of this.
29#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
30pub struct InternalConverseOutput {
31    /// <p>The result from the call to <code>Converse</code>.</p>
32    pub output: Option<ConverseOutput>,
33    /// <p>The reason why the model stopped generating output.</p>
34    pub stop_reason: StopReason,
35    /// <p>The total number of tokens used in the call to <code>Converse</code>. The total includes the tokens input to the model and the tokens generated by the model.</p>
36    pub usage: Option<TokenUsage>,
37    /// <p>Metrics for the call to <code>Converse</code>.</p>
38    pub metrics: Option<ConverseMetrics>,
39    /// <p>Additional fields in the response that are unique to the model.</p>
40    pub additional_model_response_fields: Option<serde_json::Value>,
41    /// The AWS request id, taken from the response's `x-amzn-RequestId`
42    /// header. Always present on a successful call, and the identifier AWS
43    /// support asks for.
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub request_id: Option<String>,
46    /// <p>A trace object that contains information about the Guardrail behavior.</p>
47    ///
48    /// Populated when the request carried a guardrail configuration with
49    /// tracing enabled (see
50    /// [`CompletionModel::with_guardrail`](crate::completion::CompletionModel::with_guardrail));
51    /// it is the only place Bedrock explains *why* a turn stopped with
52    /// [`StopReason::GuardrailIntervened`].
53    #[serde(skip)]
54    pub trace: Option<aws_bedrock::ConverseTrace>,
55    /// <p>Model performance settings for the request.</p>
56    #[serde(skip)]
57    pub performance_config: Option<aws_bedrock::PerformanceConfiguration>,
58    /// <p>The processing tier used to serve the request.</p>
59    #[serde(skip)]
60    pub service_tier: Option<aws_bedrock::ServiceTier>,
61}
62
63impl InternalConverseOutput {
64    pub fn usage(&self) -> Option<&TokenUsage> {
65        self.usage.as_ref()
66    }
67
68    /// Bedrock's guardrail trace for this turn, when one was requested.
69    pub fn trace(&self) -> Option<&aws_bedrock::ConverseTrace> {
70        self.trace.as_ref()
71    }
72
73    /// The AWS request id for this call.
74    pub fn request_id(&self) -> Option<&str> {
75        self.request_id.as_deref()
76    }
77}
78
79impl TryFrom<aws_sdk_bedrockruntime::operation::converse::ConverseOutput>
80    for InternalConverseOutput
81{
82    type Error = TypeConversionError;
83
84    fn try_from(
85        value: aws_sdk_bedrockruntime::operation::converse::ConverseOutput,
86    ) -> Result<Self, Self::Error> {
87        // The request id is only reachable through the trait, and only before
88        // the struct is destructured.
89        let request_id =
90            aws_sdk_bedrockruntime::operation::RequestId::request_id(&value).map(str::to_string);
91
92        // `ConverseOutput` is `#[non_exhaustive]` and hides `_request_id`, so
93        // the rest pattern is mandatory and cannot be traded for a
94        // compile-time guard: every field the SDK adds arrives here silently.
95        // That is how the guardrail trace, performance config and service tier
96        // came to be dropped, so the list below is checked against the SDK
97        // type when the dependency moves, and `converse_output_carries_every_
98        // sdk_field` pins the ones known today.
99        let aws_sdk_bedrockruntime::operation::converse::ConverseOutput {
100            output,
101            stop_reason,
102            usage,
103            metrics,
104            additional_model_response_fields,
105            trace,
106            performance_config,
107            service_tier,
108            ..
109        } = value;
110
111        Ok(Self {
112            output: output.map(|x| x.try_into()).transpose()?,
113            stop_reason: stop_reason.try_into()?,
114            usage: usage.map(|x| x.try_into()).transpose()?,
115            metrics: metrics.map(|x| x.try_into()).transpose()?,
116            additional_model_response_fields: additional_model_response_fields
117                .map(|doc| AwsDocument(doc).into()),
118            request_id,
119            trace,
120            performance_config,
121            service_tier,
122        })
123    }
124}
125
126/// Opaque struct used as inner data for the `Unknown` variant defined in enums in
127/// the crate.
128///
129/// This is not intended to be used directly.
130#[derive(Clone, Eq, Ord, PartialEq, PartialOrd, Debug, Hash, Serialize, Deserialize)]
131pub struct UnknownVariantValue(pub(crate) String);
132
133impl fmt::Display for UnknownVariantValue {
134    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
135        write!(f, "{}", self.0)
136    }
137}
138
139#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
140pub enum StopReason {
141    ContentFiltered,
142    EndTurn,
143    GuardrailIntervened,
144    MaxTokens,
145    StopSequence,
146    ToolUse,
147    Unknown(UnknownVariantValue),
148}
149
150#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
151pub struct TokenUsage {
152    pub input_tokens: i32,
153    pub output_tokens: i32,
154    pub total_tokens: i32,
155    #[serde(skip_serializing_if = "Option::is_none")]
156    pub cache_read_input_tokens: Option<i32>,
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub cache_write_input_tokens: Option<i32>,
159}
160
161#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
162pub struct ConverseMetrics {
163    pub latency_ms: i64,
164}
165
166#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
167pub enum ConverseOutput {
168    Message(Message),
169    Unknown,
170}
171
172impl ConverseOutput {
173    pub fn as_message(&self) -> Result<&Message, TypeConversionError> {
174        match self {
175            Self::Message(message) => Ok(message),
176            invalid => Err(TypeConversionError::new(
177                format!("Tried to get ConverseOutput as message but is invalid: {invalid:?}")
178                    .as_ref(),
179            )),
180        }
181    }
182}
183
184#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
185pub struct Message {
186    pub role: ConversationRole,
187    pub content: Vec<ContentBlock>,
188}
189#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
190pub enum ConversationRole {
191    Assistant,
192    User,
193    Unknown(UnknownVariantValue),
194}
195#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
196pub enum ContentBlock {
197    CachePoint(CachePointBlock),
198    CitationsContent(CitationsContentBlock),
199    Document(DocumentBlock),
200    GuardContent(GuardrailConverseContentBlock),
201    Image(ImageBlock),
202    ReasoningContent(ReasoningContentBlock),
203    Text(String),
204    ToolResult(ToolResultBlock),
205    ToolUse(ToolUseBlock),
206    Video(VideoBlock),
207    Unknown,
208}
209#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
210pub struct CachePointBlock {
211    #[serde(rename = "type")]
212    pub kind: CachePointType,
213}
214
215#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
216pub enum CachePointType {
217    Default,
218    Unknown(UnknownVariantValue),
219}
220
221#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
222pub struct CitationsContentBlock {
223    pub content: Option<Vec<CitationGeneratedContent>>,
224    pub citations: Option<Vec<Citation>>,
225}
226#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
227pub enum CitationGeneratedContent {
228    Text(String),
229    Unknown,
230}
231#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
232pub struct Citation {
233    pub title: Option<String>,
234    pub source_content: Option<Vec<CitationSourceContent>>,
235    pub location: Option<CitationLocation>,
236}
237
238#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
239pub enum CitationSourceContent {
240    Text(String),
241    Unknown,
242}
243
244#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
245pub enum CitationLocation {
246    DocumentChar(DocumentCharLocation),
247    DocumentChunk(DocumentChunkLocation),
248    DocumentPage(DocumentPageLocation),
249    Unknown,
250}
251#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
252pub struct DocumentCharLocation {
253    pub document_index: Option<i32>,
254    pub start: Option<i32>,
255    pub end: Option<i32>,
256}
257#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
258pub struct DocumentChunkLocation {
259    pub document_index: Option<i32>,
260    pub start: Option<i32>,
261    pub end: Option<i32>,
262}
263#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
264pub struct DocumentPageLocation {
265    pub document_index: Option<i32>,
266    pub start: Option<i32>,
267    pub end: Option<i32>,
268}
269
270#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
271pub struct DocumentBlock {
272    pub format: DocumentFormat,
273    pub name: String,
274    pub source: Option<DocumentSource>,
275    pub context: Option<String>,
276    pub citations: Option<CitationsConfig>,
277}
278#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
279pub enum DocumentFormat {
280    Csv,
281    Doc,
282    Docx,
283    Html,
284    Md,
285    Pdf,
286    Txt,
287    Xls,
288    Xlsx,
289    Unknown(UnknownVariantValue),
290}
291#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
292pub enum DocumentSource {
293    Bytes(Blob),
294    Content(Vec<DocumentContentBlock>),
295    S3Location(S3Location),
296    Text(String),
297    Unknown,
298}
299#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
300pub enum DocumentContentBlock {
301    Text(String),
302    Unknown,
303}
304#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
305pub struct S3Location {
306    pub uri: String,
307    pub bucket_owner: Option<String>,
308}
309
310#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
311pub struct Blob {
312    pub inner: Vec<u8>,
313}
314#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
315pub struct CitationsConfig {
316    pub enabled: bool,
317}
318#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
319pub enum GuardrailConverseContentBlock {
320    Image(GuardrailConverseImageBlock),
321    Text(GuardrailConverseTextBlock),
322    Unknown,
323}
324#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
325pub struct GuardrailConverseImageBlock {
326    pub format: GuardrailConverseImageFormat,
327    pub source: Option<GuardrailConverseImageSource>,
328}
329#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
330pub enum GuardrailConverseImageFormat {
331    Jpeg,
332    Png,
333    Unknown(UnknownVariantValue),
334}
335#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
336pub enum GuardrailConverseImageSource {
337    Bytes(Blob),
338    Unknown,
339}
340#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
341pub struct GuardrailConverseTextBlock {
342    pub text: String,
343    pub qualifiers: Option<Vec<GuardrailConverseContentQualifier>>,
344}
345#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
346pub enum GuardrailConverseContentQualifier {
347    GroundingSource,
348    GuardContent,
349    Query,
350    Unknown(UnknownVariantValue),
351}
352#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
353pub struct ImageBlock {
354    pub format: ImageFormat,
355    pub source: Option<ImageSource>,
356}
357#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
358pub enum ImageFormat {
359    Gif,
360    Jpeg,
361    Png,
362    Webp,
363    Unknown(UnknownVariantValue),
364}
365#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
366pub enum ImageSource {
367    Bytes(Blob),
368    S3Location(S3Location),
369    Unknown,
370}
371
372#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
373pub enum ReasoningContentBlock {
374    ReasoningText(ReasoningTextBlock),
375    RedactedContent(Blob),
376    Unknown,
377}
378#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
379pub struct ReasoningTextBlock {
380    pub text: String,
381    pub signature: Option<String>,
382}
383
384#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
385pub struct ToolResultBlock {
386    pub tool_use_id: String,
387    pub content: Vec<ToolResultContentBlock>,
388    pub status: Option<ToolResultStatus>,
389}
390
391#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
392pub enum ToolResultContentBlock {
393    Document(DocumentBlock),
394    Image(ImageBlock),
395    Json(serde_json::Value),
396    Text(String),
397    Video(VideoBlock),
398    Unknown,
399}
400#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
401pub struct VideoBlock {
402    pub format: VideoFormat,
403    pub source: Option<VideoSource>,
404}
405#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
406pub enum VideoFormat {
407    Flv,
408    Mkv,
409    Mov,
410    Mp4,
411    Mpeg,
412    Mpg,
413    ThreeGp,
414    Webm,
415    Wmv,
416    Unknown(UnknownVariantValue),
417}
418#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
419pub enum VideoSource {
420    Bytes(Blob),
421    S3Location(S3Location),
422    Unknown,
423}
424#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
425pub struct ToolUseBlock {
426    pub tool_use_id: String,
427    pub name: String,
428    pub input: serde_json::Value,
429}
430
431#[derive(Clone, PartialEq, Debug, Serialize, Deserialize)]
432pub enum ToolResultStatus {
433    /// Renamed due to linting
434    #[serde(rename = "Error")]
435    IsError,
436    Success,
437    Unknown(UnknownVariantValue),
438}
439
440// TryFrom<T> implementations
441//
442// The AWS SDK's "string enums" (unit variants only) and unions (one payload
443// per variant) mirror mechanically, so the impls are macro-generated. The
444// error strings match the historical hand-written impls exactly:
445// `Unknown variant for TYPE: {invalid:?}`.
446
447/// Mirror a unit-variant AWS enum: emits owned + borrowed `TryFrom<aws>`
448/// impls. Unlisted variants (including this crate's `Unknown`) fall through to
449/// a `TypeConversionError`.
450macro_rules! mirror_enum {
451    ($ours:ident, $aws:ty { $($aws_variant:ident => $ours_variant:ident),+ $(,)? }) => {
452        impl TryFrom<$aws> for $ours {
453            type Error = TypeConversionError;
454            fn try_from(value: $aws) -> Result<Self, Self::Error> {
455                <$ours>::try_from(&value)
456            }
457        }
458        impl TryFrom<&$aws> for $ours {
459            type Error = TypeConversionError;
460            fn try_from(value: &$aws) -> Result<Self, Self::Error> {
461                type Aws = $aws;
462                match value {
463                    $(Aws::$aws_variant => Ok($ours::$ours_variant),)+
464                    invalid => Err(TypeConversionError::new(&format!(
465                        concat!("Unknown variant for ", stringify!($ours), ": {:?}"),
466                        invalid
467                    ))),
468                }
469            }
470        }
471    };
472}
473
474/// Mirror an AWS union (every listed variant carries a single payload that
475/// converts via `TryInto`): emits the owned `TryFrom<aws>` impl.
476macro_rules! mirror_union {
477    ($ours:ident, $aws:ty { $($aws_variant:ident => $ours_variant:ident),+ $(,)? }) => {
478        impl TryFrom<$aws> for $ours {
479            type Error = TypeConversionError;
480            fn try_from(value: $aws) -> Result<Self, Self::Error> {
481                type Aws = $aws;
482                match value {
483                    $(Aws::$aws_variant(value) => Ok($ours::$ours_variant(value.try_into()?)),)+
484                    invalid => Err(TypeConversionError::new(&format!(
485                        concat!("Unknown variant for ", stringify!($ours), ": {:?}"),
486                        invalid
487                    ))),
488                }
489            }
490        }
491    };
492}
493
494mirror_enum!(StopReason, aws_bedrock::StopReason {
495    ContentFiltered => ContentFiltered,
496    EndTurn => EndTurn,
497    GuardrailIntervened => GuardrailIntervened,
498    MaxTokens => MaxTokens,
499    StopSequence => StopSequence,
500    ToolUse => ToolUse,
501});
502mirror_enum!(ConversationRole, aws_bedrock::ConversationRole {
503    Assistant => Assistant,
504    User => User,
505});
506mirror_enum!(CachePointType, aws_bedrock::CachePointType {
507    Default => Default,
508});
509mirror_enum!(DocumentFormat, aws_bedrock::DocumentFormat {
510    Csv => Csv,
511    Doc => Doc,
512    Docx => Docx,
513    Html => Html,
514    Md => Md,
515    Pdf => Pdf,
516    Txt => Txt,
517    Xls => Xls,
518    Xlsx => Xlsx,
519});
520mirror_enum!(ImageFormat, aws_bedrock::ImageFormat {
521    Gif => Gif,
522    Jpeg => Jpeg,
523    Png => Png,
524    Webp => Webp,
525});
526mirror_enum!(VideoFormat, aws_bedrock::VideoFormat {
527    Flv => Flv,
528    Mkv => Mkv,
529    Mov => Mov,
530    Mp4 => Mp4,
531    Mpeg => Mpeg,
532    Mpg => Mpg,
533    ThreeGp => ThreeGp,
534    Webm => Webm,
535    Wmv => Wmv,
536});
537mirror_enum!(ToolResultStatus, aws_bedrock::ToolResultStatus {
538    Error => IsError,
539    Success => Success,
540});
541mirror_enum!(GuardrailConverseImageFormat, aws_bedrock::GuardrailConverseImageFormat {
542    Jpeg => Jpeg,
543    Png => Png,
544});
545mirror_enum!(GuardrailConverseContentQualifier, aws_bedrock::GuardrailConverseContentQualifier {
546    GroundingSource => GroundingSource,
547    GuardContent => GuardContent,
548    Query => Query,
549});
550
551mirror_union!(ConverseOutput, aws_bedrock::ConverseOutput {
552    Message => Message,
553});
554mirror_union!(ContentBlock, aws_bedrock::ContentBlock {
555    CachePoint => CachePoint,
556    CitationsContent => CitationsContent,
557    Document => Document,
558    GuardContent => GuardContent,
559    Image => Image,
560    ReasoningContent => ReasoningContent,
561    Text => Text,
562    ToolResult => ToolResult,
563    ToolUse => ToolUse,
564    Video => Video,
565});
566mirror_union!(CitationGeneratedContent, aws_bedrock::CitationGeneratedContent {
567    Text => Text,
568});
569mirror_union!(CitationSourceContent, aws_bedrock::CitationSourceContent {
570    Text => Text,
571});
572mirror_union!(CitationLocation, aws_bedrock::CitationLocation {
573    DocumentChar => DocumentChar,
574    DocumentChunk => DocumentChunk,
575    DocumentPage => DocumentPage,
576});
577mirror_union!(DocumentContentBlock, aws_bedrock::DocumentContentBlock {
578    Text => Text,
579});
580mirror_union!(GuardrailConverseContentBlock, aws_bedrock::GuardrailConverseContentBlock {
581    Image => Image,
582    Text => Text,
583});
584mirror_union!(GuardrailConverseImageSource, aws_bedrock::GuardrailConverseImageSource {
585    Bytes => Bytes,
586});
587mirror_union!(ImageSource, aws_bedrock::ImageSource {
588    Bytes => Bytes,
589    S3Location => S3Location,
590});
591mirror_union!(VideoSource, aws_bedrock::VideoSource {
592    Bytes => Bytes,
593    S3Location => S3Location,
594});
595mirror_union!(ReasoningContentBlock, aws_bedrock::ReasoningContentBlock {
596    ReasoningText => ReasoningText,
597    RedactedContent => RedactedContent,
598});
599
600// `DocumentSource::Content` carries a `Vec` payload and
601// `ToolResultContentBlock::Json` bridges Smithy `Document` into
602// `serde_json::Value`, so those two unions stay hand-written.
603
604impl TryFrom<aws_bedrock::DocumentSource> for DocumentSource {
605    type Error = TypeConversionError;
606    fn try_from(value: aws_bedrock::DocumentSource) -> Result<Self, Self::Error> {
607        match value {
608            aws_bedrock::DocumentSource::Bytes(value) => {
609                Ok(DocumentSource::Bytes(value.try_into()?))
610            }
611            aws_bedrock::DocumentSource::Content(value) => Ok(DocumentSource::Content(
612                value
613                    .into_iter()
614                    .map(TryInto::try_into)
615                    .collect::<Result<_, Self::Error>>()?,
616            )),
617            aws_bedrock::DocumentSource::S3Location(value) => {
618                Ok(DocumentSource::S3Location(value.try_into()?))
619            }
620            aws_bedrock::DocumentSource::Text(value) => Ok(DocumentSource::Text(value)),
621            invalid => Err(TypeConversionError::new(&format!(
622                "Unknown variant for DocumentSource: {invalid:?}"
623            ))),
624        }
625    }
626}
627
628impl TryFrom<aws_bedrock::ToolResultContentBlock> for ToolResultContentBlock {
629    type Error = TypeConversionError;
630    fn try_from(value: aws_bedrock::ToolResultContentBlock) -> Result<Self, Self::Error> {
631        match value {
632            aws_bedrock::ToolResultContentBlock::Document(value) => {
633                Ok(ToolResultContentBlock::Document(value.try_into()?))
634            }
635            aws_bedrock::ToolResultContentBlock::Image(value) => {
636                Ok(ToolResultContentBlock::Image(value.try_into()?))
637            }
638            aws_bedrock::ToolResultContentBlock::Json(value) => {
639                Ok(ToolResultContentBlock::Json(AwsDocument(value).into()))
640            }
641            aws_bedrock::ToolResultContentBlock::Text(value) => {
642                Ok(ToolResultContentBlock::Text(value))
643            }
644            aws_bedrock::ToolResultContentBlock::Video(value) => {
645                Ok(ToolResultContentBlock::Video(value.try_into()?))
646            }
647            invalid => Err(TypeConversionError::new(&format!(
648                "Unknown variant for ToolResultContentBlock: {invalid:?}"
649            ))),
650        }
651    }
652}
653
654// Struct conversions.
655
656impl TryFrom<aws_bedrock::TokenUsage> for TokenUsage {
657    type Error = TypeConversionError;
658    fn try_from(value: aws_bedrock::TokenUsage) -> Result<Self, Self::Error> {
659        Ok(TokenUsage {
660            input_tokens: value.input_tokens,
661            output_tokens: value.output_tokens,
662            total_tokens: value.total_tokens,
663            cache_read_input_tokens: value.cache_read_input_tokens,
664            cache_write_input_tokens: value.cache_write_input_tokens,
665        })
666    }
667}
668
669impl TryFrom<aws_bedrock::ConverseMetrics> for ConverseMetrics {
670    type Error = TypeConversionError;
671    fn try_from(value: aws_bedrock::ConverseMetrics) -> Result<Self, Self::Error> {
672        Ok(ConverseMetrics {
673            latency_ms: value.latency_ms,
674        })
675    }
676}
677
678impl TryFrom<aws_bedrock::Message> for Message {
679    type Error = TypeConversionError;
680    fn try_from(value: aws_bedrock::Message) -> Result<Self, Self::Error> {
681        Ok(Message {
682            role: value.role.try_into()?,
683            content: value
684                .content
685                .into_iter()
686                .map(TryInto::try_into)
687                .collect::<Result<_, Self::Error>>()?,
688        })
689    }
690}
691
692impl TryFrom<aws_bedrock::CachePointBlock> for CachePointBlock {
693    type Error = TypeConversionError;
694    fn try_from(value: aws_bedrock::CachePointBlock) -> Result<Self, Self::Error> {
695        Ok(CachePointBlock {
696            kind: value.r#type.try_into()?,
697        })
698    }
699}
700
701impl TryFrom<aws_bedrock::CitationsContentBlock> for CitationsContentBlock {
702    type Error = TypeConversionError;
703    fn try_from(value: aws_bedrock::CitationsContentBlock) -> Result<Self, Self::Error> {
704        Ok(CitationsContentBlock {
705            content: Some(
706                value
707                    .content
708                    .unwrap_or_default()
709                    .into_iter()
710                    .map(TryInto::try_into)
711                    .collect::<Result<_, Self::Error>>()?,
712            ),
713            citations: Some(
714                value
715                    .citations
716                    .unwrap_or_default()
717                    .into_iter()
718                    .map(TryInto::try_into)
719                    .collect::<Result<_, Self::Error>>()?,
720            ),
721        })
722    }
723}
724
725impl TryFrom<aws_bedrock::Citation> for Citation {
726    type Error = TypeConversionError;
727    fn try_from(value: aws_bedrock::Citation) -> Result<Self, Self::Error> {
728        Ok(Citation {
729            title: value.title,
730            source_content: Some(
731                value
732                    .source_content
733                    .unwrap_or_default()
734                    .into_iter()
735                    .map(TryInto::try_into)
736                    .collect::<Result<_, Self::Error>>()?,
737            ),
738            location: value.location.map(TryInto::try_into).transpose()?,
739        })
740    }
741}
742
743/// The three citation-location structs are field-identical; mirror them with
744/// one macro.
745macro_rules! mirror_location {
746    ($($name:ident),+ $(,)?) => {$(
747        impl TryFrom<aws_bedrock::$name> for $name {
748            type Error = TypeConversionError;
749            fn try_from(value: aws_bedrock::$name) -> Result<Self, Self::Error> {
750                Ok($name {
751                    document_index: value.document_index,
752                    start: value.start,
753                    end: value.end,
754                })
755            }
756        }
757    )+};
758}
759mirror_location!(
760    DocumentCharLocation,
761    DocumentChunkLocation,
762    DocumentPageLocation
763);
764
765impl TryFrom<aws_bedrock::DocumentBlock> for DocumentBlock {
766    type Error = TypeConversionError;
767    fn try_from(value: aws_bedrock::DocumentBlock) -> Result<Self, Self::Error> {
768        Ok(DocumentBlock {
769            format: value.format.try_into()?,
770            name: value.name,
771            source: value.source.map(TryInto::try_into).transpose()?,
772            context: value.context,
773            citations: value.citations.map(TryInto::try_into).transpose()?,
774        })
775    }
776}
777
778impl TryFrom<aws_bedrock::S3Location> for S3Location {
779    type Error = TypeConversionError;
780    fn try_from(value: aws_bedrock::S3Location) -> Result<Self, Self::Error> {
781        Ok(S3Location {
782            uri: value.uri,
783            bucket_owner: value.bucket_owner,
784        })
785    }
786}
787
788impl TryFrom<aws_sdk_bedrockruntime::primitives::Blob> for Blob {
789    type Error = TypeConversionError;
790    fn try_from(value: aws_sdk_bedrockruntime::primitives::Blob) -> Result<Self, Self::Error> {
791        Ok(Blob {
792            inner: value.into_inner(),
793        })
794    }
795}
796
797impl TryFrom<aws_bedrock::CitationsConfig> for CitationsConfig {
798    type Error = TypeConversionError;
799    fn try_from(value: aws_bedrock::CitationsConfig) -> Result<Self, Self::Error> {
800        Ok(CitationsConfig {
801            enabled: value.enabled,
802        })
803    }
804}
805
806impl TryFrom<aws_bedrock::GuardrailConverseImageBlock> for GuardrailConverseImageBlock {
807    type Error = TypeConversionError;
808    fn try_from(value: aws_bedrock::GuardrailConverseImageBlock) -> Result<Self, Self::Error> {
809        Ok(GuardrailConverseImageBlock {
810            format: value.format.try_into()?,
811            source: value.source.map(TryInto::try_into).transpose()?,
812        })
813    }
814}
815
816impl TryFrom<aws_bedrock::GuardrailConverseTextBlock> for GuardrailConverseTextBlock {
817    type Error = TypeConversionError;
818    fn try_from(value: aws_bedrock::GuardrailConverseTextBlock) -> Result<Self, Self::Error> {
819        Ok(GuardrailConverseTextBlock {
820            text: value.text,
821            qualifiers: Some(
822                value
823                    .qualifiers
824                    .unwrap_or_default()
825                    .into_iter()
826                    .map(|v| (&v).try_into())
827                    .collect::<Result<_, Self::Error>>()?,
828            ),
829        })
830    }
831}
832
833impl TryFrom<aws_bedrock::ImageBlock> for ImageBlock {
834    type Error = TypeConversionError;
835    fn try_from(value: aws_bedrock::ImageBlock) -> Result<Self, Self::Error> {
836        Ok(ImageBlock {
837            format: value.format.try_into()?,
838            source: value.source.map(TryInto::try_into).transpose()?,
839        })
840    }
841}
842
843impl TryFrom<aws_bedrock::ReasoningTextBlock> for ReasoningTextBlock {
844    type Error = TypeConversionError;
845    fn try_from(value: aws_bedrock::ReasoningTextBlock) -> Result<Self, Self::Error> {
846        Ok(ReasoningTextBlock {
847            text: value.text,
848            signature: value.signature,
849        })
850    }
851}
852
853impl TryFrom<aws_bedrock::ToolResultBlock> for ToolResultBlock {
854    type Error = TypeConversionError;
855    fn try_from(value: aws_bedrock::ToolResultBlock) -> Result<Self, Self::Error> {
856        Ok(ToolResultBlock {
857            tool_use_id: value.tool_use_id,
858            content: value
859                .content
860                .into_iter()
861                .map(TryInto::try_into)
862                .collect::<Result<_, Self::Error>>()?,
863            status: value.status.map(|v| (&v).try_into()).transpose()?,
864        })
865    }
866}
867
868impl TryFrom<aws_bedrock::VideoBlock> for VideoBlock {
869    type Error = TypeConversionError;
870    fn try_from(value: aws_bedrock::VideoBlock) -> Result<Self, Self::Error> {
871        Ok(VideoBlock {
872            format: value.format.try_into()?,
873            source: value.source.map(TryInto::try_into).transpose()?,
874        })
875    }
876}
877
878impl TryFrom<aws_bedrock::ToolUseBlock> for ToolUseBlock {
879    type Error = TypeConversionError;
880    fn try_from(value: aws_bedrock::ToolUseBlock) -> Result<Self, Self::Error> {
881        Ok(ToolUseBlock {
882            tool_use_id: value.tool_use_id,
883            name: value.name,
884            input: AwsDocument(value.input).into(),
885        })
886    }
887}
888
889#[cfg(test)]
890mod tests {
891    use super::*;
892    use serde_json::json;
893
894    /// The escape hatch's contract is that nothing the provider sent was
895    /// dropped, and the SDK's output type is `#[non_exhaustive]`, so the
896    /// conversion's rest pattern hides every field added upstream. This pins
897    /// the ones known today — `trace`, `performance_config` and `service_tier`
898    /// were all silently discarded before.
899    #[test]
900    fn converse_output_carries_every_sdk_field() {
901        let sdk_output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
902            .stop_reason(aws_bedrock::StopReason::GuardrailIntervened)
903            .output(aws_bedrock::ConverseOutput::Message(
904                aws_bedrock::Message::builder()
905                    .role(aws_bedrock::ConversationRole::Assistant)
906                    .content(aws_bedrock::ContentBlock::Text("blocked".into()))
907                    .build()
908                    .unwrap(),
909            ))
910            .usage(
911                aws_bedrock::TokenUsage::builder()
912                    .input_tokens(1)
913                    .output_tokens(2)
914                    .total_tokens(3)
915                    .build()
916                    .unwrap(),
917            )
918            .metrics(
919                aws_bedrock::ConverseMetrics::builder()
920                    .latency_ms(4)
921                    .build()
922                    .unwrap(),
923            )
924            .trace(
925                aws_bedrock::ConverseTrace::builder()
926                    .guardrail(aws_bedrock::GuardrailTraceAssessment::builder().build())
927                    .build(),
928            )
929            .performance_config(
930                aws_bedrock::PerformanceConfiguration::builder()
931                    .latency(aws_bedrock::PerformanceConfigLatency::Standard)
932                    .build(),
933            )
934            .service_tier(
935                aws_bedrock::ServiceTier::builder()
936                    .r#type(aws_bedrock::ServiceTierType::Default)
937                    .build()
938                    .unwrap(),
939            )
940            .build()
941            .unwrap();
942
943        let mirrored = InternalConverseOutput::try_from(sdk_output).unwrap();
944
945        assert!(mirrored.output.is_some());
946        assert_eq!(mirrored.stop_reason, StopReason::GuardrailIntervened);
947        assert!(mirrored.usage.is_some());
948        assert!(mirrored.metrics.is_some());
949        assert!(
950            mirrored.trace().is_some(),
951            "the guardrail trace must survive the conversion"
952        );
953        assert!(
954            mirrored.performance_config.is_some(),
955            "the performance configuration must survive the conversion"
956        );
957        assert!(
958            mirrored.service_tier.is_some(),
959            "the service tier must survive the conversion"
960        );
961    }
962
963    /// The SDK types behind `trace`, `performance_config` and `service_tier`
964    /// are not `Serialize`, so they are `#[serde(skip)]`: serializing must
965    /// still succeed and must not invent values on the way back.
966    #[test]
967    fn skipped_provider_fields_round_trip_as_absent() {
968        let output = InternalConverseOutput {
969            output: None,
970            stop_reason: StopReason::EndTurn,
971            usage: None,
972            metrics: None,
973            additional_model_response_fields: None,
974            request_id: Some("req-1".to_string()),
975            trace: None,
976            performance_config: None,
977            service_tier: None,
978        };
979
980        let json = serde_json::to_string(&output).unwrap();
981        let restored: InternalConverseOutput = serde_json::from_str(&json).unwrap();
982
983        assert_eq!(restored.request_id(), Some("req-1"));
984        assert!(restored.trace().is_none());
985        assert!(restored.performance_config.is_none());
986        assert!(restored.service_tier.is_none());
987    }
988
989    #[test]
990    fn mirror_enum_converts_known_variants() {
991        assert_eq!(
992            StopReason::try_from(aws_bedrock::StopReason::EndTurn).unwrap(),
993            StopReason::EndTurn
994        );
995        // Borrowed impl.
996        assert_eq!(
997            StopReason::try_from(&aws_bedrock::StopReason::ToolUse).unwrap(),
998            StopReason::ToolUse
999        );
1000        // A renamed pairing (aws `Error` -> ours `IsError`).
1001        assert_eq!(
1002            ToolResultStatus::try_from(aws_bedrock::ToolResultStatus::Error).unwrap(),
1003            ToolResultStatus::IsError
1004        );
1005    }
1006
1007    #[test]
1008    fn mirror_enum_unknown_variant_preserves_error_string() {
1009        let unknown = aws_bedrock::StopReason::from("weird_stop");
1010        let err = StopReason::try_from(unknown.clone()).unwrap_err();
1011        assert_eq!(
1012            err.to_string(),
1013            format!("Unknown variant for StopReason: {unknown:?}")
1014        );
1015
1016        let err =
1017            ConversationRole::try_from(aws_bedrock::ConversationRole::from("nope")).unwrap_err();
1018        assert!(
1019            err.to_string()
1020                .starts_with("Unknown variant for ConversationRole:")
1021        );
1022    }
1023
1024    #[test]
1025    fn additional_model_response_fields_survive_as_json() {
1026        let doc: AwsDocument = json!({"reasoning_effort": "low", "depth": 3}).into();
1027        let output = aws_sdk_bedrockruntime::operation::converse::ConverseOutput::builder()
1028            .stop_reason(aws_bedrock::StopReason::EndTurn)
1029            .additional_model_response_fields(doc.0)
1030            .build()
1031            .unwrap();
1032
1033        let internal = InternalConverseOutput::try_from(output).unwrap();
1034        assert_eq!(
1035            internal.additional_model_response_fields,
1036            Some(json!({"reasoning_effort": "low", "depth": 3}))
1037        );
1038
1039        // The whole normalized output stays serializable and the extras
1040        // survive a serde round trip.
1041        let value = serde_json::to_value(&internal).unwrap();
1042        assert_eq!(
1043            value.get("additional_model_response_fields"),
1044            Some(&json!({"reasoning_effort": "low", "depth": 3}))
1045        );
1046        let back: InternalConverseOutput = serde_json::from_value(value).unwrap();
1047        assert_eq!(back, internal);
1048    }
1049
1050    #[test]
1051    fn tool_use_input_decodes_into_json_value() {
1052        let aws_block = aws_bedrock::ToolUseBlock::builder()
1053            .tool_use_id("call_1")
1054            .name("add")
1055            .input(AwsDocument::from(json!({"x": 1, "y": 2})).0)
1056            .build()
1057            .unwrap();
1058
1059        let ours = ToolUseBlock::try_from(aws_block).unwrap();
1060        assert_eq!(ours.tool_use_id, "call_1");
1061        assert_eq!(ours.name, "add");
1062        assert_eq!(ours.input, json!({"x": 1, "y": 2}));
1063    }
1064}
rig_bedrock/types/converse_output.rs

rig_bedrock/types/
converse_output.rs
