stored_maxsim/
stored_maxsim.rs

1//! A text gatherer implemented by the caller, reranked by MaxSim over a
2//! lateweave vector store that a writer keeps moving.
3//!
4//! ```bash
5//! cargo run --example stored_maxsim
6//! ```
7
8use std::collections::BTreeMap;
9use std::sync::Arc;
10use std::time::SystemTime;
11
12use lateweave::{
13    Candidate, CandidateGenerator, DocumentKey, Encoding, Feature, Gathered, MaxSimReranker,
14    MultiVectorSource, Query, Representation, Requirements, ResourceBudget, Result, SearchPipeline,
15    SearchRequest, Subset, TokenMatrix, VectorStore, VectorStoreWriter, DEFAULT_FEATURE,
16};
17
18const DIMENSION: usize = 4;
19
20/// Gathers every document containing a query term, scored by match count.
21struct TermGatherer {
22    requires: Requirements,
23    documents: Vec<(&'static str, &'static str)>,
24    indexed_at: SystemTime,
25}
26
27impl CandidateGenerator for TermGatherer {
28    fn requires(&self) -> &Requirements {
29        &self.requires
30    }
31
32    fn score_semantics(&self) -> &str {
33        "term-match-count"
34    }
35
36    fn gather(&self, query: &Query, limit: usize, subset: Option<&Subset>) -> Result<Gathered> {
37        let terms = query.text().split_whitespace().collect::<Vec<_>>();
38        let mut matches = self
39            .documents
40            .iter()
41            .map(|&(id, text)| (DocumentKey::new("fruit", id), text))
42            .filter(|(key, _)| subset.is_none_or(|subset| subset.contains(key)))
43            .map(|(key, text)| {
44                let count = text
45                    .split_whitespace()
46                    .filter(|word| terms.contains(word))
47                    .count();
48                (key, count as f32)
49            })
50            .filter(|&(_, count)| count > 0.0)
51            .collect::<Vec<_>>();
52        matches.sort_by(|left, right| right.1.total_cmp(&left.1).then(left.0.cmp(&right.0)));
53        Ok(Gathered {
54            candidates: matches
55                .into_iter()
56                .take(limit)
57                .enumerate()
58                .map(|(rank, (key, gather_score))| Candidate {
59                    key,
60                    gather_score,
61                    gather_rank: rank,
62                    provenance: "term-gatherer".to_string(),
63                })
64                .collect(),
65            as_of: self.indexed_at,
66        })
67    }
68}
69
70fn unit(axis: usize) -> [f32; DIMENSION] {
71    let mut row = [0.0; DIMENSION];
72    row[axis] = 1.0;
73    row
74}
75
76fn main() -> Result<()> {
77    let representation = Representation::new("toy-encoder", "1", DIMENSION, true)?;
78    let directory = std::env::temp_dir().join(format!("lateweave-example-{}", std::process::id()));
79    let path = directory.join("vectors");
80
81    // The indexing side: one token per word, the axis standing in for what
82    // a real encoder produces.
83    let mut writer =
84        VectorStoreWriter::create(&path, Encoding::Float32, "fruit", representation.clone())?;
85    let embeddings = [unit(0), unit(2), unit(1), unit(2), unit(0), unit(3)].concat();
86    writer.append(["a", "b", "c"], &embeddings, DIMENSION, &[2, 2, 2], None)?;
87    writer.commit()?;
88
89    // The serving side.
90    let gatherer = TermGatherer {
91        requires: BTreeMap::new(),
92        documents: vec![("a", "red apple"), ("b", "green apple"), ("c", "red car")],
93        indexed_at: SystemTime::now(),
94    };
95    let store = Arc::new(VectorStore::open(&path)?) as Arc<dyn MultiVectorSource>;
96    let reranker = MaxSimReranker::new([store], DEFAULT_FEATURE)?;
97    let pipeline = SearchPipeline::new(Arc::new(gatherer), Some(Arc::new(reranker)));
98
99    let query = Query::new("apple").with_feature(
100        DEFAULT_FEATURE,
101        Feature::lazy(representation, || {
102            TokenMatrix::new(unit(1).to_vec(), DIMENSION)
103        }),
104    );
105    let request = SearchRequest::new(10, 2).with_budget(ResourceBudget::new(4096, 64, Some(1))?);
106
107    let print = |label: &str| -> Result<()> {
108        let result = pipeline.search(&query, &request)?;
109        println!("{label}:");
110        for document in &result.documents {
111            println!(
112                "  #{} {:?} score {}",
113                document.rank, document.key, document.score
114            );
115        }
116        println!("  {:?}", result.diagnostics);
117        Ok(())
118    };
119    print("before the delete")?;
120
121    // The indexing side deletes "b" from the vectors; the gatherer still
122    // finds it, and the pipeline drops it.
123    writer.delete(["b"]);
124    writer.commit()?;
125    print("after the delete")?;
126
127    std::fs::remove_dir_all(directory)?;
128    Ok(())
129}
stored_maxsim/stored_maxsim.rs

stored_maxsim/
stored_maxsim.rs
