strands_agents/types/
media.rs

1//! Media-related type definitions for the SDK.
2
3use serde::{Deserialize, Serialize};
4
5use super::citations::CitationsConfig;
6
7/// Supported document formats.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
9#[serde(rename_all = "lowercase")]
10pub enum DocumentFormat {
11    Pdf,
12    Csv,
13    Doc,
14    Docx,
15    Xls,
16    Xlsx,
17    Html,
18    Txt,
19    Md,
20}
21
22impl DocumentFormat {
23    pub fn as_str(&self) -> &'static str {
24        match self {
25            Self::Pdf => "pdf",
26            Self::Csv => "csv",
27            Self::Doc => "doc",
28            Self::Docx => "docx",
29            Self::Xls => "xls",
30            Self::Xlsx => "xlsx",
31            Self::Html => "html",
32            Self::Txt => "txt",
33            Self::Md => "md",
34        }
35    }
36}
37
38/// Contains the content of a document.
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct DocumentSource {
41    #[serde(with = "base64_bytes")]
42    pub bytes: Vec<u8>,
43}
44
45impl DocumentSource {
46    pub fn new(bytes: Vec<u8>) -> Self {
47        Self { bytes }
48    }
49}
50
51/// A document to include in a message.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct DocumentContent {
54    pub format: DocumentFormat,
55    pub name: String,
56    pub source: DocumentSource,
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub citations: Option<CitationsConfig>,
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub context: Option<String>,
61}
62
63impl DocumentContent {
64    pub fn new(format: DocumentFormat, name: impl Into<String>, bytes: Vec<u8>) -> Self {
65        Self {
66            format,
67            name: name.into(),
68            source: DocumentSource::new(bytes),
69            citations: None,
70            context: None,
71        }
72    }
73
74    pub fn with_citations(mut self, config: CitationsConfig) -> Self {
75        self.citations = Some(config);
76        self
77    }
78
79    pub fn with_context(mut self, context: impl Into<String>) -> Self {
80        self.context = Some(context.into());
81        self
82    }
83}
84
85/// Supported image formats.
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
87#[serde(rename_all = "lowercase")]
88pub enum ImageFormat {
89    Png,
90    Jpeg,
91    Gif,
92    Webp,
93}
94
95impl ImageFormat {
96    pub fn as_str(&self) -> &'static str {
97        match self {
98            Self::Png => "png",
99            Self::Jpeg => "jpeg",
100            Self::Gif => "gif",
101            Self::Webp => "webp",
102        }
103    }
104}
105
106/// Contains the content of an image.
107#[derive(Debug, Clone, Serialize, Deserialize)]
108pub struct ImageSource {
109    #[serde(with = "base64_bytes")]
110    pub bytes: Vec<u8>,
111}
112
113impl ImageSource {
114    pub fn new(bytes: Vec<u8>) -> Self {
115        Self { bytes }
116    }
117}
118
119/// An image to include in a message.
120#[derive(Debug, Clone, Serialize, Deserialize)]
121pub struct ImageContent {
122    pub format: ImageFormat,
123    pub source: ImageSource,
124}
125
126impl ImageContent {
127    pub fn new(format: ImageFormat, bytes: Vec<u8>) -> Self {
128        Self {
129            format,
130            source: ImageSource::new(bytes),
131        }
132    }
133}
134
135/// Supported video formats.
136#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
137#[serde(rename_all = "lowercase")]
138pub enum VideoFormat {
139    Flv,
140    Mkv,
141    Mov,
142    Mpeg,
143    Mpg,
144    Mp4,
145    #[serde(rename = "three_gp")]
146    ThreeGp,
147    Webm,
148    Wmv,
149}
150
151impl VideoFormat {
152    pub fn as_str(&self) -> &'static str {
153        match self {
154            Self::Flv => "flv",
155            Self::Mkv => "mkv",
156            Self::Mov => "mov",
157            Self::Mpeg => "mpeg",
158            Self::Mpg => "mpg",
159            Self::Mp4 => "mp4",
160            Self::ThreeGp => "three_gp",
161            Self::Webm => "webm",
162            Self::Wmv => "wmv",
163        }
164    }
165}
166
167/// Contains the content of a video.
168#[derive(Debug, Clone, Serialize, Deserialize)]
169pub struct VideoSource {
170    #[serde(with = "base64_bytes")]
171    pub bytes: Vec<u8>,
172}
173
174impl VideoSource {
175    pub fn new(bytes: Vec<u8>) -> Self {
176        Self { bytes }
177    }
178}
179
180/// A video to include in a message.
181#[derive(Debug, Clone, Serialize, Deserialize)]
182pub struct VideoContent {
183    pub format: VideoFormat,
184    pub source: VideoSource,
185}
186
187impl VideoContent {
188    pub fn new(format: VideoFormat, bytes: Vec<u8>) -> Self {
189        Self {
190            format,
191            source: VideoSource::new(bytes),
192        }
193    }
194}
195
196/// Module for base64 encoding/decoding bytes in serde.
197mod base64_bytes {
198    use base64::{engine::general_purpose::STANDARD, Engine};
199    use serde::{self, Deserialize, Deserializer, Serializer};
200
201    pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
202    where
203        S: Serializer,
204    {
205        serializer.serialize_str(&STANDARD.encode(bytes))
206    }
207
208    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
209    where
210        D: Deserializer<'de>,
211    {
212        let s = String::deserialize(deserializer)?;
213        STANDARD.decode(&s).map_err(serde::de::Error::custom)
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn test_document_content() {
223        let doc = DocumentContent::new(DocumentFormat::Pdf, "test.pdf", vec![1, 2, 3]);
224        assert_eq!(doc.format, DocumentFormat::Pdf);
225        assert_eq!(doc.name, "test.pdf");
226    }
227
228    #[test]
229    fn test_image_content() {
230        let img = ImageContent::new(ImageFormat::Png, vec![1, 2, 3]);
231        assert_eq!(img.format, ImageFormat::Png);
232    }
233
234    #[test]
235    fn test_video_content() {
236        let vid = VideoContent::new(VideoFormat::Mp4, vec![1, 2, 3]);
237        assert_eq!(vid.format, VideoFormat::Mp4);
238    }
239
240    #[test]
241    fn test_document_format_as_str() {
242        assert_eq!(DocumentFormat::Pdf.as_str(), "pdf");
243        assert_eq!(DocumentFormat::Docx.as_str(), "docx");
244    }
245}
246
strands_agents/types/media.rs

strands_agents/types/
media.rs
