From e1691a8d429e178d5708f60200fad43e4786e93f Mon Sep 17 00:00:00 2001 From: alexthebuildr <116134064+ztsalexey@users.noreply.github.com> Date: Thu, 12 Mar 2026 15:49:00 -0600 Subject: [PATCH] feat: configurable hybrid search fusion strategy (#234) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: configurable hybrid search fusion strategy (#169) Add WeightedScore fusion as an alternative to the default RRF algorithm. Users can now tune search behavior via env vars (SEARCH_FUSION_STRATEGY, SEARCH_FTS_WEIGHT, SEARCH_VECTOR_WEIGHT, SEARCH_RRF_K) or by passing SearchConfig with the new fields. Default behavior (RRF, k=60) is unchanged. - Add FusionStrategy enum (Rrf/WeightedScore) to workspace::search - Add weighted_score_fusion() and fuse_results() dispatcher - Add config/search.rs with WorkspaceSearchConfig from env vars - Wire search defaults through Workspace struct - Update both postgres and libsql backends to use fuse_results() - Add 7 new tests (4 fusion + 3 config) Co-Authored-By: Claude Opus 4.6 * fix: swap default search weights to match issue #169 spec (0.7 vector / 0.3 FTS) The issue spec says "0.7/0.3 (vector/keyword) for weighted mode" but our defaults had fts_weight=0.7, vector_weight=0.3 (inverted). Also fixes the misleading docstring on weighted_score_fusion that claimed 1/rank normalizes to [0,1]. Co-Authored-By: Claude Opus 4.6 * fix: validate weight inputs and update stale doc comments - Reject NaN, infinite, and negative values for SEARCH_FTS_WEIGHT and SEARCH_VECTOR_WEIGHT with a clear ConfigError - Fix module-level docs that incorrectly claimed WeightedScore "normalizes per-method scores to [0,1]" - Update SearchResult.score doc from "Combined RRF score" to strategy-agnostic "Combined fusion score" Co-Authored-By: Claude Opus 4.6 * fix: validate weight setters against NaN/inf/negative values with_fts_weight() and with_vector_weight() now silently ignore non-finite (NaN, ±inf) and negative values, matching the env var validation already in place for SEARCH_FTS_WEIGHT / SEARCH_VECTOR_WEIGHT. Values > 1.0 remain valid since weights are normalized internally. Co-Authored-By: Claude Opus 4.6 * fix: use crate-wide ENV_MUTEX in search config tests Replace the module-local `ENV_MUTEX` in `search.rs` with a shared `crate::config::helpers::ENV_MUTEX` to prevent cross-module env races when `cargo test` runs tests in parallel. Addresses copilot review comment. Tracked in #245. Co-Authored-By: Claude Opus 4.6 * fix: per-strategy weight defaults to match issue #169 spec RRF mode now defaults to 0.5/0.5 (fts/vector) and WeightedScore defaults to 0.3/0.7, matching the acceptance criteria in #169. Previously both modes used 0.3/0.7 uniformly. Co-Authored-By: Claude Opus 4.6 * fix: reject both weights=0 in weighted fusion mode When both SEARCH_FTS_WEIGHT and SEARCH_VECTOR_WEIGHT are 0.0 under WeightedScore strategy, all scores would be 0.0, producing arbitrary ordering. RRF mode is unaffected since it ignores weights entirely. Addresses Copilot review comment. The other comment (rrf_k=0 division by zero) is a false positive — ranks are 1-based, so k=0 just gives inverse-rank scoring with no infinity. Co-Authored-By: Claude Opus 4.6 * fix: clarify weight doc comments and error key - SearchConfig field docs: clarify that Default always uses 0.5, per-strategy defaults only apply via WorkspaceSearchConfig::resolve() - WorkspaceSearchConfig field docs: same clarification - Error key for both-weights-zero now references both env vars Co-Authored-By: Claude Opus 4.6 * fix: remove broken intra-doc links to pub(crate) resolve() WorkspaceSearchConfig::resolve is pub(crate), so linking to it from public field docs triggers rustdoc private_intra_doc_links warnings. Switch to plain-text references. Co-Authored-By: Claude Opus 4.6 * fix: add document_path to weighted_score_fusion results The weighted_score_fusion function was missing the document_path field added in a recent main branch commit, causing a compile error after rebase. Co-Authored-By: Claude Opus 4.6 * chore: trigger CI re-check after rebase * fix: resolve pre-existing staging fmt and clippy issues - Fix import ordering in cli/mod.rs (cargo fmt) - Fix line wrapping in tools/mcp/auth.rs (cargo fmt) - Move path_routing_tests before MemoryTreeTool to fix clippy::items_after_test_module [skip-regression-check] * fix: remove duplicate path_routing_tests module after rebase [skip-regression-check] --------- Co-authored-by: Claude Opus 4.6 --- src/app.rs | 3 +- src/config/mod.rs | 5 + src/config/search.rs | 211 ++++++++++++++++++++++++ src/db/libsql/workspace.rs | 4 +- src/tools/builtin/memory.rs | 40 ++--- src/workspace/mod.rs | 22 ++- src/workspace/repository.rs | 4 +- src/workspace/search.rs | 321 +++++++++++++++++++++++++++++++++++- 8 files changed, 575 insertions(+), 35 deletions(-) create mode 100644 src/config/search.rs diff --git a/src/app.rs b/src/app.rs index 22416cfc..59890bfc 100644 --- a/src/app.rs +++ b/src/app.rs @@ -303,7 +303,8 @@ impl AppBuilder { // Register memory tools if database is available let workspace = if let Some(ref db) = self.db { - let mut ws = Workspace::new_with_db("default", db.clone()); + let mut ws = Workspace::new_with_db("default", db.clone()) + .with_search_config(&self.config.search); if let Some(ref emb) = embeddings { ws = ws.with_embeddings(emb.clone()); } diff --git a/src/config/mod.rs b/src/config/mod.rs index afc54372..34c34423 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -18,6 +18,7 @@ pub mod relay; mod routines; mod safety; mod sandbox; +mod search; mod secrets; mod skills; mod transcription; @@ -44,6 +45,7 @@ pub use self::routines::RoutineConfig; pub use self::safety::SafetyConfig; use self::safety::resolve_safety_config; pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig}; +pub use self::search::WorkspaceSearchConfig; pub use self::secrets::SecretsConfig; pub use self::skills::SkillsConfig; pub use self::transcription::TranscriptionConfig; @@ -91,6 +93,7 @@ pub struct Config { pub claude_code: ClaudeCodeConfig, pub skills: SkillsConfig, pub transcription: TranscriptionConfig, + pub search: WorkspaceSearchConfig, pub observability: crate::observability::ObservabilityConfig, /// Channel-relay integration (Slack via external relay service). /// Present only when both `CHANNEL_RELAY_URL` and `CHANNEL_RELAY_API_KEY` are set. @@ -166,6 +169,7 @@ impl Config { ..SkillsConfig::default() }, transcription: TranscriptionConfig::default(), + search: WorkspaceSearchConfig::default(), observability: crate::observability::ObservabilityConfig::default(), relay: None, } @@ -318,6 +322,7 @@ impl Config { claude_code: ClaudeCodeConfig::resolve()?, skills: SkillsConfig::resolve()?, transcription: TranscriptionConfig::resolve(settings)?, + search: WorkspaceSearchConfig::resolve()?, observability: crate::observability::ObservabilityConfig { backend: std::env::var("OBSERVABILITY_BACKEND").unwrap_or_else(|_| "none".into()), }, diff --git a/src/config/search.rs b/src/config/search.rs new file mode 100644 index 00000000..9555fecc --- /dev/null +++ b/src/config/search.rs @@ -0,0 +1,211 @@ +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; +use crate::workspace::FusionStrategy; + +/// Workspace search configuration resolved from environment variables. +#[derive(Debug, Clone)] +pub struct WorkspaceSearchConfig { + /// Fusion strategy: "rrf" or "weighted". + pub fusion_strategy: FusionStrategy, + /// RRF constant k (default 60). + pub rrf_k: u32, + /// FTS weight for fusion. + /// + /// [`Default`] uses 0.5. When the configuration is resolved, per-strategy + /// defaults are applied: 0.5 (RRF) or 0.3 (weighted). + pub fts_weight: f32, + /// Vector weight for fusion. + /// + /// [`Default`] uses 0.5. When the configuration is resolved, per-strategy + /// defaults are applied: 0.5 (RRF) or 0.7 (weighted). + pub vector_weight: f32, +} + +impl Default for WorkspaceSearchConfig { + fn default() -> Self { + Self { + fusion_strategy: FusionStrategy::default(), + rrf_k: 60, + fts_weight: 0.5, + vector_weight: 0.5, + } + } +} + +impl WorkspaceSearchConfig { + pub(crate) fn resolve() -> Result { + let fusion_strategy = match optional_env("SEARCH_FUSION_STRATEGY")? { + Some(s) => match s.to_lowercase().as_str() { + "rrf" => FusionStrategy::Rrf, + "weighted" => FusionStrategy::WeightedScore, + other => { + return Err(ConfigError::InvalidValue { + key: "SEARCH_FUSION_STRATEGY".to_string(), + message: format!("must be 'rrf' or 'weighted', got '{other}'"), + }); + } + }, + None => FusionStrategy::default(), + }; + + let rrf_k = parse_optional_env("SEARCH_RRF_K", 60u32)?; + + // Per-strategy weight defaults: RRF uses 0.5/0.5, weighted uses 0.3/0.7 (vector-biased). + let (default_fts, default_vec) = match fusion_strategy { + FusionStrategy::Rrf => (0.5f32, 0.5f32), + FusionStrategy::WeightedScore => (0.3f32, 0.7f32), + }; + let fts_weight = parse_optional_env("SEARCH_FTS_WEIGHT", default_fts)?; + let vector_weight = parse_optional_env("SEARCH_VECTOR_WEIGHT", default_vec)?; + + if !fts_weight.is_finite() || fts_weight < 0.0 { + return Err(ConfigError::InvalidValue { + key: "SEARCH_FTS_WEIGHT".to_string(), + message: "must be a finite, non-negative float".to_string(), + }); + } + if !vector_weight.is_finite() || vector_weight < 0.0 { + return Err(ConfigError::InvalidValue { + key: "SEARCH_VECTOR_WEIGHT".to_string(), + message: "must be a finite, non-negative float".to_string(), + }); + } + if matches!(fusion_strategy, FusionStrategy::WeightedScore) + && fts_weight == 0.0 + && vector_weight == 0.0 + { + return Err(ConfigError::InvalidValue { + key: "SEARCH_FTS_WEIGHT/SEARCH_VECTOR_WEIGHT".to_string(), + message: "weighted fusion requires at least one non-zero weight".to_string(), + }); + } + + Ok(Self { + fusion_strategy, + rrf_k, + fts_weight, + vector_weight, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::helpers::ENV_MUTEX; + + fn clear_search_env() { + // SAFETY: Only called under ENV_MUTEX in tests. + unsafe { + std::env::remove_var("SEARCH_FUSION_STRATEGY"); + std::env::remove_var("SEARCH_RRF_K"); + std::env::remove_var("SEARCH_FTS_WEIGHT"); + std::env::remove_var("SEARCH_VECTOR_WEIGHT"); + } + } + + #[test] + fn defaults_when_no_env() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::Rrf); + assert_eq!(config.rrf_k, 60); + assert!((config.fts_weight - 0.5).abs() < 0.001); + assert!((config.vector_weight - 0.5).abs() < 0.001); + } + + #[test] + fn env_overrides() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "weighted"); + std::env::set_var("SEARCH_RRF_K", "30"); + std::env::set_var("SEARCH_FTS_WEIGHT", "0.9"); + std::env::set_var("SEARCH_VECTOR_WEIGHT", "0.1"); + } + + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::WeightedScore); + assert_eq!(config.rrf_k, 30); + assert!((config.fts_weight - 0.9).abs() < 0.001); + assert!((config.vector_weight - 0.1).abs() < 0.001); + + clear_search_env(); + } + + #[test] + fn invalid_strategy_rejected() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "bm25"); + } + + let result = WorkspaceSearchConfig::resolve(); + assert!(result.is_err()); + + clear_search_env(); + } + + #[test] + fn weighted_strategy_defaults() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "weighted"); + } + + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::WeightedScore); + // Weighted mode should default to 0.3 FTS / 0.7 vector + assert!((config.fts_weight - 0.3).abs() < 0.001); + assert!((config.vector_weight - 0.7).abs() < 0.001); + + clear_search_env(); + } + + #[test] + fn weighted_both_zero_rejected() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FUSION_STRATEGY", "weighted"); + std::env::set_var("SEARCH_FTS_WEIGHT", "0.0"); + std::env::set_var("SEARCH_VECTOR_WEIGHT", "0.0"); + } + + let result = WorkspaceSearchConfig::resolve(); + assert!(result.is_err()); + + clear_search_env(); + } + + #[test] + fn rrf_both_zero_allowed() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_search_env(); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("SEARCH_FTS_WEIGHT", "0.0"); + std::env::set_var("SEARCH_VECTOR_WEIGHT", "0.0"); + } + + // RRF ignores weights, so both=0 is fine + let config = WorkspaceSearchConfig::resolve().expect("should resolve"); + assert_eq!(config.fusion_strategy, FusionStrategy::Rrf); + + clear_search_env(); + } +} diff --git a/src/db/libsql/workspace.rs b/src/db/libsql/workspace.rs index 19000404..68bd58ba 100644 --- a/src/db/libsql/workspace.rs +++ b/src/db/libsql/workspace.rs @@ -14,7 +14,7 @@ use crate::db::WorkspaceStore; use crate::error::WorkspaceError; use crate::workspace::{ MemoryChunk, MemoryDocument, RankedResult, SearchConfig, SearchResult, WorkspaceEntry, - reciprocal_rank_fusion, + fuse_results, }; use chrono::Utc; @@ -614,6 +614,6 @@ impl WorkspaceStore for LibSqlBackend { ); } - Ok(reciprocal_rank_fusion(fts_results, vector_results, config)) + Ok(fuse_results(fts_results, vector_results, config)) } } diff --git a/src/tools/builtin/memory.rs b/src/tools/builtin/memory.rs index dadab877..de04575b 100644 --- a/src/tools/builtin/memory.rs +++ b/src/tools/builtin/memory.rs @@ -539,6 +539,26 @@ impl Tool for MemoryTreeTool { } } +#[cfg(test)] +mod path_routing_tests { + use super::looks_like_filesystem_path; + + #[test] + fn detects_filesystem_paths() { + assert!(looks_like_filesystem_path("/Users/nige/file.md")); + assert!(looks_like_filesystem_path("C:\\Users\\nige\\file.md")); + assert!(looks_like_filesystem_path("D:/work/file.md")); + assert!(looks_like_filesystem_path("~/notes.md")); + } + + #[test] + fn allows_workspace_memory_paths() { + assert!(!looks_like_filesystem_path("MEMORY.md")); + assert!(!looks_like_filesystem_path("daily/2026-03-11.md")); + assert!(!looks_like_filesystem_path("projects/alpha/notes.md")); + } +} + #[cfg(all(test, feature = "postgres"))] mod tests { use super::*; @@ -616,23 +636,3 @@ mod tests { assert_eq!(schema["properties"]["depth"]["default"], 1); } } - -#[cfg(test)] -mod path_routing_tests { - use super::looks_like_filesystem_path; - - #[test] - fn detects_filesystem_paths() { - assert!(looks_like_filesystem_path("/Users/nige/file.md")); - assert!(looks_like_filesystem_path("C:\\Users\\nige\\file.md")); - assert!(looks_like_filesystem_path("D:/work/file.md")); - assert!(looks_like_filesystem_path("~/notes.md")); - } - - #[test] - fn allows_workspace_memory_paths() { - assert!(!looks_like_filesystem_path("MEMORY.md")); - assert!(!looks_like_filesystem_path("daily/2026-03-11.md")); - assert!(!looks_like_filesystem_path("projects/alpha/notes.md")); - } -} diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index fa48072b..ad233caf 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -55,7 +55,9 @@ pub use embeddings::{ }; #[cfg(feature = "postgres")] pub use repository::Repository; -pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion}; +pub use search::{ + FusionStrategy, RankedResult, SearchConfig, SearchResult, fuse_results, reciprocal_rank_fusion, +}; use std::sync::Arc; @@ -332,6 +334,8 @@ pub struct Workspace { storage: WorkspaceStorage, /// Embedding provider for semantic search. embeddings: Option>, + /// Default search configuration applied to all queries. + search_defaults: SearchConfig, } impl Workspace { @@ -343,6 +347,7 @@ impl Workspace { agent_id: None, storage: WorkspaceStorage::Repo(Repository::new(pool)), embeddings: None, + search_defaults: SearchConfig::default(), } } @@ -355,6 +360,7 @@ impl Workspace { agent_id: None, storage: WorkspaceStorage::Db(db), embeddings: None, + search_defaults: SearchConfig::default(), } } @@ -370,6 +376,16 @@ impl Workspace { self } + /// Set the default search configuration from workspace search config. + pub fn with_search_config(mut self, config: &crate::config::WorkspaceSearchConfig) -> Self { + self.search_defaults = SearchConfig::default() + .with_fusion_strategy(config.fusion_strategy) + .with_rrf_k(config.rrf_k) + .with_fts_weight(config.fts_weight) + .with_vector_weight(config.vector_weight); + self + } + /// Get the user ID. pub fn user_id(&self) -> &str { &self.user_id @@ -709,13 +725,13 @@ impl Workspace { /// Hybrid search across all memory documents. /// /// Combines full-text search (BM25) with semantic search (vector similarity) - /// using Reciprocal Rank Fusion (RRF). + /// using the configured fusion strategy. pub async fn search( &self, query: &str, limit: usize, ) -> Result, WorkspaceError> { - self.search_with_config(query, SearchConfig::default().with_limit(limit)) + self.search_with_config(query, self.search_defaults.clone().with_limit(limit)) .await } diff --git a/src/workspace/repository.rs b/src/workspace/repository.rs index de8c3169..82e4f949 100644 --- a/src/workspace/repository.rs +++ b/src/workspace/repository.rs @@ -12,7 +12,7 @@ use uuid::Uuid; use crate::error::WorkspaceError; use crate::workspace::document::{MemoryChunk, MemoryDocument, WorkspaceEntry}; -use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion}; +use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, fuse_results}; /// Database repository for workspace operations. pub struct Repository { @@ -415,7 +415,7 @@ impl Repository { Vec::new() }; - Ok(reciprocal_rank_fusion(fts_results, vector_results, config)) + Ok(fuse_results(fts_results, vector_results, config)) } /// Full-text search using PostgreSQL ts_rank_cd. diff --git a/src/workspace/search.rs b/src/workspace/search.rs index dff15298..8b78a125 100644 --- a/src/workspace/search.rs +++ b/src/workspace/search.rs @@ -1,17 +1,30 @@ //! Hybrid search combining full-text and semantic search. //! -//! Uses Reciprocal Rank Fusion (RRF) to combine results from: -//! 1. PostgreSQL full-text search (ts_rank_cd) -//! 2. pgvector cosine similarity search +//! Supports two fusion strategies: +//! 1. **RRF** (Reciprocal Rank Fusion) — the default, rank-based method. +//! `score = sum(1 / (k + rank))` for each retrieval method. +//! 2. **WeightedScore** — converts ranks to scores via `1/rank`, combines with +//! configurable weights (`fts_weight * fts_score + vector_weight * vector_score`), +//! then normalizes to \[0,1\] by dividing by the maximum combined score. //! -//! RRF formula: score = sum(1 / (k + rank)) for each retrieval method -//! This is robust to different score scales and produces better results -//! than simple score averaging. +//! Both strategies combine results from: +//! - PostgreSQL / libSQL full-text search +//! - pgvector / libsql_vector cosine similarity search use std::collections::HashMap; use uuid::Uuid; +/// Strategy used to fuse FTS and vector search results. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum FusionStrategy { + /// Reciprocal Rank Fusion (default). Ignores `fts_weight`/`vector_weight`. + #[default] + Rrf, + /// Weighted score fusion using normalized rank-derived scores. + WeightedScore, +} + /// Configuration for hybrid search. #[derive(Debug, Clone)] pub struct SearchConfig { @@ -27,6 +40,16 @@ pub struct SearchConfig { pub min_score: f32, /// Maximum results to fetch from each method before fusion. pub pre_fusion_limit: usize, + /// Fusion strategy to use when combining results. + pub fusion_strategy: FusionStrategy, + /// Weight for FTS results in `WeightedScore` fusion (default 0.5). + /// Ignored by `Rrf` fusion. For env-based config via + /// `WorkspaceSearchConfig::resolve`, defaults are per-strategy. + pub fts_weight: f32, + /// Weight for vector results in `WeightedScore` fusion (default 0.5). + /// Ignored by `Rrf` fusion. For env-based config via + /// `WorkspaceSearchConfig::resolve`, defaults are per-strategy. + pub vector_weight: f32, } impl Default for SearchConfig { @@ -38,6 +61,9 @@ impl Default for SearchConfig { use_vector: true, min_score: 0.0, pre_fusion_limit: 50, + fusion_strategy: FusionStrategy::default(), + fts_weight: 0.5, + vector_weight: 0.5, } } } @@ -74,6 +100,32 @@ impl SearchConfig { self.min_score = score.clamp(0.0, 1.0); self } + + /// Set the fusion strategy. + pub fn with_fusion_strategy(mut self, strategy: FusionStrategy) -> Self { + self.fusion_strategy = strategy; + self + } + + /// Set the FTS weight for `WeightedScore` fusion. + /// + /// Non-finite (NaN, ±inf) or negative values are ignored. + pub fn with_fts_weight(mut self, weight: f32) -> Self { + if weight.is_finite() && weight >= 0.0 { + self.fts_weight = weight; + } + self + } + + /// Set the vector weight for `WeightedScore` fusion. + /// + /// Non-finite (NaN, ±inf) or negative values are ignored. + pub fn with_vector_weight(mut self, weight: f32) -> Self { + if weight.is_finite() && weight >= 0.0 { + self.vector_weight = weight; + } + self + } } /// A search result with hybrid scoring. @@ -87,7 +139,7 @@ pub struct SearchResult { pub chunk_id: Uuid, /// Chunk content. pub content: String, - /// Combined RRF score (0.0-1.0 normalized). + /// Combined fusion score (0.0-1.0 normalized). Strategy-dependent (RRF or WeightedScore). pub score: f32, /// Rank in FTS results (1-based, None if not in FTS results). pub fts_rank: Option, @@ -123,6 +175,22 @@ pub struct RankedResult { pub rank: u32, // 1-based rank } +/// Fuse FTS and vector search results using the strategy specified in `config`. +/// +/// This is the primary entry point for result fusion. Delegates to +/// [`reciprocal_rank_fusion`] or [`weighted_score_fusion`] based on +/// `config.fusion_strategy`. +pub fn fuse_results( + fts_results: Vec, + vector_results: Vec, + config: &SearchConfig, +) -> Vec { + match config.fusion_strategy { + FusionStrategy::Rrf => reciprocal_rank_fusion(fts_results, vector_results, config), + FusionStrategy::WeightedScore => weighted_score_fusion(fts_results, vector_results, config), + } +} + /// Reciprocal Rank Fusion algorithm. /// /// Combines ranked results from multiple retrieval methods using the formula: @@ -235,6 +303,109 @@ pub fn reciprocal_rank_fusion( results } +/// Weighted score fusion. +/// +/// Converts ranks from each method into scores using `1/rank` +/// (so rank 1 → 1.0, rank N → 1/N), then combines them with +/// configurable weights: `fts_weight * fts_score + vector_weight * vector_score`. +/// +/// The combined scores are then normalized to [0,1] by dividing by the +/// maximum score; post-processing (normalization, min_score filter, sort, +/// truncate) matches RRF. +pub fn weighted_score_fusion( + fts_results: Vec, + vector_results: Vec, + config: &SearchConfig, +) -> Vec { + struct ChunkInfo { + document_id: Uuid, + document_path: String, + content: String, + score: f32, + fts_rank: Option, + vector_rank: Option, + } + + let mut chunk_scores: HashMap = HashMap::new(); + + // Process FTS results: score = fts_weight * (1 / rank) + for result in fts_results { + let score = config.fts_weight * (1.0 / result.rank as f32); + chunk_scores + .entry(result.chunk_id) + .and_modify(|info| { + info.score += score; + info.fts_rank = Some(result.rank); + }) + .or_insert(ChunkInfo { + document_id: result.document_id, + document_path: result.document_path, + content: result.content, + score, + fts_rank: Some(result.rank), + vector_rank: None, + }); + } + + // Process vector results: score = vector_weight * (1 / rank) + for result in vector_results { + let score = config.vector_weight * (1.0 / result.rank as f32); + chunk_scores + .entry(result.chunk_id) + .and_modify(|info| { + info.score += score; + info.vector_rank = Some(result.rank); + }) + .or_insert(ChunkInfo { + document_id: result.document_id, + document_path: result.document_path, + content: result.content, + score, + fts_rank: None, + vector_rank: Some(result.rank), + }); + } + + let mut results: Vec = chunk_scores + .into_iter() + .map(|(chunk_id, info)| SearchResult { + document_id: info.document_id, + document_path: info.document_path, + chunk_id, + content: info.content, + score: info.score, + fts_rank: info.fts_rank, + vector_rank: info.vector_rank, + }) + .collect(); + + // Normalize scores to 0-1 range + if let Some(max_score) = results.iter().map(|r| r.score).reduce(f32::max) + && max_score > 0.0 + { + for result in &mut results { + result.score /= max_score; + } + } + + // Filter by minimum score + if config.min_score > 0.0 { + results.retain(|r| r.score >= config.min_score); + } + + // Sort by score descending + results.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + // Limit results + results.truncate(config.limit); + + results +} + #[cfg(test)] mod tests { use super::*; @@ -457,6 +628,142 @@ mod tests { let vector_only = SearchConfig::default().vector_only(); assert!(!vector_only.use_fts); assert!(vector_only.use_vector); + + let weighted = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_fts_weight(0.8) + .with_vector_weight(0.2); + assert_eq!(weighted.fusion_strategy, FusionStrategy::WeightedScore); + assert!((weighted.fts_weight - 0.8).abs() < 0.001); + assert!((weighted.vector_weight - 0.2).abs() < 0.001); + } + + #[test] + fn test_weighted_fusion_basic() { + // With equal weights, a hybrid match should still rank highest. + let config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_fts_weight(1.0) + .with_vector_weight(1.0) + .with_limit(10); + + let chunk1 = Uuid::new_v4(); // In both + let chunk2 = Uuid::new_v4(); // FTS only + let chunk3 = Uuid::new_v4(); // Vector only + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 2)]; + let vec_results = vec![make_result(chunk1, doc, 1), make_result(chunk3, doc, 2)]; + + let results = weighted_score_fusion(fts, vec_results, &config); + + assert_eq!(results.len(), 3); + // Hybrid match (chunk1) should be first — it gets score from both + assert_eq!(results[0].chunk_id, chunk1); + assert!(results[0].is_hybrid()); + assert!(results[0].score > results[1].score); + } + + #[test] + fn test_weighted_fusion_fts_boost() { + // High FTS weight should elevate FTS-only results above vector-only. + let config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_fts_weight(2.0) + .with_vector_weight(0.5) + .with_limit(10); + + let chunk_fts = Uuid::new_v4(); // FTS only, rank 2 + let chunk_vec = Uuid::new_v4(); // Vector only, rank 2 + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk_fts, doc, 2)]; + let vec_results = vec![make_result(chunk_vec, doc, 2)]; + + let results = weighted_score_fusion(fts, vec_results, &config); + + assert_eq!(results.len(), 2); + // FTS result should rank higher because of the 2.0 weight vs 0.5 + assert_eq!(results[0].chunk_id, chunk_fts); + assert!(results[0].from_fts()); + assert!(!results[0].from_vector()); + } + + #[test] + fn test_weighted_fusion_single_source() { + // Only FTS results — should still work correctly. + let config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_limit(10); + + let chunk1 = Uuid::new_v4(); + let chunk2 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk1, doc, 1), make_result(chunk2, doc, 3)]; + + let results = weighted_score_fusion(fts, Vec::new(), &config); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].chunk_id, chunk1); + assert!(results[0].score > results[1].score); + // Top result should be normalized to 1.0 + assert!((results[0].score - 1.0).abs() < 0.001); + } + + #[test] + fn test_weight_setters_reject_invalid() { + let config = SearchConfig::default(); + let original_fts = config.fts_weight; + let original_vec = config.vector_weight; + + // NaN is ignored + let c = config.clone().with_fts_weight(f32::NAN); + assert!((c.fts_weight - original_fts).abs() < 0.001); + + // Infinity is ignored + let c = config.clone().with_vector_weight(f32::INFINITY); + assert!((c.vector_weight - original_vec).abs() < 0.001); + + // Negative is ignored + let c = config.clone().with_fts_weight(-1.0); + assert!((c.fts_weight - original_fts).abs() < 0.001); + + // Negative infinity is ignored + let c = config.clone().with_vector_weight(f32::NEG_INFINITY); + assert!((c.vector_weight - original_vec).abs() < 0.001); + + // Valid values > 1.0 are accepted (weights don't need to sum to 1.0) + let c = config.clone().with_fts_weight(2.0); + assert!((c.fts_weight - 2.0).abs() < 0.001); + + // Zero is valid + let c = config.clone().with_vector_weight(0.0); + assert!(c.vector_weight.abs() < 0.001); + } + + #[test] + fn test_fuse_results_dispatches_correctly() { + let chunk1 = Uuid::new_v4(); + let doc = Uuid::new_v4(); + + let fts = vec![make_result(chunk1, doc, 1)]; + + // RRF strategy + let rrf_config = SearchConfig::default().with_limit(10); + let rrf_results = fuse_results(fts.clone(), Vec::new(), &rrf_config); + assert_eq!(rrf_results.len(), 1); + + // Weighted strategy + let weighted_config = SearchConfig::default() + .with_fusion_strategy(FusionStrategy::WeightedScore) + .with_limit(10); + let weighted_results = fuse_results(fts, Vec::new(), &weighted_config); + assert_eq!(weighted_results.len(), 1); + + // Both should normalize single result to 1.0 + assert!((rrf_results[0].score - 1.0).abs() < 0.001); + assert!((weighted_results[0].score - 1.0).abs() < 0.001); } // --- Edge case tests ---