From f4bfc12293991f782ed7b75c852733bd24173d1f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 31 Aug 2026 17:00:42 +0300 Subject: [PATCH] ARM ONLY, NEVER MERGE: z-scored branch fusion Hybrid without a reranker keeps each database's vector and FTS branches apart and compares every candidate by per-branch z-score, summing both branches for a chunk found in each. Databases are then compared by how exceptional a hit is for them rather than by raw score, which is not comparable across indexes. Targets the measured ceiling on rank-and-score fusion: 73% of candidates at n=4 and 84% at n=8 tie on both score and rank, so no key built from those two can separate them and they fall to declaration order under every other arm. Continuous keys should barely collide. Implementation from the multi-fusion session; branch depth via HAIKU_RAG_BRANCH_DEPTH, default 20. Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc --- haiku_rag_slim/haiku/rag/client/search.py | 98 ++++++++++++++++++++--- tests/multi_db/test_search.py | 91 +++++++++++++++++++++ 2 files changed, 178 insertions(+), 11 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 9c3441a6..73b8dda7 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -1,4 +1,6 @@ import base64 +import os +import statistics from collections.abc import Sequence from typing import TYPE_CHECKING @@ -105,20 +107,24 @@ async def search_sources( fetch_limit = _fetch_limit(client, query, limit) query_vector = await _embed_query(selected[0], query, resolved) text = query if isinstance(query, str) else "" - per_source = await gather_all( - *( - c.chunk_repository.search( - query=text, - limit=fetch_limit, - search_type=resolved, - filter=filter, - query_vector=query_vector, + if resolved == "hybrid" and isinstance(query, str) and client.reranker is None: + # ARM C+D (never merge): z-scored branch fusion. + ranked = await _fuse_branches(selected, text, query_vector, filter, limit) + else: + per_source = await gather_all( + *( + c.chunk_repository.search( + query=text, + limit=fetch_limit, + search_type=resolved, + filter=filter, + query_vector=query_vector, + ) + for c in selected ) - for c in selected ) - ) - ranked = await _fuse(client, selected, query, per_source, limit) + ranked = await _fuse(client, selected, query, per_source, limit) results: list[SearchResult] = [] for owner, chunk, score in ranked: @@ -199,6 +205,76 @@ async def _fuse( return [(client, chunk, score) for score, client, chunk in scored[:limit]] +async def _fuse_branches( + clients: list["HaikuRAG"], + query: str, + query_vector: list[float] | None, + filter: str | None, + limit: int, +) -> list[tuple["HaikuRAG", Chunk, float]]: + """ARM C+D (never merge): one ranked list from every database's vector and + FTS branches, compared by per-branch z-score. + + Each branch's scores are normalized against that branch's own candidate + distribution, so databases are compared by how exceptional a hit is for + them rather than by raw score. A chunk in both branches of its database + sums both z-scores. Ties keep arrival order: client order, vector before + FTS, rank within a branch. + """ + + async def branches(client: "HaikuRAG"): + return await gather_all( + client.chunk_repository.search( + query=query, + limit=_BRANCH_DEPTH, + search_type="vector", + filter=filter, + query_vector=query_vector, + ), + client.chunk_repository.search( + query=query, + limit=_BRANCH_DEPTH, + search_type="fts", + filter=filter, + ), + ) + + per_client = await gather_all(*(branches(c) for c in clients)) + + totals: dict[tuple[int, str], float] = {} + seen: dict[tuple[int, str], tuple[HaikuRAG, Chunk]] = {} + for position, (client, (vector, fts)) in enumerate( + zip(clients, per_client, strict=True) + ): + for branch in (vector, fts): + zs = _z_scores([score for _, score in branch]) + for (chunk, _), z in zip(branch, zs, strict=True): + key = (position, chunk.id or chunk.content) + totals[key] = totals.get(key, 0.0) + z + seen.setdefault(key, (client, chunk)) + ranked = sorted(seen, key=lambda key: totals[key], reverse=True)[:limit] + return [(*seen[key], totals[key]) for key in ranked] + + +def _z_scores(scores: list[float]) -> list[float]: + """Each score as standard deviations above its own list's mean. + + A list too short or too flat to carry a distribution normalizes to zeros: + nothing in it is exceptional. + """ + if len(scores) < 2: + return [0.0] * len(scores) + mean = statistics.fmean(scores) + sd = statistics.pstdev(scores) + if sd == 0: + return [0.0] * len(scores) + return [(score - mean) / sd for score in scores] + + +# ARM C+D (never merge): candidates fetched per branch, an eval grid knob. +_BRANCH_DEPTH = int(os.environ.get("HAIKU_RAG_BRANCH_DEPTH", "20")) + + # Reciprocal rank fusion's smoothing constant, the value the literature uses. _RRF_K = 60 diff --git a/tests/multi_db/test_search.py b/tests/multi_db/test_search.py index 3332d9b3..40b13080 100644 --- a/tests/multi_db/test_search.py +++ b/tests/multi_db/test_search.py @@ -545,6 +545,97 @@ class TestReciprocalRankFusion: ] +class TestZScoredBranchFusion: + """ARM C+D (never merge): hybrid without a reranker fuses every database's + vector and FTS branches by per-branch z-score.""" + + @staticmethod + def _client(vector: list[tuple[Chunk, float]], fts: list[tuple[Chunk, float]]): + from types import SimpleNamespace + + class Repo: + async def search( + self, query, limit, search_type, filter=None, query_vector=None + ): + return vector if search_type == "vector" else fts + + return SimpleNamespace(chunk_repository=Repo()) + + @staticmethod + def _branch(source: str, scores: list[float]) -> list[tuple[Chunk, float]]: + return [ + (Chunk(id=f"{source}{i}", content=f"{source} {i}"), score) + for i, score in enumerate(scores) + ] + + @pytest.mark.asyncio + async def test_a_spike_beats_a_flat_profile(self): + """The database whose top hit stands out from its own candidates wins, + whatever the raw magnitudes.""" + from haiku.rag.client.search import _fuse_branches + + spike = self._client(self._branch("a", [0.9, 0.2, 0.19, 0.18]), []) + flat = self._client(self._branch("b", [5.0, 4.99, 4.98, 4.97]), []) + + ranked = await _fuse_branches([flat, spike], "q", None, None, 2) + + assert [chunk.id for _, chunk, _ in ranked] == ["a0", "b0"] + + @pytest.mark.asyncio + async def test_agreement_within_a_database_sums(self): + """A chunk topping both of its database's branches carries both + z-scores.""" + from haiku.rag.client.search import _fuse_branches + + scores = [0.9, 0.2, 0.19, 0.18] + both = self._client(self._branch("a", scores), self._branch("a", scores)) + one = self._client(self._branch("b", scores), []) + + ranked = await _fuse_branches([one, both], "q", None, None, 3) + + assert [chunk.id for _, chunk, _ in ranked][:2] == ["a0", "b0"] + assert ranked[0][2] == pytest.approx(2 * ranked[1][2]) + + @pytest.mark.asyncio + async def test_a_flat_or_thin_branch_is_nothing_special(self): + from haiku.rag.client.search import _z_scores + + assert _z_scores([]) == [] + assert _z_scores([0.9]) == [0.0] + assert _z_scores([0.5, 0.5, 0.5]) == [0.0, 0.0, 0.0] + + @pytest.mark.asyncio + async def test_hybrid_without_a_reranker_fetches_both_branches( + self, tmp_path, monkeypatch, query_embedding + ): + from haiku.rag.store.repositories.chunk import ChunkRepository + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + monkeypatch.setattr(HaikuRAG, "reranker", property(lambda self: None)) + + asked: list[tuple[str, int]] = [] + search = ChunkRepository.search + + async def spy(self, *args, **kwargs): + asked.append((kwargs["search_type"], kwargs["limit"])) + return await search(self, *args, **kwargs) + + monkeypatch.setattr(ChunkRepository, "search", spy) + + async with HaikuRAG(config=config) as rag: + results = await rag.search("cats", limit=3) + + assert sorted(asked) == [ + ("fts", 20), + ("fts", 20), + ("vector", 20), + ("vector", 20), + ] + assert {r.source for r in results} == {"alpha", "beta"} + + class TestFusingWhatARerankerReturns: @pytest.mark.asyncio async def test_a_reranker_returning_copies_is_named(self, tmp_path):