harness_core/image.rs
1//! Optional image-generation trait. **Strictly opt-in** — nothing in `Model`,
2//! `AgentLoop`, `Hook`, `Guide`, `Sensor`, or `Memory` references this. Code
3//! that generates illustrations holds an `Arc<dyn ImageModel>` explicitly;
4//! everything else compiles without ever touching this module.
5//!
6//! Implementations live in `harness-models` (`ChatImageModel`, `OpenAiImage`,
7//! `DashScopeImage`).
8//!
9//! # Output convention: bytes, always
10//!
11//! Providers return generated images in at least three shapes — an inline
12//! `data:image/jpeg;base64,…` URI on the chat channel, a `b64_json` field, or
13//! a remote URL that **expires**. Adapters normalise all of them to owned
14//! bytes before returning. A caller never learns which transport its provider
15//! used, and never holds a handle that can rot.
16//!
17//! That last part is the whole point: an expiring URL handed to application
18//! code is a bug that surfaces a day later, in production, as a broken image.
19
20use crate::b64::base64_encode;
21use async_trait::async_trait;
22use serde::{Deserialize, Serialize};
23use std::fmt;
24
25/// An image fed *into* generation.
26///
27/// This is the mechanism behind character consistency across an illustrated
28/// sequence: generate page one, then pass it back as a reference for pages
29/// two onward so the same character keeps the same face.
30#[derive(Debug, Clone, Serialize, Deserialize)]
31#[non_exhaustive]
32pub enum ImageRef {
33 /// Raw bytes plus a MIME type like `"image/png"`. Adapters encode as the
34 /// provider requires.
35 Bytes { media_type: String, bytes: Vec<u8> },
36 /// Already-encoded standard base64, paired with a MIME type. Saves a
37 /// re-encode when the bytes came off the wire in this shape.
38 Base64 { media_type: String, base64: String },
39 /// A URL the *provider* fetches. Only usable with providers that accept
40 /// remote references; others must reject it rather than guess.
41 Url(String),
42}
43
44impl ImageRef {
45 /// Build a reference from raw bytes.
46 pub fn bytes(media_type: impl Into<String>, bytes: impl Into<Vec<u8>>) -> Self {
47 ImageRef::Bytes {
48 media_type: media_type.into(),
49 bytes: bytes.into(),
50 }
51 }
52
53 /// Standard-base64 payload for this reference, or `None` for [`ImageRef::Url`]
54 /// (nothing local to encode). Lets an adapter render `Bytes` and `Base64`
55 /// through one arm.
56 pub fn as_base64(&self) -> Option<(&str, std::borrow::Cow<'_, str>)> {
57 match self {
58 ImageRef::Bytes { media_type, bytes } => {
59 Some((media_type, base64_encode(bytes).into()))
60 }
61 ImageRef::Base64 { media_type, base64 } => Some((media_type, base64.as_str().into())),
62 ImageRef::Url(_) => None,
63 }
64 }
65}
66
67/// What to generate.
68#[derive(Debug, Clone, Default, Serialize, Deserialize)]
69pub struct ImageRequest {
70 pub prompt: String,
71 /// Provider-native size string (`"1024x1024"`, `"1024*1024"` — they differ,
72 /// deliberately not normalised). `None` uses the provider default.
73 pub size: Option<String>,
74 /// How many images to return. `0` is treated as `1`.
75 pub n: u8,
76 /// Reference images. An adapter whose provider cannot accept references
77 /// MUST return [`ImageError::Unsupported`] rather than dropping them —
78 /// a silently ignored reference yields a plausible-looking *wrong* image,
79 /// which is the hardest kind of failure to notice.
80 pub references: Vec<ImageRef>,
81}
82
83impl ImageRequest {
84 pub fn new(prompt: impl Into<String>) -> Self {
85 ImageRequest {
86 prompt: prompt.into(),
87 size: None,
88 n: 1,
89 references: Vec::new(),
90 }
91 }
92
93 pub fn with_size(mut self, size: impl Into<String>) -> Self {
94 self.size = Some(size.into());
95 self
96 }
97
98 pub fn with_n(mut self, n: u8) -> Self {
99 self.n = n;
100 self
101 }
102
103 pub fn with_references(mut self, references: Vec<ImageRef>) -> Self {
104 self.references = references;
105 self
106 }
107
108 /// Requested count, with `0` normalised to `1`.
109 pub fn count(&self) -> usize {
110 self.n.max(1) as usize
111 }
112}
113
114/// A generated image, fully materialised.
115///
116/// `bytes` serialises as standard base64 so a `ModelOutput` carrying images
117/// round-trips through session recording and deterministic replay.
118#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
119pub struct GeneratedImage {
120 /// MIME type, e.g. `"image/jpeg"`.
121 pub media_type: String,
122 #[serde(with = "crate::b64::serde_bytes_b64")]
123 pub bytes: Vec<u8>,
124}
125
126impl GeneratedImage {
127 pub fn new(media_type: impl Into<String>, bytes: impl Into<Vec<u8>>) -> Self {
128 GeneratedImage {
129 media_type: media_type.into(),
130 bytes: bytes.into(),
131 }
132 }
133}
134
135/// Failures from [`ImageModel::generate`]. Kept separate from `ModelError` for
136/// the same reason `EmbedError` is: the surface differs (no thinking, no
137/// tools, no streaming) and adapters should not reach across modules.
138#[derive(Debug)]
139#[non_exhaustive]
140pub enum ImageError {
141 /// Network / DNS / TLS / timeout.
142 Transport(String),
143 /// Non-2xx response or a body that did not contain an image.
144 Provider(String),
145 /// Caller passed something ungeneratable (empty prompt, `n` above the
146 /// provider's ceiling).
147 BadInput(String),
148 /// Provider or model does not support a requested feature — reference
149 /// images, `n > 1`, a particular size.
150 Unsupported(String),
151 /// Rate limit or quota exhaustion. Separate from `Provider` because the
152 /// correct response is to back off and retry, not to fail the task.
153 RateLimited(String),
154}
155
156impl fmt::Display for ImageError {
157 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158 match self {
159 ImageError::Transport(s) => write!(f, "image transport: {s}"),
160 ImageError::Provider(s) => write!(f, "image provider: {s}"),
161 ImageError::BadInput(s) => write!(f, "image bad input: {s}"),
162 ImageError::Unsupported(s) => write!(f, "image unsupported: {s}"),
163 ImageError::RateLimited(s) => write!(f, "image rate limited: {s}"),
164 }
165 }
166}
167
168impl std::error::Error for ImageError {}
169
170/// Producer of images from a text prompt.
171///
172/// Adapters MUST:
173/// - Return exactly [`ImageRequest::count`] images, or an error. Never fewer,
174/// silently.
175/// - Return fully materialised bytes — decode base64, fetch remote URLs. A
176/// caller must never receive an expiring handle.
177/// - Reject rather than ignore unsupported inputs (see [`ImageRequest::references`]).
178/// - Map provider rate-limit responses to [`ImageError::RateLimited`], so
179/// retry layers engage instead of failing the task.
180#[async_trait]
181pub trait ImageModel: Send + Sync + 'static {
182 async fn generate(&self, req: &ImageRequest) -> Result<Vec<GeneratedImage>, ImageError>;
183
184 /// Human-readable identifier, e.g. `"cpa:gemini-3.1-flash-image"`. Used in
185 /// logs and to tag stored artifacts, so a model swap is detectable after
186 /// the fact.
187 fn handle(&self) -> &str;
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193
194 #[test]
195 fn count_normalises_zero_to_one() {
196 assert_eq!(ImageRequest::new("x").with_n(0).count(), 1);
197 assert_eq!(ImageRequest::new("x").count(), 1);
198 assert_eq!(ImageRequest::new("x").with_n(4).count(), 4);
199 }
200
201 #[test]
202 fn image_ref_bytes_and_base64_agree() {
203 let raw = ImageRef::bytes("image/png", b"foobar".to_vec());
204 let (mt, b64) = raw.as_base64().unwrap();
205 assert_eq!(mt, "image/png");
206 assert_eq!(b64.as_ref(), "Zm9vYmFy");
207
208 let pre = ImageRef::Base64 {
209 media_type: "image/png".into(),
210 base64: "Zm9vYmFy".into(),
211 };
212 assert_eq!(pre.as_base64().unwrap().1.as_ref(), b64.as_ref());
213 }
214
215 #[test]
216 fn image_ref_url_has_no_local_payload() {
217 assert!(
218 ImageRef::Url("https://x/a.png".into())
219 .as_base64()
220 .is_none()
221 );
222 }
223
224 #[test]
225 fn generated_image_round_trips_through_json() {
226 // Non-UTF8 bytes on purpose: a JPEG header is not valid text, and a
227 // naive String-based encoding would corrupt it here.
228 let img = GeneratedImage::new("image/jpeg", vec![0xff, 0xd8, 0xff, 0xe0, 0x00]);
229 let json = serde_json::to_string(&img).unwrap();
230 assert!(
231 json.contains("\"/9j/4AA=\""),
232 "expected base64 payload, got {json}"
233 );
234 assert_eq!(serde_json::from_str::<GeneratedImage>(&json).unwrap(), img);
235 }
236}