Add ranked lexical search

Ripgrep gave us a trustworthy exact-search baseline. Its weakness appears when a Query matches many Chunks: every match is real, but you still need to decide which one the model should read first.

I use SQLite here because it is a small database built into Python, so you do not need to run a database server. Its FTS5 feature searches all the text in our Chunks. FTS stands for full-text search.

FTS5 builds a search index, which is a stored map from words to the Chunks containing them. That map avoids reading every Chunk for every Query.

FTS5 uses a ranking formula called BM25. The formula gives more weight to useful word matches and less weight to common words or unusually long Chunks.

That is enough theory to understand the implementation. The class below creates the index, replaces its contents, and returns ranked matches.

Build a tiny index

import sqlite3
from collections.abc import Iterable
from dataclasses import dataclass


@dataclass(frozen=True)
class LexicalResult:
    chunk_id: str
    source_id: str
    text: str
    score: float


class LexicalIndex:
    def __init__(self, connection: sqlite3.Connection) -> None:
        self.connection = connection
        connection.execute(
            "CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5("
            "chunk_id UNINDEXED, source_id UNINDEXED, text)"
        )

    def replace(self, chunks: Iterable[tuple[str, str, str]]) -> None:
        with self.connection:
            self.connection.execute("DELETE FROM chunks_fts")
            self.connection.executemany(
                "INSERT INTO chunks_fts(chunk_id, source_id, text) "
                "VALUES (?, ?, ?)",
                chunks,
            )

    def search(self, query: str, *, limit: int = 8) -> list[LexicalResult]:
        rows = self.connection.execute(
            "SELECT chunk_id, source_id, text, bm25(chunks_fts) "
            "FROM chunks_fts WHERE chunks_fts MATCH ? "
            "ORDER BY bm25(chunks_fts) LIMIT ?",
            (query, limit),
        ).fetchall()
        return [
            LexicalResult(row[0], row[1], row[2], -float(row[3]))
            for row in rows
        ]

replace() clears the old rows and inserts (chunk_id, source_id, text) tuples. search() turns SQLite's BM25 distance into a score where larger means better.

The class can be used with an in-memory database:

index = LexicalIndex(sqlite3.connect(":memory:"))
index.replace((chunk.id, chunk.source_id, chunk.text) for chunk in chunks)
results = index.search('"restart limit" OR DEP-1042')

A successful code example does not tell us when the method is useful. Compare an easy word match such as deployment limit with undo a bad release, which uses different wording from the rollback guide.

A Query that word search handles poorly

BM25 rewards terms that appear in the Query, gives rarer terms more weight, and adjusts for Chunk length. It is still lexical. If the rollback guide never uses the word undo, BM25 cannot infer that undo a bad release means rollback.

That failure is useful. You now know exactly what embeddings need to improve, rather than adding them because every RAG diagram has a vector database in it.

Keep the exact method

Do not throw grep away. Error codes, command flags, plan names, and function names remain strong lexical targets. Later, the hybrid retriever will combine these lexical results with semantic results.

Before we add semantic search, there is another class of Query that neither grep nor embeddings should own: exact questions about structured records.

Retrieve structured facts with SQL