backend/src/finrag/retrieval/lexical.py

"""BM25 lexical retrieval over the same chunks indexed in Qdrant.

Why this exists: dense embeddings (Cohere v3) lose exact-match signal on
years, tickers, GAAP terminology, and dollar amounts — exactly the tokens
that matter most in financial documents. BM25 catches these.

This module is *not* a full retriever. It returns (chunk_id, score) pairs.
The fusion + hydration into full RetrievedChunk objects happens in
retrieval.hybrid (Decision 10).

Pipeline:
    build:  data/processed/*.jsonl  ──▶  tokenize  ──▶  BM25Okapi
                                                          │
                                                          ▼
                                          pickled to data/bm25_index.pkl
    load:   pickle load → ready to search
"""

from __future__ import annotations

import pickle
import re
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from rank_bm25 import BM25Okapi

from finrag.ingestion.parse import PROCESSED_DIR, Chunk

# parse.py → ingestion/ → finrag/ → src/ → backend/ → ROOT
REPO_ROOT = Path(__file__).resolve().parents[4]
INDEX_PATH = REPO_ROOT / "data" / "bm25_index.pkl"

# Token regex: any run of alphanumeric chars including underscores. Drops
# punctuation, splits on whitespace + symbols. Lowercased before splitting.
# Critical: this exact function is also called on queries — the same vocab
# must be used on both sides or no terms will match.
_TOKEN_RE = re.compile(r"\w+")


def tokenize(text: str) -> list[str]:
    """Lowercase + simple word tokenization.

    The same function runs on chunk text at index time and on user queries
    at search time. Don't tweak one side without the other.
    """
    return _TOKEN_RE.findall(text.lower())


# ── On-disk format ────────────────────────────────────────────────────────
@dataclass
class _BM25Bundle:
    """What we pickle. Separated so we can version the schema later.

    `chunks` is a list of (id, ticker, fiscal_year, chunk_type) tuples — the
    minimal payload we need for post-search filtering and lookups. The full
    chunk text/metadata lives in Qdrant; storing it twice would double disk
    use and risk drift between stores.
    """

    bm25: BM25Okapi
    # Parallel arrays — index `i` in `bm25` corresponds to chunks[i].
    # We use a tuple-list rather than a dict because BM25Okapi indexes by
    # position, not by chunk_id.
    chunk_ids: list[str]
    tickers: list[str]
    fiscal_years: list[int]
    chunk_types: list[str]


# Pin __module__ so pickle records the dotted path "finrag.retrieval.lexical"
# instead of "__main__" when this file is run via `python -m`. Without this,
# the pickle is only loadable from the same entrypoint that built it.
_BM25Bundle.__module__ = "finrag.retrieval.lexical"


# ── Build ─────────────────────────────────────────────────────────────────
def _load_all_chunks(processed_dir: Path) -> list[Chunk]:
    chunks: list[Chunk] = []
    for jsonl in sorted(processed_dir.glob("*.jsonl")):
        for line in jsonl.read_text(encoding="utf-8").splitlines():
            if line.strip():
                chunks.append(Chunk.model_validate_json(line))
    return chunks


def build_index(processed_dir: Path = PROCESSED_DIR) -> _BM25Bundle:
    """Build a BM25 index from all chunks in processed_dir and persist it."""
    chunks = _load_all_chunks(processed_dir)
    if not chunks:
        raise RuntimeError(f"No chunks found in {processed_dir}")

    print(f"Tokenizing {len(chunks)} chunks…")
    tokenized = [tokenize(c.text) for c in chunks]

    print("Building BM25Okapi index…")
    # k1=1.5, b=0.75 are BM25's standard defaults. The rank_bm25 library
    # exposes these as kwargs; leave them at defaults unless we have a
    # specific reason — these are well-calibrated for English text and
    # any tuning we'd do should be eval-driven, not guess-driven.
    bm25 = BM25Okapi(tokenized)

    bundle = _BM25Bundle(
        bm25=bm25,
        chunk_ids=[c.chunk_id for c in chunks],
        tickers=[c.ticker for c in chunks],
        fiscal_years=[c.fiscal_year for c in chunks],
        chunk_types=[c.chunk_type for c in chunks],
    )

    INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
    with INDEX_PATH.open("wb") as f:
        pickle.dump(bundle, f, protocol=pickle.HIGHEST_PROTOCOL)
    print(f"Wrote {INDEX_PATH} ({INDEX_PATH.stat().st_size / 1024:.0f} KB)")

    return bundle


# ── Load + search ────────────────────────────────────────────────────────
_cached_bundle: _BM25Bundle | None = None


def load_index() -> _BM25Bundle:
    """Load the pickled index from disk, caching in-process.

    Module-level cache (not lru_cache) because the underlying BM25Okapi
    object is heavyweight (~30 MB at our scale) — we want exactly one in
    memory regardless of how many callers ask for it.
    """
    global _cached_bundle
    if _cached_bundle is None:
        if not INDEX_PATH.exists():
            raise FileNotFoundError(
                f"BM25 index not found at {INDEX_PATH}. "
                "Run `uv run python -m finrag.retrieval.lexical` to build it."
            )
        with INDEX_PATH.open("rb") as f:
            _cached_bundle = pickle.load(f)
    return _cached_bundle


def search(
    query: str,
    top_k: int = 50,
    ticker: str | None = None,
    fiscal_year: int | None = None,
    chunk_type: str | None = None,
) -> list[tuple[str, float]]:
    """Return ranked (chunk_id, score) tuples for a query.

    Filtering is post-ranking: we ask BM25 for top-N (where N > top_k to
    leave headroom after filtering), then drop chunks that don't match.
    This is fine at our scale (~4k chunks); at 1M+ you'd want a filter-
    aware index structure or pre-shard by ticker.
    """
    bundle = load_index()
    tokens = tokenize(query)
    if not tokens:
        return []

    # get_scores returns one score per indexed document, in index order
    scores = bundle.bm25.get_scores(tokens)

    # Build candidate list — over-fetch to allow for filter attrition.
    # 4x is a heuristic; if filters are tight (e.g. one ticker × one year),
    # we may want more — but unbounded over-fetch defeats the purpose.
    candidate_count = top_k * 4 if (ticker or fiscal_year or chunk_type) else top_k
    candidate_count = min(candidate_count, len(scores))

    # argpartition is O(n) vs argsort's O(n log n) — meaningful at scale.
    # We get the top-K unordered, then sort just those K.
    top_indices = np.argpartition(-scores, candidate_count - 1)[:candidate_count]
    # Sort the candidates by descending score
    top_indices = top_indices[np.argsort(-scores[top_indices])]

    results: list[tuple[str, float]] = []
    for i in top_indices:
        if ticker and bundle.tickers[i] != ticker:
            continue
        if fiscal_year and bundle.fiscal_years[i] != fiscal_year:
            continue
        if chunk_type and bundle.chunk_types[i] != chunk_type:
            continue
        results.append((bundle.chunk_ids[i], float(scores[i])))
        if len(results) >= top_k:
            break

    return results


# ── CLI ───────────────────────────────────────────────────────────────────
def main() -> None:
    build_index()

    # Sanity check: run a couple of test queries
    print("\nSanity-check queries:")
    for q in [
        "services revenue 2023",
        "iPhone net sales",
        "SG&A expense",
        "Tesla R&D",
    ]:
        results = search(q, top_k=3)
        print(f"\n  Q: {q!r}")
        for chunk_id, score in results:
            print(f"    {chunk_id}  score={score:.3f}")


if __name__ == "__main__":
    main()