grit_core/

search.rs

//! Hybrid retrieval: three legs — FTS5 BM25, sqlite-vec cosine, graph
//! expansion from seed nodes — fused with Reciprocal Rank Fusion, returned
//! with provenance and trimmed to a caller-supplied budget so Layer 3 can ask
//! for "context that fits".

use std::collections::HashMap;

use rusqlite::{Connection, params};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::Grit;
use crate::clock::TimestampMs;
use crate::error::Result;
use crate::model::{EDGE_COLS, EPISODE_COLS, Edge, Episode, NODE_COLS, Node, collect};
use crate::query::ids_json;
use crate::vecext::f32s_as_bytes;

/// RRF constant; the standard k=60 from the original TREC paper.
const RRF_K: f64 = 60.0;
/// How many candidates each leg contributes before fusion.
const LEG_LIMIT: i64 = 50;
/// How many top fused node hits seed the graph-expansion leg.
const EXPANSION_SEEDS: usize = 5;
/// Rough chars-per-token for [`Budget::approx_tokens`].
const CHARS_PER_TOKEN: usize = 4;

/// How much context to return (Layer 3 asks for "context that fits").
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Budget {
    /// At most this many hits.
    Items(usize),
    /// Stop once the accumulated text is roughly this many tokens
    /// (chars / 4 heuristic; no tokenizer dependency in a storage engine).
    ApproxTokens(usize),
}

impl Budget {
    /// Budget of at most `n` hits.
    pub fn items(n: usize) -> Self {
        Budget::Items(n)
    }

    /// Budget of roughly `n` tokens of text.
    pub fn approx_tokens(n: usize) -> Self {
        Budget::ApproxTokens(n)
    }

    fn max_items(self) -> usize {
        match self {
            Budget::Items(n) => n,
            // One token is at least one char; never fetch more than this.
            Budget::ApproxTokens(n) => n.max(1),
        }
    }
}

/// A hybrid search request. Build with [`Query::text`], refine with the
/// builder methods, execute with [`Grit::search`].
///
/// # Example
/// ```no_run
/// # use grit_core::{Budget, Grit, Options, Query};
/// # let g = Grit::open("memory.db", Options::new("laptop"))?;
/// let hits = g.search(
///     Query::text("exactness").group("algebra").budget(Budget::items(20)),
/// )?;
/// # Ok::<(), grit_core::Error>(())
/// ```
#[derive(Debug, Clone)]
pub struct Query {
    text: String,
    vector: Option<Vec<f32>>,
    group_id: Option<String>,
    as_of: Option<TimestampMs>,
    as_at: Option<TimestampMs>,
    budget: Budget,
    targets: Option<Vec<SearchKind>>,
}

/// Result kinds a [`Query`] may be restricted to (see [`Query::targets`]).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchKind {
    /// Entity nodes.
    Node,
    /// Fact edges.
    Edge,
    /// Provenance episodes.
    Episode,
}

impl Query {
    /// Full-text query over node names/summaries, edge facts, episode content.
    pub fn text(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            vector: None,
            group_id: None,
            as_of: None,
            as_at: None,
            budget: Budget::Items(20),
            targets: None,
        }
    }

    /// Restrict results to the given kinds. The fused ranking is computed
    /// over ALL kinds (graph expansion still works), then filtered BEFORE
    /// the budget is spent — so `targets(&[SearchKind::Edge])` with a
    /// budget of 3 returns the top 3 edges, not whatever edges survive a
    /// mixed top-3 cut.
    pub fn targets(mut self, kinds: &[SearchKind]) -> Self {
        self.targets = Some(kinds.to_vec());
        self
    }

    /// Add a query embedding (computed by the caller — grit never embeds) to
    /// enable the vector legs.
    pub fn vector(mut self, embedding: Vec<f32>) -> Self {
        self.vector = Some(embedding);
        self
    }

    /// Restrict to a namespace.
    pub fn group(mut self, group_id: impl Into<String>) -> Self {
        self.group_id = Some(group_id.into());
        self
    }

    /// Event-time instant: only return facts true at `t`.
    pub fn as_of(mut self, t: TimestampMs) -> Self {
        self.as_of = Some(t);
        self
    }

    /// System-time instant: only return facts believed at `t` (time travel).
    pub fn as_at(mut self, t: TimestampMs) -> Self {
        self.as_at = Some(t);
        self
    }

    /// Cap the result size.
    pub fn budget(mut self, budget: Budget) -> Self {
        self.budget = budget;
        self
    }
}

