crates\nfarms-core\src\kb.rs

源码路径:crates\nfarms-core\src\kb.rs · 554 行

//! 本地知识库:把 `wiki\` 下的文档当作对话的参考资料。
//!
//! # 为什么是关键词检索,不是向量
//!
//! 项目里没有 embedding 模型,也不该为了这个功能去下一个(那是几百 MB 起步,
//! 而它只在桌宠对话里用)。所以用**关键词打分**:
//!
//! * 中文按 **2-gram** 切(不分词器也能用:把「知识库」切成 `知识`/`识库`);
//! * 英文/数字按非字母数字切,转小写;
//! * 每个词在每个文档里的**出现次数**加起来当分数。
//!
//! 这比向量差,但有几个好处:**零依赖、可解释、可测** —— 命中哪一段能直接
//! 看出来,而不是一个 0.83 的相似度。
//!
//! # 目录
//!
//! 读的是**程序根下的 `wiki\`**(`paths.wiki()`)。取不到的常见原因是
//! 源码里根本没有这个目录 —— 它由构建脚本创建,内容由用户自己放。
//! 所以 [`load`] 拿不到东西时返回空 `Vec`,**不是错误**:
//! 没有知识库,对话照样该能用。
//!
//! # 这个目录实测有多大
//!
//! `dist\wiki\` 下 200 多个文件、合计好几 MB,而且**有大量重复**
//! (同一份 `开发过程记录.md` / `CONVERSATION-LOG` 出现在多个深度)。
//! 所以有两条约束,都是量出来的而不是猜的:
//!
//! * [`MAX_FILE_BYTES`] —— 单篇超过 512 KB 整篇跳过(实测最大一篇 716 KB);
//! * [`context_from`] —— 文档由调用方**读一次留着**,不每条消息重读。

use std::path::Path;

/// 一篇文档。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Doc {
    /// 相对 `wiki\` 的路径,用来在提示里注明出处
    pub name: String,
    pub text: String,
}

/// 一条命中。
///
/// **没有 `Eq`** —— 里面有个 `f32`(`norm`),而浮点不满足全等。
/// `PartialEq` 够了:这个类型只用来排序和展示,不做键。
#[derive(Debug, Clone, PartialEq)]
pub struct Hit {
    pub name: String,
    /// 命中处附近的一小段原文
    pub snippet: String,
    /// 原始命中次数(各词次数之和,单词上限 [`MAX_TF`])。**人能看懂的量的**。
    pub score: u32,
    /// 按文档长度归一后的分数。**只用于排序,不对外展示** ——
    /// 它是 `score / sqrt(字符数/1000)`,报出去没人知道怎么读。
    norm: f32,
}

/// 单篇文档最多取多少字符进上下文。
///
/// 不是"越小越好":截太短会把关键句切掉。但也不能不截 ——
/// wiki 下有一份 23 KB 的记录,整篇塞进 32 K 上下文的模型里,
/// 一次就把预算吃掉四分之一。
const MAX_DOC_CHARS: usize = 2400;

/// 单个文件超过这个大就**整篇跳过**。
///
/// 不是洁癖:实测 `wiki\` 下有 716 KB 的单篇教程和 440 KB 的清单 ——
/// 它们进去只会把检索结果淹掉(几万个词命中同一篇),而且读一次要几十毫秒。
/// 512 KB 已经能装下任何"正常文档",超出的多半是导出/拼接产物。
const MAX_FILE_BYTES: u64 = 512 * 1024;

/// 递归读 `dir` 下的 `.md`/`.txt`。目录不存在或读不到都返回空 —— 见模块说明。
pub fn load(dir: &Path) -> Vec<Doc> {
    let mut out = Vec::new();
    walk(dir, dir, &mut out, 0);
    // 名字排序:同一份目录每次读出来的顺序不该变,
    // 否则同一句话两次检索出来的上下文不一样,排查时无从对比。
    out.sort_by(|a, b| a.name.cmp(&b.name));
    out
}

fn walk(root: &Path, dir: &Path, out: &mut Vec<Doc>, depth: u32) {
    // 深度上限:防止有人在 wiki 下放了个指向上级的软链接,把整块盘读进来
    if depth > 4 {
        return;
    }
    let Ok(rd) = std::fs::read_dir(dir) else {
        return;
    };
    for e in rd.flatten() {
        let p = e.path();
        if p.is_dir() {
            walk(root, &p, out, depth + 1);
            continue;
        }
        let ext = p
            .extension()
            .map(|x| x.to_string_lossy().to_ascii_lowercase())
            .unwrap_or_default();
        if ext != "md" && ext != "txt" {
            continue;
        }
        // 先看大小再读 —— 716 KB 那篇读进来再丢掉是白花的 I/O
        if e.metadata().map(|m| m.len()).unwrap_or(0) > MAX_FILE_BYTES {
            continue;
        }
        let Ok(text) = std::fs::read_to_string(&p) else {
            continue;
        };
        let name = p
            .strip_prefix(root)
            .unwrap_or(&p)
            .to_string_lossy()
            .replace('\\', "/");
        out.push(Doc { name, text });
    }
}

