sqlite_graphrag/commands/
hybrid_search.rs

1use crate::cli::MemoryType;
2use crate::errors::AppError;
3use crate::output::{self, JsonOutputFormat, RecallItem};
4use crate::paths::AppPaths;
5use crate::storage::connection::open_ro;
6use crate::storage::memories;
7
8use std::collections::HashMap;
9
10/// Arguments for the `hybrid-search` subcommand.
11///
12/// When `--namespace` is omitted the search runs against the `global` namespace,
13/// which is the default namespace used by `remember` when no `--namespace` flag
14/// is provided. Pass an explicit `--namespace` value to search a different
15/// isolated namespace.
16#[derive(clap::Args)]
17pub struct HybridSearchArgs {
18    pub query: String,
19    #[arg(short = 'k', long, default_value = "10")]
20    pub k: usize,
21    #[arg(long, default_value = "60")]
22    pub rrf_k: u32,
23    #[arg(long, default_value = "1.0")]
24    pub weight_vec: f32,
25    #[arg(long, default_value = "1.0")]
26    pub weight_fts: f32,
27    /// Filter by memory.type. Note: distinct from graph entity_type
28    /// (project/tool/person/file/concept/incident/decision/memory/dashboard/issue_tracker)
29    /// used in --entities-file.
30    #[arg(long, value_enum)]
31    pub r#type: Option<MemoryType>,
32    #[arg(long)]
33    pub namespace: Option<String>,
34    #[arg(long)]
35    pub with_graph: bool,
36    #[arg(long, default_value = "2")]
37    pub max_hops: u32,
38    #[arg(long, default_value = "0.3")]
39    pub min_weight: f64,
40    #[arg(long, value_enum, default_value_t = JsonOutputFormat::Json)]
41    pub format: JsonOutputFormat,
42    #[arg(long, env = "SQLITE_GRAPHRAG_DB_PATH")]
43    pub db: Option<String>,
44    /// Accept `--json` as a no-op because output is already JSON by default.
45    #[arg(long, help = "No-op; JSON is always emitted on stdout")]
46    pub json: bool,
47}
48
49#[derive(serde::Serialize)]
50pub struct HybridSearchItem {
51    pub memory_id: i64,
52    pub name: String,
53    pub namespace: String,
54    #[serde(rename = "type")]
55    pub memory_type: String,
56    pub description: String,
57    pub body: String,
58    pub combined_score: f64,
59    /// Alias de `combined_score` para contrato documentado em SKILL.md.
60    pub score: f64,
61    /// Fonte do match: sempre "hybrid" (RRF de vec + fts). Adicionado em v2.0.1.
62    pub source: String,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub vec_rank: Option<usize>,
65    #[serde(skip_serializing_if = "Option::is_none")]
66    pub fts_rank: Option<usize>,
67    /// Score RRF combinado — alias explícito de `combined_score` para contratos de integração.
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub rrf_score: Option<f64>,
70}
71
72/// Pesos RRF usados na busca híbrida: vec (vetorial) e fts (texto).
73#[derive(serde::Serialize)]
74pub struct Weights {
75    pub vec: f32,
76    pub fts: f32,
77}
78
79#[derive(serde::Serialize)]
80pub struct HybridSearchResponse {
81    pub query: String,
82    pub k: usize,
83    /// Parâmetro k do RRF usado no ranking combinado.
84    pub rrf_k: u32,
85    /// Pesos aplicados às fontes vec e fts no RRF.
86    pub weights: Weights,
87    pub results: Vec<HybridSearchItem>,
88    pub graph_matches: Vec<RecallItem>,
89    /// Tempo total de execução em milissegundos desde início do handler até serialização.
90    pub elapsed_ms: u64,
91}
92
93pub fn run(args: HybridSearchArgs) -> Result<(), AppError> {
94    let start = std::time::Instant::now();
95    let _ = args.format;
96
97    let namespace = crate::namespace::resolve_namespace(args.namespace.as_deref())?;
98    let paths = AppPaths::resolve(args.db.as_deref())?;
99    if !paths.db.exists() {
100        return Err(AppError::NotFound(
101            crate::i18n::erros::banco_nao_encontrado(&paths.db.display().to_string()),
102        ));
103    }
104
105    output::emit_progress_i18n(
106        "Computing query embedding...",
107        "Calculando embedding da consulta...",
108    );
109    let embedding = crate::daemon::embed_query_or_local(&paths.models, &args.query)?;
110
111    let conn = open_ro(&paths.db)?;
112
113    let memory_type_str = args.r#type.map(|t| t.as_str());
114
115    let vec_results =
116        memories::knn_search(&conn, &embedding, &namespace, memory_type_str, args.k * 2)?;
117
118    // Mapear posição de ranking vetorial por memory_id (1-indexed conforme schema)
119    let vec_rank_map: HashMap<i64, usize> = vec_results
120        .iter()
121        .enumerate()
122        .map(|(pos, (id, _))| (*id, pos + 1))
123        .collect();
124
125    let fts_results =
126        memories::fts_search(&conn, &args.query, &namespace, memory_type_str, args.k * 2)?;
127
128    // Mapear posição de ranking FTS por memory_id (1-indexed conforme schema)
129    let fts_rank_map: HashMap<i64, usize> = fts_results
130        .iter()
131        .enumerate()
132        .map(|(pos, row)| (row.id, pos + 1))
133        .collect();
134
135    let rrf_k = args.rrf_k as f64;
136
137    // Acumular scores RRF combinados
138    let mut combined_scores: HashMap<i64, f64> = HashMap::new();
139
140    for (rank, (memory_id, _)) in vec_results.iter().enumerate() {
141        let score = args.weight_vec as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
142        *combined_scores.entry(*memory_id).or_insert(0.0) += score;
143    }
144
145    for (rank, row) in fts_results.iter().enumerate() {
146        let score = args.weight_fts as f64 * (1.0 / (rrf_k + rank as f64 + 1.0));
147        *combined_scores.entry(row.id).or_insert(0.0) += score;
148    }
149
150    // Ordenar por score descendente e tomar os top-k
151    let mut ranked: Vec<(i64, f64)> = combined_scores.into_iter().collect();
152    ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
153    ranked.truncate(args.k);
154
155    // Coletar todos os IDs para busca em batch (evitar N+1)
156    let top_ids: Vec<i64> = ranked.iter().map(|(id, _)| *id).collect();
157
158    // Buscar dados completos das memórias top
159    let mut memory_data: HashMap<i64, memories::MemoryRow> = HashMap::new();
160    for id in &top_ids {
161        if let Some(row) = memories::read_full(&conn, *id)? {
162            memory_data.insert(*id, row);
163        }
164    }
165
166    // Construir resultados finais na ordem de ranking
167    let results: Vec<HybridSearchItem> = ranked
168        .into_iter()
169        .filter_map(|(memory_id, combined_score)| {
170            memory_data.remove(&memory_id).map(|row| HybridSearchItem {
171                memory_id: row.id,
172                name: row.name,
173                namespace: row.namespace,
174                memory_type: row.memory_type,
175                description: row.description,
176                body: row.body,
177                combined_score,
178                score: combined_score,
179                source: "hybrid".to_string(),
180                vec_rank: vec_rank_map.get(&memory_id).copied(),
181                fts_rank: fts_rank_map.get(&memory_id).copied(),
182                rrf_score: Some(combined_score),
183            })
184        })
185        .collect();
186
187    output::emit_json(&HybridSearchResponse {
188        query: args.query,
189        k: args.k,
190        rrf_k: args.rrf_k,
191        weights: Weights {
192            vec: args.weight_vec,
193            fts: args.weight_fts,
194        },
195        results,
196        graph_matches: vec![],
197        elapsed_ms: start.elapsed().as_millis() as u64,
198    })?;
199
200    Ok(())
201}
202
203#[cfg(test)]
204mod testes {
205    use super::*;
206
207    fn resposta_vazia(
208        k: usize,
209        rrf_k: u32,
210        weight_vec: f32,
211        weight_fts: f32,
212    ) -> HybridSearchResponse {
213        HybridSearchResponse {
214            query: "busca teste".to_string(),
215            k,
216            rrf_k,
217            weights: Weights {
218                vec: weight_vec,
219                fts: weight_fts,
220            },
221            results: vec![],
222            graph_matches: vec![],
223            elapsed_ms: 0,
224        }
225    }
226
227    #[test]
228    fn hybrid_search_response_vazia_serializa_campos_corretos() {
229        let resp = resposta_vazia(10, 60, 1.0, 1.0);
230        let json = serde_json::to_string(&resp).unwrap();
231        assert!(json.contains("\"results\""), "deve conter campo results");
232        assert!(json.contains("\"query\""), "deve conter campo query");
233        assert!(json.contains("\"k\""), "deve conter campo k");
234        assert!(
235            json.contains("\"graph_matches\""),
236            "deve conter campo graph_matches"
237        );
238        assert!(
239            !json.contains("\"combined_rank\""),
240            "NÃO deve conter combined_rank"
241        );
242        assert!(
243            !json.contains("\"vec_rank_list\""),
244            "NÃO deve conter vec_rank_list"
245        );
246        assert!(
247            !json.contains("\"fts_rank_list\""),
248            "NÃO deve conter fts_rank_list"
249        );
250    }
251
252    #[test]
253    fn hybrid_search_response_serializa_rrf_k_e_weights() {
254        let resp = resposta_vazia(5, 60, 0.7, 0.3);
255        let json = serde_json::to_string(&resp).unwrap();
256        assert!(json.contains("\"rrf_k\""), "deve conter campo rrf_k");
257        assert!(json.contains("\"weights\""), "deve conter campo weights");
258        assert!(json.contains("\"vec\""), "deve conter campo weights.vec");
259        assert!(json.contains("\"fts\""), "deve conter campo weights.fts");
260    }
261
262    #[test]
263    fn hybrid_search_response_serializa_elapsed_ms() {
264        let mut resp = resposta_vazia(5, 60, 1.0, 1.0);
265        resp.elapsed_ms = 123;
266        let json = serde_json::to_string(&resp).unwrap();
267        assert!(
268            json.contains("\"elapsed_ms\""),
269            "deve conter campo elapsed_ms"
270        );
271        assert!(json.contains("123"), "deve serializar valor de elapsed_ms");
272    }
273
274    #[test]
275    fn weights_struct_serializa_corretamente() {
276        let w = Weights { vec: 0.6, fts: 0.4 };
277        let json = serde_json::to_string(&w).unwrap();
278        assert!(json.contains("\"vec\""));
279        assert!(json.contains("\"fts\""));
280    }
281
282    #[test]
283    fn hybrid_search_item_omite_fts_rank_quando_none() {
284        let item = HybridSearchItem {
285            memory_id: 1,
286            name: "mem".to_string(),
287            namespace: "default".to_string(),
288            memory_type: "user".to_string(),
289            description: "desc".to_string(),
290            body: "conteúdo".to_string(),
291            combined_score: 0.0328,
292            score: 0.0328,
293            source: "hybrid".to_string(),
294            vec_rank: Some(1),
295            fts_rank: None,
296            rrf_score: Some(0.0328),
297        };
298        let json = serde_json::to_string(&item).unwrap();
299        assert!(
300            json.contains("\"vec_rank\""),
301            "deve conter vec_rank quando Some"
302        );
303        assert!(
304            !json.contains("\"fts_rank\""),
305            "NÃO deve conter fts_rank quando None"
306        );
307    }
308
309    #[test]
310    fn hybrid_search_item_omite_vec_rank_quando_none() {
311        let item = HybridSearchItem {
312            memory_id: 2,
313            name: "mem2".to_string(),
314            namespace: "default".to_string(),
315            memory_type: "fact".to_string(),
316            description: "desc2".to_string(),
317            body: "corpo2".to_string(),
318            combined_score: 0.016,
319            score: 0.016,
320            source: "hybrid".to_string(),
321            vec_rank: None,
322            fts_rank: Some(2),
323            rrf_score: Some(0.016),
324        };
325        let json = serde_json::to_string(&item).unwrap();
326        assert!(
327            !json.contains("\"vec_rank\""),
328            "NÃO deve conter vec_rank quando None"
329        );
330        assert!(
331            json.contains("\"fts_rank\""),
332            "deve conter fts_rank quando Some"
333        );
334    }
335
336    #[test]
337    fn hybrid_search_item_serializa_ambos_ranks_quando_some() {
338        let item = HybridSearchItem {
339            memory_id: 3,
340            name: "mem3".to_string(),
341            namespace: "ns".to_string(),
342            memory_type: "entity".to_string(),
343            description: "desc3".to_string(),
344            body: "corpo3".to_string(),
345            combined_score: 0.05,
346            score: 0.05,
347            source: "hybrid".to_string(),
348            vec_rank: Some(3),
349            fts_rank: Some(1),
350            rrf_score: Some(0.05),
351        };
352        let json = serde_json::to_string(&item).unwrap();
353        assert!(json.contains("\"vec_rank\""), "deve conter vec_rank");
354        assert!(json.contains("\"fts_rank\""), "deve conter fts_rank");
355        assert!(json.contains("\"type\""), "deve serializar type renomeado");
356        assert!(!json.contains("memory_type"), "NÃO deve expor memory_type");
357    }
358
359    #[test]
360    fn hybrid_search_response_serializa_k_corretamente() {
361        let resp = resposta_vazia(5, 60, 1.0, 1.0);
362        let json = serde_json::to_string(&resp).unwrap();
363        assert!(json.contains("\"k\":5"), "deve serializar k=5");
364    }
365}
sqlite_graphrag/commands/hybrid_search.rs

sqlite_graphrag/commands/
hybrid_search.rs
