use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Context, Result};
use rusqlite::{params, Connection};
use tokio::sync::Mutex;
#[derive(Debug, Clone)]
pub struct MessageFtsHit {
pub message_id: String,
pub conversation_id: String,
pub role: String,
pub created_at: i64,
pub score: f32,
}
#[derive(Clone)]
pub struct MessageFtsIndex {
conn: Arc<Mutex<Connection>>,
}
impl MessageFtsIndex {
pub fn open(path: PathBuf) -> Result<Self> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("creating message-fts db dir {}", parent.display()))?;
}
let conn = Connection::open(&path)
.with_context(|| format!("opening message-fts db {}", path.display()))?;
Self::init_schema(&conn)?;
Ok(Self {
conn: Arc::new(Mutex::new(conn)),
})
}
pub fn open_in_memory() -> Result<Self> {
let conn = Connection::open_in_memory().context("opening in-memory message-fts db")?;
Self::init_schema(&conn)?;
Ok(Self {
conn: Arc::new(Mutex::new(conn)),
})
}
fn init_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(
"PRAGMA journal_mode = WAL;
CREATE TABLE IF NOT EXISTS fts_messages (
rowid INTEGER PRIMARY KEY,
message_id TEXT UNIQUE NOT NULL,
conversation_id TEXT NOT NULL,
role TEXT NOT NULL,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_fts_messages_conversation
ON fts_messages(conversation_id);",
)
.context("initializing fts_messages schema")?;
conn.execute_batch(
"CREATE VIRTUAL TABLE IF NOT EXISTS message_fts USING fts5(body, content='');",
)
.context("initializing message_fts FTS5 table (is FTS5 compiled in?)")?;
Ok(())
}
pub async fn index_message(
&self,
message_id: &str,
conversation_id: &str,
role: &str,
text: &str,
created_at: i64,
) -> Result<()> {
if text.trim().is_empty() {
return Ok(());
}
let conn = self.conn.lock().await;
conn.execute(
"INSERT OR IGNORE INTO fts_messages
(message_id, conversation_id, role, created_at)
VALUES (?1, ?2, ?3, ?4)",
params![message_id, conversation_id, role, created_at],
)
.context("inserting fts_messages metadata row")?;
if conn.changes() == 0 {
return Ok(());
}
let rowid = conn.last_insert_rowid();
conn.execute(
"INSERT INTO message_fts (rowid, body) VALUES (?1, ?2)",
params![rowid, text],
)
.context("inserting message_fts body row")?;
Ok(())
}
pub async fn indexed_ids(&self) -> Result<HashSet<String>> {
let conn = self.conn.lock().await;
let mut stmt = conn.prepare("SELECT message_id FROM fts_messages")?;
let rows = stmt.query_map([], |row| row.get::<_, String>(0))?;
let mut set = HashSet::new();
for row in rows {
set.insert(row?);
}
Ok(set)
}
pub async fn search(
&self,
query: &str,
limit: usize,
conversation_ids: Option<&[String]>,
) -> Result<Vec<MessageFtsHit>> {
if limit == 0 {
return Ok(Vec::new());
}
let Some(match_expr) = sanitize_fts_query(query) else {
return Ok(Vec::new());
};
let conn = self.conn.lock().await;
let fetch = match conversation_ids {
Some(ids) if !ids.is_empty() => limit.saturating_mul(8).max(64),
_ => limit,
};
let mut stmt = conn.prepare(
"SELECT m.message_id, m.conversation_id, m.role, m.created_at, bm25(message_fts)
FROM message_fts
JOIN fts_messages m ON m.rowid = message_fts.rowid
WHERE message_fts MATCH ?1
ORDER BY bm25(message_fts)
LIMIT ?2",
)?;
let rows = stmt.query_map(params![match_expr, fetch as i64], |row| {
let bm25 = row.get::<_, f64>(4)? as f32;
Ok(MessageFtsHit {
message_id: row.get(0)?,
conversation_id: row.get(1)?,
role: row.get(2)?,
created_at: row.get(3)?,
score: bm25_to_relevance(bm25),
})
})?;
let mut out = Vec::new();
for row in rows {
let hit = row?;
if let Some(ids) = conversation_ids {
if !ids.is_empty() && !ids.iter().any(|c| c == &hit.conversation_id) {
continue;
}
}
out.push(hit);
if out.len() >= limit {
break;
}
}
Ok(out)
}
}
fn bm25_to_relevance(bm25: f32) -> f32 {
let strength = (-bm25).max(0.0);
strength / (1.0 + strength)
}
fn sanitize_fts_query(raw: &str) -> Option<String> {
let terms: Vec<String> = raw
.split(|c: char| !c.is_alphanumeric())
.filter(|t| !t.is_empty())
.map(|t| format!("\"{t}\""))
.collect();
if terms.is_empty() {
return None;
}
Some(terms.join(" OR "))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn fts_round_trip_surfaces_matching_session() {
let index = MessageFtsIndex::open_in_memory().expect("open index");
let docs = [
("m1", "c1", "assistant", "rust borrow checker lifetimes"),
("m2", "c1", "user", "favourite pizza toppings pepperoni"),
("m3", "c2", "user", "singapore travel itinerary"),
];
for (id, conv, role, text) in docs {
index
.index_message(id, conv, role, text, 0)
.await
.expect("index");
}
let hits = index.search("pizza", 5, None).await.expect("search");
assert!(!hits.is_empty(), "expected at least one hit");
assert_eq!(hits[0].message_id, "m2", "pizza query must surface m2");
assert_eq!(hits[0].conversation_id, "c1");
}
#[tokio::test]
async fn search_scopes_to_conversation_ids() {
let index = MessageFtsIndex::open_in_memory().expect("open index");
for (id, conv) in [("m1", "c1"), ("m2", "c2")] {
index
.index_message(id, conv, "user", "alpha beta gamma", 0)
.await
.expect("index");
}
let hits = index
.search("alpha beta", 5, Some(&["c2".to_owned()]))
.await
.expect("search");
assert!(!hits.is_empty());
assert!(
hits.iter().all(|h| h.conversation_id == "c2"),
"all hits must be scoped to c2"
);
}
#[tokio::test]
async fn reindex_is_idempotent() {
let index = MessageFtsIndex::open_in_memory().expect("open index");
index
.index_message("m1", "c1", "user", "hello world", 0)
.await
.expect("index");
index
.index_message("m1", "c1", "user", "hello world", 1)
.await
.expect("reindex");
let ids = index.indexed_ids().await.expect("ids");
assert_eq!(ids.len(), 1, "re-index must not duplicate the row");
let hits = index.search("hello", 5, None).await.expect("search");
assert_eq!(hits.len(), 1);
}
#[tokio::test]
async fn sanitize_handles_meta_chars() {
let index = MessageFtsIndex::open_in_memory().expect("open index");
index
.index_message("m1", "c1", "user", "the foo and the bar baz", 0)
.await
.expect("index");
let hits = index
.search("foo:bar* \"baz AND", 5, None)
.await
.expect("search must not error on meta-chars");
assert!(!hits.is_empty(), "term tokens should still match");
assert_eq!(hits[0].message_id, "m1");
}
#[test]
fn sanitize_fts_query_shapes_terms() {
assert_eq!(
sanitize_fts_query("foo:bar* \"baz"),
Some("\"foo\" OR \"bar\" OR \"baz\"".to_owned())
);
assert_eq!(sanitize_fts_query(" "), None);
assert_eq!(sanitize_fts_query("!!! @@@ ---"), None);
assert_eq!(sanitize_fts_query("solo"), Some("\"solo\"".to_owned()));
}
#[test]
fn bm25_relevance_is_bounded_and_monotonic() {
let strong = bm25_to_relevance(-5.0);
let weak = bm25_to_relevance(-0.5);
assert!(strong > weak);
assert!((0.0..=1.0).contains(&strong));
assert!((0.0..=1.0).contains(&weak));
assert_eq!(bm25_to_relevance(1.0), 0.0);
}
}