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    text.chars()
15        .map(|character| {
16            if character.is_alphanumeric() {
17                character
18            } else {
19                ' '
20            }
21        })
22        .collect::<String>()
23        .split_whitespace()
24        .map(|token| format!("\"{token}\"*"))
25        .collect::<Vec<String>>()
26        .join(" OR ")
27}
28
29impl Database {
30    pub fn search_fts(&self, query: &str, limit: usize) -> Result<Vec<FtsResult>, StorageError> {
31        let sanitized = sanitize_fts_query(query);
32        if sanitized.is_empty() {
33            return Ok(vec![]);
34        }
35
36        let mut statement = self.connection().prepare(
37            "SELECT m.*, rank
38             FROM memories m
39             JOIN memories_fts ON memories_fts.rowid = m.rowid
40             WHERE memories_fts MATCH ?1
41             ORDER BY rank
42             LIMIT ?2",
43        )?;
44        let rows = statement.query_map(params![sanitized, limit as i64], |row| {
45            let mem = memory::row_to_memory(row)?;
46            let rank: f64 = row.get("rank")?;
47            Ok(FtsResult { memory: mem, rank })
48        })?;
49        let mut results = Vec::new();
50        for row in rows {
51            results.push(row?);
52        }
53        Ok(results)
54    }
55}