/// What a search hit points at.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum SearchTarget {
    /// An entity node.
    Node(Node),
    /// A fact edge.
    Edge(Edge),
    /// A provenance episode.
    Episode(Episode),
}

/// One fused search result with provenance.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchHit {
    /// The matched node/edge/episode.
    pub target: SearchTarget,
    /// RRF fusion score (higher is better; comparable within one search only).
    pub score: f64,
    /// Episodes this fact traces to (empty for episode hits — they *are* the
    /// provenance).
    pub episodes: Vec<Uuid>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum CandidateKind {
    Node,
    Edge,
    Episode,
}

impl Grit {
    /// Run hybrid retrieval: BM25 + vector cosine + graph expansion, RRF-fused,
    /// validity-filtered at the query's `as_of`/`as_at`, trimmed to budget.
    pub fn search(&self, query: Query) -> Result<Vec<SearchHit>> {
        let now = self.now_ms();
        let as_of = query.as_of.unwrap_or(now);
        let as_at = query.as_at.unwrap_or(now);
        let group = query.group_id.as_deref();
        let match_expr = fts_match_expr(&query.text);

        let conn = self.read();
        // One read transaction: every leg, the expansion, and the fetch
        // phase see the same committed snapshot.
        let tx = conn.unchecked_transaction()?;

        let mut legs: Vec<Vec<(CandidateKind, Uuid)>> = Vec::new();
        if let Some(expr) = &match_expr {
            legs.push(tag(
                CandidateKind::Node,
                fts_nodes(&tx, expr, group, as_at)?,
            ));
            legs.push(tag(
                CandidateKind::Edge,
                fts_edges(&tx, expr, group, as_at, as_of)?,
            ));
            legs.push(tag(CandidateKind::Episode, fts_episodes(&tx, expr, group)?));
        }
        if let Some(vector) = &query.vector
            && vec_tables_exist(&tx)?
        {
            legs.push(tag(CandidateKind::Node, vec_leg(&tx, "vec_nodes", vector)?));
            legs.push(tag(CandidateKind::Edge, vec_leg(&tx, "vec_edges", vector)?));
        }

        // Leg 3: expand one hop from the best node candidates so far; related
        // entities surface even when they don't match the text themselves.
        let mut fused = rrf(&legs);
        let seed_nodes: Vec<Uuid> = fused
            .iter()
            .filter(|((kind, _), _)| *kind == CandidateKind::Node)
            .take(EXPANSION_SEEDS)
            .map(|((_, id), _)| *id)
            .collect();
        if !seed_nodes.is_empty() {
            legs.push(tag(
                CandidateKind::Node,
                expansion_leg(&tx, &seed_nodes, group, as_at, as_of)?,
            ));
            fused = rrf(&legs);
        }

        // Fetch rows in fused order, validity-filtered, until the budget is
        // spent. Provenance comes along for free via `mentions`.
        let mut hits = Vec::new();
        let mut spent_chars = 0usize;
        for ((kind, id), score) in fused {
            if let Some(targets) = &query.targets {
                let kind = match kind {
                    CandidateKind::Node => SearchKind::Node,
                    CandidateKind::Edge => SearchKind::Edge,
                    CandidateKind::Episode => SearchKind::Episode,
                };
                if !targets.contains(&kind) {
                    continue;
                }
            }
            if hits.len() >= query.budget.max_items() {
                break;
            }
            if let Budget::ApproxTokens(tokens) = query.budget
                && spent_chars / CHARS_PER_TOKEN >= tokens
            {
                break;
            }
            let target = match kind {
                CandidateKind::Node => fetch_node(&tx, id, group, as_at)?.map(SearchTarget::Node),
                CandidateKind::Edge => {
                    fetch_edge(&tx, id, group, as_at, as_of)?.map(SearchTarget::Edge)
                }
                CandidateKind::Episode => fetch_episode(&tx, id, group)?.map(SearchTarget::Episode),
            };
            let Some(target) = target else { continue };
            spent_chars += target_chars(&target);
            let episodes = match kind {
                CandidateKind::Episode => Vec::new(),
                _ => episode_ids_for(&tx, id)?,
            };
            hits.push(SearchHit {
                target,
                score,
                episodes,
            });
        }
        Ok(hits)
    }
}

/// Sanitize free text into an FTS5 MATCH expression: each whitespace token is
/// double-quoted (quotes doubled inside), joined with implicit AND. Returns
/// `None` for effectively-empty input.
fn fts_match_expr(text: &str) -> Option<String> {
    let tokens: Vec<String> = text
        .split_whitespace()
        .map(|t| format!("\"{}\"", t.replace('"', "\"\"")))
        .collect();
    if tokens.is_empty() {
        None
    } else {
        Some(tokens.join(" "))
    }
}

