agentdb/
fts.rs

1use crate::error::{AgentDbError, Result};
2use rusqlite::params;
3use rusqlite::Connection;
4use serde_json::Value;
5use std::sync::{Arc, Mutex};
6
7/// A full-text search result
8#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
9pub struct FtsResult {
10    pub id: String,
11    pub collection_id: String,
12    pub snippet: String,
13    pub rank: f64,
14    pub metadata: Option<Value>,
15}
16
17/// Manages FTS5 virtual tables per collection
18pub struct FullTextStore {
19    conn: Arc<Mutex<Connection>>,
20}
21
22impl FullTextStore {
23    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
24        Self { conn }
25    }
26
27    /// Ensure the FTS5 virtual table exists for a collection
28    pub fn ensure_fts_table(&self, collection_name: &str) -> Result<()> {
29        let table = fts_table_name(collection_name);
30        let conn = self.conn.lock().unwrap();
31        conn.execute_batch(&format!(
32            "CREATE VIRTUAL TABLE IF NOT EXISTS {table}
33             USING fts5(
34               vec_id,
35               collection_id UNINDEXED,
36               text,
37               tokenize='porter ascii'
38             );"
39        ))?;
40        Ok(())
41    }
42
43    /// Index a text document for a given collection
44    pub fn index_text(
45        &self,
46        collection_name: &str,
47        vec_id: &str,
48        collection_id: &str,
49        text: &str,
50    ) -> Result<()> {
51        self.ensure_fts_table(collection_name)?;
52        let table = fts_table_name(collection_name);
53        let conn = self.conn.lock().unwrap();
54        conn.execute(
55            &format!("DELETE FROM {table} WHERE vec_id = ?1"),
56            params![vec_id],
57        )?;
58        conn.execute(
59            &format!("INSERT INTO {table} (vec_id, collection_id, text) VALUES (?1, ?2, ?3)"),
60            params![vec_id, collection_id, text],
61        )?;
62        Ok(())
63    }
64
65    /// Full-text search over a collection
66    pub fn search(
67        &self,
68        collection_name: &str,
69        query: &str,
70        top_k: usize,
71    ) -> Result<Vec<FtsResult>> {
72        let table = fts_table_name(collection_name);
73        let conn = self.conn.lock().unwrap();
74
75        let exists: bool = conn
76            .query_row(
77                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1",
78                params![table],
79                |r| r.get::<_, i64>(0),
80            )
81            .map(|c| c > 0)
82            .unwrap_or(false);
83
84        if !exists {
85            return Ok(vec![]);
86        }
87
88        let sql = format!(
89            "SELECT f.vec_id,
90                    f.collection_id,
91                    snippet({table}, 2, '<b>', '</b>', '...', 10),
92                    bm25({table}) AS rank,
93                    v.metadata
94             FROM {table} f
95             LEFT JOIN _adb_vectors v
96               ON v.id = f.vec_id AND v.collection_id = f.collection_id
97             WHERE {table} MATCH ?1
98             ORDER BY rank
99             LIMIT ?2"
100        );
101
102        let mut stmt = conn.prepare(&sql)?;
103        let rows = stmt.query_map(params![query, top_k as i64], |row| {
104            let meta_str: Option<String> = row.get(4)?;
105            Ok(FtsResult {
106                id: row.get(0)?,
107                collection_id: row.get(1)?,
108                snippet: row.get(2)?,
109                rank: row.get(3)?,
110                metadata: meta_str
111                    .as_deref()
112                    .and_then(|s| serde_json::from_str(s).ok()),
113            })
114        })?;
115
116        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
117    }
118
119    /// Delete a document from the FTS index
120    pub fn delete_text(&self, collection_name: &str, vec_id: &str) -> Result<()> {
121        let table = fts_table_name(collection_name);
122        let conn = self.conn.lock().unwrap();
123        conn.execute(
124            &format!("DELETE FROM {table} WHERE vec_id = ?1"),
125            params![vec_id],
126        )?;
127        Ok(())
128    }
129
130    /// Merge FTS index segments for faster queries
131    pub fn optimize(&self, collection_name: &str) -> Result<()> {
132        let table = fts_table_name(collection_name);
133        let conn = self.conn.lock().unwrap();
134        conn.execute_batch(&format!("INSERT INTO {table}({table}) VALUES('optimize');"))?;
135        Ok(())
136    }
137}
138
139fn fts_table_name(name: &str) -> String {
140    let safe: String = name
141        .chars()
142        .map(|c| {
143            if c.is_alphanumeric() || c == '_' {
144                c
145            } else {
146                '_'
147            }
148        })
149        .collect();
150    format!("_adb_fts_{safe}")
151}
agentdb/fts.rs

agentdb/
fts.rs
