Skip to main content

canonic/
index.rs

1//! Tantivy-backed full-text index for search and near-duplicate detection.
2
3use crate::corpus::{walk_responses, CannedResponse};
4use anyhow::{bail, Context, Result};
5use serde::{Deserialize, Serialize};
6use std::collections::HashSet;
7use std::fs;
8use std::path::{Path, PathBuf};
9use tantivy::collector::TopDocs;
10use tantivy::query::QueryParser;
11use tantivy::schema::{
12    Field, IndexRecordOption, Schema, TextFieldIndexing, TextOptions, Value, STORED, STRING,
13};
14use tantivy::{doc, Index, IndexWriter, ReloadPolicy, TantivyDocument};
15
16/// One document stored in the search index.
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
18pub struct IndexDoc {
19    pub id: String,
20    pub title: String,
21    pub content: String,
22    pub path: String,
23    pub tags: Vec<String>,
24    pub sop: Option<String>,
25}
26
27/// A ranked hit from Tantivy BM25 search.
28#[derive(Debug, Clone, PartialEq)]
29pub struct SearchHit {
30    pub id: String,
31    pub score: f64,
32    pub title: String,
33    pub snippet: String,
34    pub path: PathBuf,
35}
36
37/// A pair of near-duplicate responses for curation.
38#[derive(Debug, Clone, Serialize, PartialEq)]
39pub struct DedupePair {
40    pub left_id: String,
41    pub right_id: String,
42    pub score: f64,
43    pub left_path: String,
44    pub right_path: String,
45    pub reason: String,
46}
47
48/// Default index directory (gitignored) under the working tree.
49pub fn default_index_dir() -> PathBuf {
50    PathBuf::from(".canonic-index")
51}
52
53struct Fields {
54    id: Field,
55    title: Field,
56    body: Field,
57    path: Field,
58    tags: Field,
59    sop: Field,
60}
61
62fn build_schema() -> (Schema, Fields) {
63    let mut builder = Schema::builder();
64    let id = builder.add_text_field("id", STRING | STORED);
65    let text_indexing = TextFieldIndexing::default()
66        .set_tokenizer("default")
67        .set_index_option(IndexRecordOption::WithFreqsAndPositions);
68    let text_opts = TextOptions::default()
69        .set_indexing_options(text_indexing)
70        .set_stored();
71    let title = builder.add_text_field("title", text_opts.clone());
72    let body = builder.add_text_field("body", text_opts.clone());
73    let path = builder.add_text_field("path", STRING | STORED);
74    let tags = builder.add_text_field("tags", text_opts);
75    let sop = builder.add_text_field("sop", STRING | STORED);
76    let schema = builder.build();
77    (
78        schema,
79        Fields {
80            id,
81            title,
82            body,
83            path,
84            tags,
85            sop,
86        },
87    )
88}
89
90impl From<&CannedResponse> for IndexDoc {
91    fn from(r: &CannedResponse) -> Self {
92        IndexDoc {
93            id: r.id.clone(),
94            title: r.title.clone(),
95            content: r.content.clone(),
96            path: r.path.display().to_string(),
97            tags: r.tags.clone(),
98            sop: r.sop.clone(),
99        }
100    }
101}
102
103/// Tokenize for pure lexical helpers (tests / jaccard).
104pub fn tokenize(text: &str) -> Vec<String> {
105    text.to_lowercase()
106        .split(|c: char| !c.is_alphanumeric())
107        .filter(|t| !t.is_empty())
108        .map(|s| s.to_string())
109        .collect()
110}
111
112/// Jaccard similarity over token sets (pure; used as a second signal for dedupe tests).
113pub fn jaccard_similarity(a: &str, b: &str) -> f64 {
114    let ta: HashSet<_> = tokenize(a).into_iter().collect();
115    let tb: HashSet<_> = tokenize(b).into_iter().collect();
116    if ta.is_empty() && tb.is_empty() {
117        return 1.0;
118    }
119    if ta.is_empty() || tb.is_empty() {
120        return 0.0;
121    }
122    let inter = ta.intersection(&tb).count() as f64;
123    let union = ta.union(&tb).count() as f64;
124    inter / union
125}
126
127/// Reindex all markdown responses under `corpus_dir` into Tantivy at `index_dir`.
128/// Returns number of documents written.
129pub fn reindex(corpus_dir: &Path, index_dir: &Path) -> Result<usize> {
130    let responses = walk_responses(corpus_dir)?;
131    if index_dir.exists() {
132        fs::remove_dir_all(index_dir)
133            .with_context(|| format!("clear old index {}", index_dir.display()))?;
134    }
135    fs::create_dir_all(index_dir)?;
136    let (schema, fields) = build_schema();
137    let index = Index::create_in_dir(index_dir, schema)?;
138    let mut writer: IndexWriter = index.writer(50_000_000)?;
139    let mut n = 0usize;
140    for r in &responses {
141        let tags = r.tags.join(" ");
142        let sop = r.sop.clone().unwrap_or_default();
143        writer.add_document(doc!(
144            fields.id => r.id.as_str(),
145            fields.title => r.title.as_str(),
146            fields.body => r.content.as_str(),
147            fields.path => r.path.display().to_string(),
148            fields.tags => tags.as_str(),
149            fields.sop => sop.as_str(),
150        ))?;
151        n += 1;
152    }
153    writer.commit()?;
154    Ok(n)
155}
156
157fn open_index(index_dir: &Path) -> Result<(Index, Fields)> {
158    if !index_dir.exists() {
159        bail!("index not found at {}", index_dir.display());
160    }
161    let index = Index::open_in_dir(index_dir)
162        .with_context(|| format!("open tantivy index {}", index_dir.display()))?;
163    let schema = index.schema();
164    let fields = Fields {
165        id: schema.get_field("id")?,
166        title: schema.get_field("title")?,
167        body: schema.get_field("body")?,
168        path: schema.get_field("path")?,
169        tags: schema.get_field("tags")?,
170        sop: schema.get_field("sop")?,
171    };
172    Ok((index, fields))
173}
174
175fn field_text(doc: &TantivyDocument, field: Field) -> String {
176    doc.get_first(field)
177        .and_then(|v| v.as_str())
178        .unwrap_or("")
179        .to_string()
180}
181
182/// Strip characters that break Tantivy's query parser; keep alphanumeric terms.
183pub fn sanitize_query(query: &str) -> String {
184    tokenize(query).join(" ")
185}
186
187/// Search the Tantivy index with BM25. Returns hits sorted by score desc.
188pub fn search(index_dir: &Path, query: &str, limit: usize) -> Result<Vec<SearchHit>> {
189    let query = sanitize_query(query);
190    if query.is_empty() || limit == 0 {
191        return Ok(vec![]);
192    }
193    let (index, fields) = open_index(index_dir)?;
194    let reader = index
195        .reader_builder()
196        .reload_policy(ReloadPolicy::Manual)
197        .try_into()?;
198    let searcher = reader.searcher();
199    let parser = QueryParser::for_index(&index, vec![fields.title, fields.body, fields.tags]);
200    let q = parser
201        .parse_query(&query)
202        .with_context(|| format!("parse query {query:?}"))?;
203    let top = searcher.search(&q, &TopDocs::with_limit(limit))?;
204    let mut hits = Vec::new();
205    for (score, addr) in top {
206        let retrieved: TantivyDocument = searcher.doc(addr)?;
207        let id = field_text(&retrieved, fields.id);
208        let title = field_text(&retrieved, fields.title);
209        let body = field_text(&retrieved, fields.body);
210        let path = field_text(&retrieved, fields.path);
211        hits.push(SearchHit {
212            id,
213            score: score as f64,
214            title,
215            snippet: snippet_for(&body, &query),
216            path: PathBuf::from(path),
217        });
218    }
219    Ok(hits)
220}
221
222fn snippet_for(text: &str, query: &str) -> String {
223    let flat = text.replace('\n', " ");
224    if flat.is_empty() {
225        return String::new();
226    }
227    let lower = flat.to_lowercase();
228    let mut pos = None;
229    for qt in tokenize(query) {
230        if let Some(i) = lower.find(qt.as_str()) {
231            pos = Some(i);
232            break;
233        }
234    }
235    let match_char = pos.map(|i| lower[..i].chars().count()).unwrap_or(0);
236    let start_char = match_char.saturating_sub(40);
237    let total_chars = flat.chars().count();
238    let end_char = (start_char + 120).min(total_chars);
239    let mut s: String = flat
240        .chars()
241        .skip(start_char)
242        .take(end_char - start_char)
243        .collect();
244    if start_char > 0 {
245        s = format!("...{s}");
246    }
247    if end_char < total_chars {
248        s.push_str("...");
249    }
250    s
251}
252
253/// Build a short self-query for near-duplicate search from a response.
254pub fn self_query_for(doc: &CannedResponse) -> String {
255    let mut parts = vec![doc.title.clone()];
256    let words: Vec<_> = doc
257        .content
258        .split_whitespace()
259        .filter(|w| w.len() > 3)
260        .take(40)
261        .map(|s| s.trim_matches(|c: char| !c.is_alphanumeric() && c != '-'))
262        .filter(|s| !s.is_empty())
263        .collect();
264    parts.push(words.join(" "));
265    parts.join(" ")
266}
267
268/// Find near-duplicate pairs using Tantivy self-queries plus optional Jaccard floor.
269///
270/// For each document, query the index with a content-derived query and flag other
271/// documents that rank above `score_threshold` (excluding self). Pairs are unique
272/// and ordered by score desc.
273pub fn find_duplicates(
274    index_dir: &Path,
275    corpus_dir: &Path,
276    score_threshold: f64,
277    per_doc_limit: usize,
278) -> Result<Vec<DedupePair>> {
279    let docs = walk_responses(corpus_dir)?;
280    let mut pairs: Vec<DedupePair> = Vec::new();
281    let mut seen_keys: HashSet<(String, String)> = HashSet::new();
282
283    for doc in &docs {
284        let q = self_query_for(doc);
285        if q.trim().is_empty() {
286            continue;
287        }
288        let hits = search(index_dir, &q, per_doc_limit.max(2))?;
289        for hit in hits {
290            if hit.id == doc.id {
291                continue;
292            }
293            if hit.score < score_threshold {
294                continue;
295            }
296            let (a, b) = if doc.id <= hit.id {
297                (doc.id.clone(), hit.id.clone())
298            } else {
299                (hit.id.clone(), doc.id.clone())
300            };
301            if !seen_keys.insert((a.clone(), b.clone())) {
302                continue;
303            }
304            let other = docs.iter().find(|d| d.id == hit.id);
305            let jacc = other
306                .map(|o| jaccard_similarity(&doc.content, &o.content))
307                .unwrap_or(0.0);
308            pairs.push(DedupePair {
309                left_id: a,
310                right_id: b,
311                score: hit.score,
312                left_path: doc.path.display().to_string(),
313                right_path: hit.path.display().to_string(),
314                reason: format!(
315                    "tantivy self-query hit score={:.3}; content jaccard={jacc:.3}",
316                    hit.score
317                ),
318            });
319        }
320    }
321    pairs.sort_by(|x, y| {
322        y.score
323            .partial_cmp(&x.score)
324            .unwrap_or(std::cmp::Ordering::Equal)
325    });
326    Ok(pairs)
327}
328
329/// Pure near-dup detection on an in-memory list via Jaccard (no Tantivy I/O).
330pub fn find_duplicates_jaccard(docs: &[CannedResponse], threshold: f64) -> Vec<DedupePair> {
331    let mut pairs = Vec::new();
332    for i in 0..docs.len() {
333        for j in (i + 1)..docs.len() {
334            let score = jaccard_similarity(&docs[i].content, &docs[j].content);
335            if score >= threshold {
336                pairs.push(DedupePair {
337                    left_id: docs[i].id.clone(),
338                    right_id: docs[j].id.clone(),
339                    score,
340                    left_path: docs[i].path.display().to_string(),
341                    right_path: docs[j].path.display().to_string(),
342                    reason: format!("content jaccard={score:.3}"),
343                });
344            }
345        }
346    }
347    pairs.sort_by(|a, b| {
348        b.score
349            .partial_cmp(&a.score)
350            .unwrap_or(std::cmp::Ordering::Equal)
351    });
352    pairs
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358    use std::io::Write;
359
360    #[test]
361    fn tokenize_splits_and_lowercases() {
362        assert_eq!(
363            tokenize("Hello, Project-Space!"),
364            vec!["hello", "project", "space"]
365        );
366    }
367
368    #[test]
369    fn jaccard_high_for_near_duplicates() {
370        let a = "project space is not a backup or archive on the cluster";
371        let b = "project space is not a backup archive for cluster users";
372        let c = "small compute request needs an sbu calculation for gpu nodes";
373        assert!(jaccard_similarity(a, b) > jaccard_similarity(a, c));
374        assert!(jaccard_similarity(a, b) > 0.3);
375    }
376
377    #[test]
378    fn reindex_and_search_ranks_project_space_query() {
379        let dir = tempfile::tempdir().unwrap();
380        let corpus = dir.path().join("responses");
381        fs::create_dir_all(&corpus).unwrap();
382        fs::write(
383            corpus.join("snell-project-space-not-backup.md"),
384            "---\nid: snell-project-space-not-backup\ntitle: Project space is not a backup\nprefix: snell\nsop: none\n---\n\nProject space is not a backup or archive. Use tape for long-term retention.\n",
385        )
386        .unwrap();
387        fs::write(
388            corpus.join("snell-small-compute-sbu-calculation.md"),
389            "---\nid: snell-small-compute-sbu-calculation\ntitle: SBU calculation\nprefix: snell\nsop: none\n---\n\nSmall compute needs an SBU calculation for GPU and CPU hours.\n",
390        )
391        .unwrap();
392
393        let idx = dir.path().join("index");
394        let n = reindex(&corpus, &idx).unwrap();
395        assert_eq!(n, 2);
396        let hits = search(&idx, "project space backup archive tape", 5).unwrap();
397        assert!(!hits.is_empty());
398        assert_eq!(hits[0].id, "snell-project-space-not-backup");
399        assert!(hits[0].score > 0.0);
400    }
401
402    #[test]
403    fn find_duplicates_flags_near_copy() {
404        let dir = tempfile::tempdir().unwrap();
405        let corpus = dir.path().join("responses");
406        fs::create_dir_all(&corpus).unwrap();
407        let body = "Project space on the cluster is user-managed working storage and not a backup system for research data that must be archived offline.";
408        fs::write(
409            corpus.join("snell-a.md"),
410            format!("---\nid: snell-a\ntitle: Project space note A\nprefix: snell\nsop: none\n---\n\n{body}\n"),
411        )
412        .unwrap();
413        fs::write(
414            corpus.join("snell-b.md"),
415            format!("---\nid: snell-b\ntitle: Project space note B\nprefix: snell\nsop: none\n---\n\n{body} Please confirm your backup plan.\n"),
416        )
417        .unwrap();
418        fs::write(
419            corpus.join("snell-c.md"),
420            "---\nid: snell-c\ntitle: Unrelated SBU math\nprefix: snell\nsop: none\n---\n\nGPU hours and thin node SBU rates for small compute grants.\n",
421        )
422        .unwrap();
423
424        let idx = dir.path().join("index");
425        reindex(&corpus, &idx).unwrap();
426        // Low threshold so near-copies surface; unrelated should not pair with them at high score.
427        let pairs = find_duplicates(&idx, &corpus, 0.1, 5).unwrap();
428        assert!(
429            pairs.iter().any(|p| {
430                (p.left_id == "snell-a" && p.right_id == "snell-b")
431                    || (p.left_id == "snell-b" && p.right_id == "snell-a")
432            }),
433            "expected a/b pair, got {pairs:?}"
434        );
435    }
436
437    #[test]
438    fn jaccard_dedupe_pure_path() {
439        let docs = vec![
440            CannedResponse {
441                id: "snell-a".into(),
442                title: "A".into(),
443                prefix: Some("snell".into()),
444                sop: Some("none".into()),
445                body: String::new(),
446                content: "alpha beta gamma delta shared tokens for test".into(),
447                path: PathBuf::from("a.md"),
448                tags: vec![],
449            },
450            CannedResponse {
451                id: "snell-b".into(),
452                title: "B".into(),
453                prefix: Some("snell".into()),
454                sop: Some("none".into()),
455                body: String::new(),
456                content: "alpha beta gamma delta shared tokens for test again".into(),
457                path: PathBuf::from("b.md"),
458                tags: vec![],
459            },
460            CannedResponse {
461                id: "snell-c".into(),
462                title: "C".into(),
463                prefix: Some("snell".into()),
464                sop: Some("none".into()),
465                body: String::new(),
466                content: "completely different vocabulary about licensing".into(),
467                path: PathBuf::from("c.md"),
468                tags: vec![],
469            },
470        ];
471        let pairs = find_duplicates_jaccard(&docs, 0.5);
472        assert_eq!(pairs.len(), 1);
473        assert!(pairs[0].left_id == "snell-a" || pairs[0].right_id == "snell-a");
474    }
475
476    #[test]
477    fn empty_query_returns_no_hits() {
478        let dir = tempfile::tempdir().unwrap();
479        let corpus = dir.path().join("responses");
480        fs::create_dir_all(&corpus).unwrap();
481        let mut f = fs::File::create(corpus.join("snell-x.md")).unwrap();
482        writeln!(
483            f,
484            "---\nid: snell-x\ntitle: X\nprefix: snell\nsop: none\n---\n\nhello searchable\n"
485        )
486        .unwrap();
487        let idx = dir.path().join("index");
488        reindex(&corpus, &idx).unwrap();
489        assert!(search(&idx, "   ", 5).unwrap().is_empty());
490        assert!(search(&idx, "searchable", 0).unwrap().is_empty());
491    }
492}