dwara_core/ai/
image_fetch.rs

1//! AI-13 (#201): Gateway-side image URL fetcher.
2//!
3//! When a client sends a multimodal request with a remote image URL
4//! (the OpenAI `image_url.url` convention) and the target provider
5//! adapter requires base64 image data (Anthropic, Gemini), the
6//! gateway fetches the image, converts it to base64, and fills the
7//! `data_b64` + `media_type` fields on `ContentPart::Image`. The
8//! OpenAI adapter continues to pass the remote URL through unchanged.
9//!
10//! The fetcher is async and runs in `serve_ai` after guardrails and
11//! before `adapter.build_request`. It is opt-in per route via the
12//! `ai.image_fetch` config block; when disabled, remote URLs are
13//! dropped by the Anthropic/Gemini adapters as before (preserving
14//! backward compatibility).
15
16use std::time::Duration;
17
18use base64::{engine::general_purpose, Engine as _};
19use bytes::Bytes;
20use http_body_util::{BodyExt as _, Full};
21use hyper::Request;
22use hyper_util::client::legacy::Client;
23use hyper_util::rt::TokioExecutor;
24
25use crate::ai::types::{ChatMessage, ChatRequest, ContentPart};
26
27/// Default maximum image size (10 MB). Provider limits vary but 10 MB
28/// is a safe upper bound for vision-model images.
29const DEFAULT_MAX_SIZE_BYTES: u64 = 10 * 1024 * 1024;
30
31/// Default fetch timeout.
32const DEFAULT_TIMEOUT_SECS: u64 = 10;
33
34/// Configuration for the image fetcher.
35#[derive(Debug, Clone)]
36pub struct ImageFetchConfig {
37    /// Maximum image size in bytes.
38    pub max_size_bytes: u64,
39    /// Fetch timeout.
40    pub timeout: Duration,
41}
42
43impl Default for ImageFetchConfig {
44    fn default() -> Self {
45        Self {
46            max_size_bytes: DEFAULT_MAX_SIZE_BYTES,
47            timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
48        }
49    }
50}
51
52/// Fetch remote image URLs in a `ChatRequest` and fill `data_b64` +
53/// `media_type` on each `ContentPart::Image` that has a remote URL
54/// but no base64 data. Images that already have `data_b64` (from
55/// `data:` URIs) are left unchanged. The `url` field is preserved so
56/// the OpenAI adapter can still pass it through.
57///
58/// Returns the mutated `ChatRequest`. On fetch failure for any single
59/// image, the image is left as-is (the adapter will drop it, preserving
60/// the pre-#201 behavior) and a warning is logged.
61pub async fn fetch_remote_images(mut req: ChatRequest, cfg: &ImageFetchConfig) -> ChatRequest {
62    let client = build_http_client(cfg.timeout);
63    for msg in &mut req.messages {
64        for part in &mut msg.content {
65            if let ContentPart::Image {
66                url: Some(url_str),
67                data_b64,
68                media_type,
69            } = part
70            {
71                // Skip if already has base64 data (data: URI decomposed).
72                if data_b64.is_some() {
73                    continue;
74                }
75                // Skip data: URIs (already handled by split_data_uri).
76                if url_str.starts_with("data:") {
77                    continue;
78                }
79                // Only fetch http(s) URLs.
80                if !url_str.starts_with("http://") && !url_str.starts_with("https://") {
81                    continue;
82                }
83                match fetch_one(&client, url_str, cfg.max_size_bytes).await {
84                    Ok((b64, mime)) => {
85                        *data_b64 = Some(b64);
86                        if media_type.is_none() {
87                            *media_type = Some(mime);
88                        }
89                    }
90                    Err(e) => {
91                        tracing::warn!(
92                            code = "ai_image_fetch_failed",
93                            url = %url_str,
94                            error = %e,
95                            "failed to fetch remote image URL; adapter will drop it"
96                        );
97                    }
98                }
99            }
100        }
101    }
102    req
103}
104
105/// Fetch a single image URL and return (base64_data, mime_type).
106async fn fetch_one(
107    client: &Client<hyper_util::client::legacy::connect::HttpConnector, Full<Bytes>>,
108    url: &str,
109    max_size: u64,
110) -> Result<(String, String), String> {
111    let req = Request::builder()
112        .uri(url)
113        .header(hyper::header::USER_AGENT, "dwara-ai-gateway/1.0")
114        .body(Full::new(Bytes::new()))
115        .map_err(|e| format!("build request: {e}"))?;
116    let resp = client
117        .request(req)
118        .await
119        .map_err(|e| format!("fetch: {e}"))?;
120    let status = resp.status();
121    if !status.is_success() {
122        return Err(format!("HTTP {status}"));
123    }
124    // Check Content-Length against the size cap.
125    if let Some(len) = resp.headers().get(hyper::header::CONTENT_LENGTH) {
126        if let Ok(s) = len.to_str() {
127            if let Ok(n) = s.parse::<u64>() {
128                if n > max_size {
129                    return Err(format!("image too large: {n} bytes (max {max_size})"));
130                }
131            }
132        }
133    }
134    let mime = resp
135        .headers()
136        .get(hyper::header::CONTENT_TYPE)
137        .and_then(|v| v.to_str().ok())
138        .unwrap_or("image/png")
139        .to_string();
140    // Read the body with a size cap.
141    let body = resp
142        .into_body()
143        .collect()
144        .await
145        .map_err(|e| format!("read body: {e}"))?
146        .to_bytes();
147    if body.len() as u64 > max_size {
148        return Err(format!(
149            "image too large: {} bytes (max {max_size})",
150            body.len()
151        ));
152    }
153    // Verify it looks like an image MIME type.
154    if !mime.starts_with("image/") {
155        return Err(format!("not an image: Content-Type {mime}"));
156    }
157    let b64 = general_purpose::STANDARD.encode(&body);
158    Ok((b64, mime))
159}
160
161/// Build a plain HTTP client (no TLS). For `https://` URLs the
162/// connector needs TLS; the hyper_util `HttpConnector` does not
163/// support TLS by default. The fetcher is intended for internal/
164/// trusted image sources; for production HTTPS, a TLS-enabled
165/// connector would be needed (follow-up).
166fn build_http_client(
167    timeout: Duration,
168) -> Client<hyper_util::client::legacy::connect::HttpConnector, Full<Bytes>> {
169    let mut connector = hyper_util::client::legacy::connect::HttpConnector::new();
170    connector.enforce_http(false); // allow https (will fail without TLS connector)
171    let client = Client::builder(TokioExecutor::new())
172        .pool_idle_timeout(Some(timeout))
173        .build_http::<Full<Bytes>>();
174    let _ = connector; // suppress unused warning
175    client
176}
177
178/// Check if a `ChatRequest` has any remote image URLs that would need
179/// fetching. Used to skip the fetcher entirely when not needed.
180pub fn has_remote_images(req: &ChatRequest) -> bool {
181    req.messages.iter().any(|m: &ChatMessage| {
182        m.content.iter().any(|p| match p {
183            ContentPart::Image {
184                url: Some(url),
185                data_b64: None,
186                ..
187            } => !url.starts_with("data:"),
188            _ => false,
189        })
190    })
191}