fn tag(kind: CandidateKind, ids: Vec<Uuid>) -> Vec<(CandidateKind, Uuid)> {
    ids.into_iter().map(|id| (kind, id)).collect()
}

/// Reciprocal Rank Fusion across legs; stable order (score desc, then id).
fn rrf(legs: &[Vec<(CandidateKind, Uuid)>]) -> Vec<((CandidateKind, Uuid), f64)> {
    let mut scores: HashMap<(CandidateKind, Uuid), f64> = HashMap::new();
    for leg in legs {
        for (rank, key) in leg.iter().enumerate() {
            *scores.entry(*key).or_insert(0.0) += 1.0 / (RRF_K + rank as f64 + 1.0);
        }
    }
    let mut fused: Vec<_> = scores.into_iter().collect();
    fused.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.1.cmp(&b.0.1)));
    fused
}

/// Malformed uuids in our own tables are corruption, surfaced loudly — a
/// silent drop would make rows invisible with no signal.
fn parse_ids(rows: Vec<String>) -> Result<Vec<Uuid>> {
    rows.iter()
        .map(|s| {
            Uuid::parse_str(s).map_err(|_| crate::Error::Corrupt(format!("bad uuid in index: {s}")))
        })
        .collect()
}

fn fts_nodes(
    conn: &Connection,
    expr: &str,
    group: Option<&str>,
    as_at: TimestampMs,
) -> Result<Vec<Uuid>> {
    let mut stmt = conn.prepare_cached(
        "SELECT n.id FROM nodes_fts AS f
         JOIN nodes AS n ON n.rowid = f.rowid
         WHERE nodes_fts MATCH ?1
           AND (?2 IS NULL OR n.group_id = ?2)
           AND n.created_at <= ?3 AND (n.expired_at IS NULL OR n.expired_at > ?3)
         ORDER BY f.rank LIMIT ?4",
    )?;
    let rows = stmt.query_map(params![expr, group, as_at, LEG_LIMIT], |r| r.get(0))?;
    parse_ids(collect(rows)?)
}

fn fts_edges(
    conn: &Connection,
    expr: &str,
    group: Option<&str>,
    as_at: TimestampMs,
    as_of: TimestampMs,
) -> Result<Vec<Uuid>> {
    let mut stmt = conn.prepare_cached(
        "SELECT e.id FROM edges_fts AS f
         JOIN edges AS e ON e.rowid = f.rowid
         WHERE edges_fts MATCH ?1
           AND (?2 IS NULL OR e.group_id = ?2)
           AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
           AND (e.valid_at IS NULL OR e.valid_at <= ?4)
           AND NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
                           WHERE iv.edge_id = e.id
                             AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)
         ORDER BY f.rank LIMIT ?5",
    )?;
    let rows = stmt.query_map(params![expr, group, as_at, as_of, LEG_LIMIT], |r| r.get(0))?;
    parse_ids(collect(rows)?)
}

fn fts_episodes(conn: &Connection, expr: &str, group: Option<&str>) -> Result<Vec<Uuid>> {
    let mut stmt = conn.prepare_cached(
        "SELECT e.id FROM episodes_fts AS f
         JOIN episodes AS e ON e.rowid = f.rowid
         WHERE episodes_fts MATCH ?1
           AND (?2 IS NULL OR e.group_id = ?2)
         ORDER BY f.rank LIMIT ?3",
    )?;
    let rows = stmt.query_map(params![expr, group, LEG_LIMIT], |r| r.get(0))?;
    parse_ids(collect(rows)?)
}

fn vec_tables_exist(conn: &Connection) -> Result<bool> {
    let n: i64 = conn.query_row(
        "SELECT COUNT(*) FROM sqlite_master WHERE name IN ('vec_nodes', 'vec_edges')",
        [],
        |r| r.get(0),
    )?;
    Ok(n == 2)
}

fn vec_leg(conn: &Connection, table: &str, vector: &[f32]) -> Result<Vec<Uuid>> {
    let registered_dim: Option<i64> = conn
        .query_row("SELECT dim FROM embedding_meta WHERE id = 1", [], |r| {
            r.get(0)
        })
        .map(Some)
        .or_else(|e| match e {
            rusqlite::Error::QueryReturnedNoRows => Ok(None),
            other => Err(other),
        })?;
    if registered_dim != Some(vector.len() as i64) {
        // Wrong-dim query vectors contribute nothing rather than erroring the
        // whole search; the caller may be mid-migration between models.
        return Ok(Vec::new());
    }
    // `table` is one of two compile-time names; the vector is a bound blob.
    let mut stmt = conn.prepare_cached(&format!(
        "SELECT id FROM {table} WHERE embedding MATCH ?1 AND k = ?2 ORDER BY distance"
    ))?;
    let rows = stmt.query_map(params![f32s_as_bytes(vector), LEG_LIMIT], |r| {
        r.get::<_, String>(0)
    })?;
    parse_ids(collect(rows)?)
}

