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:
@@ -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())
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
Reference in New Issue
Block a user