Menu
Public repository / master
AkurAI-Framework / crates/cli/src/images.rs
Latest change 40e7ccf643510ab1faf1693f1491727fbfac8767 - File copies for any stored file (AKURAI_FILE_COPY_DIR/COLLECTIONS) replace the image API's save dir by Ólafur Búi Ólafsson
//! `POST /api/images/generations` — generate an image through an upstream image
//! API and store it in a collection's `file` field.
//!
//! Request body:
//!
//! ```json
//! { "collection": "images", "field": "image",
//! "record": { "parent": 3 },
//! "params": { "prompt": "a fox", "steps": 40, "init_image": "blob:<id>" } }
//! ```
//!
//! - `params` is POSTed as JSON to `AKURAI_IMAGE_URL` (bearer
//! `AKURAI_IMAGE_API_KEY` when set). A string `"blob:<id>"` (top level or in an
//! array) is replaced by a data URL of that stored blob, so edits of stored
//! images never round-trip through the browser.
//! - The first image of the response's `data[]` or `images[]` is used: bytes
//! from `b64_json`, else a GET of its `url` (relative to the upstream origin)
//! with the same bearer.
//! - The record is `record`, plus every declared field found in `params`, plus
//! every declared field on the upstream image item (so the seed actually used
//! is stored), plus the file, plus `duration_ms` (upstream round-trip) when the
//! collection declares it. Unknown keys are ignored.
//!
//! Plain HTTP only, like the LLM and embedding clients; TLS terminates at the edge.
use std::io::{Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use std::time::Duration;
use akurai_collections::FieldKind;
use akurai_http::{Method, Request, Response};
use akurai_json::Value;
use crate::collections_api::CollectionsApi;
/// Default upstream timeout: a large diffusion render on a small GPU takes minutes.
const DEFAULT_TIMEOUT_SECS: u64 = 900;
/// Where images come from: `AKURAI_IMAGE_URL`, `AKURAI_IMAGE_API_KEY`,
/// `AKURAI_IMAGE_TIMEOUT` (seconds). Copies to disk are the collections layer's job
/// (`AKURAI_FILE_COPY_DIR`).
pub struct Upstream {
pub host: String,
pub port: u16,
pub path: String,
pub api_key: Option<String>,
pub timeout: Duration,
}
impl Upstream {
/// `None` when `AKURAI_IMAGE_URL` is unset; `Err` when it is not plain HTTP.
pub fn from_env() -> Option<Result<Upstream, String>> {
let url = std::env::var("AKURAI_IMAGE_URL")
.ok()
.filter(|u| !u.trim().is_empty())?;
let timeout = std::env::var("AKURAI_IMAGE_TIMEOUT")
.ok()
.and_then(|s| s.trim().parse().ok())
.unwrap_or(DEFAULT_TIMEOUT_SECS);
let api_key = std::env::var("AKURAI_IMAGE_API_KEY")
.ok()
.filter(|k| !k.is_empty());
Some(Upstream::parse(&url, api_key, Duration::from_secs(timeout)))
}
fn parse(url: &str, api_key: Option<String>, timeout: Duration) -> Result<Upstream, String> {
let url = url.trim();
if akurai_http::endpoint::is_https_endpoint(url) {
return Err(
"AKURAI_IMAGE_URL must be plain http:// (TLS terminates at the edge)".into(),
);
}
let (host, port) = akurai_http::endpoint::parse_plain_endpoint(url)?;
let after_scheme = url.split_once("://").map_or(url, |(_, rest)| rest);
let path = after_scheme
.find('/')
.map_or("/", |i| &after_scheme[i..])
.to_string();
Ok(Upstream {
host,
port,
path,
api_key,
timeout,
})
}
}
/// One upstream round-trip: `(method, path, body)` → `(status, body)`.
pub type Send<'a> = dyn FnMut(&str, &str, Option<&[u8]>) -> Result<(u16, Vec<u8>), String> + 'a;
/// Handle the endpoint with the environment's upstream and a real socket.
pub fn handle(api: &CollectionsApi, req: &Request) -> Response {
let upstream = match Upstream::from_env() {
None => {
return error(
503,
"image generation not configured (set AKURAI_IMAGE_URL)",
)
}
Some(Err(e)) => return error(503, &e),
Some(Ok(u)) => u,
};
generate(api, req, &upstream, &mut |method, path, body| {
socket_send(&upstream, method, path, body)
})
}
/// The transport-free core of [`handle`], so it can be tested with a fake upstream.
pub fn generate(
api: &CollectionsApi,
req: &Request,
upstream: &Upstream,
send: &mut Send,
) -> Response {
if req.method != Method::Post {
return error(405, "use POST");
}
let body = match std::str::from_utf8(&req.body)
.ok()
.and_then(|s| akurai_json::parse(s).ok())
{
Some(v) => v,
None => return error(400, "body must be JSON"),
};
let (Some(collection), Some(field), Some(params)) = (
body.get("collection").and_then(Value::as_str),
body.get("field").and_then(Value::as_str),
body.get("params").filter(|p| matches!(p, Value::Object(_))),
) else {
return error(
400,
"body needs \"collection\", \"field\" and an object \"params\"",
);
};
let Some(coll) = api.collections().iter().find(|c| c.name == collection) else {
return error(404, "unknown collection");
};
let params = inline_blobs(api, params);
let started = std::time::Instant::now();
let (status, raw) = match send("POST", &upstream.path, Some(params.to_json().as_bytes())) {
Ok(r) => r,
Err(e) => return error(502, &format!("upstream unreachable: {e}")),
};
if !(200..300).contains(&status) {
return error(502, &format!("upstream {status}: {}", snippet(&raw)));
}
let parsed = std::str::from_utf8(&raw)
.ok()
.and_then(|s| akurai_json::parse(s).ok());
let Some(item) = parsed.as_ref().and_then(first_image) else {
return error(
502,
&format!("upstream returned no image: {}", snippet(&raw)),
);
};
let bytes = if let Some(b64) = item.get("b64_json").and_then(Value::as_str) {
match akurai_ws::base64::decode(b64) {
Ok(b) => b,
Err(e) => return error(502, &format!("upstream b64_json: {e}")),
}
} else if let Some(url) = item.get("url").and_then(Value::as_str) {
match send("GET", &same_origin_path(url), None) {
Ok((200..=299, b)) => b,
Ok((s, b)) => return error(502, &format!("upstream image {s}: {}", snippet(&b))),
Err(e) => return error(502, &format!("upstream image unreachable: {e}")),
}
} else {
return error(502, "upstream image has neither b64_json nor url");
};
let mut pairs = match body.get("record") {
Some(Value::Object(p)) => p.clone(),
_ => Vec::new(),
};
for f in coll
.fields
.iter()
.filter(|f| f.name != field && !matches!(f.kind, FieldKind::File))
{
for source in [params.get(&f.name), item.get(&f.name)] {
if let Some(v) = source.filter(|v| !matches!(v, Value::Null)) {
set(&mut pairs, &f.name, coerce(&f.kind, v.clone()));
}
}
}
// A declared `duration_ms` int field records the upstream round-trip time.
if matches!(
coll.field("duration_ms").map(|f| &f.kind),
Some(FieldKind::Int)
) {
set(
&mut pairs,
"duration_ms",
Value::Int(started.elapsed().as_millis() as i64),
);
}
let content_type = sniff(&bytes);
let filename = format!("image.{}", content_type.trim_start_matches("image/"));
let (status, record) = api.create_with_file(
collection,
Value::Object(pairs),
field,
&bytes,
&filename,
content_type,
);
crate::json_api::response(status, &record)
}
/// The first entry of an OpenAI `data[]` or an Unsloth-style `images[]`.
fn first_image(v: &Value) -> Option<&Value> {
["data", "images"].iter().find_map(|k| match v.get(k) {
Some(Value::Array(items)) => items.first(),
_ => None,
})
}
/// Replace `"blob:<id>"` strings (top level or inside arrays) with data URLs.
fn inline_blobs(api: &CollectionsApi, params: &Value) -> Value {
let swap = |v: &Value| match v.as_str().and_then(|s| s.strip_prefix("blob:")) {
Some(id) => match api.blob(id) {
Some(bytes) => Value::Str(format!(
"data:{};base64,{}",
sniff(&bytes),
akurai_ws::base64::encode(&bytes)
)),
None => v.clone(),
},
None => v.clone(),
};
match params {
Value::Object(pairs) => Value::Object(
pairs
.iter()
.map(|(k, v)| {
let v = match v {
Value::Array(items) => Value::Array(items.iter().map(swap).collect()),
other => swap(other),
};
(k.clone(), v)
})
.collect(),
),
other => other.clone(),
}
}
/// An int sent for a float field (JSON `4` for guidance 4.0) is stored as a float.
fn coerce(kind: &FieldKind, v: Value) -> Value {
match (kind, v) {
(FieldKind::Float, Value::Int(i)) => Value::Float(i as f64),
(_, v) => v,
}
}
fn set(pairs: &mut Vec<(String, Value)>, key: &str, value: Value) {
match pairs.iter_mut().find(|(k, _)| k == key) {
Some(slot) => slot.1 = value,
None => pairs.push((key.to_string(), value)),
}
}
/// An absolute `http://host:port/x` becomes `/x`; a relative path is kept.
// ponytail: images are fetched from the upstream origin; a different-host url is not followed.
fn same_origin_path(url: &str) -> String {
match url.split_once("://") {
Some((_, rest)) => rest.find('/').map_or("/".into(), |i| rest[i..].to_string()),
None => url.to_string(),
}
}
fn sniff(bytes: &[u8]) -> &'static str {
match bytes {
[0x89, b'P', b'N', b'G', ..] => "image/png",
[0xff, 0xd8, 0xff, ..] => "image/jpeg",
[b'R', b'I', b'F', b'F', _, _, _, _, b'W', b'E', b'B', b'P', ..] => "image/webp",
_ => "image/png",
}
}
fn snippet(raw: &[u8]) -> String {
String::from_utf8_lossy(raw).chars().take(300).collect()
}
fn error(status: u16, msg: &str) -> Response {
crate::json_api::response(
status,
&Value::Object(vec![("error".into(), Value::Str(msg.into()))]),
)
}
/// One HTTP/1.1 exchange over a fresh connection (`Connection: close`).
fn socket_send(
up: &Upstream,
method: &str,
path: &str,
body: Option<&[u8]>,
) -> Result<(u16, Vec<u8>), String> {
let addr = (up.host.as_str(), up.port)
.to_socket_addrs()
.map_err(|e| e.to_string())?
.next()
.ok_or("upstream host did not resolve")?;
let mut stream =
TcpStream::connect_timeout(&addr, Duration::from_secs(10)).map_err(|e| e.to_string())?;
stream
.set_read_timeout(Some(up.timeout))
.map_err(|e| e.to_string())?;
stream
.set_write_timeout(Some(up.timeout))
.map_err(|e| e.to_string())?;
stream
.write_all(&build_request(up, method, path, body))
.map_err(|e| e.to_string())?;
let mut raw = Vec::new();
stream.read_to_end(&mut raw).map_err(|e| e.to_string())?;
split_response(&raw)
}
fn build_request(up: &Upstream, method: &str, path: &str, body: Option<&[u8]>) -> Vec<u8> {
let mut head = format!(
"{method} {path} HTTP/1.1\r\nHost: {}:{}\r\nAccept: */*\r\nConnection: close\r\n",
up.host, up.port
);
// CR/LF in an env-sourced key must not inject headers.
if let Some(key) = up.api_key.as_deref().filter(|k| !k.contains(['\r', '\n'])) {
head.push_str(&format!("Authorization: Bearer {key}\r\n"));
}
if let Some(b) = body {
head.push_str(&format!(
"Content-Type: application/json\r\nContent-Length: {}\r\n",
b.len()
));
}
head.push_str("\r\n");
let mut out = head.into_bytes();
out.extend_from_slice(body.unwrap_or_default());
out
}
/// Split a raw HTTP response into `(status, body)`.
// ponytail: no chunked decoding; JSON/PNG upstreams send Content-Length. Add it if one streams.
fn split_response(raw: &[u8]) -> Result<(u16, Vec<u8>), String> {
let end = raw
.windows(4)
.position(|w| w == b"\r\n\r\n")
.ok_or("malformed upstream response")?;
let head = String::from_utf8_lossy(&raw[..end]);
if head
.to_ascii_lowercase()
.contains("transfer-encoding: chunked")
{
return Err("chunked upstream responses are not supported".into());
}
let status = head
.split_whitespace()
.nth(1)
.and_then(|s| s.parse().ok())
.ok_or("malformed upstream status line")?;
Ok((status, raw[end + 4..].to_vec()))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::collections_api::tests_support::api_with;
use akurai_collections::{Collection, Field};
const PNG: &[u8] = b"\x89PNG\r\n\x1a\nrest-of-image";
fn api(label: &str) -> CollectionsApi {
api_with(
label,
vec![Collection::new(
"images",
vec![
Field::new("prompt", FieldKind::Text).required(),
Field::new("seed", FieldKind::Int),
Field::new("guidance", FieldKind::Float),
Field::new("parent", FieldKind::Int),
Field::new("duration_ms", FieldKind::Int),
Field::new("image", FieldKind::File),
],
)],
)
}
fn upstream() -> Upstream {
Upstream::parse(
"http://127.0.0.1:8888/api/gen",
Some("k".into()),
Duration::from_secs(1),
)
.unwrap()
}
fn post(body: &str) -> Request {
let head = format!(
"POST /api/images/generations HTTP/1.1\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n",
body.len()
);
let mut req = Request::parse_head(head.as_bytes()).unwrap();
req.body = body.as_bytes().to_vec();
req
}
fn json(r: &Response) -> Value {
akurai_json::parse(std::str::from_utf8(&r.body).unwrap()).unwrap()
}
#[test]
fn stores_a_b64_image_with_declared_params_and_strips_routing_keys() {
let api = api("img-b64");
let mut sent = Vec::new();
let b64 = akurai_ws::base64::encode(PNG);
let r = generate(
&api,
&post(
r#"{"collection":"images","field":"image","params":{"prompt":"a fox","guidance":4,"steps":40}}"#,
),
&upstream(),
&mut |m, p, b| {
sent.push((
m.to_string(),
p.to_string(),
String::from_utf8(b.unwrap_or_default().to_vec()).unwrap(),
));
Ok((
200,
format!(r#"{{"data":[{{"b64_json":"{b64}"}}]}}"#).into_bytes(),
))
},
);
assert_eq!(r.status, 201, "{}", String::from_utf8_lossy(&r.body));
let rec = json(&r);
assert_eq!(rec.get("prompt").and_then(Value::as_str), Some("a fox"));
assert_eq!(rec.get("guidance").and_then(Value::as_f64), Some(4.0));
let file = rec.get("image").unwrap();
assert_eq!(
file.get("content_type").and_then(Value::as_str),
Some("image/png")
);
assert_eq!(
api.blob(file.get("blob").and_then(Value::as_str).unwrap())
.unwrap(),
PNG
);
assert_eq!(sent.len(), 1);
assert_eq!(
(sent[0].0.as_str(), sent[0].1.as_str()),
("POST", "/api/gen")
);
assert!(
sent[0].2.contains("\"steps\":40") && !sent[0].2.contains("collection"),
"{}",
sent[0].2
);
}
#[test]
fn fetches_a_url_image_and_records_upstream_values_and_record_fields() {
let api = api("img-url");
let mut calls = 0;
let r = generate(
&api,
&post(
r#"{"collection":"images","field":"image","record":{"parent":7},"params":{"prompt":"edit it"}}"#,
),
&upstream(),
&mut |m, p, _| {
calls += 1;
match (m, p) {
("POST", _) => Ok((
200,
br#"{"images":[{"url":"/g/1/file","seed":42}]}"#.to_vec(),
)),
("GET", "/g/1/file") => Ok((200, PNG.to_vec())),
other => panic!("unexpected call {other:?}"),
}
},
);
assert_eq!(r.status, 201, "{}", String::from_utf8_lossy(&r.body));
let rec = json(&r);
assert_eq!(rec.get("seed").and_then(Value::as_i64), Some(42));
assert!(
rec.get("duration_ms").and_then(Value::as_i64).is_some(),
"duration_ms recorded"
);
assert_eq!(rec.get("parent").and_then(Value::as_i64), Some(7));
assert_eq!(calls, 2);
}
#[test]
fn inlines_stored_blobs_as_data_urls() {
let api = api("img-blob");
let (_, src) = api.create_with_file(
"images",
Value::Object(vec![("prompt".into(), Value::Str("src".into()))]),
"image",
PNG,
"s.png",
"image/png",
);
let id = src
.get("image")
.and_then(|f| f.get("blob"))
.and_then(Value::as_str)
.unwrap()
.to_string();
let mut body_sent = String::new();
let b64 = akurai_ws::base64::encode(PNG);
generate(
&api,
&post(&format!(
r#"{{"collection":"images","field":"image","params":{{"prompt":"x","init_image":"blob:{id}","reference_images":["blob:{id}"]}}}}"#
)),
&upstream(),
&mut |_, _, b| {
body_sent = String::from_utf8(b.unwrap().to_vec()).unwrap();
Ok((
200,
format!(r#"{{"data":[{{"b64_json":"{b64}"}}]}}"#).into_bytes(),
))
},
);
let data_url = format!("data:image/png;base64,{b64}");
assert_eq!(body_sent.matches(&data_url).count(), 2, "{body_sent}");
}
#[test]
fn reports_bad_bodies_and_upstream_failures() {
let api = api("img-err");
let mut never = |_: &str, _: &str, _: Option<&[u8]>| -> Result<(u16, Vec<u8>), String> {
panic!("no call")
};
assert_eq!(
generate(&api, &post("{"), &upstream(), &mut never).status,
400
);
assert_eq!(
generate(
&api,
&post(r#"{"collection":"images"}"#),
&upstream(),
&mut never
)
.status,
400
);
let unknown = r#"{"collection":"nope","field":"image","params":{}}"#;
assert_eq!(
generate(&api, &post(unknown), &upstream(), &mut never).status,
404
);
let ok = r#"{"collection":"images","field":"image","params":{"prompt":"x"}}"#;
let r = generate(&api, &post(ok), &upstream(), &mut |_, _, _| {
Ok((503, b"no model".to_vec()))
});
assert_eq!(r.status, 502);
assert!(String::from_utf8_lossy(&r.body).contains("no model"));
let r = generate(&api, &post(ok), &upstream(), &mut |_, _, _| {
Err("refused".into())
});
assert_eq!(r.status, 502);
}
#[test]
fn parses_upstream_urls_and_http_responses() {
let u = upstream();
assert_eq!(
(u.host.as_str(), u.port, u.path.as_str()),
("127.0.0.1", 8888, "/api/gen")
);
assert!(Upstream::parse("https://x/y", None, Duration::ZERO).is_err());
assert_eq!(same_origin_path("http://h:1/a/b?t=1"), "/a/b?t=1");
assert_eq!(
split_response(b"HTTP/1.1 201 Created\r\nA: b\r\n\r\nbody").unwrap(),
(201, b"body".to_vec())
);
assert!(split_response(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n").is_err());
let req = String::from_utf8(build_request(&u, "POST", "/p", Some(b"{}"))).unwrap();
assert!(
req.starts_with("POST /p HTTP/1.1\r\n")
&& req.contains("Authorization: Bearer k\r\n")
&& req.ends_with("\r\n\r\n{}")
);
}
}