gemini_adk_rs/tool/
media.rs

1//! Tool media returns — the ADK pattern where a function response carries
2//! images or other media to the model alongside its JSON payload, so
3//! vision tools (screenshots, chart renderers, document croppers) feed
4//! the model something it can actually look at.
5//!
6//! The mechanism is a convention, so it works with every existing tool
7//! shape unchanged: a tool embeds media under the reserved `"_media"` key
8//! of its JSON result (via [`attach`]), and the text-agent loop lifts it
9//! out of the function response and delivers it to the model as
10//! `inline_data` parts in the same turn. Tools that never touch `_media`
11//! behave exactly as before.
12//!
13//! ```ignore
14//! T::simple("chart", "Render a chart", |args| async move {
15//!     let png: Vec<u8> = render(&args)?;
16//!     let mut result = json!({"rendered": true});
17//!     media::attach(&mut result, "image/png", &png);
18//!     Ok(result)
19//! })
20//! ```
21
22use base64::Engine as _;
23
24/// Reserved result key holding media attachments.
25pub const MEDIA_KEY: &str = "_media";
26
27/// One media attachment lifted from a tool result.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct MediaAttachment {
30    /// MIME type, e.g. `"image/png"`.
31    pub mime_type: String,
32    /// Base64-encoded payload, ready for an `inline_data` part.
33    pub data_base64: String,
34}
35
36/// Attach raw media bytes to a tool result. Repeated calls accumulate.
37pub fn attach(result: &mut serde_json::Value, mime_type: &str, data: &[u8]) {
38    let encoded = base64::engine::general_purpose::STANDARD.encode(data);
39    let entry = serde_json::json!({"mime_type": mime_type, "data": encoded});
40    match result {
41        serde_json::Value::Object(map) => {
42            if let Some(list) = map
43                .entry(MEDIA_KEY)
44                .or_insert_with(|| serde_json::Value::Array(Vec::new()))
45                .as_array_mut()
46            {
47                list.push(entry);
48            }
49        }
50        other => {
51            // Non-object results are wrapped so the attachment has a home.
52            let wrapped = serde_json::json!({"result": other.take(), MEDIA_KEY: [entry]});
53            *other = wrapped;
54        }
55    }
56}
57
58/// Remove and return any media attachments from a tool result — called by
59/// the agent loop so the model receives media as parts, not as base64 JSON
60/// noise inside the function response.
61pub fn extract(result: &mut serde_json::Value) -> Vec<MediaAttachment> {
62    let Some(map) = result.as_object_mut() else {
63        return Vec::new();
64    };
65    let Some(raw) = map.remove(MEDIA_KEY) else {
66        return Vec::new();
67    };
68    raw.as_array()
69        .into_iter()
70        .flatten()
71        .filter_map(|entry| {
72            Some(MediaAttachment {
73                mime_type: entry.get("mime_type")?.as_str()?.to_string(),
74                data_base64: entry.get("data")?.as_str()?.to_string(),
75            })
76        })
77        .collect()
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use serde_json::json;
84
85    #[test]
86    fn attach_then_extract_round_trips() {
87        let mut result = json!({"ok": true});
88        attach(&mut result, "image/png", b"\x89PNG");
89        attach(&mut result, "image/jpeg", b"\xff\xd8");
90        let media = extract(&mut result);
91        assert_eq!(media.len(), 2);
92        assert_eq!(media[0].mime_type, "image/png");
93        assert_eq!(
94            media[0].data_base64,
95            base64::engine::general_purpose::STANDARD.encode(b"\x89PNG")
96        );
97        // The reserved key is gone; the payload is untouched.
98        assert_eq!(result, json!({"ok": true}));
99    }
100
101    #[test]
102    fn non_object_results_are_wrapped() {
103        let mut result = json!("plain text");
104        attach(&mut result, "image/png", b"x");
105        assert_eq!(result["result"], json!("plain text"));
106        assert_eq!(extract(&mut result).len(), 1);
107    }
108
109    #[test]
110    fn extract_without_media_is_a_no_op() {
111        let mut result = json!({"ok": true});
112        assert!(extract(&mut result).is_empty());
113        assert_eq!(result, json!({"ok": true}));
114    }
115}
gemini_adk_rs/tool/media.rs

gemini_adk_rs/tool/
media.rs
