jameshtwose/dm-pal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

building out the components, updating the UI in overview and adding a little sqlite viewer

Browse Source
This commit is contained in:
itsamejms
2026-07-17 09:31:51 +01:00
parent 82f2dcf3df
commit b89a39ec9d
30 changed files with 2664 additions and 635 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/commands/image_commands.rs
+212 -42
View File
@@ -1,8 +1,9 @@
use crate::llm::AppState;
use serde::Deserialize;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use tauri::Manager;
use tauri::ipc::Channel;
// ponytail: DefaultHasher is fine for a cache filename — not crypto, just a stable key.
@@ -14,14 +15,57 @@ pub struct ImageRequest {
}
/// One line of Ollama's NDJSON image-generation response.
#[derive(Debug, Deserialize)]
/// `step`/`total` carry progress; the final `done: true` line carries the
/// singular `image` base64 field. All fields optional per-line.
#[derive(Debug, Default, Deserialize)]
struct OllamaImageLine {
#[serde(default)]
done: bool,
#[serde(default)]
step: Option<u32>,
#[serde(default)]
total: Option<u32>,
#[serde(default)]
image: Option<String>,
}
/// Parsed view of one NDJSON line used by both the streaming and the
/// buffered paths. Extracted so the parsing logic is unit-testable without
/// touching the network.
#[derive(Debug, PartialEq)]
struct ImageProgress {
done: bool,
step: Option<u32>,
total: Option<u32>,
image: Option<String>,
}
fn parse_image_line(line: &str) -> Option<ImageProgress> {
let line = line.trim();
if line.is_empty() {
return None;
}
let parsed: OllamaImageLine = serde_json::from_str(line).ok()?;
Some(ImageProgress {
done: parsed.done,
step: parsed.step,
total: parsed.total,
image: parsed.image,
})
}
/// Channel events for streaming image generation.
#[derive(Clone, Serialize)]
#[serde(tag = "type", content = "data", rename_all = "camelCase")]
pub enum ImageEvent {
/// Progress update: (step, total). Either may be None if Ollama omits it.
Progress { step: Option<u32>, total: Option<u32> },
/// Final result: a `data:image/png;base64,...` URL ready for `<img src>`.
Done(String),
/// Fatal error.
Error(String),
}
/// Generate (or fetch from disk cache) an image for `prompt` via the configured
/// Ollama image model. Returns a `data:image/png;base64,...` URL ready for `<img src>`.
///
@@ -29,7 +73,6 @@ struct OllamaImageLine {
/// so the front-end can fall back to a placeholder instead of a confusing timeout.
#[tauri::command]
pub async fn generate_image(
app: tauri::AppHandle,
state: tauri::State<'_, AppState>,
req: ImageRequest,
) -> Result<String, String> {
@@ -40,34 +83,92 @@ pub async fn generate_image(
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
let model = req.model.unwrap_or(config.image_model.clone());
let cache_dir = app
.path()
.app_data_dir()
.map_err(|e| e.to_string())?
.join("dm-toolkit")
.join("images");
std::fs::create_dir_all(&cache_dir).map_err(|e| e.to_string())?;
// Stable cache key over model + prompt. Regenerating with a tweaked prompt
// produces a new file; identical prompt reuses the cached PNG.
let mut hasher = DefaultHasher::new();
model.hash(&mut hasher);
req.prompt.hash(&mut hasher);
let cache_path = cache_dir.join(format!("{:016x}.png", hasher.finish()));
// Cache hit: return the stored PNG without calling the model.
if cache_path.exists() {
let (cache_path, cache_hit) = prepare_cache(&state.data_dir, &model, &req.prompt)?;
if cache_hit {
let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
return Ok(data_url(&bytes));
}
let png_bytes = request_image_bytes(&config, &model, &req.prompt, None).await?;
std::fs::write(&cache_path, &png_bytes).map_err(|e| e.to_string())?;
Ok(data_url(&png_bytes))
}
/// Streaming variant: emits `ImageEvent::Progress` as Ollama reports `step`/`total`,
/// then `ImageEvent::Done` with the data URL (or `Error`). Reuses the same disk cache
/// as `generate_image`. The DM gets a real progress bar for the multi-second wait.
#[tauri::command]
pub async fn generate_image_stream(
state: tauri::State<'_, AppState>,
req: ImageRequest,
channel: Channel<ImageEvent>,
) -> Result<(), String> {
if cfg!(not(target_os = "macos")) {
let _ = channel.send(ImageEvent::Error(
"image generation is macOS-only via Ollama (for now)".into(),
));
return Ok(());
}
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
let model = req.model.unwrap_or(config.image_model.clone());
let (cache_path, cache_hit) = prepare_cache(&state.data_dir, &model, &req.prompt)?;
if cache_hit {
let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
let _ = channel.send(ImageEvent::Done(data_url(&bytes)));
return Ok(());
}
// Drive the request on a background task so the command returns immediately
// and progress flows through the channel. Errors become ImageEvent::Error.
let channel = std::sync::Arc::new(channel);
let ch = channel.clone();
tauri::async_runtime::spawn(async move {
match request_image_bytes(&config, &model, &req.prompt, Some(ch)).await {
Ok(png_bytes) => {
let _ = std::fs::write(&cache_path, &png_bytes);
let _ = channel.send(ImageEvent::Done(data_url(&png_bytes)));
}
Err(e) => {
let _ = channel.send(ImageEvent::Error(e));
}
}
});
Ok(())
}
/// Resolve the on-disk cache path for (model, prompt) and report a cache hit.
/// `base` is the configured data dir (AppState.data_dir).
fn prepare_cache(
base: &std::path::Path,
model: &str,
prompt: &str,
) -> Result<(std::path::PathBuf, bool), String> {
let cache_dir = base.join("images");
std::fs::create_dir_all(&cache_dir).map_err(|e| e.to_string())?;
let mut hasher = DefaultHasher::new();
model.hash(&mut hasher);
prompt.hash(&mut hasher);
let cache_path = cache_dir.join(format!("{:016x}.png", hasher.finish()));
let hit = cache_path.exists();
Ok((cache_path, hit))
}
/// POST the image request to Ollama and collect the PNG bytes. When a channel
/// is given, parse the NDJSON body incrementally and emit progress; otherwise
/// read the whole body at once (legacy buffered path).
async fn request_image_bytes(
config: &crate::llm::LlmConfig,
model: &str,
prompt: &str,
channel: Option<std::sync::Arc<Channel<ImageEvent>>>,
) -> Result<Vec<u8>, String> {
let client = reqwest::Client::new();
let url = format!("{}/api/generate", config.api_url.trim_end_matches('/'));
let body = serde_json::json!({
"model": model,
"prompt": req.prompt,
"stream": false,
});
let body = serde_json::json!({ "model": model, "prompt": prompt, "stream": false });
let res = client
.post(&url)
@@ -82,19 +183,47 @@ pub async fn generate_image(
return Err(format!("image error {status}: {text}"));
}
// Ollama returns newline-delimited JSON even with stream:false for image models;
// the final line with `done: true` carries the singular `image` base64 field.
let body_text = res.text().await.map_err(|e| format!("image read error: {e}"))?;
// ponytail: Ollama image models emit NDJSON even with stream:false —
// progress lines (step/total) then a final done:true carrying the image.
// The buffered path reads it all at once; the streaming path splits lines
// as they arrive so the bar animates.
let mut png_b64: Option<String> = None;
for line in body_text.lines() {
let line = line.trim();
if line.is_empty() {
continue;
let mut buffer = String::new();
let apply_line = |line: &str, png: &mut Option<String>, ch: Option<&Channel<ImageEvent>>| {
if let Some(p) = parse_image_line(line) {
if let Some(c) = ch {
let _ = c.send(ImageEvent::Progress { step: p.step, total: p.total });
}
if let Some(b64) = p.image {
*png = Some(b64);
}
}
if let Ok(parsed) = serde_json::from_str::<OllamaImageLine>(line) {
if let Some(b64) = parsed.image {
png_b64 = Some(b64);
if parsed.done {
};
if let Some(ch) = channel {
let mut stream = res.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| format!("image read error: {e}"))?;
buffer.push_str(&String::from_utf8_lossy(&chunk));
// Process complete lines; keep the trailing partial line in buffer.
while let Some(idx) = buffer.find('\n') {
let line = buffer.split_off(idx + 1);
let complete = std::mem::replace(&mut buffer, line);
apply_line(&complete, &mut png_b64, Some(&ch));
}
}
if !buffer.trim().is_empty() {
apply_line(&buffer, &mut png_b64, Some(&ch));
}
} else {
let body_text = res.text().await.map_err(|e| format!("image read error: {e}"))?;
for line in body_text.lines() {
apply_line(line, &mut png_b64, None);
// ponytail: stop after the done line in the buffered path — the
// final image is the one we want.
if let Some(p) = parse_image_line(line) {
if p.done && p.image.is_some() {
break;
}
}
@@ -102,11 +231,8 @@ pub async fn generate_image(
}
let b64 = png_b64.ok_or_else(|| "no image data in Ollama response".to_string())?;
// Decode + persist to cache, then return a data URL.
let png_bytes = base64_decode(&b64)?;
std::fs::write(&cache_path, &png_bytes).map_err(|e| e.to_string())?;
Ok(data_url(&png_bytes))
Ok(png_bytes)
}
fn data_url(png: &[u8]) -> String {
@@ -115,7 +241,7 @@ fn data_url(png: &[u8]) -> String {
let mut out = String::with_capacity((png.len() + 2) / 3 * 4);
let mut chunks = png.chunks_exact(3);
for c in &mut chunks {
let n = (c[0] as usize) << 16 | (c[1] as usize) << 8 | c[2] as usize;
let n = (c[0] as usize) << 16 | (c[1] as usize) << 8 | (c[2] as usize);
out.push(table[(n >> 18) & 63] as char);
out.push(table[(n >> 12) & 63] as char);
out.push(table[(n >> 6) & 63] as char);
@@ -200,4 +326,48 @@ mod tests {
let s = format!("{:016x}", h.finish());
assert_eq!(s.len(), 16);
}
#[test]
fn parse_progress_and_image_lines() {
// intermediate progress line, no image
let p = parse_image_line(r#"{"model":"x/flux2-klein:4b","done":false,"total":4,"step":1}"#).unwrap();
assert_eq!(p, ImageProgress { done: false, step: Some(1), total: Some(4), image: None });
// final done line carries the image
let p = parse_image_line(r#"{"model":"x/flux2-klein:4b","done":true,"image":"iVBORw0KGgoAAAANSUhEUgAA"}"#).unwrap();
assert!(p.done);
assert_eq!(p.image.as_deref(), Some("iVBORw0KGgoAAAANSUhEUgAA"));
assert!(p.total.is_none() && p.step.is_none());
// blank/garbage lines are ignored, not errors
assert!(parse_image_line("").is_none());
assert!(parse_image_line("not json").is_none());
}
#[test]
fn split_ndjson_buffer_keeps_trailing_partial() {
// Simulate two chunks arriving separately where the split falls mid-line.
let chunk1 = "{\"done\":false,\"step\":1,\"total\":4}\n{\"done\":tru";
let chunk2 = "e,\"image\":\"abc\"}\n";
let mut buffer = String::new();
let mut png: Option<String> = None;
let whole = format!("{chunk1}{chunk2}");
// emulate the streaming loop over the concatenated body
buffer.push_str(&whole);
let mut lines = Vec::new();
while let Some(idx) = buffer.find('\n') {
let line = buffer.split_off(idx + 1);
let complete = std::mem::replace(&mut buffer, line);
lines.push(complete);
}
if !buffer.trim().is_empty() {
lines.push(std::mem::take(&mut buffer));
}
for line in &lines {
if let Some(p) = parse_image_line(line) {
if let Some(b) = p.image { png = Some(b); }
}
}
assert_eq!(png.as_deref(), Some("abc"));
}
}