kibble 0.1.0

chew through any source into clean datasets — a fast ingestion, RAG & fine-tuning toolkit
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

//! Query the retrieval index: BM25 lexical ranking + semantic cosine ranking,
//! fused by Reciprocal Rank Fusion.

use crate::bm25::{tokenize, Bm25Index};
use crate::embed::Embedder;
use crate::index::Chunk;
use crate::vectors::cosine;
use std::collections::HashMap;
use std::io;
use std::path::Path;

/// One search result.
#[derive(Debug)]
pub struct Hit {
    pub score: f32,
    pub chunk: Chunk,
}

/// Reciprocal Rank Fusion. Each `rankings[r]` is chunk indices best-first.
/// `fused(c) = Σ_r 1/(rrf_k + rank_r(c))`; returns top-`k` (idx, score) desc.
pub fn rrf_fuse(rankings: &[Vec<usize>], rrf_k: f32, k: usize) -> Vec<(usize, f32)> {
    let mut fused: HashMap<usize, f32> = HashMap::new();
    for ranking in rankings {
        for (rank, &idx) in ranking.iter().enumerate() {
            *fused.entry(idx).or_insert(0.0) += 1.0 / (rrf_k + rank as f32);
        }
    }
    let mut out: Vec<(usize, f32)> = fused.into_iter().collect();
    out.sort_by(|a, b| {
        b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0))
    });
    out.truncate(k);
    out
}

/// Chunk indices ordered by cosine similarity of their vectors to the query.
pub async fn semantic_ranks<E: Embedder>(
    embedder: &E,
    store: &Path,
    model: &str,
    query: &str,
    chunk_texts: &[String],
    batch_size: usize,
) -> io::Result<Vec<usize>> {
    let qv = crate::vectors::get_or_embed(embedder, store, model, &[query.to_string()], 1).await?;
    let cvs = crate::vectors::get_or_embed(embedder, store, model, chunk_texts, batch_size).await?;
    let q = &qv[0];
    let mut sims: Vec<(usize, f32)> = cvs.iter().enumerate().map(|(i, cv)| (i, cosine(q, cv))).collect();
    sims.sort_by(|a, b| {
        b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0))
    });
    Ok(sims.into_iter().map(|(i, _)| i).collect())
}

fn load_index(dir: &Path) -> io::Result<(Vec<Chunk>, Bm25Index, serde_json::Value)> {
    let chunks_path = dir.join("chunks.jsonl");
    if !chunks_path.is_file() {
        return Err(io::Error::other("no index found — run `kibble index` first"));
    }
    let mut chunks = Vec::new();
    for line in std::fs::read_to_string(&chunks_path)?.lines() {
        if line.trim().is_empty() {
            continue;
        }
        chunks.push(serde_json::from_str::<Chunk>(line).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?);
    }
    if chunks.is_empty() {
        return Err(io::Error::other("index is empty — run `kibble index` first"));
    }
    let bm25: Bm25Index = serde_json::from_str(&std::fs::read_to_string(dir.join("bm25.json"))?)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
    let meta: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(dir.join("meta.json"))?)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
    Ok((chunks, bm25, meta))
}