fn expansion_leg(
    conn: &Connection,
    seeds: &[Uuid],
    group: Option<&str>,
    as_at: TimestampMs,
    as_of: TimestampMs,
) -> Result<Vec<Uuid>> {
    let mut stmt = conn.prepare_cached(
        "SELECT DISTINCT CASE WHEN e.src = s.value THEN e.dst ELSE e.src END
         FROM json_each(?1) AS s
         JOIN edges AS e ON e.src = s.value OR e.dst = s.value
         WHERE (?2 IS NULL OR e.group_id = ?2)
           AND e.created_at <= ?3 AND (e.expired_at IS NULL OR e.expired_at > ?3)
           AND (e.valid_at IS NULL OR e.valid_at <= ?4)
           AND NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
                           WHERE iv.edge_id = e.id
                             AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)
         LIMIT ?5",
    )?;
    let rows = stmt.query_map(
        params![ids_json(seeds.iter()), group, as_at, as_of, LEG_LIMIT],
        |r| r.get(0),
    )?;
    parse_ids(collect(rows)?)
}

fn fetch_node(
    conn: &Connection,
    id: Uuid,
    group: Option<&str>,
    as_at: TimestampMs,
) -> Result<Option<Node>> {
    let mut stmt = conn.prepare_cached(&format!(
        "SELECT {NODE_COLS} FROM nodes
         WHERE id = ?1 AND (?2 IS NULL OR group_id = ?2)
           AND created_at <= ?3 AND (expired_at IS NULL OR expired_at > ?3)"
    ))?;
    let mut rows = collect(stmt.query_map(params![id.to_string(), group, as_at], Node::from_row)?)?;
    Ok(rows.pop())
}

fn fetch_edge(
    conn: &Connection,
    id: Uuid,
    group: Option<&str>,
    as_at: TimestampMs,
    as_of: TimestampMs,
) -> Result<Option<Edge>> {
    // The EXISTS guards keep ghost edges (an endpoint purged or merged away)
    // out of results.
    let mut stmt = conn.prepare_cached(&format!(
        "SELECT {EDGE_COLS} FROM edges
         WHERE id = ?1 AND (?2 IS NULL OR group_id = ?2)
           AND created_at <= ?3 AND (expired_at IS NULL OR expired_at > ?3)
           AND (valid_at IS NULL OR valid_at <= ?4)
           AND NOT EXISTS (SELECT 1 FROM edge_invalidations AS iv
                           WHERE iv.edge_id = edges.id
                             AND iv.recorded_at <= ?3 AND iv.invalid_at <= ?4)
           AND EXISTS (SELECT 1 FROM nodes WHERE id = edges.src
                       AND created_at <= ?3 AND (expired_at IS NULL OR expired_at > ?3))
           AND EXISTS (SELECT 1 FROM nodes WHERE id = edges.dst
                       AND created_at <= ?3 AND (expired_at IS NULL OR expired_at > ?3))"
    ))?;
    let mut rows =
        collect(stmt.query_map(params![id.to_string(), group, as_at, as_of], Edge::from_row)?)?;
    Ok(rows.pop())
}

fn fetch_episode(conn: &Connection, id: Uuid, group: Option<&str>) -> Result<Option<Episode>> {
    let mut stmt = conn.prepare_cached(&format!(
        "SELECT {EPISODE_COLS} FROM episodes
         WHERE id = ?1 AND (?2 IS NULL OR group_id = ?2)"
    ))?;
    let mut rows = collect(stmt.query_map(params![id.to_string(), group], Episode::from_row)?)?;
    Ok(rows.pop())
}

fn episode_ids_for(conn: &Connection, target: Uuid) -> Result<Vec<Uuid>> {
    // Alias-aware: mention rows keep original target ids across merges.
    crate::query::episodes_mentioning(conn, &target.to_string())
}

fn target_chars(target: &SearchTarget) -> usize {
    match target {
        SearchTarget::Node(n) => n.name.len() + n.summary.len(),
        SearchTarget::Edge(e) => e.fact.len(),
        SearchTarget::Episode(e) => e.content.len(),
    }
}