/// 把查询切成检索词。
///
/// 中文走 2-gram:**不分词器也能用**。代价是会有一些无意义的组合
/// (「的我」),但它们在任何文档里出现次数都差不多,排序时自然被摊平。
pub fn terms(q: &str) -> Vec<String> {
    let mut out = Vec::new();
    let mut cjk: Vec<char> = Vec::new();
    let mut word = String::new();

    fn flush_word(word: &mut String, out: &mut Vec<String>) {
        if word.chars().count() >= 2 {
            out.push(word.to_lowercase());
        }
        word.clear();
    }
    fn flush_cjk(cjk: &mut Vec<char>, out: &mut Vec<String>) {
        if cjk.len() >= 2 {
            for w in cjk.windows(2) {
                out.push(w.iter().collect());
            }
        }
        else if cjk.len() == 1 {
            out.push(cjk[0].to_string());
        }
        cjk.clear();
    }

    for ch in q.chars() {
        if is_cjk(ch) {
            flush_word(&mut word, &mut out);
            cjk.push(ch);
        }
        else if ch.is_alphanumeric() {
            flush_cjk(&mut cjk, &mut out);
            word.push(ch);
        }
        else {
            flush_word(&mut word, &mut out);
            flush_cjk(&mut cjk, &mut out);
        }
    }
    flush_word(&mut word, &mut out);
    flush_cjk(&mut cjk, &mut out);
    out
}

fn is_cjk(c: char) -> bool {
    matches!(c as u32, 0x4E00..=0x9FFF | 0x3400..=0x4DBF)
}

/// 检索。返回按分数降序的前 `k` 条。
pub fn search(docs: &[Doc], query: &str, k: usize) -> Vec<Hit> {
    let ts = terms(query);
    if ts.is_empty() || docs.is_empty() || k == 0 {
        return Vec::new();
    }
    let mut hits: Vec<Hit> = Vec::new();
    for d in docs {
        let lower = d.text.to_lowercase();
        let mut score = 0u32;
        let mut best_at = 0usize;
        let mut best_n = 0u32;
        // **文件名也参与检索** —— 理由见 [`NAME_BONUS`]
        let name_lower = d.name.to_lowercase();
        for t in &ts {
            if name_lower.contains(t.as_str()) {
                score = score.saturating_add(NAME_BONUS);
            }
            // 每个词只记一次出现位置,但**次数要全算** —— 出现 20 次的
            // 那个词比出现 1 次的更能说明这篇文档在讲它。
            let mut n = 0u32;
            let mut from = 0usize;
            while let Some(rel) = lower.get(from..).and_then(|s| s.find(t.as_str())) {
                let at = from + rel;
                n += 1;
                if n > best_n {
                    best_n = n;
                    best_at = at;
                }
                from = at + t.len().max(1);
                // 到上限就不用再数了,**顺便省掉长文档里的全文扫描**
                if n >= MAX_TF {
                    break;
                }
            }
            score = score.saturating_add(n.min(MAX_TF));
        }
        if score == 0 {
            continue;
        }
        hits.push(Hit {
            name: d.name.clone(),
            snippet: snippet(&d.text, best_at),
            score,
            // 长度归一后的分数**只用来排序**,不往外报 ——
            // 报出去的是原始命中次数,那是人能看懂的量
            norm: score as f32 / ((d.text.chars().count() as f32 / LEN_UNIT).max(1.0)).sqrt(),
        });
    }
    // 分数相同的按名字排:结果**可复现**,不然同一句话两次检索顺序不同
    hits.sort_by(|a, b| {
        b.norm
            .partial_cmp(&a.norm)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| b.score.cmp(&a.score))
            .then_with(|| a.name.cmp(&b.name))
    });
    hits.truncate(k);
    hits
}

