Richard ZinkandClaude Opus 5.5
b99f71526a
Agents tidy up folders after generating images (sort into subfolders, rename revisions); the image records kept the old path, so /api/images/<id>/file answered 404 and only the requesting device still showed them from its cache. - Generated images get a copy in the image store (images.blob, kept by the blob GC); existing ones are copied once on start-up. - Records follow: after bash commands with mv/rename (missing images of the owner from the last two weeks, found by name or by content hash), when moved with the file manager (file or folder), and on access (served from the copy, relocated in the background; without a copy found by file name). - Chat gallery falls back to the preview the model saw when the file is gone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
528 lines
25 KiB
Rust
528 lines
25 KiB
Rust
//! Image generation via fal.ai – as an agent tool and for the manual image page.
|
||
|
||
use crate::app::App;
|
||
use crate::config::{ImageModel, expand_env};
|
||
use crate::store::{ImageRecord, UsageRecord};
|
||
use crate::tools::{Tool, ToolCtx, ToolOutput, arg_str};
|
||
use crate::types::{ImageData, new_id, now_ms};
|
||
use anyhow::{Context, Result, anyhow, bail};
|
||
use async_trait::async_trait;
|
||
use base64::Engine;
|
||
use serde_json::{Value, json};
|
||
use std::path::{Path, PathBuf};
|
||
use std::time::Duration;
|
||
|
||
/// Size presets offered in the UI and to the model (FLUX-style names).
|
||
pub const SIZES: &[&str] = &["square_hd", "square", "landscape_4_3", "landscape_16_9", "portrait_4_3", "portrait_16_9"];
|
||
|
||
fn aspect_ratio(size: &str) -> &'static str {
|
||
match size {
|
||
"landscape_4_3" => "4:3",
|
||
"landscape_16_9" => "16:9",
|
||
"portrait_4_3" => "3:4",
|
||
"portrait_16_9" => "9:16",
|
||
_ => "1:1",
|
||
}
|
||
}
|
||
|
||
pub struct GenerateRequest {
|
||
pub prompt: String,
|
||
pub model: Option<String>,
|
||
pub size: Option<String>,
|
||
pub count: u32,
|
||
/// Directory the files are written to.
|
||
pub dir: PathBuf,
|
||
pub user_id: Option<String>,
|
||
pub session_id: Option<String>,
|
||
/// Images to refine (sent to models with a `reference_param`).
|
||
pub references: Vec<PathBuf>,
|
||
}
|
||
|
||
fn slug(s: &str) -> String {
|
||
let mut out: String = s
|
||
.chars()
|
||
.map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' })
|
||
.collect();
|
||
while out.contains("--") {
|
||
out = out.replace("--", "-");
|
||
}
|
||
out.trim_matches('-').chars().take(40).collect::<String>().trim_end_matches('-').to_string()
|
||
}
|
||
|
||
/// Purposes the agent can ask for; models list theirs in `good_for`.
|
||
pub const PURPOSES: &[&str] = &["photo", "illustration", "text", "design"];
|
||
|
||
/// Defaults for well-known fal.ai models: (id, stage, good_for, reference_param, size_param, price).
|
||
const KNOWN: &[(&str, &str, &[&str], Option<&str>, Option<&str>, f64)] = &[
|
||
("fal-ai/flux/schnell", "preview", &["photo", "illustration"], None, None, 0.003),
|
||
("fal-ai/nano-banana-pro", "final", &["photo", "illustration", "text"], None, Some("aspect_ratio"), 0.15),
|
||
("fal-ai/nano-banana-pro/edit", "final", &["photo", "illustration", "text", "design"], Some("image_urls"), Some("aspect_ratio"), 0.15),
|
||
("fal-ai/flux-2-pro", "final", &["photo"], None, None, 0.03),
|
||
("fal-ai/ideogram/v3", "final", &["text", "design"], None, Some("aspect_ratio"), 0.06),
|
||
("fal-ai/recraft/v4/pro/text-to-image", "final", &["design", "illustration"], None, None, 0.05),
|
||
];
|
||
|
||
/// A configured model with the known defaults filled in where the config leaves them out.
|
||
fn with_defaults(mut m: ImageModel) -> ImageModel {
|
||
if let Some((_, stage, good_for, reference, size, price)) = KNOWN.iter().find(|k| k.0 == m.id) {
|
||
m.stage.get_or_insert_with(|| stage.to_string());
|
||
if m.good_for.is_empty() {
|
||
m.good_for = good_for.iter().map(|s| s.to_string()).collect();
|
||
}
|
||
if m.reference_param.is_none() {
|
||
m.reference_param = reference.map(String::from);
|
||
}
|
||
if m.size_param.is_none() {
|
||
m.size_param = size.map(String::from);
|
||
}
|
||
m.price.get_or_insert(*price);
|
||
}
|
||
m
|
||
}
|
||
|
||
fn models(app: &App) -> Vec<ImageModel> {
|
||
app.config().images.models.iter().cloned().map(with_defaults).collect()
|
||
}
|
||
|
||
fn model_config(app: &App, id: &str) -> ImageModel {
|
||
models(app).into_iter().find(|m| m.id == id).unwrap_or_else(|| with_defaults(ImageModel { id: id.to_string(), ..Default::default() }))
|
||
}
|
||
|
||
/// Pick the model for a request: an explicit model wins; previews use the first preview model
|
||
/// (else the default); final images use the best-matching final model, preferring one that can
|
||
/// take the reference image when there is one.
|
||
pub fn choose_model(app: &App, stage: &str, purpose: Option<&str>, with_reference: bool) -> ImageModel {
|
||
let default = app.config().images.default_model.clone();
|
||
pick(&models(app), stage, purpose, with_reference).unwrap_or_else(|| model_config(app, &default))
|
||
}
|
||
|
||
fn pick(all: &[ImageModel], stage: &str, purpose: Option<&str>, with_reference: bool) -> Option<ImageModel> {
|
||
if stage == "preview" && !with_reference {
|
||
return all.iter().find(|m| m.stage.as_deref() == Some("preview")).cloned();
|
||
}
|
||
let score = |m: &ImageModel| {
|
||
let mut s = 0;
|
||
if m.stage.as_deref() == Some("final") {
|
||
s += 4;
|
||
}
|
||
match (with_reference, m.reference_param.is_some()) {
|
||
(true, true) => s += 8,
|
||
// Without a reference, plain text-to-image endpoints come first.
|
||
(false, true) => s -= 1,
|
||
_ => {}
|
||
}
|
||
if purpose.is_some_and(|p| m.good_for.iter().any(|g| g == p)) {
|
||
s += 2;
|
||
}
|
||
s
|
||
};
|
||
// max_by_key keeps the last of equal scores; reverse so the config order wins ties.
|
||
all.iter().rev().max_by_key(|m| score(m)).cloned()
|
||
}
|
||
|
||
/// Read local reference images as data URIs (fal.ai accepts them in place of URLs).
|
||
async fn reference_uris(paths: &[PathBuf]) -> Result<Vec<String>> {
|
||
let mut out = Vec::new();
|
||
for p in paths {
|
||
let bytes = tokio::fs::read(p).await.with_context(|| format!("reference image {}", p.display()))?;
|
||
let mime = match p.extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase()).as_deref() {
|
||
Some("jpg" | "jpeg") => "image/jpeg",
|
||
Some("webp") => "image/webp",
|
||
_ => "image/png",
|
||
};
|
||
out.push(format!("data:{mime};base64,{}", base64::engine::general_purpose::STANDARD.encode(bytes)));
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
async fn download(url: &str) -> Result<(Vec<u8>, String)> {
|
||
if let Some(rest) = url.strip_prefix("data:") {
|
||
let (meta, data) = rest.split_once(',').context("malformed data URI")?;
|
||
let ctype = meta.split(';').next().unwrap_or("image/png").to_string();
|
||
return Ok((base64::engine::general_purpose::STANDARD.decode(data)?, ctype));
|
||
}
|
||
let resp = crate::provider::http_client().get(url).timeout(Duration::from_secs(120)).send().await?.error_for_status()?;
|
||
let ctype = resp.headers().get(reqwest::header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).unwrap_or("image/png").to_string();
|
||
Ok((resp.bytes().await?.to_vec(), ctype))
|
||
}
|
||
|
||
/// Generate images, store them in `req.dir` and record them (history + cost).
|
||
pub async fn generate(app: &App, req: GenerateRequest) -> Result<Vec<ImageRecord>> {
|
||
let cfg = app.config();
|
||
let ic = &cfg.images;
|
||
if !ic.enabled {
|
||
bail!("image generation is disabled – enable [images] in the settings");
|
||
}
|
||
let key = expand_env(&ic.api_key);
|
||
if key.is_empty() {
|
||
bail!("no fal.ai API key configured (images.api_key, e.g. ${{FAL_KEY}})");
|
||
}
|
||
crate::budget::check(app, "fal", req.user_id.as_deref())?;
|
||
let model_id = req.model.clone().filter(|m| !m.is_empty()).unwrap_or_else(|| ic.default_model.clone());
|
||
let model = model_config(app, &model_id);
|
||
let explicit_size = req.size.clone().filter(|s| SIZES.contains(&s.as_str()));
|
||
let size = explicit_size.clone().unwrap_or_else(|| "landscape_4_3".into());
|
||
let count = req.count.clamp(1, 4);
|
||
|
||
let mut body = json!({"prompt": req.prompt, "num_images": count, "output_format": "png"});
|
||
let refining = !req.references.is_empty() && model.reference_param.is_some();
|
||
match model.size_param.as_deref() {
|
||
// When refining without an explicit size, keep the reference's aspect ratio.
|
||
Some("aspect_ratio") if refining && explicit_size.is_none() => body["aspect_ratio"] = json!("auto"),
|
||
Some("aspect_ratio") => body["aspect_ratio"] = json!(aspect_ratio(&size)),
|
||
Some("none") => {}
|
||
_ => body["image_size"] = json!(size),
|
||
}
|
||
if refining && let Some(param) = &model.reference_param {
|
||
body[param.as_str()] = json!(reference_uris(&req.references).await?);
|
||
}
|
||
for (k, v) in &model.options {
|
||
body[k] = v.clone();
|
||
}
|
||
|
||
let url = format!("{}/{}", expand_env(&ic.base_url).trim_end_matches('/'), model_id.trim_matches('/'));
|
||
let resp = crate::provider::http_client()
|
||
.post(&url)
|
||
.header("Authorization", format!("Key {key}"))
|
||
.timeout(Duration::from_secs(300))
|
||
.json(&body)
|
||
.send()
|
||
.await
|
||
.context("fal.ai request failed")?;
|
||
if !resp.status().is_success() {
|
||
return Err(crate::provider::http_error(resp).await);
|
||
}
|
||
let out: Value = resp.json().await?;
|
||
let images = out["images"].as_array().cloned().or_else(|| out["image"].as_object().map(|o| vec![Value::Object(o.clone())])).unwrap_or_default();
|
||
if images.is_empty() {
|
||
bail!("fal.ai returned no images: {}", out.to_string().chars().take(300).collect::<String>());
|
||
}
|
||
|
||
tokio::fs::create_dir_all(&req.dir).await?;
|
||
let stamp = chrono::Local::now().format("%Y%m%d-%H%M%S");
|
||
let base = slug(&req.prompt);
|
||
let mut records = Vec::new();
|
||
for (i, img) in images.iter().enumerate() {
|
||
let src = img["url"].as_str().context("image without url")?;
|
||
let (bytes, ctype) = download(src).await?;
|
||
let ext = match ctype.as_str() {
|
||
"image/jpeg" => "jpg",
|
||
"image/webp" => "webp",
|
||
_ => "png",
|
||
};
|
||
let name = if images.len() > 1 { format!("{stamp}-{base}-{}.{ext}", i + 1) } else { format!("{stamp}-{base}.{ext}") };
|
||
let path = req.dir.join(name);
|
||
tokio::fs::write(&path, &bytes).await?;
|
||
// A copy in the image store keeps it viewable when the file is moved or deleted.
|
||
let blob = app.store.put_blob(&ctype, &bytes).map_err(|e| tracing::warn!("image copy: {e:#}")).ok();
|
||
let rec = ImageRecord {
|
||
id: new_id(),
|
||
user_id: req.user_id.clone(),
|
||
session_id: req.session_id.clone(),
|
||
prompt: req.prompt.clone(),
|
||
model: model_id.clone(),
|
||
path: path.display().to_string(),
|
||
width: img["width"].as_i64().unwrap_or(0),
|
||
height: img["height"].as_i64().unwrap_or(0),
|
||
created_at: now_ms(),
|
||
blob,
|
||
};
|
||
app.store.add_image(&rec)?;
|
||
records.push(rec);
|
||
}
|
||
|
||
// Costs go into the normal statistics (provider "fal").
|
||
let cost = model.price.unwrap_or(0.0) * records.len() as f64;
|
||
let session = req.session_id.clone().unwrap_or_default();
|
||
let _ = app.store.add_usage(&UsageRecord {
|
||
ts: now_ms(),
|
||
session_id: session.clone(),
|
||
root_session_id: if session.is_empty() { String::new() } else { app.root_session_id(&session) },
|
||
provider: "fal".into(),
|
||
model: model_id,
|
||
agent: "image".into(),
|
||
input_tokens: 0,
|
||
output_tokens: 0,
|
||
cache_read_tokens: 0,
|
||
cache_write_tokens: 0,
|
||
ttft_ms: 0,
|
||
duration_ms: 0,
|
||
cost,
|
||
error: false,
|
||
user_id: None,
|
||
allowance: None,
|
||
});
|
||
Ok(records)
|
||
}
|
||
|
||
/// A generated image whose file is gone: the same file name below its original folder
|
||
/// (or the folder above it), e.g. after an agent sorted the images into subfolders.
|
||
/// Only inside `roots`; looks at a few thousand entries at most.
|
||
pub fn find_moved(path: &Path, roots: &[PathBuf]) -> Option<PathBuf> {
|
||
let name = path.file_name()?;
|
||
search_near(path, roots, |p, _| p.file_name() == Some(name))
|
||
}
|
||
|
||
/// Like `find_moved`, but by content: an image file with the same bytes (renamed).
|
||
pub fn find_by_content(path: &Path, roots: &[PathBuf], hash: &str, size: u64) -> Option<PathBuf> {
|
||
search_near(path, roots, |p, len| {
|
||
let ext = p.extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase());
|
||
len == size && matches!(ext.as_deref(), Some("png" | "jpg" | "jpeg" | "webp")) && std::fs::read(p).is_ok_and(|b| crate::blobs::hash(&b) == hash)
|
||
})
|
||
}
|
||
|
||
/// Files below the folder of `path` and the folder above it (3 levels, hidden folders
|
||
/// skipped, inside `roots`) for which `matches(path, size)` holds.
|
||
fn search_near(path: &Path, roots: &[PathBuf], matches: impl Fn(&Path, u64) -> bool) -> Option<PathBuf> {
|
||
let dir = path.parent()?;
|
||
let mut budget = 5000usize;
|
||
for base in [Some(dir), dir.parent()].into_iter().flatten() {
|
||
if !roots.iter().any(|r| base.starts_with(r)) {
|
||
continue;
|
||
}
|
||
let mut stack = vec![(base.to_path_buf(), 0usize)];
|
||
while let Some((d, depth)) = stack.pop() {
|
||
let Ok(entries) = std::fs::read_dir(&d) else { continue };
|
||
for e in entries.flatten() {
|
||
if budget == 0 {
|
||
return None;
|
||
}
|
||
budget -= 1;
|
||
let p = e.path();
|
||
let Ok(ft) = e.file_type() else { continue };
|
||
if ft.is_file() && p != path && matches(&p, e.metadata().map(|m| m.len()).unwrap_or(0)) {
|
||
return Some(p);
|
||
}
|
||
let hidden = p.file_name().and_then(|n| n.to_str()).is_some_and(|n| n.starts_with('.') || n == "node_modules");
|
||
if ft.is_dir() && depth < 3 && !hidden {
|
||
stack.push((p, depth + 1));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// Where the file of a generated image is now: by name, then (with a copy in the image
|
||
/// store) by content. Corrects the stored path when found.
|
||
pub fn relocate(app: &App, rec: &ImageRecord) -> Option<PathBuf> {
|
||
let path = Path::new(&rec.path);
|
||
let roots = app.user_roots(rec.user_id.as_deref());
|
||
let found = find_moved(path, &roots).or_else(|| {
|
||
let hash = rec.blob.as_deref()?;
|
||
let (_, bytes) = app.store.blob(hash).ok()??;
|
||
find_by_content(path, &roots, hash, bytes.len() as u64)
|
||
})?;
|
||
let _ = app.store.set_image_path(&rec.id, &found.display().to_string());
|
||
Some(found)
|
||
}
|
||
|
||
/// After an agent moved or renamed files: find the generated images of `owner` from the
|
||
/// last two weeks whose files are gone and correct their paths. Returns how many.
|
||
pub fn relocate_missing(app: &App, owner: Option<&str>) -> usize {
|
||
let since = now_ms() - 14 * 86_400_000;
|
||
let Ok(list) = app.store.images(owner, 500) else { return 0 };
|
||
list.iter().filter(|r| r.created_at >= since && !Path::new(&r.path).exists()).filter(|r| relocate(app, r).is_some()).count()
|
||
}
|
||
|
||
/// Copies in the image store for generated images from before they were kept (once per
|
||
/// image, files up to 25 MB): later moves and renames are then found by content.
|
||
pub fn backfill_copies(app: &App) -> Result<usize> {
|
||
let mut n = 0;
|
||
for (id, path) in app.store.images_without_copy()? {
|
||
let Ok(meta) = std::fs::metadata(&path) else { continue };
|
||
if !meta.is_file() || meta.len() > 25 * 1024 * 1024 {
|
||
continue;
|
||
}
|
||
let Ok(bytes) = std::fs::read(&path) else { continue };
|
||
let mime = match Path::new(&path).extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase()).as_deref() {
|
||
Some("jpg" | "jpeg") => "image/jpeg",
|
||
Some("webp") => "image/webp",
|
||
_ => "image/png",
|
||
};
|
||
let hash = app.store.put_blob(mime, &bytes)?;
|
||
app.store.set_image_blob(&id, &hash)?;
|
||
n += 1;
|
||
}
|
||
Ok(n)
|
||
}
|
||
|
||
/// Downscaled copy for the model's context.
|
||
fn preview_for_model(path: &Path) -> Option<ImageData> {
|
||
let img = image::open(path).ok()?;
|
||
let img = img.resize(1024, 1024, image::imageops::FilterType::Triangle);
|
||
let mut buf = Vec::new();
|
||
img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png).ok()?;
|
||
Some(ImageData { blob: None, size: None, media_type: "image/png".into(), data: base64::engine::general_purpose::STANDARD.encode(buf) })
|
||
}
|
||
|
||
pub struct GenerateImage;
|
||
|
||
#[async_trait]
|
||
impl Tool for GenerateImage {
|
||
fn name(&self) -> &str {
|
||
"generate_image"
|
||
}
|
||
fn description(&self) -> String {
|
||
"Generate images from a text prompt (fal.ai) and save them in the project. Use for \
|
||
illustrations, icons, backgrounds or photos the user needs. Write a detailed visual prompt \
|
||
in English (subject, style, composition, colours, lighting). The server picks the model. \
|
||
Work in two steps: first stage \"preview\" (cheap drafts, count 2-4); show them and ask the \
|
||
user with ask_user which one to refine – include the final model and price the preview \
|
||
result names. Then call again with stage \"final\", the chosen file in \"reference\" and \
|
||
the same or an improved prompt. Skip the preview only if the user asks for a final image \
|
||
directly. Each image costs money – generate only what is needed."
|
||
.into()
|
||
}
|
||
fn schema(&self) -> Value {
|
||
json!({
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt": {"type": "string"},
|
||
"stage": {"type": "string", "enum": ["preview", "final"], "description": "preview = cheap drafts (default), final = best model for the purpose"},
|
||
"purpose": {"type": "string", "enum": PURPOSES, "description": "What the image is: photo (realistic), illustration, text (lettering, posters, signs), design (logos, icons, UI, vector style)"},
|
||
"reference": {"type": "array", "items": {"type": "string"}, "maxItems": 4, "description": "Image files (relative paths) to refine, e.g. the preview the user chose"},
|
||
"size": {"type": "string", "enum": SIZES, "description": "Default: landscape_4_3; when refining, the reference's aspect ratio"},
|
||
"count": {"type": "integer", "minimum": 1, "maximum": 4},
|
||
"directory": {"type": "string", "description": "Target directory relative to the working directory (default: images.output_dir)"},
|
||
"model": {"type": "string", "description": "fal.ai model id – only if the user asks for a specific model"}
|
||
},
|
||
"required": ["prompt"]
|
||
})
|
||
}
|
||
fn title(&self, input: &Value) -> String {
|
||
let stage = if input["stage"].as_str() == Some("final") { "Generate final image" } else { "Generate image" };
|
||
format!("{stage}: {}", input["prompt"].as_str().unwrap_or("").chars().take(80).collect::<String>())
|
||
}
|
||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||
let prompt = arg_str(&input, "prompt")?.to_string();
|
||
let stage = if input["stage"].as_str() == Some("final") { "final" } else { "preview" };
|
||
let purpose = input["purpose"].as_str().filter(|p| PURPOSES.contains(p));
|
||
let dir_arg = input["directory"].as_str().map(String::from).unwrap_or_else(|| ctx.app.config().images.output_dir.clone());
|
||
let dir = ctx.resolve(&dir_arg)?;
|
||
let mut references = Vec::new();
|
||
for r in input["reference"].as_array().into_iter().flatten().filter_map(|v| v.as_str()).take(4) {
|
||
references.push(ctx.resolve(r)?);
|
||
}
|
||
let model = match input["model"].as_str().filter(|m| !m.is_empty()) {
|
||
Some(m) => model_config(&ctx.app, m),
|
||
None => choose_model(&ctx.app, stage, purpose, !references.is_empty()),
|
||
};
|
||
let owner = ctx.app.session(&ctx.session_id).ok().and_then(|s| s.owner_id);
|
||
let refined = !references.is_empty() && model.reference_param.is_some();
|
||
let records = generate(
|
||
&ctx.app,
|
||
GenerateRequest {
|
||
prompt,
|
||
model: Some(model.id.clone()),
|
||
size: input["size"].as_str().map(String::from),
|
||
count: input["count"].as_u64().unwrap_or(if stage == "preview" { 2 } else { 1 }) as u32,
|
||
dir,
|
||
user_id: owner,
|
||
session_id: Some(ctx.session_id.clone()),
|
||
references,
|
||
},
|
||
)
|
||
.await
|
||
.map_err(|e| anyhow!("{e:#}"))?;
|
||
let price = |m: &ImageModel| m.price.map(|p| format!(" (${p:.3} per image)")).unwrap_or_default();
|
||
let label = |m: &ImageModel| m.name.clone().unwrap_or_else(|| m.id.clone());
|
||
let mut text = format!("Model: {}{}\nGenerated:\n", label(&model), price(&model));
|
||
let mut images = Vec::new();
|
||
for r in &records {
|
||
let rel = Path::new(&r.path).strip_prefix(&ctx.cwd).map(|p| p.display().to_string()).unwrap_or_else(|_| r.path.clone());
|
||
text.push_str(&format!("- {rel} ({}×{})\n", r.width, r.height));
|
||
if let Some(p) = preview_for_model(Path::new(&r.path)) {
|
||
images.push(p);
|
||
}
|
||
}
|
||
if stage == "preview" && !refined {
|
||
let fin = choose_model(&ctx.app, "final", purpose, true);
|
||
if fin.id != model.id {
|
||
let how = if fin.reference_param.is_some() { "refines the chosen preview" } else { "generates it again from the prompt" };
|
||
text.push_str(&format!(
|
||
"\nThese are previews. Ask the user (ask_user) which one to refine. The final image would use {}{} and {how}.",
|
||
label(&fin),
|
||
price(&fin)
|
||
));
|
||
}
|
||
} else if !input["reference"].as_array().is_none_or(|a| a.is_empty()) && !refined {
|
||
text.push_str("\nNote: this model takes no reference image, so it was generated from the prompt only.");
|
||
}
|
||
let ids: Vec<&str> = records.iter().map(|r| r.id.as_str()).collect();
|
||
Ok(ToolOutput { content: text, images, meta: Some(json!({"image_ids": ids})), ..Default::default() })
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn finds_moved_files_and_keeps_their_copies() {
|
||
let dir = std::env::temp_dir().join(format!("cc-img-moved-{}", std::process::id()));
|
||
let art = dir.join("art");
|
||
std::fs::create_dir_all(art.join("older")).unwrap();
|
||
std::fs::write(art.join("older/20261110-a.png"), b"x").unwrap();
|
||
// Moved into a subfolder: found; renamed or outside the roots: not.
|
||
assert_eq!(find_moved(&art.join("20261110-a.png"), std::slice::from_ref(&dir)), Some(art.join("older/20261110-a.png")));
|
||
assert_eq!(find_moved(&art.join("20261110-b.png"), std::slice::from_ref(&dir)), None);
|
||
assert_eq!(find_moved(&art.join("20261110-a.png"), &[dir.join("elsewhere")]), None);
|
||
// Renamed: found by content.
|
||
std::fs::write(art.join("older/level-5.png"), b"renamed bytes").unwrap();
|
||
let h = crate::blobs::hash(b"renamed bytes");
|
||
assert_eq!(find_by_content(&art.join("20261110-c.png"), std::slice::from_ref(&dir), &h, 13), Some(art.join("older/level-5.png")));
|
||
assert_eq!(find_by_content(&art.join("20261110-c.png"), std::slice::from_ref(&dir), &crate::blobs::hash(b"other"), 13), None);
|
||
// The image store keeps the copy of a generated image.
|
||
let store = crate::store::Store::open(&dir.join("t.db")).unwrap();
|
||
let hash = store.put_blob("image/png", b"png bytes").unwrap();
|
||
store
|
||
.add_image(&ImageRecord { id: "i".into(), user_id: None, session_id: None, prompt: "p".into(), model: "m".into(), path: "/gone.png".into(), width: 1, height: 1, created_at: 0, blob: Some(hash.clone()) })
|
||
.unwrap();
|
||
assert_eq!(store.gc_blobs().unwrap(), 0);
|
||
// Moved with the file manager: the entry follows (file or folder, not look-alikes).
|
||
assert_eq!(store.move_image_paths("/gone.png", "/art/new.png").unwrap(), 1);
|
||
assert_eq!(store.move_image_paths("/art", "/bilder/").unwrap(), 1);
|
||
assert_eq!(store.move_image_paths("/bild", "/x").unwrap(), 0);
|
||
assert_eq!(store.image("i").unwrap().unwrap().path, "/bilder/new.png");
|
||
assert_eq!(store.blob(&hash).unwrap().unwrap().1, b"png bytes");
|
||
std::fs::remove_dir_all(&dir).unwrap();
|
||
}
|
||
|
||
fn m(id: &str) -> ImageModel {
|
||
with_defaults(ImageModel { id: id.into(), ..Default::default() })
|
||
}
|
||
|
||
#[test]
|
||
fn picks_by_stage_purpose_and_reference() {
|
||
let all = vec![
|
||
m("fal-ai/flux/schnell"),
|
||
m("fal-ai/nano-banana-pro/edit"),
|
||
m("fal-ai/nano-banana-pro"),
|
||
m("fal-ai/flux-2-pro"),
|
||
m("fal-ai/ideogram/v3"),
|
||
m("fal-ai/recraft/v4/pro/text-to-image"),
|
||
];
|
||
let id = |stage, purpose, r| pick(&all, stage, purpose, r).unwrap().id;
|
||
assert_eq!(id("preview", Some("photo"), false), "fal-ai/flux/schnell");
|
||
// Refining a preview: the edit model, whatever the purpose.
|
||
assert_eq!(id("final", Some("photo"), true), "fal-ai/nano-banana-pro/edit");
|
||
assert_eq!(id("final", Some("design"), true), "fal-ai/nano-banana-pro/edit");
|
||
// From the prompt only: the first fitting text-to-image model (config order).
|
||
assert_eq!(id("final", Some("photo"), false), "fal-ai/nano-banana-pro");
|
||
assert_eq!(id("final", Some("text"), false), "fal-ai/nano-banana-pro");
|
||
assert_eq!(id("final", Some("design"), false), "fal-ai/ideogram/v3");
|
||
let no_edit: Vec<_> = all.iter().filter(|m| m.reference_param.is_none()).cloned().collect();
|
||
assert_eq!(pick(&no_edit, "final", Some("design"), true).unwrap().id, "fal-ai/ideogram/v3");
|
||
assert_eq!(pick(&no_edit, "final", Some("photo"), false).unwrap().id, "fal-ai/nano-banana-pro");
|
||
}
|
||
|
||
#[test]
|
||
fn config_overrides_known_defaults() {
|
||
let custom = with_defaults(ImageModel { id: "fal-ai/flux/schnell".into(), stage: Some("final".into()), price: Some(0.01), ..Default::default() });
|
||
assert_eq!(custom.stage.as_deref(), Some("final"));
|
||
assert_eq!(custom.price, Some(0.01));
|
||
assert!(m("some/unknown-model").stage.is_none());
|
||
}
|
||
}
|