context_harness/
sqlite_store.rs

1//! SQLite-backed [`Store`] implementation.
2//!
3//! Maps each [`Store`] operation to the existing SQLite queries used by
4//! the ingestion, search, and retrieval modules.
5
6use anyhow::Result;
7use async_trait::async_trait;
8use sqlx::{Row, SqlitePool};
9
10use context_harness_core::embedding::{blob_to_vec, cosine_similarity, vec_to_blob};
11use context_harness_core::models::{Chunk, Document};
12use context_harness_core::store::{
13    ChunkCandidate, ChunkResponse, DocumentMetadata, DocumentResponse, Store,
14};
15
16/// SQLite implementation of the [`Store`] trait.
17///
18/// Wraps a [`SqlitePool`] and translates every `Store` method into one
19/// or more SQL statements against the existing schema (documents, chunks,
20/// chunks_fts, chunk_vectors, embeddings).
21pub struct SqliteStore {
22    pool: SqlitePool,
23}
24
25impl SqliteStore {
26    pub fn new(pool: SqlitePool) -> Self {
27        Self { pool }
28    }
29
30    #[allow(dead_code)]
31    pub fn pool(&self) -> &SqlitePool {
32        &self.pool
33    }
34}
35
36fn fts_query_from_user_text(query: &str) -> String {
37    // Join terms with OR, not a bare space. In FTS5 a space between terms is an
38    // implicit AND, so a natural-language query ("the post about saving tokens…")
39    // requires every word to appear in one document and usually matches nothing.
40    // OR lets BM25 rank by how many (and how rare) the matched terms are.
41    query
42        .split(|c: char| !(c.is_alphanumeric() || c == '_'))
43        .filter(|term| !term.is_empty())
44        .collect::<Vec<_>>()
45        .join(" OR ")
46}
47
48fn format_ts_iso(ts: i64) -> String {
49    chrono::DateTime::from_timestamp(ts, 0)
50        .map(|dt| dt.format("%Y-%m-%dT%H:%M:%SZ").to_string())
51        .unwrap_or_else(|| ts.to_string())
52}
53
54#[async_trait]
55impl Store for SqliteStore {
56    async fn upsert_document(&self, doc: &Document) -> Result<String> {
57        sqlx::query(
58            r#"
59            INSERT INTO documents (id, source, source_id, source_url, title, author,
60                                   created_at, updated_at, content_type, body,
61                                   metadata_json, raw_json, dedup_hash)
62            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
63            ON CONFLICT(source, source_id) DO UPDATE SET
64                source_url = excluded.source_url,
65                title = excluded.title,
66                author = excluded.author,
67                updated_at = excluded.updated_at,
68                content_type = excluded.content_type,
69                body = excluded.body,
70                metadata_json = excluded.metadata_json,
71                raw_json = excluded.raw_json,
72                dedup_hash = excluded.dedup_hash
73            "#,
74        )
75        .bind(&doc.id)
76        .bind(&doc.source)
77        .bind(&doc.source_id)
78        .bind(&doc.source_url)
79        .bind(&doc.title)
80        .bind(&doc.author)
81        .bind(doc.created_at)
82        .bind(doc.updated_at)
83        .bind(&doc.content_type)
84        .bind(&doc.body)
85        .bind(&doc.metadata_json)
86        .bind(&doc.raw_json)
87        .bind(&doc.dedup_hash)
88        .execute(&self.pool)
89        .await?;
90
91        Ok(doc.id.clone())
92    }
93
94    async fn replace_chunks(
95        &self,
96        doc_id: &str,
97        chunks: &[Chunk],
98        vectors: Option<&[Vec<f32>]>,
99    ) -> Result<()> {
100        let mut tx = self.pool.begin().await?;
101
102        sqlx::query(
103            "DELETE FROM chunk_vectors WHERE chunk_id IN (SELECT id FROM chunks WHERE document_id = ?)",
104        )
105        .bind(doc_id)
106        .execute(&mut *tx)
107        .await?;
108
109        sqlx::query(
110            "DELETE FROM embeddings WHERE chunk_id IN (SELECT id FROM chunks WHERE document_id = ?)",
111        )
112        .bind(doc_id)
113        .execute(&mut *tx)
114        .await?;
115
116        sqlx::query("DELETE FROM chunks_fts WHERE document_id = ?")
117            .bind(doc_id)
118            .execute(&mut *tx)
119            .await?;
120
121        sqlx::query("DELETE FROM chunks WHERE document_id = ?")
122            .bind(doc_id)
123            .execute(&mut *tx)
124            .await?;
125
126        for (i, chunk) in chunks.iter().enumerate() {
127            sqlx::query(
128                "INSERT INTO chunks (id, document_id, chunk_index, text, hash) VALUES (?, ?, ?, ?, ?)",
129            )
130            .bind(&chunk.id)
131            .bind(&chunk.document_id)
132            .bind(chunk.chunk_index)
133            .bind(&chunk.text)
134            .bind(&chunk.hash)
135            .execute(&mut *tx)
136            .await?;
137
138            sqlx::query("INSERT INTO chunks_fts (chunk_id, document_id, text) VALUES (?, ?, ?)")
139                .bind(&chunk.id)
140                .bind(&chunk.document_id)
141                .bind(&chunk.text)
142                .execute(&mut *tx)
143                .await?;
144
145            if let Some(vecs) = vectors {
146                if let Some(vec) = vecs.get(i) {
147                    let blob = vec_to_blob(vec);
148                    sqlx::query(
149                        r#"
150                        INSERT INTO chunk_vectors (chunk_id, document_id, embedding)
151                        VALUES (?, ?, ?)
152                        ON CONFLICT(chunk_id) DO UPDATE SET
153                            document_id = excluded.document_id,
154                            embedding = excluded.embedding
155                        "#,
156                    )
157                    .bind(&chunk.id)
158                    .bind(doc_id)
159                    .bind(&blob)
160                    .execute(&mut *tx)
161                    .await?;
162                }
163            }
164        }
165
166        tx.commit().await?;
167        Ok(())
168    }
169
170    async fn upsert_embedding(
171        &self,
172        chunk_id: &str,
173        doc_id: &str,
174        vector: &[f32],
175        model: &str,
176        dims: usize,
177        content_hash: &str,
178    ) -> Result<()> {
179        let now = chrono::Utc::now().timestamp();
180        let blob = vec_to_blob(vector);
181
182        sqlx::query(
183            r#"
184            INSERT INTO embeddings (chunk_id, model, dims, created_at, hash)
185            VALUES (?, ?, ?, ?, ?)
186            ON CONFLICT(chunk_id) DO UPDATE SET
187                model = excluded.model,
188                dims = excluded.dims,
189                created_at = excluded.created_at,
190                hash = excluded.hash
191            "#,
192        )
193        .bind(chunk_id)
194        .bind(model)
195        .bind(dims as i64)
196        .bind(now)
197        .bind(content_hash)
198        .execute(&self.pool)
199        .await?;
200
201        sqlx::query(
202            r#"
203            INSERT INTO chunk_vectors (chunk_id, document_id, embedding)
204            VALUES (?, ?, ?)
205            ON CONFLICT(chunk_id) DO UPDATE SET
206                document_id = excluded.document_id,
207                embedding = excluded.embedding
208            "#,
209        )
210        .bind(chunk_id)
211        .bind(doc_id)
212        .bind(&blob)
213        .execute(&self.pool)
214        .await?;
215
216        Ok(())
217    }
218
219    async fn get_document(&self, id: &str) -> Result<Option<DocumentResponse>> {
220        let doc_row = sqlx::query(
221            "SELECT id, source, source_id, source_url, title, author, created_at, updated_at, content_type, body, metadata_json FROM documents WHERE id = ?",
222        )
223        .bind(id)
224        .fetch_optional(&self.pool)
225        .await?;
226
227        let doc_row = match doc_row {
228            Some(row) => row,
229            None => return Ok(None),
230        };
231
232        let created_at: i64 = doc_row.get("created_at");
233        let updated_at: i64 = doc_row.get("updated_at");
234        let metadata_json: String = doc_row.get("metadata_json");
235
236        let metadata: serde_json::Value =
237            serde_json::from_str(&metadata_json).unwrap_or(serde_json::json!({}));
238
239        let chunk_rows = sqlx::query(
240            "SELECT chunk_index, text FROM chunks WHERE document_id = ? ORDER BY chunk_index ASC",
241        )
242        .bind(id)
243        .fetch_all(&self.pool)
244        .await?;
245
246        let chunks: Vec<ChunkResponse> = chunk_rows
247            .iter()
248            .map(|row| ChunkResponse {
249                index: row.get("chunk_index"),
250                text: row.get("text"),
251            })
252            .collect();
253
254        Ok(Some(DocumentResponse {
255            id: doc_row.get("id"),
256            source: doc_row.get("source"),
257            source_id: doc_row.get("source_id"),
258            source_url: doc_row.get("source_url"),
259            title: doc_row.get("title"),
260            author: doc_row.get("author"),
261            created_at: format_ts_iso(created_at),
262            updated_at: format_ts_iso(updated_at),
263            content_type: doc_row.get("content_type"),
264            body: doc_row.get("body"),
265            metadata,
266            chunks,
267        }))
268    }
269
270    async fn get_document_metadata(&self, id: &str) -> Result<Option<DocumentMetadata>> {
271        let row = sqlx::query(
272            "SELECT id, title, source, source_id, updated_at, source_url FROM documents WHERE id = ?",
273        )
274        .bind(id)
275        .fetch_optional(&self.pool)
276        .await?;
277
278        Ok(row.map(|r| DocumentMetadata {
279            id: r.get("id"),
280            title: r.get("title"),
281            source: r.get("source"),
282            source_id: r.get("source_id"),
283            source_url: r.get("source_url"),
284            updated_at: r.get("updated_at"),
285        }))
286    }
287
288    async fn keyword_search(
289        &self,
290        query: &str,
291        limit: i64,
292        _source: Option<&str>,
293        _since: Option<&str>,
294    ) -> Result<Vec<ChunkCandidate>> {
295        let fts_query = fts_query_from_user_text(query);
296        if fts_query.is_empty() {
297            return Ok(Vec::new());
298        }
299
300        let rows = sqlx::query(
301            r#"
302            SELECT chunk_id, document_id, rank,
303                   snippet(chunks_fts, 2, '>>>', '<<<', '...', 48) AS snippet
304            FROM chunks_fts
305            WHERE chunks_fts MATCH ?
306            ORDER BY rank
307            LIMIT ?
308            "#,
309        )
310        .bind(fts_query)
311        .bind(limit)
312        .fetch_all(&self.pool)
313        .await?;
314
315        let candidates: Vec<ChunkCandidate> = rows
316            .iter()
317            .map(|row| {
318                let rank: f64 = row.get("rank");
319                ChunkCandidate {
320                    chunk_id: row.get("chunk_id"),
321                    document_id: row.get("document_id"),
322                    raw_score: -rank,
323                    snippet: row.get("snippet"),
324                }
325            })
326            .collect();
327
328        Ok(candidates)
329    }
330
331    async fn vector_search(
332        &self,
333        query_vec: &[f32],
334        limit: i64,
335        _source: Option<&str>,
336        _since: Option<&str>,
337    ) -> Result<Vec<ChunkCandidate>> {
338        let rows = sqlx::query(
339            r#"
340            SELECT cv.chunk_id, cv.document_id, cv.embedding,
341                   COALESCE(substr(c.text, 1, 240), '') AS snippet
342            FROM chunk_vectors cv
343            JOIN chunks c ON c.id = cv.chunk_id
344            "#,
345        )
346        .fetch_all(&self.pool)
347        .await?;
348
349        let mut candidates: Vec<ChunkCandidate> = rows
350            .iter()
351            .map(|row| {
352                let blob: Vec<u8> = row.get("embedding");
353                let vec = blob_to_vec(&blob);
354                let similarity = cosine_similarity(query_vec, &vec) as f64;
355                ChunkCandidate {
356                    chunk_id: row.get("chunk_id"),
357                    document_id: row.get("document_id"),
358                    raw_score: similarity,
359                    snippet: row.get("snippet"),
360                }
361            })
362            .collect();
363
364        candidates.sort_by(|a, b| {
365            b.raw_score
366                .partial_cmp(&a.raw_score)
367                .unwrap_or(std::cmp::Ordering::Equal)
368        });
369        candidates.truncate(limit as usize);
370
371        Ok(candidates)
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use super::fts_query_from_user_text;
378
379    #[test]
380    fn multi_word_query_joins_terms_with_or() {
381        // The fix: FTS5 treats a space between terms as an implicit AND, so a
382        // natural-language query used to require every word in one document and
383        // usually matched nothing. Terms must be OR-joined instead.
384        assert_eq!(
385            fts_query_from_user_text("the post about saving tokens"),
386            "the OR post OR about OR saving OR tokens"
387        );
388    }
389
390    #[test]
391    fn single_term_query_is_unchanged() {
392        assert_eq!(fts_query_from_user_text("tokens"), "tokens");
393    }
394
395    #[test]
396    fn punctuation_and_hyphens_split_into_or_terms() {
397        assert_eq!(
398            fts_query_from_user_text("saving-tokens, now!"),
399            "saving OR tokens OR now"
400        );
401    }
402
403    #[test]
404    fn underscores_are_kept_within_a_term() {
405        assert_eq!(
406            fts_query_from_user_text("save_tokens here"),
407            "save_tokens OR here"
408        );
409    }
410
411    #[test]
412    fn empty_or_punctuation_only_query_is_empty() {
413        assert_eq!(fts_query_from_user_text(""), "");
414        assert_eq!(fts_query_from_user_text("   "), "");
415        assert_eq!(fts_query_from_user_text("!!! ,.-"), "");
416    }
417}
context_harness/sqlite_store.rs

context_harness/
sqlite_store.rs
