From d11c9f257585125946ecbc51b0fa32ac5e433285 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Fri, 10 Jul 2026 19:26:21 +0300 Subject: [PATCH] feat: add weighted RRF to warm search --- src/brainlayer/mcp/search_handler.py | 5 +++ src/brainlayer/search_repo.py | 64 ++++++++++++++++++++++------ tests/test_hybrid_search.py | 53 +++++++++++++++++++++++ tests/test_search_handler.py | 9 ++++ 4 files changed, 119 insertions(+), 12 deletions(-) diff --git a/src/brainlayer/mcp/search_handler.py b/src/brainlayer/mcp/search_handler.py index cd9f6ab4..c256b618 100644 --- a/src/brainlayer/mcp/search_handler.py +++ b/src/brainlayer/mcp/search_handler.py @@ -1264,6 +1264,7 @@ def _emit_kg_degrade(reason: str) -> None: agent_id=agent_id, brainbar_helper_fast_profile=brainbar_helper_fast_profile, consumer_scope=consumer_scope, + warm_rrf=True, ) chunk_results = kg_results.get("chunks", {}) if order == "origin" and chunk_results.get("ids") and chunk_results["ids"][0]: @@ -1927,6 +1928,7 @@ def embed_query_guarded() -> list[float] | None: profile_scope=profile_scope, brainbar_helper_fast_profile=brainbar_helper_fast_profile, consumer_scope=consumer_scope, + warm_rrf=True, ) break except Exception as e: @@ -1960,6 +1962,7 @@ def embed_query_guarded() -> list[float] | None: text = "No results found." if fallback_reason: empty["fallback_reason"] = fallback_reason + empty["degraded"] = True text = f"No results found. Search mode: FTS fallback ({fallback_reason}); vector embedding was skipped." return ([TextContent(type="text", text=text)], empty) @@ -2000,6 +2003,7 @@ def embed_query_guarded() -> list[float] | None: } if fallback_reason: structured["fallback_reason"] = fallback_reason + structured["degraded"] = True if order == "origin": structured["order"] = order structured["order_scope"] = _ORIGIN_ORDER_SCOPE @@ -2096,6 +2100,7 @@ def embed_query_guarded() -> list[float] | None: } if fallback_reason: structured["fallback_reason"] = fallback_reason + structured["degraded"] = True if order == "origin": structured["order"] = order structured["order_scope"] = _ORIGIN_ORDER_SCOPE diff --git a/src/brainlayer/search_repo.py b/src/brainlayer/search_repo.py index 881382e6..192a7dc6 100644 --- a/src/brainlayer/search_repo.py +++ b/src/brainlayer/search_repo.py @@ -125,6 +125,34 @@ def _env_flag_enabled(name: str) -> bool: return os.environ.get(name, "").strip().casefold() in {"1", "true", "yes", "on"} +def _rrf_vector_alpha() -> float: + raw = os.environ.get("BRAINLAYER_RRF_ALPHA") + if raw is None: + return RRF_VECTOR_ALPHA + try: + alpha = float(raw) + except (TypeError, ValueError): + return RRF_VECTOR_ALPHA + if not math.isfinite(alpha): + return RRF_VECTOR_ALPHA + return max(0.0, min(1.0, alpha)) + + +def _weighted_rrf_score( + *, + semantic_rank: int | None, + fts_rank: int | None, + k: int, + alpha: float, +) -> float: + score = 0.0 + if semantic_rank is not None: + score += alpha / (k + semantic_rank) + if fts_rank is not None: + score += (1.0 - alpha) / (k + fts_rank) + return score + + def _has_recency_intent(query_text: str) -> bool: """Return true when recency words appear as terms, not substrings.""" return bool(_RECENCY_SINGLE_TERM_RE.search(query_text) or _RECENCY_THIS_WEEK_RE.search(query_text)) @@ -311,6 +339,8 @@ def _hybrid_cache_key( content_class_filter: Optional[str] = None, recency_rerank: bool = False, importance_rerank: bool = False, + warm_rrf: bool = False, + rrf_vector_alpha: float = RRF_VECTOR_ALPHA, ) -> tuple: return ( store_key, @@ -338,6 +368,8 @@ def _hybrid_cache_key( normalize_content_class(content_class_filter) if content_class_filter else None, recency_rerank, importance_rerank, + warm_rrf, + rrf_vector_alpha, ) @@ -1756,6 +1788,7 @@ def hybrid_search( *, recency_rerank: bool = False, importance_rerank: bool = False, + warm_rrf: bool = False, ) -> Dict[str, List]: """Hybrid search combining semantic (vector) + keyword (FTS5) via Reciprocal Rank Fusion. @@ -1770,6 +1803,7 @@ def hybrid_search( recency_rerank = recency_rerank or _env_flag_enabled("BRAINLAYER_SCORE_RECENCY_DECAY") importance_rerank = importance_rerank or _env_flag_enabled("BRAINLAYER_SCORE_IMPORTANCE_BOOST") + rrf_vector_alpha = _rrf_vector_alpha() if warm_rrf else RRF_VECTOR_ALPHA project_filter = _effective_project_filter(project_filter, consumer_scope) source_filter = _effective_source_filter(source_filter, consumer_scope) include_checkpoints = _effective_include_checkpoints(include_checkpoints, consumer_scope) @@ -1804,6 +1838,8 @@ def hybrid_search( content_class_filter, recency_rerank, importance_rerank, + warm_rrf, + rrf_vector_alpha, ) + ( fts_query_override, kg_boost, @@ -2081,7 +2117,7 @@ def _fetch_fts_rows(table_name: str, timeout_ms: float | None = None) -> list[tu fts_results = _fetch_fts_rows("chunks_fts", timeout_ms=fts_timeout_ms) if include_operational: fts_results.extend(_fetch_fts_rows("chunks_fts_operational", timeout_ms=fts_timeout_ms)) - if getattr(self, "_trigram_fts_available", False) and not brainbar_helper_fast_profile: + if getattr(self, "_trigram_fts_available", False) and not brainbar_helper_fast_profile and not warm_rrf: trigram_fts_results = _fetch_fts_rows("chunks_fts_trigram") search_profile.emit( profile_scope, @@ -2278,8 +2314,11 @@ def _ingest_keyword_rows(rows: list[tuple], ranks: dict[str, int]) -> None: "dist": semantic["distances"][0][i], } - # Union of all chunk_ids from both sources - all_chunk_ids = set(semantic_by_id.keys()) | set(fts_ranks.keys()) | set(trigram_ranks.keys()) + # Warm interactive RRF has exactly two legs. Other consumers keep the + # legacy trigram candidate leg until they migrate independently. + all_chunk_ids = set(semantic_by_id.keys()) | set(fts_ranks.keys()) + if not warm_rrf: + all_chunk_ids |= set(trigram_ranks.keys()) scored = [] match_features_by_id: dict[str, dict[str, bool]] = {} @@ -2288,14 +2327,14 @@ def _ingest_keyword_rows(rows: list[tuple], ranks: dict[str, int]) -> None: sem_entry = semantic_by_id.get(cid) fts_rank = fts_ranks.get(cid) trigram_rank = trigram_ranks.get(cid) - lexical_leg_weight = 1.0 - RRF_VECTOR_ALPHA - - if sem_entry is not None: - score += RRF_VECTOR_ALPHA / (k + sem_entry["rank"]) - if fts_rank is not None: - score += lexical_leg_weight / (k + fts_rank) - if trigram_rank is not None: - score += lexical_leg_weight / (k + trigram_rank) + score = _weighted_rrf_score( + semantic_rank=sem_entry["rank"] if sem_entry is not None else None, + fts_rank=fts_rank, + k=k, + alpha=rrf_vector_alpha, + ) + if not warm_rrf and trigram_rank is not None: + score += (1.0 - RRF_VECTOR_ALPHA) / (k + trigram_rank) # Get data — prefer semantic (has distance) if sem_entry is not None: @@ -2368,8 +2407,9 @@ def _ingest_keyword_rows(rows: list[tuple], ranks: dict[str, int]) -> None: match_features_by_id[cid] = { "vector": sem_entry is not None, "fts": fts_rank is not None, - "trigram": trigram_rank is not None, } + if not warm_rrf: + match_features_by_id[cid]["trigram"] = trigram_rank is not None scored.append((score, cid, doc, meta, dist)) agent_profile = None diff --git a/tests/test_hybrid_search.py b/tests/test_hybrid_search.py index 2d2005bb..5dc99afe 100644 --- a/tests/test_hybrid_search.py +++ b/tests/test_hybrid_search.py @@ -206,6 +206,7 @@ def test_hybrid_search_rrf_scoring(self, store): assert set(ids) == {"both", "fts-only", "vec-only"} def test_hybrid_search_weighted_rrf_uses_vector_alpha(self, store, monkeypatch): + monkeypatch.delenv("BRAINLAYER_RRF_ALPHA", raising=False) query_embedding = _embed("weighted rrf alpha control") _insert_chunk( store, @@ -226,11 +227,63 @@ def test_hybrid_search_weighted_rrf_uses_vector_alpha(self, store, monkeypatch): query_embedding=query_embedding, query_text="weighted rrf alpha control", n_results=2, + warm_rrf=True, ) assert 0.0 < RRF_VECTOR_ALPHA < 0.5 assert results["ids"][0] == ["lexical-only", "vector-only"] + monkeypatch.setenv("BRAINLAYER_RRF_ALPHA", "1.0") + overridden = store.hybrid_search( + query_embedding=query_embedding, + query_text="weighted rrf alpha control", + n_results=2, + warm_rrf=True, + ) + + assert overridden["ids"][0] == ["vector-only", "lexical-only"] + + def test_weighted_rrf_two_leg_math_is_deterministic(self): + score = search_repo._weighted_rrf_score( + semantic_rank=2, + fts_rank=5, + k=60, + alpha=0.25, + ) + + assert score == pytest.approx(0.25 / 62 + 0.75 / 65) + + def test_rrf_alpha_env_override(self, monkeypatch): + monkeypatch.setenv("BRAINLAYER_RRF_ALPHA", "0.4") + + assert search_repo._rrf_vector_alpha() == pytest.approx(0.4) + + def test_rrf_alpha_nonfinite_override_uses_default(self, monkeypatch): + monkeypatch.setenv("BRAINLAYER_RRF_ALPHA", "nan") + + assert search_repo._rrf_vector_alpha() == pytest.approx(RRF_VECTOR_ALPHA) + + def test_warm_weighted_rrf_drops_trigram_fusion_leg(self, store): + _insert_chunk( + store, + chunk_id="trigram-only", + content="stalker-golem queue note", + embedding=_embed("distant vector"), + ) + store.build_binary_index() + cursor = store.conn.cursor() + cursor.execute("DELETE FROM chunk_vectors") + cursor.execute("DELETE FROM chunk_vectors_binary") + + results = store.hybrid_search( + query_embedding=_embed("nothing close"), + query_text="alker-go", + n_results=5, + warm_rrf=True, + ) + + assert results["ids"][0] == [] + def test_hybrid_search_default_does_not_apply_recency_or_importance_boost(self, store, monkeypatch): monkeypatch.delenv("BRAINLAYER_SCORE_IMPORTANCE_BOOST", raising=False) monkeypatch.delenv("BRAINLAYER_SCORE_RECENCY_DECAY", raising=False) diff --git a/tests/test_search_handler.py b/tests/test_search_handler.py index 8db19202..d482d8d1 100644 --- a/tests/test_search_handler.py +++ b/tests/test_search_handler.py @@ -217,8 +217,10 @@ async def test_brain_search_falls_back_to_fts_when_embedding_times_out(monkeypat assert elapsed < 0.5 assert store.hybrid_kwargs["query_embedding"] is None + assert store.hybrid_kwargs["warm_rrf"] is True assert structured["search_mode"] == "fts_fallback" assert structured["fallback_reason"] == "embed_timeout" + assert structured["degraded"] is True assert [item["chunk_id"] for item in structured["results"]] == ["fts-fallback-hit"] assert "FTS fallback" in content[0].text finally: @@ -240,8 +242,10 @@ async def test_brain_search_falls_back_to_fts_when_embedding_raises(monkeypatch) ) assert store.hybrid_kwargs["query_embedding"] is None + assert store.hybrid_kwargs["warm_rrf"] is True assert structured["search_mode"] == "fts_fallback" assert structured["fallback_reason"] == "embed_error:RuntimeError" + assert structured["degraded"] is True assert [item["chunk_id"] for item in structured["results"]] == ["fts-fallback-hit"] @@ -261,8 +265,10 @@ async def test_brain_search_uses_hybrid_when_embedding_is_fast(monkeypatch): ) assert store.hybrid_kwargs["query_embedding"] == [0.1, 0.2, 0.3] + assert store.hybrid_kwargs["warm_rrf"] is True assert structured["search_mode"] == "hybrid" assert "fallback_reason" not in structured + assert "degraded" not in structured assert [item["chunk_id"] for item in structured["results"]] == ["hybrid-hit"] @@ -408,8 +414,10 @@ async def test_brain_search_empty_fts_fallback_includes_fallback_metadata(monkey ) assert store.hybrid_kwargs["query_embedding"] is None + assert store.hybrid_kwargs["warm_rrf"] is True assert structured["search_mode"] == "fts_fallback" assert structured["fallback_reason"] == "embed_error:RuntimeError" + assert structured["degraded"] is True assert structured["total"] == 0 assert structured["results"] == [] assert "Search mode: FTS fallback (embed_error:RuntimeError)" in content[0].text @@ -783,6 +791,7 @@ async def test_brain_search_entity_route_threads_agent_id_to_kg_hybrid_search(mo assert store.kg_hybrid_kwargs is not None assert store.kg_hybrid_kwargs["agent_id"] == "codex-test-agent" + assert store.kg_hybrid_kwargs["warm_rrf"] is True @pytest.mark.asyncio