AnyBot/src/docindex.rs

//! 文件索引 (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);
    }
}