Metafazer/finrag-backend

1"""BM25 sparse retrieval index for SEC filings.2 3Provides keyword-based retrieval using the Okapi BM25 algorithm.4Complements dense vector search by catching exact financial terms5that embedding models often miss (ticker symbols, accounting terms,6specific numerical references like "diluted EPS").7 8Design decisions:9- rank-bm25 for lightweight, dependency-free BM25 implementation10- Custom financial-aware tokenizer (preserves ticker symbols, numbers)11- In-memory index with serialization for persistence12- Metadata-filtered search to narrow by ticker/section/form_type13- Returns scored results compatible with Day 5 hybrid fusion14 15Debt: DAY-4-001 — Tokenizer uses simple regex. A finance-specific16      tokenizer (handling $, %, B/M/K suffixes) would improve recall.17"""18 19import pickle20import re21from pathlib import Path22 23import structlog24from rank_bm25 import BM25Okapi25 26from finrag.ingestion.chunker import Chunk27 28logger = structlog.get_logger(__name__)29 30# --------------------------------------------------------------------------- #31# Constants32# --------------------------------------------------------------------------- #33 34# BM25 parameters (Okapi defaults are well-studied)35# k1: term frequency saturation. Higher = more weight on repeated terms.36# b: document length normalization. 0 = no normalization, 1 = full.37DEFAULT_K1 = 1.538DEFAULT_B = 0.7539 40# Pattern for tokenization: split on non-alphanumeric, keep numbers/decimals41TOKEN_PATTERN = re.compile(r"[a-zA-Z0-9]+(?:\.[0-9]+)*")42 43 44# --------------------------------------------------------------------------- #45# Tokenizer46# --------------------------------------------------------------------------- #47 48 49def tokenize(text: str) -> list[str]:50    """Tokenize text for BM25 indexing.51 52    Financial-aware tokenizer that:53    - Lowercases for case-insensitive matching54    - Preserves decimal numbers (e.g., "46.2" stays as one token)55    - Preserves ticker-like tokens (e.g., "AAPL")56    - Strips punctuation but keeps alphanumeric content57 58    Args:59        text: Raw text to tokenize.60 61    Returns:62        List of lowercase tokens.63    """64    return [t.lower() for t in TOKEN_PATTERN.findall(text)]65 66 67# --------------------------------------------------------------------------- #68# BM25Index69# --------------------------------------------------------------------------- #70 71 72class BM25Index:73    """BM25 sparse retrieval index with metadata filtering.74 75    Indexes chunked documents for keyword-based retrieval.76    Supports filtered search by metadata fields and serialization77    to disk for persistence.78 79    Args:80        k1: BM25 term frequency saturation parameter.81        b: BM25 document length normalization parameter.82    """83 84    def __init__(85        self,86        k1: float = DEFAULT_K1,87        b: float = DEFAULT_B,88    ) -> None:89        """Initialize an empty BM25 index.90 91        Args:92            k1: Term frequency saturation (default 1.5).93            b: Length normalization (default 0.75).94        """95        self._k1 = k196        self._b = b97 98        # Storage for indexed documents99        self._chunks: list[Chunk] = []100        self._tokenized_corpus: list[list[str]] = []101        self._bm25: BM25Okapi | None = None102 103        logger.debug("bm25_index_initialized", k1=k1, b=b)104 105    @property106    def count(self) -> int:107        """Number of documents in the index."""108        return len(self._chunks)109 110    @property111    def is_built(self) -> bool:112        """Whether the index has been built."""113        return self._bm25 is not None114 115    def add_chunks(self, chunks: list[Chunk]) -> int:116        """Add chunks to the index and rebuild.117 118        Tokenizes each chunk's text and rebuilds the BM25 index.119        Can be called multiple times to add more documents.120 121        Args:122            chunks: List of Chunk objects to index.123 124        Returns:125            Total number of documents in the index after adding.126        """127        if not chunks:128            return self.count129 130        for chunk in chunks:131            tokens = tokenize(chunk.text)132            self._chunks.append(chunk)133            self._tokenized_corpus.append(tokens)134 135        # Rebuild BM25 index with all documents136        self._bm25 = BM25Okapi(137            self._tokenized_corpus,138            k1=self._k1,139            b=self._b,140        )141 142        logger.info(143            "bm25_index_built",144            total_documents=len(self._chunks),145            new_documents=len(chunks),146            avg_tokens=sum(len(t) for t in self._tokenized_corpus) // max(len(self._tokenized_corpus), 1),147        )148        return self.count149 150    def query(151        self,152        query_text: str,153        n_results: int = 10,154        where: dict | None = None,155    ) -> list[dict]:156        """Search the BM25 index with optional metadata filtering.157 158        Scores all documents against the query, optionally filters159        by metadata, and returns the top-n results sorted by score.160 161        Args:162            query_text: Natural language query string.163            n_results: Maximum number of results to return.164            where: Optional metadata filter dict.165                Supports flat key-value: {"ticker": "AAPL"}166                and compound: {"$and": [{"ticker": "AAPL"}, {"form_type": "10-K"}]}167 168        Returns:169            List of result dicts sorted by BM25 score (descending),170            each containing:171            - text: The chunk text172            - metadata: Dict of chunk metadata173            - score: BM25 relevance score174            - chunk_id: The deterministic chunk ID175        """176        if not self.is_built:177            logger.warning("bm25_query_on_empty_index")178            return []179 180        query_tokens = tokenize(query_text)181        if not query_tokens:182            return []183 184        # Score all documents185        scores = self._bm25.get_scores(query_tokens)186 187        # Build (score, index) pairs and apply metadata filter188        scored_indices: list[tuple[float, int]] = []189        for idx, score in enumerate(scores):190            if score <= 0:191                continue192 193            chunk = self._chunks[idx]194 195            # Apply metadata filter196            if where and not self._matches_filter(chunk, where):197                continue198 199            scored_indices.append((score, idx))200 201        # Sort by score descending, take top n202        scored_indices.sort(key=lambda x: x[0], reverse=True)203        top_results = scored_indices[:n_results]204 205        # Build output206        output: list[dict] = []207        for score, idx in top_results:208            chunk = self._chunks[idx]209            output.append(210                {211                    "chunk_id": chunk.metadata.chunk_id,212                    "text": chunk.text,213                    "metadata": {214                        "ticker": chunk.metadata.ticker,215                        "company_name": chunk.metadata.company_name,216                        "form_type": chunk.metadata.form_type,217                        "filing_date": chunk.metadata.filing_date,218                        "section_name": chunk.metadata.section_name,219                        "chunk_index": chunk.metadata.chunk_index,220                        "total_chunks_in_section": chunk.metadata.total_chunks_in_section,221                        "token_count": chunk.metadata.token_count,222                    },223                    "score": float(score),224                }225            )226 227        logger.info(228            "bm25_query_executed",229            query_preview=query_text[:80],230            query_tokens=len(query_tokens),231            candidates=len(scored_indices),232            returned=len(output),233            filter=where,234        )235        return output236 237    def _matches_filter(self, chunk: Chunk, where: dict) -> bool:238        """Check if a chunk matches a metadata filter.239 240        Supports:241        - Simple: {"ticker": "AAPL"}242        - Compound: {"$and": [{"ticker": "AAPL"}, {"form_type": "10-K"}]}243 244        Args:245            chunk: Chunk to check.246            where: Filter dict.247 248        Returns:249            True if the chunk matches the filter.250        """251        if "$and" in where:252            return all(self._matches_filter(chunk, sub_filter) for sub_filter in where["$and"])253 254        meta = chunk.metadata255        for key, value in where.items():256            actual = getattr(meta, key, None)257            if actual != value:258                return False259        return True260 261    def save(self, path: Path) -> None:262        """Serialize the index to disk.263 264        Saves the chunks and tokenized corpus. The BM25 index265        is rebuilt on load (it's fast).266 267        Args:268            path: File path to save to (.pkl).269        """270        path.parent.mkdir(parents=True, exist_ok=True)271        data = {272            "chunks": [c.model_dump() for c in self._chunks],273            "tokenized_corpus": self._tokenized_corpus,274            "k1": self._k1,275            "b": self._b,276        }277        with open(path, "wb") as f:278            pickle.dump(data, f)279 280        logger.info(281            "bm25_index_saved",282            path=str(path),283            documents=len(self._chunks),284        )285 286    @classmethod287    def load(cls, path: Path) -> "BM25Index":288        """Load a serialized index from disk.289 290        Args:291            path: File path to load from (.pkl).292 293        Returns:294            Reconstructed BM25Index.295 296        Raises:297            FileNotFoundError: If the index file doesn't exist.298        """299        if not path.exists():300            msg = f"BM25 index file not found: {path}"301            raise FileNotFoundError(msg)302 303        with open(path, "rb") as f:304            data = pickle.load(f)  # noqa: S301305 306        index = cls(k1=data["k1"], b=data["b"])307        index._chunks = [Chunk(**c) for c in data["chunks"]]308        index._tokenized_corpus = data["tokenized_corpus"]309 310        if index._tokenized_corpus:311            index._bm25 = BM25Okapi(312                index._tokenized_corpus,313                k1=index._k1,314                b=index._b,315            )316 317        logger.info(318            "bm25_index_loaded",319            path=str(path),320            documents=len(index._chunks),321        )322        return index323 324    def get_stats(self) -> dict:325        """Return index statistics.326 327        Returns:328            Dict with index metadata and counts.329        """330        stats: dict = {331            "total_documents": self.count,332            "is_built": self.is_built,333            "k1": self._k1,334            "b": self._b,335        }336 337        if self._chunks:338            tickers = {c.metadata.ticker for c in self._chunks}339            form_types = {c.metadata.form_type for c in self._chunks}340            stats["unique_tickers"] = sorted(tickers)341            stats["unique_form_types"] = sorted(form_types)342            avg_tokens = sum(len(t) for t in self._tokenized_corpus) // len(self._tokenized_corpus)343            stats["avg_tokens_per_doc"] = avg_tokens344 345        return stats346