/// 取命中位置附近的一段原文,并按字符边界截断。
///
/// 用 `chars()` 而不是字节切片 —— 中文文档里按字节切几乎必然切在
/// 多字节字符中间,那会 panic。
fn snippet(text: &str, at: usize) -> String {
    let at = at.min(text.len());
    let at = (0..=at)
        .rev()
        .find(|k| text.is_char_boundary(*k))
        .unwrap_or(0);
    let head: String = text[..at]
        .chars()
        .rev()
        .take(40)
        .collect::<String>()
        .chars()
        .rev()
        .collect();
    let tail: String = text[at..].chars().take(160).collect();
    format!("{head}{tail}").replace(['\n', '\r'], " ")
}

/// 把命中拼成给模型的上下文块。
///
/// 每段都写明**来自哪个文件** —— 模型据此可以说「按 wiki 里的记录」,
/// 而不是把内容当自己的记忆说出来。出处对排查很重要。
pub fn context_block(hits: &[Hit]) -> String {
    if hits.is_empty() {
        return String::new();
    }
    let mut s = String::from(
        "以下是本机 wiki 目录里与问题相关的片段,请优先依据它们回答;\
如果里面没有答案,就直说不知道,不要编。\n",
    );
    for h in hits {
        s.push_str(&format!("\n--- {} ---\n{}\n", h.name, h.snippet));
    }
    s
}

/// 单个检索词在一篇文档里最多计几次。
///
/// **这是 TF 饱和,不是省事。** 不设上限的话,一篇把某个词重复 500 次的
/// 长文档会压倒性地赢 —— 而"重复 500 次"和"重复 20 次"在相关性上
/// 几乎没有区别。实测:既有账本 244 KB,新写的 15 KB,不设上限时
/// 新文档完全浮不上来。
const MAX_TF: u32 = 5;

/// 检索词出现在**文件名**里时的加成(每个词)。
///
/// **文件名原来根本没参与检索** —— 于是"按名字问一份文档"反而找不到它。
/// 实测:本人刚写进 wiki 的 `变更账本-2026-10-14.md`,问「变更账本」
/// 命中不了自己。这很反直觉:用户说一个词,往往**就是那份文档的名字**。
///
/// 给固定加成而不是按次数:名字只有一份,出现一次就是"全中",
/// 没有"出现多次"这回事。
const NAME_BONUS: u32 = 20;

/// 长度归一的基准(字符)。分数会除以 `sqrt(字符数 / 这个数)`。
///
/// 理由和 [`MAX_TF`] 是同一个:**长文档不该因为字多就赢**。
/// 200 KB 的文档是 15 KB 的 13 倍,开方后仍差 3.6 倍 —— 只靠 TF 饱和
/// 还不够,得再按长度摊一次。
const LEN_UNIT: f32 = 1000.0;

/// 一次检索最多带几篇进上下文。
pub const TOP_K: usize = 3;

/// 给对话用的一步到位:读目录 → 检索 → 拼块。拿不到就返回空串。
pub fn context_for(dir: &Path, query: &str) -> String {
    context_from(&load(dir), query)
}

/// 和 [`context_for`] 一样,但**文档由调用方持有**。
///
/// 分开的理由是性能:`wiki\` 下有 200 多个文件、合计好几 MB。每条消息
/// 都重读一遍是几百毫秒的无谓 I/O,而知识库**几乎不变** ——
/// 所以调用方读一次、留着,每条消息只做检索。
pub fn context_from(docs: &[Doc], query: &str) -> String {
    if docs.is_empty() {
        return String::new();
    }
    let hits = search(docs, query, TOP_K);
    let mut block = context_block(&hits);
    // 总长度也要兜一下:3 篇 × 2400 字符对 32 K 上下文还好,
    // 但 MAX_DOC_CHARS 要是被人调大,这里能挡住。
    let cap = MAX_DOC_CHARS * TOP_K;
    if block.chars().count() > cap {
        block = block.chars().take(cap).collect();
    }
    block
}

#[cfg(test)]
mod tests {
    use super::*;

    fn doc(name: &str, text: &str) -> Doc {
        Doc {
            name: name.to_string(),
            text: text.to_string(),
        }
    }

    #[test]
    fn cjk_is_split_into_bigrams_and_ascii_into_words() {
        let t = terms("本地知识库 llama.cpp");
        assert!(t.contains(&"本地".to_string()), "实际: {t:?}");
        assert!(t.contains(&"知识".to_string()), "实际: {t:?}");
        assert!(t.contains(&"识库".to_string()), "实际: {t:?}");
        assert!(t.contains(&"llama".to_string()), "实际: {t:?}");
        assert!(t.contains(&"cpp".to_string()), "实际: {t:?}");
        // 单字母不该成为检索词(噪音太大)
        assert!(!t.iter().any(|x| x == "c"), "实际: {t:?}");
    }

