diff --git a/src/cli/commands/search.rs b/src/cli/commands/search.rs index 56f358c..af02838 100644 --- a/src/cli/commands/search.rs +++ b/src/cli/commands/search.rs @@ -1,3 +1,5 @@ +use std::collections::HashMap; + use console::style; use serde::Serialize; @@ -8,9 +10,10 @@ use crate::core::paths::get_db_path; use crate::core::project::resolve_project; use crate::core::time::{ms_to_iso, parse_since}; use crate::documents::SourceType; +use crate::embedding::ollama::{OllamaClient, OllamaConfig}; use crate::search::{ - FtsQueryMode, PathFilter, SearchFilters, apply_filters, get_result_snippet, rank_rrf, - search_fts, + FtsQueryMode, HybridResult, PathFilter, SearchFilters, SearchMode, get_result_snippet, + search_fts, search_hybrid, }; #[derive(Debug, Serialize)] @@ -58,7 +61,7 @@ pub struct SearchCliFilters { pub limit: usize, } -pub fn run_search( +pub async fn run_search( config: &Config, query: &str, cli_filters: SearchCliFilters, @@ -71,15 +74,18 @@ pub fn run_search( let mut warnings: Vec = Vec::new(); - // Determine actual mode: vector search requires embeddings, which need async + Ollama. - // Until hybrid/semantic are wired up, we run lexical and warn if the user asked for more. - let actual_mode = "lexical"; - if requested_mode != "lexical" { - warnings.push(format!( - "Requested mode '{}' is not yet available; falling back to lexical search.", - requested_mode - )); - } + let actual_mode = SearchMode::parse(requested_mode).unwrap_or(SearchMode::Hybrid); + + let client = if actual_mode != SearchMode::Lexical { + let ollama_cfg = &config.embedding; + Some(OllamaClient::new(OllamaConfig { + base_url: ollama_cfg.base_url.clone(), + model: ollama_cfg.model.clone(), + ..OllamaConfig::default() + })) + } else { + None + }; let doc_count: i64 = conn .query_row("SELECT COUNT(*) FROM documents", [], |row| row.get(0)) @@ -89,7 +95,7 @@ pub fn run_search( warnings.push("No documents indexed. Run 'lore generate-docs' first.".to_string()); return Ok(SearchResponse { query: query.to_string(), - mode: actual_mode.to_string(), + mode: actual_mode.as_str().to_string(), total_results: 0, results: vec![], warnings, @@ -151,52 +157,54 @@ pub fn run_search( limit: cli_filters.limit, }; - let requested = filters.clamp_limit(); - let top_k = if filters.has_any_filter() { - (requested * 50).clamp(200, 1500) - } else { - (requested * 10).clamp(50, 1500) - }; - - let fts_results = search_fts(&conn, query, top_k, fts_mode)?; - let fts_tuples: Vec<(i64, f64)> = fts_results - .iter() - .map(|r| (r.document_id, r.bm25_score)) - .collect(); - - let snippet_map: std::collections::HashMap = fts_results + // Run FTS separately for snippet extraction (search_hybrid doesn't return snippets). + let snippet_top_k = filters + .clamp_limit() + .checked_mul(10) + .unwrap_or(500) + .clamp(50, 1500); + let fts_results = search_fts(&conn, query, snippet_top_k, fts_mode)?; + let snippet_map: HashMap = fts_results .iter() .map(|r| (r.document_id, r.snippet.clone())) .collect(); - let ranked = rank_rrf(&[], &fts_tuples); - let ranked_ids: Vec = ranked.iter().map(|r| r.document_id).collect(); + // search_hybrid handles recall sizing, RRF ranking, and filter application internally. + let (hybrid_results, mut hybrid_warnings) = search_hybrid( + &conn, + client.as_ref(), + query, + actual_mode, + &filters, + fts_mode, + ) + .await?; + warnings.append(&mut hybrid_warnings); - let filtered_ids = apply_filters(&conn, &ranked_ids, &filters)?; - - if filtered_ids.is_empty() { + if hybrid_results.is_empty() { return Ok(SearchResponse { query: query.to_string(), - mode: actual_mode.to_string(), + mode: actual_mode.as_str().to_string(), total_results: 0, results: vec![], warnings, }); } - let hydrated = hydrate_results(&conn, &filtered_ids)?; + let ranked_ids: Vec = hybrid_results.iter().map(|r| r.document_id).collect(); + let hydrated = hydrate_results(&conn, &ranked_ids)?; - let rrf_map: std::collections::HashMap = - ranked.iter().map(|r| (r.document_id, r)).collect(); + let hybrid_map: HashMap = + hybrid_results.iter().map(|r| (r.document_id, r)).collect(); let mut results: Vec = Vec::with_capacity(hydrated.len()); for row in &hydrated { - let rrf = rrf_map.get(&row.document_id); + let hr = hybrid_map.get(&row.document_id); let fts_snippet = snippet_map.get(&row.document_id).map(|s| s.as_str()); let snippet = get_result_snippet(fts_snippet, &row.content_text); let explain_data = if explain { - rrf.map(|r| ExplainData { + hr.map(|r| ExplainData { vector_rank: r.vector_rank, fts_rank: r.fts_rank, rrf_score: r.rrf_score, @@ -217,14 +225,14 @@ pub fn run_search( labels: row.labels.clone(), paths: row.paths.clone(), snippet, - score: rrf.map(|r| r.normalized_score).unwrap_or(0.0), + score: hr.map(|r| r.score).unwrap_or(0.0), explain: explain_data, }); } Ok(SearchResponse { query: query.to_string(), - mode: actual_mode.to_string(), + mode: actual_mode.as_str().to_string(), total_results: results.len(), results, warnings, @@ -360,8 +368,12 @@ pub fn print_search_results(response: &SearchResponse) { if let Some(ref explain) = result.explain { println!( - " {} fts_rank={} rrf_score={:.6}", + " {} vector_rank={} fts_rank={} rrf_score={:.6}", style("[explain]").magenta(), + explain + .vector_rank + .map(|r| r.to_string()) + .unwrap_or_else(|| "-".into()), explain .fts_rank .map(|r| r.to_string()) diff --git a/src/main.rs b/src/main.rs index 8f7ad96..fa5ec2e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1778,7 +1778,8 @@ async fn handle_search( fts_mode, &args.mode, explain, - )?; + ) + .await?; let elapsed_ms = start.elapsed().as_millis() as u64; if robot_mode {