mcp_server_sqlite/tools/
search_fts_tool.rs1use rmcp::model::{Content, IntoContents};
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8use super::ToolError;
9use crate::{mcp::McpServerSqlite, traits::SqliteServerTool};
10
11#[derive(
12 Clone,
13 Copy,
14 Debug,
15 PartialEq,
16 Eq,
17 PartialOrd,
18 Ord,
19 Hash,
20 Default,
21 Serialize,
22 Deserialize,
23 JsonSchema,
24)]
25pub struct SearchFtsTool;
30
31impl SqliteServerTool for SearchFtsTool {
32 const NAME: &str = "search_fts";
33
34 type Context = McpServerSqlite;
35 type Error = ToolError<SearchFtsError>;
36
37 type Input = SearchFtsInput;
38 type Output = SearchFtsOutput;
39
40 fn handle(
41 ctx: &Self::Context,
42 input: Self::Input,
43 ) -> Result<Self::Output, Self::Error> {
44 let conn = ctx
45 .connection()
46 .map_err(|source| ToolError::Connection { source })?;
47
48 let limit = input.limit.unwrap_or(10);
49 let snippet_tokens = input.snippet_tokens.unwrap_or(32);
50 let hl_start = input.highlight_start.as_deref().unwrap_or("<b>");
51 let hl_end = input.highlight_end.as_deref().unwrap_or("</b>");
52
53 let column_count =
54 fts_column_count(&conn, &input.fts_table).map_err(|source| {
55 ToolError::Tool(SearchFtsError::Query { source })
56 })?;
57
58 let snippet_exprs = (0..column_count)
59 .map(|i| {
60 format!(
61 "snippet({tbl}, {i}, '{hl_start}', '{hl_end}', '...', {tokens})",
62 tbl = input.fts_table,
63 tokens = snippet_tokens,
64 )
65 })
66 .collect::<Vec<_>>()
67 .join(", ");
68
69 let sql = format!(
70 "SELECT rowid, rank, {snippets} \
71 FROM [{tbl}] \
72 WHERE [{tbl}] MATCH ?1 \
73 ORDER BY rank \
74 LIMIT ?2",
75 tbl = input.fts_table,
76 snippets = snippet_exprs,
77 );
78
79 let mut stmt = conn.prepare(&sql).map_err(|source| {
80 ToolError::Tool(SearchFtsError::Query { source })
81 })?;
82
83 let results = stmt
84 .query_map(rusqlite::params![input.query, limit], |row| {
85 let rowid = row.get::<_, i64>(0)?;
86 let rank = row.get::<_, f64>(1)?;
87 let snippets = (0..column_count)
88 .map(|i| row.get::<_, String>(2 + i))
89 .collect::<Result<Vec<_>, _>>()?;
90 Ok(FtsMatch {
91 rowid,
92 rank,
93 snippets,
94 })
95 })
96 .map_err(|source| {
97 ToolError::Tool(SearchFtsError::Query { source })
98 })?
99 .collect::<Result<Vec<_>, _>>()
100 .map_err(|source| {
101 ToolError::Tool(SearchFtsError::Query { source })
102 })?;
103
104 Ok(SearchFtsOutput { results })
105 }
106}
107
108fn fts_column_count(
110 conn: &rusqlite::Connection,
111 fts_table: &str,
112) -> Result<usize, rusqlite::Error> {
113 let mut stmt =
114 conn.prepare(&format!("PRAGMA table_info([{}])", fts_table))?;
115 let count = stmt.query_map([], |_| Ok(()))?.count();
116 Ok(count)
117}
118
119#[derive(
121 Clone,
122 Debug,
123 PartialEq,
124 Eq,
125 PartialOrd,
126 Ord,
127 Hash,
128 Serialize,
129 Deserialize,
130 schemars::JsonSchema,
131)]
132pub struct SearchFtsInput {
133 #[schemars(description = "The FTS5 virtual table to search")]
135 pub fts_table: String,
136 #[schemars(description = "The FTS5 MATCH query string")]
138 pub query: String,
139 #[schemars(description = "Maximum number of results (default 10)")]
141 pub limit: Option<i64>,
142 #[schemars(description = "Max tokens per snippet (default 32)")]
144 pub snippet_tokens: Option<i32>,
145 #[schemars(
148 description = "String before matched terms in snippets (default <b>)"
149 )]
150 pub highlight_start: Option<String>,
151 #[schemars(
154 description = "String after matched terms in snippets (default </b>)"
155 )]
156 pub highlight_end: Option<String>,
157}
158
159#[derive(
161 Clone,
162 Debug,
163 PartialEq,
164 PartialOrd,
165 Serialize,
166 Deserialize,
167 schemars::JsonSchema,
168)]
169pub struct SearchFtsOutput {
170 pub results: Vec<FtsMatch>,
172}
173
174#[derive(
176 Clone,
177 Debug,
178 PartialEq,
179 PartialOrd,
180 Serialize,
181 Deserialize,
182 schemars::JsonSchema,
183)]
184pub struct FtsMatch {
185 pub rowid: i64,
187 pub rank: f64,
189 pub snippets: Vec<String>,
192}
193
194#[derive(Debug, thiserror::Error)]
196pub enum SearchFtsError {
197 #[error("FTS search failed: {source}")]
199 Query {
200 source: rusqlite::Error,
202 },
203}
204
205impl IntoContents for SearchFtsError {
207 fn into_contents(self) -> Vec<Content> {
208 vec![Content::text(self.to_string())]
209 }
210}