Merge p6/embed: local embeddings + document shelf (P6 step 2)
This commit is contained in:
daniel wang
commit
9d14a369e2
|
|
@ -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"
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
76
src/app.rs
76
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<Vec<String>> = std::sync::RwLock::new(Vec::new());
|
||||
|
||||
pub fn doc_folders() -> Vec<String> {
|
||||
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<String>) {
|
||||
*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<crate::oplog::Interrupted>,
|
||||
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() {
|
||||
|
|
|
|||
|
|
@ -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<String> {
|
||||
let mut out: Vec<String> = 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<char> = 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::<usize>() >= 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Option<Option<Cfg>>> = Mutex::new(None);
|
||||
/// Last probe/embed failure: retry only after a while, the turn must not pay for it twice.
|
||||
static DOWN: Mutex<Option<Instant>> = 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<String> {
|
||||
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<Vec<Vec<f32>>> {
|
||||
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<serde_json::Value, String> = 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::<Vec<_>>()) {
|
||||
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::<Vec<_>>()) {
|
||||
for ((id, _), v) in rows.iter().zip(vs) {
|
||||
crate::store::doc_set_vec(*id, &v);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
pub fn f32s_to_blob(v: &[f32]) -> Vec<u8> {
|
||||
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<f32> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -328,6 +328,12 @@ pub struct State {
|
|||
pub providers: std::collections::BTreeMap<String, ProviderSettings>,
|
||||
#[serde(default)]
|
||||
pub sidebar_collapsed: bool,
|
||||
/// 語意搜尋: "" = 找 Ollama 裡裝好的 embedding 模型, "off" = 關。
|
||||
#[serde(default)]
|
||||
pub embed_model: String,
|
||||
/// 文件索引的資料夾; 空 = 用這台電腦的「文件」資料夾。
|
||||
#[serde(default)]
|
||||
pub doc_folders: Vec<String>,
|
||||
#[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![],
|
||||
|
|
|
|||
228
src/store.rs
228
src/store.rs
|
|
@ -70,9 +70,31 @@ fn open() -> Option<Store> {
|
|||
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<String> {
|
|||
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<String> {
|
|||
}
|
||||
}
|
||||
}
|
||||
// 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<String> = 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<String> {
|
|||
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<String>], limit: usize) -> Vec<String> {
|
||||
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<f32>)> {
|
||||
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<u8>>(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::<String>().split_whitespace().collect::<Vec<_>>().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<String> = vec![];
|
||||
for w in q.split_whitespace() {
|
||||
let ch: Vec<char> = 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<u8>>(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<String> = words.iter().map(|(p, t)| format!("{p}\u{1}{t}")).collect();
|
||||
let vr: Vec<String> = 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
|
||||
}
|
||||
|
|
|
|||
24
src/tools.rs
24
src/tools.rs
|
|
@ -39,6 +39,10 @@ pub fn specs() -> Vec<Value> {
|
|||
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<String, String> {
|
||||
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::<String>()));
|
||||
}
|
||||
Ok(out.trim().to_string())
|
||||
}
|
||||
|
||||
fn list_dir(p: &Path) -> Result<String, String> {
|
||||
let rd = std::fs::read_dir(p).map_err(|e| format!("無法開啟資料夾 {}：{e}", p.display()))?;
|
||||
let mut rows: Vec<(bool, String, u64)> = rd
|
||||
|
|
|
|||
Loading…
Reference in New Issue