building out the components, updating the UI in overview and adding a little sqlite viewer
This commit is contained in:
itsamejms
@@ -0,0 +1,200 @@
|
||||
use crate::llm::AppState;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tauri::Manager;
|
||||
use tauri_plugin_store::StoreExt;
|
||||
|
||||
// ponytail: the data-dir pref store + key are defined in lib.rs; mirror them
|
||||
// here so this module is self-contained for reads/writes.
|
||||
const PREFS_STORE: &str = "dm-pal-prefs.json";
|
||||
const DATA_DIR_KEY: &str = "dataDir";
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct DataDirInfo {
|
||||
/// Absolute path currently in use (may differ from `configured` if it was
|
||||
/// unset/fell back to the default).
|
||||
pub current: String,
|
||||
/// Absolute path that will be used after the next restart, if different.
|
||||
pub configured: Option<String>,
|
||||
/// The OS-default location ($APPDATA/dm-toolkit).
|
||||
pub default: String,
|
||||
}
|
||||
|
||||
/// Return the current data dir + the configured (pending-restart) value, so
|
||||
/// Settings can show "using X — restart to apply Y".
|
||||
#[tauri::command]
|
||||
pub fn get_data_dir(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, AppState>,
|
||||
) -> Result<DataDirInfo, String> {
|
||||
let default = app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| e.to_string())?
|
||||
.join("dm-toolkit");
|
||||
let configured = app
|
||||
.store(PREFS_STORE)
|
||||
.map_err(|e| e.to_string())?
|
||||
.get(DATA_DIR_KEY)
|
||||
.and_then(|v| v.as_str().map(|s| s.to_string()));
|
||||
Ok(DataDirInfo {
|
||||
current: state.data_dir.display().to_string(),
|
||||
configured,
|
||||
default: default.display().to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Subdirs/files that make up the campaign data and must move together.
|
||||
const DATA_RELATIVE: &[&str] = &["lore/lore.db", "generations.db", "images"];
|
||||
|
||||
/// Recursively copy `src` to `dst`, but never overwrite an existing file —
|
||||
/// a re-locate must not clobber data already at the destination, and never
|
||||
/// deletes the source (the old location stays as a backup until the DM deletes it).
|
||||
/// When both src and dst are directories, recurse so new files inside an
|
||||
/// existing destination tree still get copied.
|
||||
fn copy_if_absent(src: &Path, dst: &Path) -> std::io::Result<()> {
|
||||
if !src.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
if src.is_dir() {
|
||||
std::fs::create_dir_all(dst)?;
|
||||
for entry in std::fs::read_dir(src)? {
|
||||
let entry = entry?;
|
||||
copy_if_absent(&entry.path(), &dst.join(entry.file_name()))?;
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
// leaf file: skip if the destination already has it (no overwrite).
|
||||
if dst.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(parent) = dst.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::copy(src, dst)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set the data dir to `new_dir`: create it, copy existing campaign data
|
||||
/// (lore.db, generations.db, images/) into it, and persist the preference so
|
||||
/// the next restart uses it. The currently-open SQLite connections stay on the
|
||||
/// old paths until restart — the caller must prompt the user to relaunch.
|
||||
#[tauri::command]
|
||||
pub fn set_data_dir(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, AppState>,
|
||||
new_dir: String,
|
||||
) -> Result<DataDirInfo, String> {
|
||||
let new_path = PathBuf::from(new_dir.trim());
|
||||
if !new_path.is_absolute() {
|
||||
return Err("Choose an absolute folder path".into());
|
||||
}
|
||||
std::fs::create_dir_all(&new_path).map_err(|e| format!("can't create dir: {e}"))?;
|
||||
|
||||
// Copy existing data into the new location (no overwrite; old data stays put).
|
||||
for rel in DATA_RELATIVE {
|
||||
let src = state.data_dir.join(rel);
|
||||
let dst = new_path.join(rel);
|
||||
if src.exists() {
|
||||
copy_if_absent(&src, &dst).map_err(|e| format!("migrate {rel}: {e}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Persist the preference for the next startup.
|
||||
let store = app.store(PREFS_STORE).map_err(|e| e.to_string())?;
|
||||
store.set(DATA_DIR_KEY, serde_json::Value::String(new_path.display().to_string()));
|
||||
store.save().map_err(|e| e.to_string())?;
|
||||
|
||||
Ok(DataDirInfo {
|
||||
current: state.data_dir.display().to_string(),
|
||||
configured: Some(new_path.display().to_string()),
|
||||
default: app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| e.to_string())?
|
||||
.join("dm-toolkit")
|
||||
.display()
|
||||
.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Clear the preference so the next restart uses the OS default. Does not
|
||||
/// delete any files — the DM re-points, then restarts.
|
||||
#[tauri::command]
|
||||
pub fn reset_data_dir(app: tauri::AppHandle) -> Result<(), String> {
|
||||
let store = app.store(PREFS_STORE).map_err(|e| e.to_string())?;
|
||||
store.delete(DATA_DIR_KEY);
|
||||
store.save().map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ─── SQLite viewer (read-only) ─────────────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SqlQueryReq {
|
||||
/// "lore" or "generations".
|
||||
pub db: String,
|
||||
pub sql: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SqlQueryResult {
|
||||
pub columns: Vec<String>,
|
||||
pub rows: Vec<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
/// Run a read-only SQL query against one of the two SQLite DBs so a DM can
|
||||
/// inspect the backend from Settings → Advanced. Only SELECT/PRAGMA/WITH.
|
||||
#[tauri::command]
|
||||
pub fn sql_query(state: tauri::State<'_, AppState>, req: SqlQueryReq) -> Result<SqlQueryResult, String> {
|
||||
let res = match req.db.as_str() {
|
||||
"lore" => state.rag.raw_query(&req.sql),
|
||||
"generations" => state.gen.raw_query(&req.sql),
|
||||
other => Err(anyhow::anyhow!("unknown db: {other}")),
|
||||
}
|
||||
.map_err(|e| e.to_string())?;
|
||||
Ok(SqlQueryResult { columns: res.0, rows: res.1 })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn copy_if_absent_does_not_overwrite() {
|
||||
let tmp = std::env::temp_dir().join(format!("dm-pal-data-test-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let src = tmp.join("src");
|
||||
let dst = tmp.join("dst");
|
||||
std::fs::create_dir_all(src.join("images")).unwrap();
|
||||
std::fs::write(src.join("images/a.png"), b"png").unwrap();
|
||||
std::fs::create_dir_all(src.join("lore")).unwrap();
|
||||
std::fs::write(src.join("lore/lore.db"), b"db").unwrap();
|
||||
|
||||
// first copy creates dst with the files
|
||||
copy_if_absent(&src, &dst).unwrap();
|
||||
assert!(dst.join("images/a.png").exists());
|
||||
assert!(dst.join("lore/lore.db").exists());
|
||||
|
||||
// mutate src; second copy must NOT overwrite existing dst files
|
||||
std::fs::write(src.join("images/a.png"), b"changed").unwrap();
|
||||
copy_if_absent(&src, &dst).unwrap();
|
||||
assert_eq!(std::fs::read(dst.join("images/a.png")).unwrap(), b"png");
|
||||
|
||||
// a new file in src does get copied (absent in dst)
|
||||
std::fs::write(src.join("images/b.png"), b"png2").unwrap();
|
||||
copy_if_absent(&src, &dst).unwrap();
|
||||
assert!(dst.join("images/b.png").exists());
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn copy_if_absent_missing_src_is_noop() {
|
||||
let tmp = std::env::temp_dir().join(format!("dm-pal-data-test2-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
copy_if_absent(&tmp.join("nope"), &tmp.join("dst")).unwrap();
|
||||
assert!(!tmp.join("dst").exists());
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod llm_commands;
|
||||
pub mod image_commands;
|
||||
pub mod rag_commands;
|
||||
pub mod generation_commands;
|
||||
pub mod generation_commands;
|
||||
pub mod data_commands;
|
||||
@@ -55,4 +55,23 @@ pub fn rag_list(state: tauri::State<'_, AppState>) -> Result<Vec<RagSource>, Str
|
||||
#[tauri::command]
|
||||
pub fn rag_clear(state: tauri::State<'_, AppState>, source: Option<String>) -> Result<(), String> {
|
||||
state.rag.clear(source.as_deref()).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RagChunk {
|
||||
pub id: i64,
|
||||
pub preview: String,
|
||||
}
|
||||
|
||||
/// Preview the chunks stored under `source` (first 200 chars each). Lets the
|
||||
/// DM verify what actually got embedded without dumping the whole corpus.
|
||||
#[tauri::command]
|
||||
pub fn rag_chunks(state: tauri::State<'_, AppState>, source: String) -> Result<Vec<RagChunk>, String> {
|
||||
state
|
||||
.rag
|
||||
.list_chunks(&source)
|
||||
.map_err(|e| e.to_string())?
|
||||
.into_iter()
|
||||
.map(|(id, preview)| Ok(RagChunk { id, preview }))
|
||||
.collect()
|
||||
}
|
||||
Reference in New Issue
Block a user