mcp-server-sqlite 1.0.0

//! The `search_fts` tool: runs a full-text search query against an FTS5 virtual
//! table, returning ranked results with highlighted snippets.

use rmcp::model::{Content, IntoContents};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use super::ToolError;
use crate::{mcp::McpServerSqlite, traits::SqliteServerTool};

#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Default,
    Serialize,
    Deserialize,
    JsonSchema,
)]
/// Search a full-text search index using SQLite's FTS5 MATCH syntax. Returns
/// results ranked by BM25 relevance with highlighted snippets showing where the
/// query matched. The FTS table must have been created previously (e.g. via
/// `create_fts_index`).
pub struct SearchFtsTool;

impl SqliteServerTool for SearchFtsTool {
    const NAME: &str = "search_fts";

    type Context = McpServerSqlite;
    type Error = ToolError<SearchFtsError>;

    type Input = SearchFtsInput;
    type Output = SearchFtsOutput;

    fn handle(
        ctx: &Self::Context,
        input: Self::Input,
    ) -> Result<Self::Output, Self::Error> {
        let conn = ctx
            .connection()
            .map_err(|source| ToolError::Connection { source })?;

        let limit = input.limit.unwrap_or(10);
        let snippet_tokens = input.snippet_tokens.unwrap_or(32);
        let hl_start = input.highlight_start.as_deref().unwrap_or("<b>");
        let hl_end = input.highlight_end.as_deref().unwrap_or("</b>");

        let column_count =
            fts_column_count(&conn, &input.fts_table).map_err(|source| {
                ToolError::Tool(SearchFtsError::Query { source })
            })?;

        let snippet_exprs = (0..column_count)
            .map(|i| {
                format!(
                    "snippet({tbl}, {i}, '{hl_start}', '{hl_end}', '...', {tokens})",
                    tbl = input.fts_table,
                    tokens = snippet_tokens,
                )
            })
            .collect::<Vec<_>>()
            .join(", ");

        let sql = format!(
            "SELECT rowid, rank, {snippets} \
             FROM [{tbl}] \
             WHERE [{tbl}] MATCH ?1 \
             ORDER BY rank \
             LIMIT ?2",
            tbl = input.fts_table,
            snippets = snippet_exprs,
        );

        let mut stmt = conn.prepare(&sql).map_err(|source| {
            ToolError::Tool(SearchFtsError::Query { source })
        })?;

        let results = stmt
            .query_map(rusqlite::params![input.query, limit], |row| {
                let rowid = row.get::<_, i64>(0)?;
                let rank = row.get::<_, f64>(1)?;
                let snippets = (0..column_count)
                    .map(|i| row.get::<_, String>(2 + i))
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(FtsMatch {
                    rowid,
                    rank,
                    snippets,
                })
            })
            .map_err(|source| {
                ToolError::Tool(SearchFtsError::Query { source })
            })?
            .collect::<Result<Vec<_>, _>>()
            .map_err(|source| {
                ToolError::Tool(SearchFtsError::Query { source })
            })?;

        Ok(SearchFtsOutput { results })
    }
}

/// Queries the FTS table's column count by inspecting the table schema.
fn fts_column_count(
    conn: &rusqlite::Connection,
    fts_table: &str,
) -> Result<usize, rusqlite::Error> {
    let mut stmt =
        conn.prepare(&format!("PRAGMA table_info([{}])", fts_table))?;
    let count = stmt.query_map([], |_| Ok(()))?.count();
    Ok(count)
}

/// The input parameters for the `search_fts` tool.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub struct SearchFtsInput {
    /// The name of the FTS5 virtual table to search.
    #[schemars(description = "The FTS5 virtual table to search")]
    pub fts_table: String,
    /// The FTS5 MATCH query string (e.g. `"sqlite AND indexing"`).
    #[schemars(description = "The FTS5 MATCH query string")]
    pub query: String,
    /// Maximum number of results to return. Defaults to 10.
    #[schemars(description = "Maximum number of results (default 10)")]
    pub limit: Option<i64>,
    /// Maximum number of tokens per snippet. Defaults to 32.
    #[schemars(description = "Max tokens per snippet (default 32)")]
    pub snippet_tokens: Option<i32>,
    /// The string inserted before a matching term in snippets. Defaults to
    /// `<b>`.
    #[schemars(
        description = "String before matched terms in snippets (default <b>)"
    )]
    pub highlight_start: Option<String>,
    /// The string inserted after a matching term in snippets. Defaults to
    /// `</b>`.
    #[schemars(
        description = "String after matched terms in snippets (default </b>)"
    )]
    pub highlight_end: Option<String>,
}

/// The results of a full-text search query.
#[derive(
    Clone,
    Debug,
    PartialEq,
    PartialOrd,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub struct SearchFtsOutput {
    /// The matching rows, ranked by BM25 relevance (best first).
    pub results: Vec<FtsMatch>,
}

/// A single full-text search result.
#[derive(
    Clone,
    Debug,
    PartialEq,
    PartialOrd,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub struct FtsMatch {
    /// The rowid of the matching row in the FTS table.
    pub rowid: i64,
    /// The BM25 relevance score. Lower (more negative) is more relevant.
    pub rank: f64,
    /// Highlighted text snippets for each indexed column, in the order the
    /// columns were defined in the FTS table.
    pub snippets: Vec<String>,
}

/// Errors specific to the `search_fts` tool.
#[derive(Debug, thiserror::Error)]
pub enum SearchFtsError {
    /// The FTS search query failed.
    #[error("FTS search failed: {source}")]
    Query {
        /// The underlying rusqlite error.
        source: rusqlite::Error,
    },
}

/// Converts the error into MCP content by rendering the display string as text.
impl IntoContents for SearchFtsError {
    fn into_contents(self) -> Vec<Content> {
        vec![Content::text(self.to_string())]
    }
}