    #[test]
    fn search_finds_the_document_that_talks_about_it() {
        let docs = vec![
            doc("a.md", "这篇讲的是天气和温度"),
            doc("b.md", "这篇讲的是本地知识库怎么接进对话"),
        ];
        let hits = search(&docs, "知识库怎么用", 3);
        assert_eq!(hits.len(), 1, "只有 b 该命中: {hits:?}");
        assert_eq!(hits[0].name, "b.md");
    }

    /// 出现次数多的文档应当排在前面 —— 这是打分的核心假设,
    /// 退化成"命中就同分"的话排序就没意义了。
    #[test]
    fn more_occurrences_rank_higher() {
        let docs = vec![
            doc("once.md", "油价 其他别的内容"),
            doc("many.md", "油价 油价 油价 还有别的"),
        ];
        let hits = search(&docs, "油价", 3);
        assert_eq!(hits.len(), 2);
        assert_eq!(hits[0].name, "many.md", "实际: {hits:?}");
        assert!(hits[0].score > hits[1].score);
    }

    #[test]
    fn no_query_or_no_docs_yields_nothing() {
        let docs = vec![doc("a.md", "内容")];
        assert!(search(&docs, "", 3).is_empty());
        assert!(search(&docs, "   ", 3).is_empty());
        assert!(search(&[], "内容", 3).is_empty());
        // k = 0 也要能挡住,别 panic
        assert!(search(&docs, "内容", 0).is_empty());
    }

    /// 中文文档里取片段**不能 panic** —— 按字节切会切在字符中间。
    #[test]
    fn snippet_does_not_panic_on_multibyte_text() {
        let docs = vec![doc(
            "中文.md",
            "这是一段很长的中文内容用来测试片段提取不会崩掉真的不会崩",
        )];
        let hits = search(&docs, "中文内容", 3);
        assert!(!hits.is_empty());
        assert!(hits[0].snippet.contains("中文"));
    }

    #[test]
    fn context_block_names_each_source() {
        let hits = vec![Hit {
            name: "设计.md".to_string(),
            snippet: "某段内容".to_string(),
            score: 1,
            norm: 1.0,
        }];
        let b = context_block(&hits);
        assert!(b.contains("设计.md"), "出处必须写进去: {b}");
        assert!(b.contains("不要编"), "要明确禁止编造: {b}");
        assert!(context_block(&[]).is_empty());
    }

    /// 目录不存在 = **没有知识库**,不是错误。对话照样要能用。
    #[test]
    fn a_missing_directory_is_not_an_error() {
        let d = std::env::temp_dir().join("nfarms-kb-肯定不存在");
        assert!(load(&d).is_empty());
        assert!(context_for(&d, "随便问").is_empty());
        assert!(context_from(&[], "随便问").is_empty());
    }

