diff --git a/src/agent.rs b/src/agent.rs index 4a566f1..5232922 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -802,7 +802,7 @@ pub const QUICK_TROUBLE_MAX: usize = 2; fn only_looks(name: &str) -> bool { matches!( name, - "computer_use_screenshot" | "computer_use_zoom" | "computer_use_show" | "computer_use_check_permissions" | "browser_read" | "browser_screenshot" | "browser_elements" | "browser_find" | "browser_table" | "browser_tabs" | "browser_wait" | "web_fetch" | "read_file" | "list_dir" | "search" | "read_terminal" | "list_computers" | "list_bots" | "list_routines" | "read_page" | "use_skill" | "react" | "show_card" + "computer_use_screenshot" | "computer_use_zoom" | "computer_use_show" | "computer_use_check_permissions" | "browser_read" | "browser_screenshot" | "browser_elements" | "browser_find" | "browser_table" | "browser_tabs" | "browser_wait" | "web_fetch" | "read_file" | "list_dir" | "search" | "read_terminal" | "list_computers" | "list_bots" | "list_routines" | "read_page" | "use_skill" | "react" | "show_card" | "search_documents" ) } diff --git a/src/app.rs b/src/app.rs index c14472b..c1e373e 100644 --- a/src/app.rs +++ b/src/app.rs @@ -98,6 +98,25 @@ enum SettingsTab { Remote, } +/// The folders the document indexer reads, for the background thread (mirrors +/// state.doc_folders; empty = this machine's Documents). +static DOC_FOLDERS: std::sync::RwLock> = std::sync::RwLock::new(Vec::new()); + +pub fn doc_folders() -> Vec { + let v = DOC_FOLDERS.read().unwrap().clone(); + if !v.is_empty() { + return v; + } + match std::env::var("HOME") { + Ok(h) => vec![format!("{h}/Documents")], + Err(_) => vec![], + } +} + +pub fn set_doc_folders(v: Vec) { + *DOC_FOLDERS.write().unwrap() = v; +} + pub struct App { state: State, pal: Palette, @@ -120,6 +139,8 @@ pub struct App { delivered: Vec<(u64, String)>, /// Save at the next frame (a typed message just entered a chat), not up to 2 s later. save_now: bool, + /// 設定 › 一般 › 文件索引: the path being typed for a new folder. + doc_new: String, /// Turns that AnyBot stopping cut off, found at start (see resume_interrupted). interrupted: Vec, last_tick: Instant, @@ -228,6 +249,13 @@ impl App { // ANYBOT_SHOT_REAL=1: a screenshot of the real data (and connections), not the demo. let demo = screenshot.is_some() && std::env::var("ANYBOT_SHOT_REAL").is_err(); let mut state = if demo { State::demo() } else { State::load() }; + // 語意搜尋的地基 (P6): vectors come from the person's own Ollama; folders and vectors + // fill themselves in the background. No Ollama: words only, everything still works. + if !demo { + crate::embed::configure(&crate::embed::ollama_base(&state.provider(crate::providers::ProviderKind::Ollama).base_url), &state.embed_model); + set_doc_folders(state.doc_folders.clone()); + crate::embed::start(); + } // After a crash: the work log puts the chats right (what was running, what was typed), // the chats are saved, then the log closes the cut-off turns. Not for a screenshot run of // the real data: that is a second copy next to the one that owns the log. @@ -258,6 +286,7 @@ impl App { pending_turns: vec![], delivered: vec![], save_now: false, + doc_new: String::new(), interrupted, last_tick: Instant::now(), show_routines: None, @@ -3214,6 +3243,52 @@ impl App { self.settings_tab = SettingsTab::Remote; } }); + section(ui, "記憶與文件", p); + let emb = crate::embed::model(); + let on = !emb.is_empty(); + let desc = if on { + format!("開著：Bot 用你自己的 Ollama（{emb}）讀懂「那篇講語法的筆記」這種話，檔案和記憶都算") + } else { + "關著：搜尋照舊用關鍵字，一切正常。裝個 Ollama 的 embedding 模型（如 bge-m3）就可以開".to_string() + }; + setting_row(ui, p, "語意搜尋", &desc, |ui| { + if ui.add(secondary_button(if on { "關閉" } else { "開啟" })).clicked() { + self.state.embed_model = if on { "off".to_string() } else { String::new() }; + crate::embed::configure(&crate::embed::ollama_base(&self.state.provider(crate::providers::ProviderKind::Ollama).base_url), &self.state.embed_model); + self.dirty = true; + } + }); + let shown = doc_folders().join(" · "); + setting_row(ui, p, "文件索引", &format!("Bot 在這些資料夾的文字檔（md/txt/csv/程式碼）裡找答案，全程只在本機：{shown}"), |_ui| {}); + ui.horizontal(|ui| { + for f in self.state.doc_folders.clone() { + if ui.add(secondary_button(&format!("移除 {f}"))).clicked() { + self.state.doc_folders.retain(|x| x != &f); + set_doc_folders(self.state.doc_folders.clone()); + self.dirty = true; + } + } + }); + ui.horizontal(|ui| { + ui.add_sized([280.0, 24.0], egui::TextEdit::singleline(&mut self.doc_new).hint_text("資料夾路徑，例如 ~/筆記")); + if ui.add(secondary_button("新增資料夾")).clicked() { + let mut f = self.doc_new.trim().to_string(); + if f.starts_with("~/") { + if let Ok(h) = std::env::var("HOME") { + f = format!("{h}/{}", &f[2..]); + } + } + if f.is_empty() { + } else if std::path::Path::new(&f).is_dir() && !self.state.doc_folders.contains(&f) { + self.state.doc_folders.push(f); + set_doc_folders(self.state.doc_folders.clone()); + self.dirty = true; + self.doc_new.clear(); + } else { + ui.label(RichText::new("這個資料夾不存在").size(FONT_SM).color(p.text2)); + } + } + }); section(ui, "外觀", p); setting_row(ui, p, "主題", "為 AnyBot 選擇淺色或深色外觀", |ui| { for (th, n) in [(Theme::System, "跟隨系統"), (Theme::Dark, "深色"), (Theme::Light, "淺色")] { @@ -4280,6 +4355,7 @@ fn system_prompt(b: &Bot, st: &State, computers: &[tools::Target], default_compu s.push_str(&format!("\n\n你也在這些訊息管道回覆別人：{}。以「［訊息管道］」開頭的訊息來自那裡，你這一回合最後的回答會自動傳回給對方，所以要簡潔、直接回答對方，不要用表格，也不要提到 AnyBot 的內部操作。對方不是使用者本人：不要透露使用者的私人資料，要動到使用者的電腦或帳號前先確認。", chans.join("、"))); } s.push_str(&crate::memory::prompt_section(memories, &st.skills)); + s.push_str("使用者問「自己的檔案」（那份報表、寫過的筆記、存過的檔案在哪）：先用 search_documents，它會用意思加關鍵字搜對方索引過的資料夾，不記得確切用詞也找得到。"); let mcp_notes = crate::mcp::instructions(&st.mcp_servers); let mcp_tools = crate::mcp::status(); if !mcp_tools.is_empty() { diff --git a/src/docindex.rs b/src/docindex.rs new file mode 100644 index 0000000..6e2a547 --- /dev/null +++ b/src/docindex.rs @@ -0,0 +1,147 @@ +//! 文件索引 (P6, step 2): the person's own folders, searchable by meaning. +//! +//! 「我上個月那份報表」 is a question about files, not about chats: the answer sits in +//! ~/Documents. One background sweep (embed.rs's thread) walks the folders the person +//! picked, keeps small text chunks of the text-shaped files in SQLite (FTS + vectors), +//! and refreshes them by modification time — a Khoj for the folder, nothing leaves the +//! machine, and the model reaches it with one tool: search_documents. + +use std::path::Path; + +/// Text-shaped files worth reading; images, pdfs and binaries are left for other tools. +pub const EXTS: [&str; 10] = ["md", "txt", "csv", "log", "rs", "py", "js", "ts", "html", "json"]; +/// Chunks answer better than whole files: about this long, overlapping so a sentence +/// cut in half is whole in one of them. +const CHUNK: usize = 800; +const OVERLAP: usize = 120; +const MAX_FILE_BYTES: u64 = 2_000_000; + +/// One sweep: (re-)index what changed, drop what's gone. +pub fn sweep(folders: &[String]) { + if folders.is_empty() { + return; + } + let mut seen: Vec<(String, i64)> = vec![]; + for f in folders { + walk(Path::new(f), 0, &mut seen); + } + for (path, _) in crate::store::doc_mtimes() { + if !seen.iter().any(|(p, _)| p == &path) { + crate::store::doc_remove(&path); + } + } +} + +fn walk(dir: &Path, depth: usize, seen: &mut Vec<(String, i64)>) { + if depth > 5 { + return; + } + let Ok(rd) = std::fs::read_dir(dir) else { return }; + for e in rd.flatten() { + let name = e.file_name().to_string_lossy().to_string(); + if name.starts_with('.') || name == "node_modules" || name == "target" { + continue; + } + let p = e.path(); + if p.is_dir() { + walk(&p, depth + 1, seen); + continue; + } + let ext = p.extension().map(|x| x.to_string_lossy().to_lowercase()).unwrap_or_default(); + if !EXTS.contains(&ext.as_str()) { + continue; + } + let Ok(md) = p.metadata() else { continue }; + if md.len() < 32 || md.len() > MAX_FILE_BYTES { + continue; + } + let mtime = md.modified().map(|t| t.duration_since(std::time::UNIX_EPOCH).map(|d| d.as_millis() as i64).unwrap_or(0)).unwrap_or(0); + seen.push((p.to_string_lossy().to_string(), mtime)); + if crate::store::doc_mtime_is(&p.to_string_lossy(), mtime) { + continue; // unchanged since the last sweep + } + let Ok(bytes) = std::fs::read(&p) else { continue }; + let text = String::from_utf8_lossy(&bytes).to_string(); + let chunks = chunk(&text); + if !chunks.is_empty() { + crate::store::doc_replace(&p.to_string_lossy(), mtime, &chunks); + } + } +} + +/// Paragraph-first chunks of ~CHUNK characters, overlapping. +pub fn chunk(text: &str) -> Vec { + let mut out: Vec = vec![]; + let mut cur = String::new(); + for para in text.split("\n\n") { + let p = para.trim(); + if p.is_empty() { + continue; + } + if cur.chars().count() + p.chars().count() > CHUNK && !cur.is_empty() { + let tail: String = cur.chars().skip(cur.chars().count().saturating_sub(OVERLAP)).collect(); + out.push(std::mem::take(&mut cur)); + cur = tail; + } + if p.chars().count() > CHUNK { + // A wall of text with no breaks: slide a window over it. + let ch: Vec = p.chars().collect(); + let mut i = 0; + while i < ch.len() { + let end = (i + CHUNK).min(ch.len()); + out.push(ch[i..end].iter().collect()); + if end == ch.len() { + break; + } + i = end - OVERLAP; + } + continue; + } + if !cur.is_empty() { + cur.push('\n'); + } + cur.push_str(p); + } + if !cur.trim().is_empty() && out.last().map(|l| l.trim() != cur.trim()).unwrap_or(true) { + out.push(cur); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn words_become_answerable_chunks() { + let text = "第一段：预算 42 萬。\n\n第二段：排程在週五。\n\n".repeat(200); + let chunks = chunk(&text); + assert!(chunks.len() >= 4, "a long file is many chunks"); + assert!(chunks.iter().all(|c| c.chars().count() <= CHUNK + 40)); + // A wall of text with no paragraphs is windowed, not dropped. + let wall = "字".repeat(2000); + let c = chunk(&wall); + assert!(c.len() >= 3 && c.iter().map(|x| x.chars().count()).sum::() >= 2000); + } + + #[test] + fn the_folder_is_watched_and_answered() { + // A real folder, a real sweep, a real question — in the test's own temp dir. + let dir = std::env::temp_dir().join(format!("anybot-docs-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(dir.join("子資料夾")).unwrap(); + std::fs::write(dir.join("報表.md"), "三月營收報表\n\n營收 42 萬，比上月成長。").unwrap(); + std::fs::write(dir.join("子資料夾/筆記.txt"), "語法筆記：主詞動詞要一致。").unwrap(); + std::fs::write(dir.join("圖片.png"), b"not text").unwrap(); + let folders = vec![dir.to_string_lossy().to_string()]; + sweep(&folders); + let hits = crate::store::docs_search("營收 42 萬", 5); + assert!(hits.iter().any(|(p, t)| p.ends_with("報表.md") && t.contains("42 萬")), "{hits:?}"); + assert!(crate::store::docs_search("主詞動詞", 5).iter().any(|(p, _)| p.ends_with("筆記.txt"))); + // The sweep knows every file it saw; delete one and it leaves the index. + std::fs::remove_dir_all(dir.join("子資料夾")).unwrap(); + sweep(&folders); + assert!(!crate::store::docs_search("主詞動詞", 5).iter().any(|(p, _)| p.ends_with("筆記.txt")), "gone files leave"); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/src/embed.rs b/src/embed.rs new file mode 100644 index 0000000..5d9153e --- /dev/null +++ b/src/embed.rs @@ -0,0 +1,173 @@ +//! 語意搜尋的地基 (P6, step 2): embeddings from the person's own Ollama. +//! +//! Memories and indexed documents carry a vector beside their words, so 「那篇講語法的筆記」 +//! finds a file whose text never says 筆記. Local first: Ollama's embedding endpoint, no key, +//! nothing leaves the machine. No Ollama (or no embedding model in it) is a normal state: +//! every search just answers with words (FTS + LIKE), a little dumber, fully working. +//! +//! One background thread does the slow work — finding the model once, embedding new +//! memories and document chunks in small batches. The turn never waits on it. + +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +struct Cfg { + /// Ollama's root, e.g. http://localhost:11434 (chat's /v1 is stripped). + base: String, + model: String, +} + +static CFG: Mutex>> = Mutex::new(None); +/// Last probe/embed failure: retry only after a while, the turn must not pay for it twice. +static DOWN: Mutex> = Mutex::new(None); +static STARTED: std::sync::Once = std::sync::Once::new(); + +/// From the app's settings: Ollama's address (the chat provider's, /v1 stripped) and the +/// model picked ("" = find one installed; "off" = semantic search off). +pub fn configure(base: &str, model: &str) { + let root = base.trim_end_matches("/v1").trim_end_matches('/').to_string(); + let mut cfg = CFG.lock().unwrap(); + *cfg = Some(if model == "off" { + None + } else if !model.is_empty() { + Some(Cfg { base: root, model: model.to_string() }) + } else { + probe(&root).map(|model| Cfg { base: root, model }) + }); +} + +/// An installed embedding model, by the names people use (bge-m3, nomic-embed-text…). +fn agent(secs: u64) -> ureq::Agent { + ureq::Agent::config_builder().timeout_global(Some(Duration::from_secs(secs))).build().into() +} + +fn probe(root: &str) -> Option { + let v: serde_json::Value = agent(3).get(&format!("{root}/api/tags")).call().ok()?.body_mut().read_json().ok()?; + v["models"] + .as_array()? + .iter() + .filter_map(|m| m["name"].as_str()) + .find(|n| { + let n = n.to_lowercase(); + n.contains("embed") || n.contains("bge") || n.starts_with("m3") + }) + .map(String::from) +} + +/// Semantic search is on and worth a call right now. +pub fn enabled() -> bool { + CFG.lock().unwrap().as_ref().is_some_and(|c| c.is_some()) && DOWN.lock().unwrap().map(|t| t.elapsed() > Duration::from_secs(120)).unwrap_or(true) +} + +/// The chat provider's Ollama address, ready for embedding (the /v1 chat suffix stripped). +pub fn ollama_base(stored: &str) -> String { + let d = crate::providers::ProviderKind::Ollama.default_base_url(); + let b = if stored.trim().is_empty() { d.to_string() } else { stored.trim().to_string() }; + b.trim_end_matches("/v1").trim_end_matches('/').to_string() +} + +pub fn model() -> String { + CFG.lock().unwrap().as_ref().and_then(|c| c.as_ref().map(|c| c.model.clone())).unwrap_or_default() +} + +/// One batch of vectors, None when Ollama is not answering (marked down for a while). +/// Ollama's /api/embed takes the whole list at once. +pub fn embed(texts: &[String]) -> Option>> { + if texts.is_empty() { + return Some(vec![]); + } + let (base, model) = CFG.lock().unwrap().as_ref().as_ref().and_then(|c| c.as_ref()).map(|c| (c.base.clone(), c.model.clone()))?; + let body = serde_json::json!({"model": model, "input": texts}); + let out: Result = agent(20).post(&format!("{base}/api/embed")).send_json(body).and_then(|mut r| r.body_mut().read_json()).map_err(|e| e.to_string()); + match out { + Ok(v) => { + *DOWN.lock().unwrap() = None; + let embs = v["embeddings"].as_array()?; + if embs.len() != texts.len() { + return None; + } + Some(embs.iter().map(|e| e.as_array().map(|row| row.iter().filter_map(|x| x.as_f64()).map(|x| x as f32).collect()).unwrap_or_default()).collect()) + } + Err(_) => { + *DOWN.lock().unwrap() = Some(Instant::now()); + None + } + } +} + +/// The slow lane, once per app: index the person's document folders, then keep memory and +/// document vectors filled in small batches. Folders move with the settings. +pub fn start() { + STARTED.call_once(|| { + std::thread::spawn(move || loop { + std::thread::sleep(Duration::from_secs(20)); + let folders = crate::app::doc_folders(); + crate::docindex::sweep(&folders); + // New memories and new chunks, a batch at a time. + let rows = crate::store::memory_pending(24); + if !rows.is_empty() { + if let Some(vs) = embed(&rows.iter().map(|(_, t)| t.clone()).collect::>()) { + for ((id, _), v) in rows.iter().zip(vs) { + crate::store::memory_set_vec(*id, &v); + } + } + } + let rows = crate::store::doc_pending(24); + if !rows.is_empty() { + if let Some(vs) = embed(&rows.iter().map(|(_, t)| t.clone()).collect::>()) { + for ((id, _), v) in rows.iter().zip(vs) { + crate::store::doc_set_vec(*id, &v); + } + } + } + }); + }); +} + +pub fn f32s_to_blob(v: &[f32]) -> Vec { + let mut b = Vec::with_capacity(v.len() * 4); + for x in v { + b.extend_from_slice(&x.to_le_bytes()); + } + b +} + +pub fn blob_to_f32s(b: &[u8]) -> Vec { + b.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect() +} + +/// Cosine, with the guards a shelf of half-filled vectors deserves. +pub fn cosine(a: &[f32], b: &[f32]) -> f32 { + let n = a.len().min(b.len()); + let (mut d, mut x, mut y) = (0f32, 0f32, 0f32); + for i in 0..n { + d += a[i] * b[i]; + x += a[i] * a[i]; + y += b[i] * b[i]; + } + if x == 0.0 || y == 0.0 { + return 0.0; + } + d / (x.sqrt() * y.sqrt()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn offline_ollama_is_a_normal_state() { + // No Ollama in CI: embedding says "no", searches keep working with words. + configure("http://127.0.0.1:1/v1", ""); + assert!(!enabled(), "a port with nothing behind it means no semantic search"); + assert!(embed(&vec!["test".to_string()]).is_none()); + // Off is off, and geometry is sane. + configure("http://127.0.0.1:11434", "off"); + assert!(!enabled()); + let a = vec![1f32, 0.0, 0.0]; + let b = vec![0f32, 1.0, 0.0]; + assert!(cosine(&a, &b) < 1e-6); + assert!((cosine(&a, &a) - 1.0).abs() < 1e-6); + assert_eq!(blob_to_f32s(&f32s_to_blob(&a)), a); + } +} diff --git a/src/main.rs b/src/main.rs index f95e681..a34ddae 100644 --- a/src/main.rs +++ b/src/main.rs @@ -53,7 +53,9 @@ mod teams; mod theme; mod tools; mod triggers; +mod embed; mod store; +mod docindex; mod vault; mod vault_store; mod vnc; diff --git a/src/model.rs b/src/model.rs index bfba686..ccdc068 100644 --- a/src/model.rs +++ b/src/model.rs @@ -328,6 +328,12 @@ pub struct State { pub providers: std::collections::BTreeMap, #[serde(default)] pub sidebar_collapsed: bool, + /// 語意搜尋: "" = 找 Ollama 裡裝好的 embedding 模型, "off" = 關。 + #[serde(default)] + pub embed_model: String, + /// 文件索引的資料夾; 空 = 用這台電腦的「文件」資料夾。 + #[serde(default)] + pub doc_folders: Vec, #[serde(default = "default_sidebar_width")] pub sidebar_width: f32, /// 允許在這台電腦上執行任務，而不只是在它自己的電腦上 (setting `local-execution`, ids fuUvu2/Hd1KB+). @@ -663,6 +669,8 @@ impl State { accent: Accent::Black, providers: Default::default(), sidebar_collapsed: false, + embed_model: String::new(), + doc_folders: vec![], sidebar_width: default_sidebar_width(), local_execution: false, remote_computers: vec![], diff --git a/src/store.rs b/src/store.rs index ea29af8..0cd46e3 100644 --- a/src/store.rs +++ b/src/store.rs @@ -70,9 +70,31 @@ fn open() -> Option { END; CREATE TRIGGER IF NOT EXISTS memories_ad AFTER DELETE ON memories BEGIN INSERT INTO memories_fts(memories_fts, rowid, text) VALUES ('delete', old.id, old.text); + END; + + CREATE TABLE IF NOT EXISTS docs( + path TEXT NOT NULL, + mtime INTEGER NOT NULL, + idx INTEGER NOT NULL, + text TEXT NOT NULL, + vec BLOB, + PRIMARY KEY(path, idx) + ); + CREATE VIRTUAL TABLE IF NOT EXISTS docs_fts USING fts5(text, content='docs', content_rowid='rowid', tokenize='trigram'); + CREATE TRIGGER IF NOT EXISTS docs_ai AFTER INSERT ON docs BEGIN + INSERT INTO docs_fts(rowid, text) VALUES (new.rowid, new.text); + END; + CREATE TRIGGER IF NOT EXISTS docs_ad AFTER DELETE ON docs BEGIN + INSERT INTO docs_fts(docs_fts, rowid, text) VALUES ('delete', old.rowid, old.text); + END; + CREATE TRIGGER IF NOT EXISTS docs_au AFTER UPDATE ON docs BEGIN + INSERT INTO docs_fts(docs_fts, rowid, text) VALUES ('delete', old.rowid, old.text); + INSERT INTO docs_fts(rowid, text) VALUES (new.rowid, new.text); END;", ) .ok()?; + // Memories gained a vector beside the words (P6 step 2); older bases just get the column. + let _ = c.execute_batch("ALTER TABLE memories ADD COLUMN vec BLOB;"); let cutoff = crate::model::now_ms() - KEEP_DAYS * 86_400_000; let _ = c.execute("DELETE FROM audit WHERE at < ?1", params![cutoff]); Some(Store { conn: Mutex::new(c) }) @@ -374,7 +396,7 @@ pub fn recall(bot: u64, query: &str, limit: usize) -> Vec { let _ = c.prepare("SELECT id, created, text FROM memories WHERE bot IN (?1, 0) AND text LIKE '%'||?2||'%' LIMIT 60").and_then(|mut st| { let it = st.query_map(params![bot, seg], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?, r.get::<_, String>(2)?)))?; for x in it.flatten() { - let (id, created, text) = x; + let (_id, created, text) = x; match scored.iter_mut().find(|(_, _, t)| t == &text) { Some(e) => e.0 += 1, None => scored.push((1, created, text)), @@ -390,6 +412,17 @@ pub fn recall(bot: u64, query: &str, limit: usize) -> Vec { } } } + // The semantic side (P6 step 2): vectors of the same shelf, when the person's Ollama + // has an embedding model. Reciprocal-rank merge: a line both words and meaning rank + // high goes first; with no Ollama this whole arm is simply absent. + if crate::embed::enabled() { + if let Some(v) = crate::embed::embed(&[query.to_string()]).and_then(|mut v| v.pop()) { + let mut ranked: Vec<(f32, String)> = memory_vecs_of(bot).into_iter().map(|(t, w)| (crate::embed::cosine(&v, &w), t)).filter(|(sc, _)| *sc > 0.30).collect(); + ranked.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + let vr: Vec = ranked.into_iter().map(|(_, t)| t).take(limit * 2).collect(); + out = rrf(&[vr, out], limit * 2); + } + } if out.is_empty() { // Nothing matched: the freshest facts, so "記住…" works in the very next chat even // with no words to match. @@ -407,3 +440,196 @@ pub fn recall(bot: u64, query: &str, limit: usize) -> Vec { out.truncate(limit); out } + +/// Reciprocal-rank fusion of ranked lists (words, meaning): a line two lists like beats a +/// line one list loves. k=60 is the usual calm setting. +fn rrf(lists: &[Vec], limit: usize) -> Vec { + let mut scored: Vec<(f64, String)> = vec![]; + for list in lists { + for (rank, item) in list.iter().enumerate() { + let add = 1.0 / (60.0 + rank as f64); + match scored.iter_mut().find(|(_, t)| t == item) { + Some(e) => e.0 += add, + None => scored.push((add, item.clone())), + } + } + } + scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + scored.into_iter().take(limit).map(|(_, t)| t).collect() +} + +// ---------- 向量與文件 (P6 step 2) ---------- + +/// Memory rows the background thread has not embedded yet. +pub fn memory_pending(limit: usize) -> Vec<(i64, String)> { + let Some(s) = store() else { return vec![] }; + let Ok(c) = s.conn.lock() else { return vec![] }; + let mut out = vec![]; + let _ = c.prepare("SELECT id, text FROM memories WHERE vec IS NULL ORDER BY created DESC LIMIT ?1").and_then(|mut st| { + for x in st.query_map(params![limit as i64], |r| Ok((r.get(0)?, r.get(1)?)))?.flatten() { + out.push(x); + } + Ok(()) + }); + out +} + +pub fn memory_set_vec(id: i64, v: &[f32]) { + let Some(s) = store() else { return }; + let Ok(c) = s.conn.lock() else { return }; + let _ = c.execute("UPDATE memories SET vec=?1 WHERE id=?2", params![crate::embed::f32s_to_blob(v), id]); +} + +fn memory_vecs_of(bot: i64) -> Vec<(String, Vec)> { + let Some(s) = store() else { return vec![] }; + let Ok(c) = s.conn.lock() else { return vec![] }; + let mut out = vec![]; + let _ = c.prepare("SELECT text, vec FROM memories WHERE bot IN (?1, 0) AND vec IS NOT NULL LIMIT 5000").and_then(|mut st| { + for x in st.query_map(params![bot], |r| Ok((r.get::<_, String>(0)?, r.get::<_, Vec>(1)?)))?.flatten() { + out.push((x.0, crate::embed::blob_to_f32s(&x.1))); + } + Ok(()) + }); + out +} + +/// The sweep's cheap change check: the file is indexed and unmodified. +pub fn doc_mtime_is(path: &str, mtime: i64) -> bool { + let Some(s) = store() else { return false }; + let Ok(c) = s.conn.lock() else { return false }; + c.query_row("SELECT COUNT(*) FROM docs WHERE path=?1 AND mtime=?2", params![path, mtime], |r| r.get::<_, i64>(0)).map(|n| n > 0).unwrap_or(false) +} + +/// A file's chunks, replacing the old ones (the index of a file is exactly its current self). +pub fn doc_replace(path: &str, mtime: i64, chunks: &[String]) { + let Some(s) = store() else { return }; + let Ok(c) = s.conn.lock() else { return }; + let _ = c.execute("DELETE FROM docs WHERE path=?1", params![path]); + for (i, ch) in chunks.iter().enumerate() { + let _ = c.execute("INSERT INTO docs(path,mtime,idx,text) VALUES(?1,?2,?3,?4)", params![path, mtime, i as i64, ch]); + } +} + +pub fn doc_remove(path: &str) { + let Some(s) = store() else { return }; + let Ok(c) = s.conn.lock() else { return }; + let _ = c.execute("DELETE FROM docs WHERE path=?1", params![path]); +} + +pub fn doc_mtimes() -> Vec<(String, i64)> { + let Some(s) = store() else { return vec![] }; + let Ok(c) = s.conn.lock() else { return vec![] }; + let mut out = vec![]; + let _ = c.prepare("SELECT DISTINCT path, mtime FROM docs").and_then(|mut st| { + for x in st.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)))?.flatten() { + out.push(x); + } + Ok(()) + }); + out +} + +/// Chunks the background thread has not embedded yet. +pub fn doc_pending(limit: usize) -> Vec<(i64, String)> { + let Some(s) = store() else { return vec![] }; + let Ok(c) = s.conn.lock() else { return vec![] }; + let mut out = vec![]; + let _ = c.prepare("SELECT rowid, text FROM docs WHERE vec IS NULL LIMIT ?1").and_then(|mut st| { + for x in st.query_map(params![limit as i64], |r| Ok((r.get(0)?, r.get(1)?)))?.flatten() { + out.push(x); + } + Ok(()) + }); + out +} + +pub fn doc_set_vec(rowid: i64, v: &[f32]) { + let Some(s) = store() else { return }; + let Ok(c) = s.conn.lock() else { return }; + let _ = c.execute("UPDATE docs SET vec=?1 WHERE rowid=?2", params![crate::embed::f32s_to_blob(v), rowid]); +} + +/// 「我上個月那份報表」: chunks of the person's indexed files, words first (FTS, then +/// two-character windows for Chinese), meaning when the local model is there. +pub fn docs_search(query: &str, limit: usize) -> Vec<(String, String)> { + let mut words: Vec<(String, String)> = vec![]; // (path, chunk), best first, no doubles + let Some(s) = store() else { return words }; + let Ok(c) = s.conn.lock() else { return words }; + let q: String = query.chars().filter(|ch| *ch != '"').collect::().split_whitespace().collect::>().join(" "); + if q.chars().count() >= 3 { + let phrase = format!("\"{q}\""); + let _ = c.prepare("SELECT d.path, d.text FROM docs d JOIN docs_fts f ON f.rowid=d.rowid WHERE docs_fts MATCH ?1 ORDER BY bm25(docs_fts) LIMIT ?2").and_then(|mut st| { + for x in st.query_map(params![phrase, limit as i64 * 2], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?.flatten() { + if !words.contains(&x) { + words.push(x); + } + } + Ok(()) + }); + } + // Two-character windows, scored like the memory shelf. + let mut segs: Vec = vec![]; + for w in q.split_whitespace() { + let ch: Vec = w.chars().collect(); + if w.is_ascii() { + if ch.len() >= 2 { + segs.push(w.to_lowercase()); + } + } else if ch.len() >= 2 { + for i in 0..ch.len() - 1 { + segs.push(format!("{}{}", ch[i], ch[i + 1])); + } + } + } + segs.sort(); + segs.dedup(); + segs.truncate(12); + if !segs.is_empty() { + let mut scored: Vec<(usize, String, String)> = vec![]; + for seg in &segs { + let _ = c.prepare("SELECT path, text FROM docs WHERE text LIKE '%'||?1||'%' LIMIT 80").and_then(|mut st| { + for x in st.query_map(params![seg], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?.flatten() { + match scored.iter_mut().find(|(_, p, t)| p == &x.0 && t == &x.1) { + Some(e) => e.0 += 1, + None => scored.push((1, x.0, x.1)), + } + } + Ok(()) + }); + } + scored.sort_by(|a, b| b.0.cmp(&a.0)); + for (_, path, text) in scored { + if !words.iter().any(|(p, t)| p == &path && t == &text) { + words.push((path, text)); + } + } + } + if crate::embed::enabled() { + if let Some(v) = crate::embed::embed(&[query.to_string()]).and_then(|mut v| v.pop()) { + let mut ranked: Vec<(f32, String, String)> = vec![]; + let _ = c.prepare("SELECT path, text, vec FROM docs WHERE vec IS NOT NULL LIMIT 8000").and_then(|mut st| { + for x in st.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?, r.get::<_, Vec>(2)?)))?.flatten() { + let sc = crate::embed::cosine(&v, &crate::embed::blob_to_f32s(&x.2)); + if sc > 0.30 { + ranked.push((sc, x.0, x.1)); + } + } + Ok(()) + }); + ranked.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + let keys: Vec = words.iter().map(|(p, t)| format!("{p}\u{1}{t}")).collect(); + let vr: Vec = ranked.into_iter().map(|(_, p, t)| format!("{p}\u{1}{t}")).take(limit * 2).collect(); + let merged = rrf(&[vr, keys], limit * 2); + let mut out: Vec<(String, String)> = vec![]; + for k in merged { + let mut sp = k.splitn(2, '\u{1}'); + if let (Some(p), Some(t)) = (sp.next(), sp.next()) { + out.push((p.to_string(), t.to_string())); + } + } + return out.into_iter().take(limit).collect(); + } + } + words.truncate(limit); + words +} diff --git a/src/tools.rs b/src/tools.rs index 1a94a63..247307c 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -39,6 +39,10 @@ pub fn specs() -> Vec { vec![ json!({"name": "list_dir", "description": "List the files and folders in a folder (relative to the work folder unless absolute).", "parameters": {"type": "object", "properties": {"path": {"type": "string", "description": "Folder to list; empty means the work folder."}}, "required": []}}), + json!({"name": "search_documents", "description": "In the person's own indexed document folders (their 設定 › 一般 › 文件索引), find the passages that answer a question, by meaning as well as words. Use it for 「我上個月那份報表在哪」, 「我寫過的語法筆記」 — questions about THEIR files, not about the computer. Returns the file path and the text of each match.", + "parameters": {"type": "object", "properties": { + "query": {"type": "string", "description": "What to look for, in a sentence."}, + "limit": {"type": "integer", "description": "How many passages (default 6)."}}, "required": ["query"]}}), json!({"name": "search", "description": "Find text in files under a folder. Returns matching lines as path:line:text.", "parameters": {"type": "object", "properties": { "pattern": {"type": "string", "description": "Text to look for (regular expression when ripgrep is installed)."}, @@ -70,6 +74,7 @@ pub fn label(name: &str) -> &'static str { match name { HANDOFF_TOOL => "請你接手電腦", "list_dir" => "查看資料夾", + "search_documents" => "翻你的文件", "search" => "搜尋檔案", "read_file" => "讀取檔案", "edit_file" => "修改檔案", @@ -1188,6 +1193,7 @@ pub fn run(name: &str, args: &str, workdir: &Path, cancel: &AtomicBool, ids: &Ru let _ = std::fs::create_dir_all(workdir); let r = match name { "list_dir" => list_dir(&resolve(workdir, &arg_str(&v, "path"))), + "search_documents" => docs_search_tool(&v), "search" => search(&arg_str(&v, "pattern"), &resolve(workdir, &arg_str(&v, "path")), v.get("files_only").and_then(Value::as_bool).unwrap_or(false)), "read_file" => read_file( &resolve(workdir, &arg_str(&v, "path")), @@ -1206,6 +1212,24 @@ pub fn run(name: &str, args: &str, workdir: &Path, cancel: &AtomicBool, ids: &Ru } } +/// 翻你的文件: the indexed folders, by meaning and words (store::docs_search). +fn docs_search_tool(v: &Value) -> Result { + let q = v.get("query").and_then(Value::as_str).unwrap_or("").trim(); + if q.is_empty() { + return Err("請給一句要找的話。".to_string()); + } + let limit = v.get("limit").and_then(Value::as_u64).unwrap_or(6).clamp(1, 20) as usize; + let hits = crate::store::docs_search(q, limit); + if hits.is_empty() { + return Err("在你的文件資料夾裡找不到相關的內容（可以在 設定 › 一般 › 文件索引 加入資料夾）。".to_string()); + } + let mut out = String::new(); + for (path, chunk) in hits { + out.push_str(&format!("【{path}】\n{}\n\n", chunk.chars().take(600).collect::())); + } + Ok(out.trim().to_string()) +} + fn list_dir(p: &Path) -> Result { let rd = std::fs::read_dir(p).map_err(|e| format!("無法開啟資料夾 {}：{e}", p.display()))?; let mut rows: Vec<(bool, String, u64)> = rd