claude_agent/types/content/
image.rs

1//! Image source types for content blocks.
2
3use std::path::Path;
4
5use base64::prelude::*;
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
9#[serde(tag = "type", rename_all = "snake_case")]
10pub enum ImageSource {
11    Base64 { media_type: String, data: String },
12    Url { url: String },
13    File { file_id: String },
14}
15
16impl ImageSource {
17    pub fn base64(media_type: impl Into<String>, data: impl Into<String>) -> Self {
18        Self::Base64 {
19            media_type: media_type.into(),
20            data: data.into(),
21        }
22    }
23
24    pub fn from_url(url: impl Into<String>) -> Self {
25        Self::Url { url: url.into() }
26    }
27
28    pub fn from_file(file_id: impl Into<String>) -> Self {
29        Self::File {
30            file_id: file_id.into(),
31        }
32    }
33
34    pub fn jpeg(data: impl Into<String>) -> Self {
35        Self::Base64 {
36            media_type: "image/jpeg".into(),
37            data: data.into(),
38        }
39    }
40
41    pub fn png(data: impl Into<String>) -> Self {
42        Self::Base64 {
43            media_type: "image/png".into(),
44            data: data.into(),
45        }
46    }
47
48    pub fn gif(data: impl Into<String>) -> Self {
49        Self::Base64 {
50            media_type: "image/gif".into(),
51            data: data.into(),
52        }
53    }
54
55    pub fn webp(data: impl Into<String>) -> Self {
56        Self::Base64 {
57            media_type: "image/webp".into(),
58            data: data.into(),
59        }
60    }
61
62    pub fn is_base64(&self) -> bool {
63        matches!(self, Self::Base64 { .. })
64    }
65
66    pub fn is_url(&self) -> bool {
67        matches!(self, Self::Url { .. })
68    }
69
70    pub fn is_file(&self) -> bool {
71        matches!(self, Self::File { .. })
72    }
73
74    pub fn file_id(&self) -> Option<&str> {
75        match self {
76            Self::File { file_id } => Some(file_id),
77            _ => None,
78        }
79    }
80
81    pub fn media_type(&self) -> Option<&str> {
82        match self {
83            Self::Base64 { media_type, .. } => Some(media_type),
84            _ => None,
85        }
86    }
87
88    pub async fn from_path(path: impl AsRef<Path>) -> crate::Result<Self> {
89        let path = path.as_ref();
90        let data = tokio::fs::read(path).await.map_err(crate::Error::Io)?;
91        let media_type = mime_guess::from_path(path)
92            .first_or_octet_stream()
93            .to_string();
94        Ok(Self::Base64 {
95            media_type,
96            data: BASE64_STANDARD.encode(&data),
97        })
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_image_source_variants() {
107        let base64 = ImageSource::jpeg("data123");
108        assert!(base64.is_base64());
109        assert_eq!(base64.media_type(), Some("image/jpeg"));
110
111        let url = ImageSource::from_url("https://example.com/img.png");
112        assert!(url.is_url());
113
114        let file = ImageSource::from_file("file_abc123");
115        assert!(file.is_file());
116        assert_eq!(file.file_id(), Some("file_abc123"));
117    }
118
119    #[test]
120    fn test_image_source_serialization() {
121        let file = ImageSource::from_file("file_xyz");
122        let json = serde_json::to_string(&file).unwrap();
123        assert!(json.contains("\"type\":\"file\""));
124        assert!(json.contains("\"file_id\":\"file_xyz\""));
125
126        let url = ImageSource::from_url("https://example.com/img.png");
127        let json = serde_json::to_string(&url).unwrap();
128        assert!(json.contains("\"type\":\"url\""));
129    }
130
131    #[tokio::test]
132    async fn test_image_source_from_path() {
133        let dir = tempfile::tempdir().unwrap();
134        let png_path = dir.path().join("test.png");
135
136        let png_data: [u8; 67] = [
137            0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
138            0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00,
139            0x00, 0x1F, 0x15, 0xC4, 0x89, 0x00, 0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x78,
140            0x9C, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00, 0x01, 0x0D, 0x0A, 0x2D, 0xB4, 0x00,
141            0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82,
142        ];
143        tokio::fs::write(&png_path, &png_data).await.unwrap();
144
145        let source = ImageSource::from_path(&png_path).await.unwrap();
146        assert!(source.is_base64());
147        assert_eq!(source.media_type(), Some("image/png"));
148
149        if let ImageSource::Base64 { data, .. } = &source {
150            let decoded = base64::prelude::BASE64_STANDARD.decode(data).unwrap();
151            assert_eq!(decoded, png_data);
152        } else {
153            panic!("Expected Base64 source");
154        }
155    }
156
157    #[tokio::test]
158    async fn test_image_source_from_path_not_found() {
159        let result = ImageSource::from_path("/nonexistent/path/image.png").await;
160        assert!(result.is_err());
161    }
162}
claude_agent/types/content/image.rs

claude_agent/types/content/
image.rs