    /// 真的读一次磁盘:临时目录放两份 .md,应当都能读到。
    #[test]
    fn loads_markdown_from_a_real_directory() {
        let d = std::env::temp_dir().join(format!("nfarms-kb-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        std::fs::create_dir_all(d.join("sub")).unwrap();
        std::fs::write(d.join("一.md"), "关于端口统一管理的记录").unwrap();
        std::fs::write(d.join("sub/二.md"), "关于桌宠透明背景的记录").unwrap();
        // 非 md/txt 不该被读
        std::fs::write(d.join("跳过.png"), "不是文本").unwrap();

        let docs = load(&d);
        assert_eq!(
            docs.len(),
            2,
            "实际: {:?}",
            docs.iter().map(|x| &x.name).collect::<Vec<_>>()
        );
        assert!(docs.iter().any(|x| x.name == "sub/二.md"), "子目录也要读");

        let ctx = context_for(&d, "透明背景怎么做的");
        assert!(ctx.contains("二.md"), "应当命中讲透明背景那篇: {ctx}");
        let _ = std::fs::remove_dir_all(&d);
    }

/// **对真实 `dist\wiki` 跑一次** —— 手动执行。
    ///
    /// 单元测试用的是临时目录里造的几份小文档,验的是**逻辑**;
    /// 这一条验的是**真实语料**:文件数、跳过的大文件、以及
    /// 「问一个这一轮刚写进 wiki 的问题,能不能命中对应文档」。
    ///
    ///     cargo test -p nfarms-core kb::tests::probe_real_wiki -- --ignored --nocapture
    ///
    /// 为什么值得单独留一条:`wiki\` 有几个 GB 的潜力(它是用户自己放的),
    /// 而"读到了但检索不出来"和"根本没读到"在界面上长得一模一样 ——
    /// **都只是没有知识库片段**。
    /// **按文件名问,要能找到那份文档。**
    ///
    /// 原来文件名不参与检索,于是 `变更账本-2026-10-14.md` 被问「变更账本」
    /// 时命中不了自己 —— 而用户说一个词,往往**就是那份文档的名字**。
    #[test]
    fn a_documents_own_name_finds_it() {
        let docs = vec![
            doc("变更账本-2026-10-14.md", "这一轮改了打印代理页和知识库"),
            doc("别的.md", "变更 变更 变更 变更 变更 变更 变更 变更"),
        ];
        let hits = search(&docs, "变更账本", 3);
        assert_eq!(hits[0].name, "变更账本-2026-10-14.md", "实际: {hits:?}");
    }

    /// **长文档不该因为字多就赢。** 同样的命中次数,短的那篇该排前面。
    ///
    /// 这条是拿真实语料量出来的:`dist\wiki` 里既有 244 KB 的账本,
    /// 也有 15 KB 的 —— 不做长度归一,新的那份**永远浮不上来**。
    #[test]
    fn a_short_document_beats_a_long_one_with_the_same_hits() {
        let mut long = String::from("桌宠透明背景");
        long.push_str(&"无关内容".repeat(20000));
        let docs = vec![
            doc("短.md", "桌宠透明背景就是这么做出来的"),
            doc("长.md", &long),
        ];
        let hits = search(&docs, "桌宠透明背景", 3);
        assert_eq!(hits[0].name, "短.md", "实际: {hits:?}");
    }

    #[test]
    #[ignore = "读真实 dist\\wiki,手动跑"]
    fn probe_real_wiki() {
        let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("..")
            .join("..")
            .join("dist")
            .join("wiki");
        if !dir.exists() {
            println!("[跳过] {} 不存在", dir.display());
            return;
        }
        let t0 = std::time::Instant::now();
        let docs = load(&dir);
        let took = t0.elapsed();
        let chars: usize = docs.iter().map(|d| d.text.chars().count()).sum();
        println!(
            "读到 {} 篇,合计 {} 字符({} KB),用时 {:?}",
            docs.len(),
            chars,
            chars / 1024,
            took
        );
        // 512 KB 上限挡掉了哪些,值得看一眼 —— 它们是真的读不到
        let skipped: Vec<String> = std::fs::read_dir(&dir)
            .into_iter()
            .flatten()
            .flatten()
            .filter(|e| e.path().is_dir())
            .flat_map(|e| std::fs::read_dir(e.path()).into_iter().flatten().flatten())
            .filter(|e| e.metadata().map(|m| m.len()).unwrap_or(0) > MAX_FILE_BYTES)
            .map(|e| e.file_name().to_string_lossy().to_string())
            .collect();
        println!("跳过(超过 {} KB)的文件: {:?}", MAX_FILE_BYTES / 1024, skipped);

        // **这一轮刚写进去的文档,问得出来吗**
        //
        // 最后一问是**知识库新鲜度探针**:`分片累积解码` 这个词只出现在移动端
        // 那一轮的代码注释与开发记录里。只要它命中,就说明 `dist\wiki\code`
        // 已经被重新生成过;如果它没命中而那几份老文档都命中,则说明知识库是旧的 ——
        // 这种"能读、但读的是旧的那一版"在界面上和"一切正常"完全一样。
        for q in [
            "打印代理页为什么重复",
            "桌宠透明背景怎么做的",
            "变更账本",
            "分片累积解码",
        ] {
            let block = context_from(&docs, q);
            let first = block
                .lines()
                .find(|l| l.starts_with("--- "))
                .unwrap_or("(没有命中)");
            println!("问「{q}」→ 首条命中: {first}");
        }
    }
    /// **单篇过大的要整篇跳过。** 实测 wiki 下有 716 KB 的单篇教程 ——
    /// 它进去只会把检索结果淹掉,而读一次要几十毫秒。
    #[test]
    fn an_oversized_file_is_skipped_whole() {
        let d = std::env::temp_dir().join(format!("nfarms-kb-big-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        std::fs::create_dir_all(&d).unwrap();
        std::fs::write(d.join("小.md"), "正常内容").unwrap();
        // 造一个刚过线的:512 KB + 1 字节
        let big = "x".repeat((MAX_FILE_BYTES + 1) as usize);
        std::fs::write(d.join("巨大.md"), big).unwrap();

        let docs = load(&d);
        assert_eq!(docs.len(), 1, "超限的那篇不该进来");
        assert_eq!(docs[0].name, "小.md");
        let _ = std::fs::remove_dir_all(&d);
    }
}