Skip to main content

forge_primitives/
image.rs

1//! Image-attachment wire-shape types + validation, shared between
2//! forge-agent and a view.
3
4use serde::{Deserialize, Serialize};
5
6use crate::content::ContentBlock;
7use crate::messages::Extras;
8
9/// MIME types the Anthropic Vision API accepts as image attachments.
10pub const SUPPORTED_IMAGE_MIME_TYPES: &[&str] =
11    &["image/png", "image/jpeg", "image/gif", "image/webp"];
12
13/// A pending image attachment: base64-encoded data and its MIME type.
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15pub struct ImageAttachment {
16    pub data: String,
17    pub mime_type: String,
18}
19
20impl ImageAttachment {
21    /// The source object both block shapes share: the mime the row names, and
22    /// the payload only where the model must receive it.
23    fn source(&self, with_payload: bool) -> serde_json::Value {
24        let mut source = serde_json::Map::new();
25        source.insert("type".to_owned(), serde_json::Value::String("base64".to_owned()));
26        source.insert("media_type".to_owned(), serde_json::Value::String(self.mime_type.clone()));
27        if with_payload {
28            source.insert("data".to_owned(), serde_json::Value::String(self.data.clone()));
29        }
30        serde_json::Value::Object(source)
31    }
32
33    /// The content block the CLI receives, payload included.
34    pub fn wire_block(&self) -> serde_json::Value {
35        serde_json::json!({
36            "type": "image",
37            "source": self.source(true),
38        })
39    }
40
41    /// The content block a view draws: the mime the row names and the size it
42    /// prints, never the payload - a frame carrying the base64 is re-served by
43    /// every read of the row.
44    pub fn display_block(&self) -> ContentBlock {
45        let mut extras = Extras::new();
46        extras.insert("forge_bytes".to_owned(), serde_json::json!(payload_bytes(&self.data)));
47        ContentBlock::Image { source: self.source(false), extras }
48    }
49}
50
51/// The image blocks a prompt's attachments draw as: the ones the sender would
52/// keep, in the order they were sent, each carrying the mime and the measured
53/// size rather than its payload.
54pub fn display_blocks(images: &[ImageAttachment]) -> Vec<ContentBlock> {
55    images
56        .iter()
57        .filter(|image| validate_image(&image.data, &image.mime_type).is_ok())
58        .map(ImageAttachment::display_block)
59        .collect()
60}
61
62/// The bytes a base64 payload stands for: four characters carry three, and
63/// padding is not a byte.
64pub fn payload_bytes(data: &str) -> u64 {
65    let padding = if data.ends_with("==") { 2 } else { u64::from(data.ends_with('=')) };
66    (data.len() as u64 * 3 / 4).saturating_sub(padding)
67}
68
69/// Returns `true` if `mime_type` is a supported image MIME type.
70pub fn is_supported_image_type(mime_type: &str) -> bool {
71    SUPPORTED_IMAGE_MIME_TYPES.contains(&mime_type)
72}
73
74/// Returns `true` if `data` is non-empty, correctly padded, and decodes
75/// as valid standard base64.
76pub fn is_valid_base64(data: &str) -> bool {
77    use base64::Engine as _;
78    if data.is_empty() {
79        return false;
80    }
81    let clean = data.trim();
82    if !clean.len().is_multiple_of(4) {
83        return false;
84    }
85    base64::engine::general_purpose::STANDARD.decode(clean).is_ok()
86}
87
88/// Validate an image attachment before sending to the API.
89///
90/// # Errors
91///
92/// Returns a human-readable description of the first failure: unsupported
93/// MIME type or invalid base64.
94pub fn validate_image(data: &str, mime_type: &str) -> Result<(), String> {
95    if !is_supported_image_type(mime_type) {
96        return Err(format!(
97            "unsupported image type \"{mime_type}\"; expected one of: {}",
98            SUPPORTED_IMAGE_MIME_TYPES.join(", ")
99        ));
100    }
101    if !is_valid_base64(data) {
102        return Err("image data is not valid base64".to_owned());
103    }
104    Ok(())
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    /// **The sender's block carries the payload.** A block without `data`
112    /// still names an image, so the model reads a picture it was sent no bytes
113    /// for - and nothing else in the suite reads this shape.
114    #[test]
115    fn the_wire_block_carries_the_base64_payload() {
116        let image =
117            ImageAttachment { data: "iVBORw0KGgo=".to_owned(), mime_type: "image/png".to_owned() };
118        assert_eq!(
119            image.wire_block(),
120            serde_json::json!({
121                "type": "image",
122                "source": {
123                    "type": "base64",
124                    "media_type": "image/png",
125                    "data": "iVBORw0KGgo=",
126                },
127            }),
128            "the block the CLI receives keeps the payload it was sent",
129        );
130    }
131}