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::QueryParser;
7use tantivy::schema::OwnedValue;
8use tantivy::TantivyDocument;
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 (tantivy_query, _errors) =
49            QueryParser::for_index(searcher.index(), vec![self.fts.content_field])
50                .parse_query_lenient(query);
51
52        let limit_for_tantivy = if scope.is_some() {
53            // Over-fetch when scope filtering in Rust; cap at a reasonable ceiling.
54            (limit * 10).max(100)
55        } else {
56            limit
57        };
58
59        let top_docs = searcher
60            .search(&tantivy_query, &TopDocs::with_limit(limit_for_tantivy))
61            .map_err(|e| crate::Error::Migration(e.to_string()))?;
62
63        if top_docs.is_empty() {
64            return Ok(Vec::new());
65        }
66
67        // Collect (score, id) pairs from the tantivy index.
68        let mut scored_ids: Vec<(f32, String)> = Vec::with_capacity(top_docs.len());
69        for (score, doc_address) in top_docs {
70            let doc: TantivyDocument = searcher
71                .doc(doc_address)
72                .map_err(|e| crate::Error::Migration(e.to_string()))?;
73            if let Some(OwnedValue::Str(id_str)) = doc.get_first(self.fts.id_field) {
74                scored_ids.push((score, id_str.to_owned()));
75            }
76        }
77
78        if scored_ids.is_empty() {
79            return Ok(Vec::new());
80        }
81
82        // Fetch full rows from turso by ID.
83        // IDs are UUID v7 strings generated internally — safe to inline directly.
84        let id_list: String = scored_ids
85            .iter()
86            .map(|(_, id)| format!("'{}'", id.replace('\'', "''")))
87            .collect::<Vec<_>>()
88            .join(", ");
89        let sql = format!(
90            "SELECT id, content, timestamp, kind, scope, session_id, token_count, metadata \
91             FROM entries WHERE id IN ({id_list})"
92        );
93
94        let conn = self.db.connect()?;
95        conn.busy_timeout(Duration::from_secs(5))?;
96
97        let mut rows = conn.query(&sql, ()).await?;
98        let mut id_to_entry = std::collections::HashMap::new();
99        while let Some(row) = rows.next().await? {
100            let entry = turso_row_to_entry(&row)?;
101            id_to_entry.insert(entry.id.clone(), entry);
102        }
103
104        // Zip scores back, apply scope filter, respect limit.
105        let scope_owned = scope.map(str::to_owned);
106        let mut result: Vec<ScoredEntry> = scored_ids
107            .into_iter()
108            .filter_map(|(score, id)| {
109                let entry = id_to_entry.remove(&id)?;
110                if scope_owned.is_some() && entry.scope.as_deref() != scope_owned.as_deref() {
111                    return None;
112                }
113                Some(ScoredEntry {
114                    score: f64::from(score),
115                    entry,
116                })
117            })
118            .take(limit)
119            .collect();
120
121        // tantivy already sorted by score DESC; preserve that order.
122        result.sort_by(|a, b| {
123            b.score
124                .partial_cmp(&a.score)
125                .unwrap_or(std::cmp::Ordering::Equal)
126        });
127
128        Ok(result)
129    }
130}
131
132impl TursoSearcher {
133    async fn search_all(
134        &self,
135        scope: Option<&str>,
136        limit: usize,
137    ) -> crate::Result<Vec<ScoredEntry>> {
138        let conn = self.db.connect()?;
139        conn.busy_timeout(Duration::from_secs(5))?;
140        let scope_owned = scope.map(str::to_owned);
141
142        let mut rows = conn
143            .query(
144                "SELECT id, content, timestamp, kind, scope, session_id, token_count, metadata \
145                 FROM entries \
146                 WHERE (?1 IS NULL OR scope = ?1) \
147                 ORDER BY timestamp DESC \
148                 LIMIT ?2",
149                (scope_owned, i64::try_from(limit).unwrap_or(i64::MAX)),
150            )
151            .await?;
152
153        let mut result = Vec::new();
154        while let Some(row) = rows.next().await? {
155            let entry = turso_row_to_entry(&row)?;
156            result.push(ScoredEntry { score: 1.0, entry });
157        }
158
159        Ok(result)
160    }
161}
context_forge/storage/turso_searcher.rs

context_forge/storage/
turso_searcher.rs
