context_forge/storage/
turso_searcher.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use tantivy::collector::TopDocs;
6use tantivy::query::{BooleanQuery, Occur, QueryParser, TermQuery};
7use tantivy::schema::{IndexRecordOption, OwnedValue};
8use tantivy::{TantivyDocument, Term};
9
10use crate::engine::MATCH_ALL_QUERY;
11use crate::entry::ScoredEntry;
12use crate::storage::fts_index::FtsIndex;
13use crate::storage::turso_storage::turso_row_to_entry;
14use crate::traits::Searcher;
15
16/// Turso-backed searcher using a standalone tantivy index for real BM25 scoring.
17///
18/// On FTS queries: tantivy returns scored IDs → turso fetches full rows by ID.
19/// On match-all: SQL ORDER BY timestamp DESC (no scoring needed).
20pub struct TursoSearcher {
21    db: Arc<turso::Database>,
22    fts: Arc<FtsIndex>,
23}
24
25impl TursoSearcher {
26    /// Create a new searcher sharing the given turso database and tantivy index.
27    #[must_use]
28    pub(crate) fn new(db: Arc<turso::Database>, fts: Arc<FtsIndex>) -> Self {
29        Self { db, fts }
30    }
31}
32
33#[async_trait]
34impl Searcher for TursoSearcher {
35    async fn search(
36        &self,
37        query: &str,
38        scope: Option<&str>,
39        limit: usize,
40    ) -> crate::Result<Vec<ScoredEntry>> {
41        if query == MATCH_ALL_QUERY {
42            return self.search_all(scope, limit).await;
43        }
44
45        // --- tantivy BM25 search ---
46        let searcher = self.fts.reader.searcher();
47        // parse_query_lenient: OR-joins terms by default and never errors on bad syntax.
48        let (content_query, _errors) =
49            QueryParser::for_index(searcher.index(), vec![self.fts.content_field])
50                .parse_query_lenient(query);
51
52        // When a scope is requested, combine the content query with an exact-match
53        // scope term so TopDocs is already scoped — no overfetch or Rust-side filter needed.
54        let tantivy_query: Box<dyn tantivy::query::Query> = if let Some(s) = scope {
55            let scope_term = Term::from_field_text(self.fts.scope_field, s);
56            Box::new(BooleanQuery::new(vec![
57                (Occur::Must, content_query),
58                (
59                    Occur::Must,
60                    Box::new(TermQuery::new(scope_term, IndexRecordOption::Basic)),
61                ),
62            ]))
63        } else {
64            content_query
65        };
66
67        let top_docs = searcher
68            .search(&tantivy_query, &TopDocs::with_limit(limit))
69            .map_err(|e| crate::Error::Migration(e.to_string()))?;
70
71        if top_docs.is_empty() {
72            return Ok(Vec::new());
73        }
74
75        // Collect (score, id) pairs from the tantivy index.
76        let mut scored_ids: Vec<(f32, String)> = Vec::with_capacity(top_docs.len());
77        for (score, doc_address) in top_docs {
78            let doc: TantivyDocument = searcher
79                .doc(doc_address)
80                .map_err(|e| crate::Error::Migration(e.to_string()))?;
81            if let Some(OwnedValue::Str(id_str)) = doc.get_first(self.fts.id_field) {
82                scored_ids.push((score, id_str.to_owned()));
83            }
84        }
85
86        if scored_ids.is_empty() {
87            return Ok(Vec::new());
88        }
89
90        // Fetch full rows from turso by ID.
91        // IDs are UUID v7 strings generated internally — safe to inline directly.
92        let id_list: String = scored_ids
93            .iter()
94            .map(|(_, id)| format!("'{}'", id.replace('\'', "''")))
95            .collect::<Vec<_>>()
96            .join(", ");
97        let sql = format!(
98            "SELECT id, content, timestamp, kind, scope, session_id, token_count, metadata \
99             FROM entries WHERE id IN ({id_list})"
100        );
101
102        let conn = self.db.connect()?;
103        conn.busy_timeout(Duration::from_secs(5))?;
104
105        let mut rows = conn.query(&sql, ()).await?;
106        let mut id_to_entry = std::collections::HashMap::new();
107        while let Some(row) = rows.next().await? {
108            let entry = turso_row_to_entry(&row)?;
109            id_to_entry.insert(entry.id.clone(), entry);
110        }
111
112        // Zip scores back in tantivy's ranked order.
113        let mut result: Vec<ScoredEntry> = scored_ids
114            .into_iter()
115            .filter_map(|(score, id)| {
116                let entry = id_to_entry.remove(&id)?;
117                Some(ScoredEntry {
118                    score: f64::from(score),
119                    entry,
120                })
121            })
122            .collect();
123
124        // tantivy already sorted by score DESC; preserve that order.
125        result.sort_by(|a, b| {
126            b.score
127                .partial_cmp(&a.score)
128                .unwrap_or(std::cmp::Ordering::Equal)
129        });
130
131        Ok(result)
132    }
133}
134
135impl TursoSearcher {
136    async fn search_all(
137        &self,
138        scope: Option<&str>,
139        limit: usize,
140    ) -> crate::Result<Vec<ScoredEntry>> {
141        let conn = self.db.connect()?;
142        conn.busy_timeout(Duration::from_secs(5))?;
143        let scope_owned = scope.map(str::to_owned);
144
145        let mut rows = conn
146            .query(
147                "SELECT id, content, timestamp, kind, scope, session_id, token_count, metadata \
148                 FROM entries \
149                 WHERE (?1 IS NULL OR scope = ?1) \
150                 ORDER BY timestamp DESC \
151                 LIMIT ?2",
152                (scope_owned, i64::try_from(limit).unwrap_or(i64::MAX)),
153            )
154            .await?;
155
156        let mut result = Vec::new();
157        while let Some(row) = rows.next().await? {
158            let entry = turso_row_to_entry(&row)?;
159            result.push(ScoredEntry { score: 1.0, entry });
160        }
161
162        Ok(result)
163    }
164}
context_forge/storage/turso_searcher.rs

context_forge/storage/
turso_searcher.rs
