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

working through the plan + UI/ UX

Browse Source
This commit is contained in:
itsamejms
2026-07-12 22:13:43 +01:00
parent 5b256242be
commit a24f3615e0
38 changed files with 5295 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/Cargo.lock
Generated
+70 -1
View File
@@ -19,6 +19,18 @@ dependencies = [
"version_check",
]
[[package]]
name = "ahash"
version = "0.8.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
dependencies = [
"cfg-if",
"once_cell",
"version_check",
"zerocopy",
]
[[package]]
name = "aho-corasick"
version = "1.1.4"
@@ -739,8 +751,10 @@ dependencies = [
name = "dm-pal"
version = "0.1.0"
dependencies = [
"anyhow",
"log",
"reqwest 0.12.28",
"rusqlite",
"serde",
"serde_json",
"tauri",
@@ -883,6 +897,18 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "fallible-iterator"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
[[package]]
name = "fallible-streaming-iterator"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
[[package]]
name = "fastrand"
version = "2.4.1"
@@ -1389,7 +1415,16 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
dependencies = [
"ahash",
"ahash 0.7.8",
]
[[package]]
name = "hashbrown"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
dependencies = [
"ahash 0.8.12",
]
[[package]]
@@ -1398,6 +1433,15 @@ version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
[[package]]
name = "hashlink"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af"
dependencies = [
"hashbrown 0.14.5",
]
[[package]]
name = "heck"
version = "0.4.1"
@@ -1898,6 +1942,17 @@ dependencies = [
"libc",
]
[[package]]
name = "libsqlite3-sys"
version = "0.28.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c10584274047cb335c23d3e61bcef8e323adae7c5c8c760540f73610177fc3f"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@@ -2883,6 +2938,20 @@ dependencies = [
"syn 1.0.109",
]
[[package]]
name = "rusqlite"
version = "0.31.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b838eba278d213a8beaf485bd313fd580ca4505a00d5871caeb1457c55322cae"
dependencies = [
"bitflags 2.13.0",
"fallible-iterator",
"fallible-streaming-iterator",
"hashlink",
"libsqlite3-sys",
"smallvec",
]
[[package]]
name = "rust_decimal"
version = "1.42.1"
src-tauri/Cargo.toml
+2
View File
@@ -18,6 +18,7 @@ tauri-build = { version = "2.6.3", features = [] }
[dependencies]
serde_json = "1.0"
serde = { version = "1.0", features = ["derive"] }
anyhow = "1.0"
log = "0.4"
tauri = { version = "2.11.3", features = [] }
tauri-plugin-log = "2"
@@ -25,3 +26,4 @@ tauri-plugin-store = "2"
tauri-plugin-fs = "2"
reqwest = { version = "0.12", features = ["json"] }
tokio = { version = "1", features = ["full"] }
rusqlite = { version = "0.31", features = ["bundled"] }
src-tauri/src/commands/generation_commands.rs
+59
View File
@@ -0,0 +1,59 @@
use crate::generations::{Generation, GenerationStore, GenerationSummary};
use serde::Deserialize;
#[derive(Debug, Deserialize)]
pub struct GenerationAddRequest {
pub kind: String,
pub title: String,
pub data: String,
/// Optional input/source (theme, name, level, etc.) for the history list.
pub source: Option<String>,
}
#[tauri::command]
pub fn generation_add(
state: tauri::State<'_, crate::llm::AppState>,
req: GenerationAddRequest,
) -> Result<i64, String> {
state
.gen
.add(&req.kind, &req.title, &req.data, req.source.as_deref())
.map_err(|e| e.to_string())
}
#[tauri::command]
pub fn generation_list(
state: tauri::State<'_, crate::llm::AppState>,
kind: Option<String>,
) -> Result<Vec<GenerationSummary>, String> {
state.gen.list(kind.as_deref()).map_err(|e| e.to_string())
}
#[tauri::command]
pub fn generation_get(
state: tauri::State<'_, crate::llm::AppState>,
id: i64,
) -> Result<Option<Generation>, String> {
state.gen.get(id).map_err(|e| e.to_string())
}
#[tauri::command]
pub fn generation_delete(
state: tauri::State<'_, crate::llm::AppState>,
id: Option<i64>,
) -> Result<usize, String> {
state.gen.delete(id).map_err(|e| e.to_string())
}
#[tauri::command]
pub fn generation_counts(
state: tauri::State<'_, crate::llm::AppState>,
) -> Result<Vec<(String, i64)>, String> {
state.gen.counts().map_err(|e| e.to_string())
}
// Reference `GenerationStore` so the type is considered used by the module
// (the public functions go through `state.gen`, which already uses it,
// but this silences a future-proofing warning if all wrappers get removed).
#[allow(dead_code)]
fn _store_type_used(_: &GenerationStore) {}
src-tauri/src/commands/image_commands.rs
+203
View File
@@ -0,0 +1,203 @@
use crate::llm::AppState;
use serde::Deserialize;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use tauri::Manager;
// ponytail: DefaultHasher is fine for a cache filename — not crypto, just a stable key.
#[derive(Debug, Deserialize)]
pub struct ImageRequest {
pub prompt: String,
/// Override the configured image model for this call.
pub model: Option<String>,
}
/// One line of Ollama's NDJSON image-generation response.
#[derive(Debug, Deserialize)]
struct OllamaImageLine {
#[serde(default)]
done: bool,
#[serde(default)]
image: Option<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>`.
///
/// Ollama image models are macOS-only today; on other platforms we return an error
/// 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> {
if cfg!(not(target_os = "macos")) {
return Err("image generation is macOS-only via Ollama (for now)".into());
}
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 bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
return Ok(data_url(&bytes));
}
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 res = client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| format!("image request failed: {e}"))?;
if !res.status().is_success() {
let status = res.status();
let text = res.text().await.unwrap_or_default();
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}"))?;
let mut png_b64: Option<String> = None;
for line in body_text.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
if let Ok(parsed) = serde_json::from_str::<OllamaImageLine>(line) {
if let Some(b64) = parsed.image {
png_b64 = Some(b64);
if parsed.done {
break;
}
}
}
}
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))
}
fn data_url(png: &[u8]) -> String {
// ponytail: no base64 crate dep — a 30-line encoder.
let table: [u8; 64] = *b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
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;
out.push(table[(n >> 18) & 63] as char);
out.push(table[(n >> 12) & 63] as char);
out.push(table[(n >> 6) & 63] as char);
out.push(table[n & 63] as char);
}
let rem = chunks.remainder();
match rem.len() {
1 => {
let n = (rem[0] as usize) << 16;
out.push(table[(n >> 18) & 63] as char);
out.push(table[(n >> 12) & 63] as char);
out.push('=');
out.push('=');
}
2 => {
let n = (rem[0] as usize) << 16 | (rem[1] as usize) << 8;
out.push(table[(n >> 18) & 63] as char);
out.push(table[(n >> 12) & 63] as char);
out.push(table[(n >> 6) & 63] as char);
out.push('=');
}
_ => {}
}
format!("data:image/png;base64,{out}")
}
/// Minimal standard base64 decoder (no extra dependency).
fn base64_decode(input: &str) -> Result<Vec<u8>, String> {
fn val(c: u8) -> Option<u8> {
match c {
b'A'..=b'Z' => Some(c - b'A'),
b'a'..=b'z' => Some(c - b'a' + 26),
b'0'..=b'9' => Some(c - b'0' + 52),
b'+' => Some(62),
b'/' => Some(63),
_ => None,
}
}
let input = input.trim();
let bytes: Vec<u8> = input.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
if bytes.is_empty() {
return Ok(Vec::new());
}
let mut out = Vec::with_capacity(bytes.len() * 3 / 4);
let mut buf: u32 = 0;
let mut bits: u32 = 0;
for &b in &bytes {
if b == b'=' {
break;
}
let v = val(b).ok_or_else(|| format!("invalid base64 char: {b}"))? as u32;
buf = (buf << 6) | v;
bits += 6;
if bits >= 8 {
bits -= 8;
out.push((buf >> bits) as u8);
buf &= (1 << bits) - 1;
}
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn roundtrip_base64() {
let data = b"hello world \x00\xff\x10";
let url = data_url(data);
let b64 = url.strip_prefix("data:image/png;base64,").unwrap();
let decoded = base64_decode(b64).unwrap();
assert_eq!(decoded, data);
}
#[test]
fn cache_key_is_stable() {
// sanity: same inputs → same filename shape (16 hex digits)
let mut h = DefaultHasher::new();
"x/flux2-klein:4b".hash(&mut h);
"prompt".hash(&mut h);
let s = format!("{:016x}", h.finish());
assert_eq!(s.len(), 16);
}
}
src-tauri/src/commands/llm_commands.rs
+42 -5
View File
@@ -16,6 +16,10 @@ pub async fn generate(state: tauri::State<'_, AppState>, req: GenerateRequest) -
let temperature = req.temperature.unwrap_or(config.temperature);
let max_tokens = req.max_tokens.unwrap_or(config.max_tokens);
// Inject retrieved lore into the system prompt so every generator stays
// consistent with the user's world bible. Shared path = every caller.
let messages = inject_lore(&state, &client, &config, messages, &req.rag_query).await;
if llm::is_ollama(&config.api_url) {
call_ollama(&client, &config, &messages, temperature, max_tokens).await
} else {
@@ -33,12 +37,13 @@ pub async fn generate_stream(
) -> Result<(), String> {
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
tauri::async_runtime::spawn(async move {
let client = reqwest::Client::new();
let messages = llm::build_messages(req.system.as_deref(), &req.prompt);
let temperature = req.temperature.unwrap_or(config.temperature);
let max_tokens = req.max_tokens.unwrap_or(config.max_tokens);
let client = reqwest::Client::new();
let messages = llm::build_messages(req.system.as_deref(), &req.prompt);
let temperature = req.temperature.unwrap_or(config.temperature);
let max_tokens = req.max_tokens.unwrap_or(config.max_tokens);
let messages = inject_lore(&state, &client, &config, messages, &req.rag_query).await;
tauri::async_runtime::spawn(async move {
// For now, we do a non-streaming call and emit the full response as one token
// Real SSE streaming from Ollama/OpenAI can be added later
let result = if llm::is_ollama(&config.api_url) {
@@ -76,6 +81,38 @@ pub fn set_llm_config(state: tauri::State<'_, AppState>, config: crate::llm::Llm
Ok(())
}
// ─── Lore RAG injection (shared by generate + generate_stream) ─
/// Best-effort: if `rag_query` is set, retrieve top lore chunks and fold them
/// into the system message so generation is grounded in the user's world bible.
/// Silently degrades to the original messages on any retrieval error.
async fn inject_lore(
state: &tauri::State<'_, AppState>,
_client: &reqwest::Client,
config: &crate::llm::LlmConfig,
mut messages: Vec<ChatMessage>,
rag_query: &Option<String>,
) -> Vec<ChatMessage> {
let Some(query) = rag_query.as_ref() else { return messages; };
if query.trim().is_empty() { return messages; }
let Ok(hits) = state.rag.search(config, query, 4).await else { return messages; };
if hits.is_empty() { return messages; }
let lore = hits.iter().map(|(t, _s, _sc)| t.clone()).collect::<Vec<_>>().join("\n\n---\n\n");
let injection = format!(
"Relevant lore from the campaign bible — stay consistent with it:\n\n{lore}"
);
if let Some(first) = messages.first_mut() {
if first.role == "system" {
first.content = format!("{}\n\n{}", first.content, injection);
return messages;
}
}
messages.insert(0, ChatMessage { role: "system".into(), content: injection });
messages
}
// ─── Ollama API ───────────────────────────────────────────────
async fn call_ollama(
src-tauri/src/commands/mod.rs
+4 -1
View File
@@ -1 +1,4 @@
pub mod llm_commands;
pub mod llm_commands;
pub mod image_commands;
pub mod rag_commands;
pub mod generation_commands;
src-tauri/src/commands/rag_commands.rs
+58
View File
@@ -0,0 +1,58 @@
use crate::llm::AppState;
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize)]
pub struct RagAddRequest {
pub source: String,
pub text: String,
}
#[derive(Debug, Serialize)]
pub struct RagHit {
pub text: String,
pub source: String,
pub score: f32,
}
#[derive(Debug, Serialize)]
pub struct RagSource {
pub source: String,
pub chunks: i64,
}
#[tauri::command]
pub async fn rag_add(state: tauri::State<'_, AppState>, req: RagAddRequest) -> Result<usize, String> {
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
state.rag.add_document(&config, &req.source, &req.text).await.map_err(|e| e.to_string())
}
#[tauri::command]
pub async fn rag_search(
state: tauri::State<'_, AppState>,
query: String,
top_k: Option<usize>,
) -> Result<Vec<RagHit>, String> {
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
let hits = state
.rag
.search(&config, &query, top_k.unwrap_or(4))
.await
.map_err(|e| e.to_string())?;
Ok(hits.into_iter().map(|(text, source, score)| RagHit { text, source, score }).collect())
}
#[tauri::command]
pub fn rag_list(state: tauri::State<'_, AppState>) -> Result<Vec<RagSource>, String> {
state
.rag
.list_sources()
.map_err(|e| e.to_string())?
.into_iter()
.map(|(source, chunks)| Ok(RagSource { source, chunks }))
.collect()
}
#[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())
}
src-tauri/src/generations/mod.rs
+197
View File
@@ -0,0 +1,197 @@
use rusqlite::{params, params_from_iter, types::Value, Connection};
use serde::Serialize;
use std::sync::Mutex;
// ponytail: a separate SQLite file from lore.db so a corrupt generations
// store can't lose the user's world bible. Same `bundled` rusqlite feature.
// Generations are append-only in v1: no schema migrations planned.
const SCHEMA: &str = "CREATE TABLE IF NOT EXISTS generations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
kind TEXT NOT NULL,
title TEXT NOT NULL,
data TEXT NOT NULL,
source TEXT,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_gen_kind ON generations(kind);
CREATE INDEX IF NOT EXISTS idx_gen_created ON generations(created_at DESC);";
#[derive(Debug, Serialize, Clone)]
pub struct Generation {
pub id: i64,
pub kind: String,
pub title: String,
pub data: String,
pub source: Option<String>,
pub created_at: i64,
}
#[derive(Debug, Serialize)]
pub struct GenerationSummary {
pub id: i64,
pub kind: String,
pub title: String,
/// Epoch ms — for relative-time display in the UI.
pub created_at: i64,
}
pub struct GenerationStore {
db: Mutex<Connection>,
}
impl GenerationStore {
pub fn open(dir: &std::path::Path) -> anyhow::Result<Self> {
std::fs::create_dir_all(dir)?;
let path = dir.join("generations.db");
let db = Connection::open(path)?;
db.execute_batch(SCHEMA)?;
Ok(Self { db: Mutex::new(db) })
}
/// Append a generation. `data` is the raw JSON string from the LLM
/// (already-validated by the caller), `source` is the human-readable
/// input (theme, name, etc.) so the history list can show context.
pub fn add(
&self,
kind: &str,
title: &str,
data: &str,
source: Option<&str>,
) -> anyhow::Result<i64> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
db.execute(
"INSERT INTO generations (kind, title, data, source, created_at) VALUES (?, ?, ?, ?, ?)",
params![kind, title, data, source, now],
)?;
Ok(db.last_insert_rowid())
}
/// All generations, newest first, optionally filtered by kind.
pub fn list(&self, kind: Option<&str>) -> anyhow::Result<Vec<GenerationSummary>> {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let (sql, params): (&str, Vec<Value>) = match kind {
Some(k) => (
"SELECT id, kind, title, created_at FROM generations WHERE kind = ?1 ORDER BY created_at DESC",
vec![Value::from(k.to_string())],
),
None => (
"SELECT id, kind, title, created_at FROM generations ORDER BY created_at DESC",
vec![],
),
};
let mut stmt = db.prepare(sql)?;
let rows = stmt.query_map(params_from_iter(params.iter()), |r| {
Ok(GenerationSummary {
id: r.get(0)?,
kind: r.get(1)?,
title: r.get(2)?,
created_at: r.get(3)?,
})
})?;
let mut out = Vec::new();
for r in rows {
out.push(r?);
}
Ok(out)
}
/// Full row including data + source. Used when rehydrating into a tool.
pub fn get(&self, id: i64) -> anyhow::Result<Option<Generation>> {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let mut stmt = db.prepare(
"SELECT id, kind, title, data, source, created_at FROM generations WHERE id = ?",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
Ok(Some(Generation {
id: r.get(0)?,
kind: r.get(1)?,
title: r.get(2)?,
data: r.get(3)?,
source: r.get(4)?,
created_at: r.get(5)?,
}))
} else {
Ok(None)
}
}
/// Delete one row, or all rows if id is None.
pub fn delete(&self, id: Option<i64>) -> anyhow::Result<usize> {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let n = match id {
Some(i) => db.execute("DELETE FROM generations WHERE id = ?", params![i])?,
None => db.execute("DELETE FROM generations", [])?,
};
Ok(n)
}
/// Count by kind — for the dashboard "X NPCs, Y encounters" pill.
pub fn counts(&self) -> anyhow::Result<Vec<(String, i64)>> {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let mut stmt = db.prepare(
"SELECT kind, COUNT(*) FROM generations GROUP BY kind ORDER BY kind",
)?;
let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)))?;
let mut out = Vec::new();
for r in rows {
out.push(r?);
}
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tmp() -> std::path::PathBuf {
let dir = std::env::temp_dir().join(format!("dm-pal-test-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
dir
}
#[test]
fn round_trip() {
let dir = tmp();
let store = GenerationStore::open(&dir).unwrap();
let id = store
.add("npc", "Thorin Stonefist", r#"{"name":"Thorin"}"#, Some("Dwarf Fighter"))
.unwrap();
let g = store.get(id).unwrap().unwrap();
assert_eq!(g.kind, "npc");
assert_eq!(g.title, "Thorin Stonefist");
assert_eq!(g.source.as_deref(), Some("Dwarf Fighter"));
assert!(g.created_at > 0);
let list = store.list(None).unwrap();
assert_eq!(list.len(), 1);
assert_eq!(list[0].id, id);
let npcs = store.list(Some("npc")).unwrap();
assert_eq!(npcs.len(), 1);
let items = store.list(Some("item")).unwrap();
assert_eq!(items.len(), 0);
let counts = store.counts().unwrap();
assert_eq!(counts, vec![("npc".to_string(), 1)]);
assert_eq!(store.delete(Some(id)).unwrap(), 1);
assert!(store.get(id).unwrap().is_none());
}
#[test]
fn clear_all() {
let dir = tmp();
let store = GenerationStore::open(&dir).unwrap();
store.add("a", "x", "{}", None).unwrap();
store.add("b", "y", "{}", None).unwrap();
assert_eq!(store.delete(None).unwrap(), 2);
assert_eq!(store.list(None).unwrap().len(), 0);
}
}
src-tauri/src/lib.rs
+29 -2
View File
@@ -1,8 +1,11 @@
mod llm;
mod rag;
mod generations;
mod commands;
use llm::AppState;
use std::sync::Mutex;
use tauri::Manager;
#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
@@ -10,8 +13,22 @@ pub fn run() {
.plugin(tauri_plugin_log::Builder::default().build())
.plugin(tauri_plugin_store::Builder::default().build())
.plugin(tauri_plugin_fs::init())
.manage(AppState {
config: Mutex::new(llm::LlmConfig::default()),
.setup(|app| {
let data_dir = app
.path()
.app_data_dir()
.expect("app data dir")
.join("dm-toolkit");
let rag = rag::RagStore::open(&data_dir.join("lore"))
.expect("open lore db");
let gen = generations::GenerationStore::open(&data_dir)
.expect("open generations db");
app.manage(AppState {
config: Mutex::new(llm::LlmConfig::default()),
rag,
gen,
});
Ok(())
})
.invoke_handler(tauri::generate_handler![
greet,
@@ -19,6 +36,16 @@ pub fn run() {
commands::llm_commands::generate_stream,
commands::llm_commands::get_llm_config,
commands::llm_commands::set_llm_config,
commands::image_commands::generate_image,
commands::rag_commands::rag_add,
commands::rag_commands::rag_search,
commands::rag_commands::rag_list,
commands::rag_commands::rag_clear,
commands::generation_commands::generation_add,
commands::generation_commands::generation_list,
commands::generation_commands::generation_get,
commands::generation_commands::generation_delete,
commands::generation_commands::generation_counts,
])
.run(tauri::generate_context!())
.expect("error while running tauri application");
src-tauri/src/llm/mod.rs
+12 -1
View File
@@ -1,6 +1,5 @@
use serde::{Deserialize, Serialize};
use std::sync::Mutex;
use tauri::ipc::Channel;
// ─── LLM Event (for streaming) ───────────────────────────────
@@ -16,6 +15,8 @@ pub enum LlmEvent {
pub struct AppState {
pub config: Mutex<LlmConfig>,
pub rag: crate::rag::RagStore,
pub gen: crate::generations::GenerationStore,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -26,6 +27,10 @@ pub struct LlmConfig {
pub temperature: f32,
pub max_tokens: u32,
pub top_p: f32,
/// Ollama image-generation model (e.g. `x/flux2-klein:4b`). macOS-only via Ollama.
pub image_model: String,
/// Ollama embedding model for lore RAG (e.g. `nomic-embed-text`).
pub embed_model: String,
}
impl Default for LlmConfig {
@@ -38,6 +43,8 @@ impl Default for LlmConfig {
temperature: 0.7,
max_tokens: 512,
top_p: 0.9,
image_model: "x/flux2-klein:4b".to_string(),
embed_model: "nomic-embed-text".to_string(),
}
}
}
@@ -50,6 +57,10 @@ pub struct GenerateRequest {
pub system: Option<String>,
pub temperature: Option<f32>,
pub max_tokens: Option<u32>,
/// Optional RAG query: when set, the top lore chunks for this query are
/// retrieved and prepended to the system prompt so generation stays
/// consistent with the user's world bible.
pub rag_query: Option<String>,
}
// ─── OpenAI-Compatible Chat Response ─────────────────────────
src-tauri/src/rag/mod.rs
+212
View File
@@ -0,0 +1,212 @@
use crate::llm::LlmConfig;
use rusqlite::{params, Connection};
use serde::Deserialize;
use std::sync::Mutex;
// ponytail: brute-force cosine over normalized vectors, not sqlite-vec.
// Ceiling: ~10k chunks stays sub-millisecond on one core. Upgrade path:
// swap search() for a sqlite-vec virtual table when chunk count grows past
// that and scan time shows up in profiles. Avoids native extension loading.
/// One embedding vector, stored as little-endian f32 bytes.
const MAX_CHUNK_CHARS: usize = 1000;
pub struct RagStore {
db: Mutex<Connection>,
}
#[derive(Debug, Deserialize)]
struct EmbedResponse {
embeddings: Vec<Vec<f32>>,
}
impl RagStore {
pub fn open(dir: &std::path::Path) -> anyhow::Result<Self> {
std::fs::create_dir_all(dir)?;
let path = dir.join("lore.db");
let db = Connection::open(path)?;
db.execute_batch(
"CREATE TABLE IF NOT EXISTS chunks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
source TEXT NOT NULL,
text TEXT NOT NULL,
emb BLOB NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_chunks_source ON chunks(source);",
)?;
Ok(Self { db: Mutex::new(db) })
}
/// Chunk `text` on paragraph boundaries, capping each chunk's length.
/// ponytail: no overlap in v1 — fine for retrieval at this scale; add
/// a sliding window if recall on boundary-spanning facts drops.
fn chunk(text: &str) -> Vec<String> {
let mut out = Vec::new();
for para in text.split("\n\n") {
let para = para.trim();
if para.is_empty() {
continue;
}
if para.len() <= MAX_CHUNK_CHARS {
out.push(para.to_string());
} else {
// Hard-cap long paragraphs on char boundaries.
for chunk in para.as_bytes().chunks(MAX_CHUNK_CHARS) {
if let Ok(s) = std::str::from_utf8(chunk) {
out.push(s.trim().to_string());
}
}
}
}
out
}
/// Embed a batch of texts via Ollama `/api/embed`. Vectors come back
/// L2-normalized from the server, so cosine similarity = dot product.
async fn embed(client: &reqwest::Client, config: &LlmConfig, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
let url = format!("{}/api/embed", config.api_url.trim_end_matches('/'));
let body = serde_json::json!({ "model": config.embed_model, "input": texts });
let res = client.post(&url).json(&body).send().await?;
if !res.status().is_success() {
let status = res.status();
let text = res.text().await.unwrap_or_default();
anyhow::bail!("embed error {status}: {text}");
}
let parsed: EmbedResponse = res.json().await?;
Ok(parsed.embeddings)
}
fn vec_to_blob(v: &[f32]) -> Vec<u8> {
let mut bytes = Vec::with_capacity(v.len() * 4);
for f in v {
bytes.extend_from_slice(&f.to_le_bytes());
}
bytes
}
fn blob_to_vec(b: &[u8]) -> Vec<f32> {
b.chunks_exact(4)
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
.collect()
}
/// Add a lore document: chunk it, embed, and store. Returns chunk count.
pub async fn add_document(
&self,
config: &LlmConfig,
source: &str,
text: &str,
) -> anyhow::Result<usize> {
let chunks = Self::chunk(text);
if chunks.is_empty() {
return Ok(0);
}
let client = reqwest::Client::new();
let embeddings = Self::embed(&client, config, &chunks).await?;
if embeddings.len() != chunks.len() {
anyhow::bail!("embedding count mismatch");
}
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let mut stmt = db.prepare("INSERT INTO chunks (source, text, emb) VALUES (?, ?, ?)")?;
for (text, emb) in chunks.iter().zip(embeddings.iter()) {
stmt.execute(params![source, text, Self::vec_to_blob(emb)])?;
}
Ok(chunks.len())
}
/// Brute-force cosine search. Returns (text, source, score) for top_k.
pub async fn search(
&self,
config: &LlmConfig,
query: &str,
top_k: usize,
) -> anyhow::Result<Vec<(String, String, f32)>> {
let client = reqwest::Client::new();
let q_emb = Self::embed(&client, config, &[query.to_string()])
.await?
.into_iter()
.next()
.ok_or_else(|| anyhow::anyhow!("no query embedding"))?;
let rows: Vec<(String, String, Vec<u8>)> = {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let mut stmt = db.prepare("SELECT text, source, emb FROM chunks")?;
let rows = stmt.query_map([], |r| {
Ok((
r.get::<_, String>(0)?,
r.get::<_, String>(1)?,
r.get::<_, Vec<u8>>(2)?,
))
})?;
rows.filter_map(|r| r.ok()).collect()
};
if rows.is_empty() {
return Ok(Vec::new());
}
// ponytail: naive O(n) scan. Fine to ~10k chunks; see module note.
let mut scored: Vec<(String, String, f32)> = rows
.into_iter()
.map(|(text, source, blob)| {
let v = Self::blob_to_vec(&blob);
let dot: f32 = v.iter().zip(q_emb.iter()).map(|(a, b)| a * b).sum();
(text, source, dot)
})
.collect();
scored.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(top_k);
Ok(scored)
}
/// (source, chunk_count) for every distinct source.
pub fn list_sources(&self) -> anyhow::Result<Vec<(String, i64)>> {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
let mut stmt = db.prepare("SELECT source, COUNT(*) FROM chunks GROUP BY source ORDER BY source")?;
let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)))?;
let mut out = Vec::new();
for r in rows {
out.push(r?);
}
Ok(out)
}
/// Clear all chunks, or just one source if given.
pub fn clear(&self, source: Option<&str>) -> anyhow::Result<()> {
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
match source {
Some(s) => { db.execute("DELETE FROM chunks WHERE source = ?", params![s])?; }
None => { db.execute("DELETE FROM chunks", [])?; }
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chunk_splits_paragraphs_and_caps_long_ones() {
let long = "a".repeat(MAX_CHUNK_CHARS * 2 + 50);
let text = format!("short para\n\n{long}\n\nanother");
let chunks = RagStore::chunk(&text);
// "short para", (2 or 3 long pieces), "another"
assert!(chunks.len() >= 4);
assert_eq!(chunks.first().unwrap(), "short para");
assert!(chunks.last().unwrap() == "another");
for c in &chunks {
assert!(c.len() <= MAX_CHUNK_CHARS);
}
}
#[test]
fn blob_roundtrip() {
let v = vec![0.0, 1.5, -2.25, 3.33];
let blob = RagStore::vec_to_blob(&v);
assert_eq!(blob.len(), v.len() * 4);
let back = RagStore::blob_to_vec(&blob);
for (a, b) in v.iter().zip(back.iter()) {
assert!((a - b).abs() < 1e-6);
}
}
}