Zalgorithm

Querying for tmux keybinding with reciprocal rank fusion

GitHub: search_rrf.py Related:

Reciprocal rank fusion (RRF) is a method for combining multiple result sets with different relevance indicators into a single result set. RRF requires no tuning, and the different relevance indicators do not have to be related to each other to achieve high-quality results.

"""Combine OR-based FTS5 and semantic rankings with Reciprocal Rank Fusion."""

import argparse
import re
import sqlite3
from dataclasses import dataclass
from pathlib import Path

from chromadb.errors import ChromaError

from search_bm25 import search as search_bm25
from search_chroma import search as search_chroma


@dataclass
class Result:
    key: str
    description: str
    score: float = 0.0
    fts_rank: int | None = None
    semantic_rank: int | None = None


def fts_query(query: str) -> str:
    """Join input words with OR, retaining common words such as 'the'."""
    words = re.findall(r"[^\W_]+", query, flags=re.UNICODE)
    if not words:
        raise ValueError("Enter at least one word to search for.")
    # Quotes make each word literal, even if the user types an operator.
    # For ordinary words, '"close" OR "the" OR "window"' is equivalent
    # to 'close OR the OR window'. Punctuation separates words.
    return " OR ".join(f'"{word}"' for word in words)


def fuse(
    lexical: list[tuple[str, str, float]],
    semantic: list[tuple[str, str, float]],
    k: int = 60,
) -> list[Result]:
    """Fuse ranked lists by key; use positions, not the original scores."""
    if k < 0:
        raise ValueError("RRF k must be nonnegative.")
    combined: dict[str, Result] = {}
    for rank, (key, description, _) in enumerate(lexical, start=1):
        result = combined.setdefault(key, Result(key, description))
        result.fts_rank = rank
        result.score += 1 / (k + rank)
    for rank, (key, description, _) in enumerate(semantic, start=1):
        result = combined.setdefault(key, Result(key, description))
        result.semantic_rank = rank
        result.score += 1 / (k + rank)
    # Missing from a list means zero contribution. Higher fused scores win.
    # Case-sensitive key ordering makes exact score ties deterministic.
    return sorted(combined.values(), key=lambda result: (-result.score, result.key))


def search(
    database: Path,
    chroma_path: Path,
    query: str,
    collection_name: str = "tmux-key-bindings",
    results: int = 5,
    candidates: int = 20,
    k: int = 60,
) -> list[Result]:
    """Retrieve independently, fuse their top candidates, then limit output."""
    if results < 1 or candidates < 1:
        raise ValueError("Results and candidates must be positive.")
    if k < 0:
        raise ValueError("RRF k must be nonnegative.")
    lexical = search_bm25(database, fts_query(query))[:candidates]
    # The original query goes to the embedding model, without OR rewriting.
    # Semantic retrieval runs even when FTS5 returns no matches.
    semantic = search_chroma(chroma_path, query, collection_name, candidates)
    return fuse(lexical, semantic, k)[:results]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("query", help="Text to search with both retrieval methods")
    parser.add_argument(
        "--database", type=Path, default=Path("tmux_key_bindings.sqlite3")
    )
    parser.add_argument("--chroma", type=Path, default=Path("data/chroma"))
    parser.add_argument("--collection", default="tmux-key-bindings")
    parser.add_argument("--results", type=int, default=5, help="Final results (default: 5)")
    parser.add_argument(
        "--candidates", type=int, default=20,
        help="Maximum entries from each ranked list to fuse (default: 20)",
    )
    parser.add_argument("--k", type=int, default=60, help="RRF constant (default: 60)")
    args = parser.parse_args()
    try:
        rows = search(
            args.database, args.chroma, args.query, args.collection,
            args.results, args.candidates, args.k,
        )
    except (sqlite3.Error, ChromaError, ValueError) as error:
        parser.exit(1, f"Search failed: {error}\n")
    print(f"FTS5 query: {fts_query(args.query)}")
    if not rows:
        print("No results found.")
        return
    print("RRF score\tFTS rank\tSemantic rank\tKey\tDescription")
    for row in rows:
        print(
            f"{row.score:.9g}\t{row.fts_rank or '-'}\t"
            f"{row.semantic_rank or '-'}\t{row.key}\t{row.description}"
        )


if __name__ == "__main__":
    main()

Implementing Reciprocal Rank Fusion (RRF) #

The fuse function implements Reciprocal Rank Fusion:

