Skip to main content

engram_storage/
fts.rs

1use rusqlite::params;
2
3use crate::database::Database;
4use crate::error::StorageError;
5use crate::memory::{self, Memory};
6
7#[derive(Debug, Clone)]
8pub struct FtsResult {
9    pub memory: Memory,
10    pub rank: f64,
11}
12
13fn sanitize_fts_query(text: &str) -> String {
14    let alphanumeric_only: String = text
15        .chars()
16        .filter(|character| character.is_alphanumeric() || character.is_whitespace())
17        .collect();
18
19    alphanumeric_only
20        .split_whitespace()
21        .map(|token| format!("\"{token}\""))
22        .collect::<Vec<String>>()
23        .join(" ")
24}
25
26impl Database {
27    pub fn search_fts(&self, query: &str, limit: usize) -> Result<Vec<FtsResult>, StorageError> {
28        let sanitized = sanitize_fts_query(query);
29        if sanitized.is_empty() {
30            return Ok(vec![]);
31        }
32
33        let mut statement = self.connection().prepare(
34            "SELECT m.*, rank
35             FROM memories m
36             JOIN memories_fts ON memories_fts.rowid = m.rowid
37             WHERE memories_fts MATCH ?1
38             ORDER BY rank
39             LIMIT ?2",
40        )?;
41        let rows = statement.query_map(params![sanitized, limit as i64], |row| {
42            let mem = memory::row_to_memory(row)?;
43            let rank: f64 = row.get("rank")?;
44            Ok(FtsResult { memory: mem, rank })
45        })?;
46        let mut results = Vec::new();
47        for row in rows {
48            results.push(row?);
49        }
50        Ok(results)
51    }
52}