feat(bd-1ksf): wire hybrid search (FTS5 + vector + RRF) to CLI

Make run_search async, replace hardcoded lexical mode with SearchMode::parse(),
wire search_hybrid() with OllamaClient for semantic/hybrid modes, graceful
degradation when Ollama unavailable.

Closes: bd-1ksf
This commit is contained in:
teernisse
2026-02-12 11:34:10 -05:00
parent 47eecce8e9
commit ecbfef537a
2 changed files with 55 additions and 42 deletions

View File

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

View File

@@ -1778,7 +1778,8 @@ async fn handle_search(
fts_mode, fts_mode,
&args.mode, &args.mode,
explain, explain,
)?; )
.await?;
let elapsed_ms = start.elapsed().as_millis() as u64; let elapsed_ms = start.elapsed().as_millis() as u64;
if robot_mode { if robot_mode {