def fuse(
    lexical: list[tuple[str, str, float]],
    semantic: list[tuple[str, str, float]],
    k: int = 60,
) -> list[Result]:
    """Fuse ranked lists by key; use positions, not the original scores."""
    if k < 0:
        raise ValueError("RRF k must be nonnegative.")
    combined: dict[str, Result] = {}
    for rank, (key, description, _) in enumerate(lexical, start=1):
        result = combined.setdefault(key, Result(key, description))
        result.fts_rank = rank
        result.score += 1 / (k + rank)
    for rank, (key, description, _) in enumerate(semantic, start=1):
        result = combined.setdefault(key, Result(key, description))
        result.semantic_rank = rank
        result.score += 1 / (k + rank)
    # Missing from a list means zero contribution. Higher fused scores win.
    # Case-sensitive key ordering makes exact score ties deterministic.
    return sorted(combined.values(), key=lambda result: (-result.score, result.key))

For a result at rank r, the contribution is:

1k+r \frac{1}{k + r}

With k = 60, the contributions look like:

rank 1 → 1 / 61 = 0.01639344262295082
rank 2 ⇒ 1 / 62 = 0.016129032258064516
...

The scores from the returned lexical and semantic results are summed, then sorted.

Using k = 60 is (mostly) a convention. The k value controls how strongly rank differences matter. Using k = 60 a fairly high k value means that a single rank-1 result doesn’t completely dominate the combined ranks.

Note: the use of += in the score calculation in the lexical enumeration block isn’t technically necessary, but it doesn’t break anything.

Does it work? #

“Work” as in give better results…

Sort of:

find-command ❯ uv run python search_rrf.py "close window"
FTS5 query: "close" OR "window"
RRF score	FTS rank	Semantic rank	Key	Description
0.0327868852	1	1	&	Kill the current window.
0.0315136476	2	5	,	Rename the current window.
0.0315136476	5	2	n	Change to the next window.
0.0314980159	3	4	c	Create a new window.
0.031024531	6	3	p	Change to the previous window.

find_command master*​
find-command ❯ uv run python search_rrf.py "go back to previous window"
FTS5 query: "go" OR "back" OR "to" OR "previous" OR "window"
RRF score	FTS rank	Semantic rank	Key	Description
0.0327868852	1	1	p	Change to the previous window.
0.0314980159	3	4	M-p	Move to the previous window with a bell or activity marker.
0.0308349146	8	2	n	Change to the next window.
0.0303657695	9	3	l	Move to the previously selected window.
0.0288501453	6	13	{	Swap the current pane with the previous pane.

find_command master*​
find-command ❯ uv run python search_rrf.py "close pane"
FTS5 query: "close" OR "pane"
RRF score	FTS rank	Semantic rank	Key	Description
0.0317540323	4	2	M	Clear the marked pane.
0.0315449578	6	1	x	Kill the current pane.
0.031024531	3	6	}	Swap the current pane with the next pane.
0.0306217859	2	9	{	Swap the current pane with the previous pane.
0.0300920728	1	13	m	Mark the current pane (see select-pane -m).

find_command master*​
find-command ❯ uv run python search_rrf.py "close current pane"
FTS5 query: "close" OR "current" OR "pane"
RRF score	FTS rank	Semantic rank	Key	Description
0.0327868852	1	1	x	Kill the current pane.
0.0317540323	4	2	}	Swap the current pane with the next pane.
0.0312576313	3	5	{	Swap the current pane with the previous pane.
0.030798389	7	3	!	Break the current pane out of the window.
0.0306217859	2	9	m	Mark the current pane (see select-pane -m).

find_command master*​
find-command ❯ uv run python search_rrf.py "go to the last pane"
FTS5 query: "go" OR "to" OR "the" OR "last" OR "pane"
RRF score	FTS rank	Semantic rank	Key	Description
0.0314980159	3	4	;	Move to the previously active pane.
0.0307765152	4	6	{	Swap the current pane with the previous pane.
0.0307692308	5	5	}	Swap the current pane with the next pane.
0.0302823315	12	1	x	Kill the current pane.
0.0300179211	2	12	Up, Down, Left, Right	Change to the pane above, below, to the left, or to the right of the current pane.

Where the top results both have the rank of 1, either approach would work. It’s more interesting where the results differ. Counterintuitively, semantic search isn’t always better:

find-command ❯ uv run python search_rrf.py "go to window 3"
FTS5 query: "go" OR "to" OR "window" OR "3"
RRF score	FTS rank	Semantic rank	Key	Description
0.0322580645	2	2	n	Change to the next window.
0.0320184426	1	4	0 to 9	Select windows 0 to 9.
0.0317460317	3	3	p	Change to the previous window.
0.0303308824	4	8	l	Move to the previously selected window.
0.0300904977	8	5	M-n	Move to the next window with a bell or activity marker.

“Go to window three” doesn’t have much semantic similarity to the “Select windows 0 to 9” text. For the FTS5 search, “to” is playing a big role in the ranking.