Richard ZinkandClaude Opus 5.5
48621f21d6
- New tool: image-to-3D (project file or an image attached to the chat, optional back/left/right views) or text-to-3D; image + prompt refines the image with the prompt first (the 3D models take no text with an image). Styles textured / geometry / lowpoly, optional PBR, face count. Runs through fal.ai's queue API with progress, saves GLB + preview, books the cost (published prices, overridable per model). Request bodies for the Hunyuan3D, Tripo and Trellis families; defaults [images] model_3d / model_3d_text. - Chat card: "Open in the 3D view"; media overview: kind "model" (also STL/3MF/OBJ written by agents) with previews. 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>,
|
||
}
|
||
|
||
pub(crate) 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).
|
||
pub(crate) 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)
|
||
}
|
||
|
||
pub(crate) 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.
|
||
pub(crate) 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());
|
||
}
|
||
}
|