dwara_core/ai/
image_fetch.rs1use 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
27const DEFAULT_MAX_SIZE_BYTES: u64 = 10 * 1024 * 1024;
30
31const DEFAULT_TIMEOUT_SECS: u64 = 10;
33
34#[derive(Debug, Clone)]
36pub struct ImageFetchConfig {
37 pub max_size_bytes: u64,
39 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
52pub 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 if data_b64.is_some() {
73 continue;
74 }
75 if url_str.starts_with("data:") {
77 continue;
78 }
79 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
105async 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 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 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 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
161fn 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); let client = Client::builder(TokioExecutor::new())
172 .pool_idle_timeout(Some(timeout))
173 .build_http::<Full<Bytes>>();
174 let _ = connector; client
176}
177
178pub 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}