/// Search the index: BM25 + (optional) semantic, fused by RRF → top-`k` hits.
pub async fn search(repo_root: &Path, query: &str, k: usize) -> io::Result<Vec<Hit>> {
    let cfg = crate::config::load_config(&repo_root.join(crate::config::CONFIG_FILE));
    let dir = repo_root.join(&cfg.index.dir);
    // load_index does synchronous fs reads + per-line JSON parse (seconds on a large
    // index); run it on the blocking pool so we don't freeze the caller's async task (#64).
    let (chunks, bm25, meta) = tokio::task::spawn_blocking(move || load_index(&dir))
        .await
        .map_err(|e| io::Error::other(format!("index load task panicked: {e}")))??;

    // Lexical ranking.
    let lex = bm25.scores(&tokenize(query));
    let lex_rank: Vec<usize> = lex.into_iter().map(|(i, _)| i).collect();
    let mut rankings: Vec<Vec<usize>> = vec![lex_rank];

    // Semantic ranking (fail-soft).
    let has_vectors = meta.get("has_vectors").and_then(|v| v.as_bool()).unwrap_or(false);
    let meta_model = meta.get("embed_model").and_then(|v| v.as_str()).unwrap_or("");
    let embed = &cfg.understand.embed;
    if has_vectors && !embed.base_url.is_empty() {
        if meta_model != embed.model {
            eprintln!("kibble: index model '{meta_model}' != configured '{}' — lexical only", embed.model);
        } else {
            match crate::embed::EndpointEmbedder::new(embed, cfg.network.proxy.as_deref()) {
                Ok(e) => {
                    let store = repo_root.join(&embed.store);
                    let texts: Vec<String> = chunks.iter().map(|c| c.text.clone()).collect();
                    match semantic_ranks(&e, &store, &embed.model, query, &texts, embed.batch_size).await {
                        Ok(sem) => rankings.push(sem),
                        Err(err) => eprintln!("kibble: semantic search skipped (embed failed): {err}"),
                    }
                }
                Err(err) => eprintln!("kibble: semantic search skipped (embed init failed): {err}"),
            }
        }
    }

    let fused = rrf_fuse(&rankings, cfg.index.rrf_k, k);
    Ok(fused.into_iter().map(|(idx, score)| Hit { score, chunk: chunks[idx].clone() }).collect())
}

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

    #[test]
    fn rrf_prefers_chunk_high_in_both() {
        // chunk 2 is top of list A and second of list B; chunk 0 is top of B only.
        let a = vec![2, 1, 0];
        let b = vec![0, 2, 1];
        let fused = rrf_fuse(&[a, b], 60.0, 3);
        assert_eq!(fused[0].0, 2); // present near top of BOTH → wins
    }

    #[tokio::test]
    async fn semantic_ranks_orders_by_similarity() {
        let dir = std::env::temp_dir().join(format!("kibble_ret_sem_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        let e = StubEmbedder::new();
        // query text equals chunk 1 exactly → identical vector → cosine 1.0 → rank 0.
        let chunks = vec!["alpha beta".to_string(), "the query text".to_string(), "gamma delta".to_string()];
        let ranks = semantic_ranks(&e, &dir, "m", "the query text", &chunks, 8).await.unwrap();
        assert_eq!(ranks[0], 1);
        std::fs::remove_dir_all(&dir).ok();
    }

    #[tokio::test]
    async fn search_lexical_only_finds_planted_chunk() {
        let dir = std::env::temp_dir().join(format!("kibble_ret_search_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(dir.join("data/notes")).unwrap();
        std::fs::write(dir.join("data/notes/fox.md"), "the quick brown fox jumps over").unwrap();
        std::fs::write(dir.join("data/notes/dog.md"), "a sleepy dog lies in the sun").unwrap();
        std::fs::write(dir.join(crate::config::CONFIG_FILE),
            "[index]\nsources=[\"data/notes\"]\nchunk_chars=1000\n[understand.embed]\nbase_url=\"\"\n").unwrap();
        crate::index::build_index(&dir, None).await.unwrap();

        let hits = search(&dir, "fox", 5).await.unwrap();
        assert!(!hits.is_empty());
        assert!(hits[0].chunk.source.contains("fox.md")); // lexical match wins
        std::fs::remove_dir_all(&dir).ok();
    }

    #[tokio::test]
    async fn search_errors_without_index() {
        let dir = std::env::temp_dir().join(format!("kibble_ret_noidx_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join(crate::config::CONFIG_FILE), "[index]\n").unwrap();
        assert!(search(&dir, "anything", 5).await.is_err());
        std::fs::remove_dir_all(&dir).ok();
    }
}