src/rag/clinical/loader.py

"""Load (or rebuild) the TF-IDF clinical RAG index.

The user's rag.py builds the index from `__main__.Chunk` (a frozen
dataclass). Picking up that pickle from a different module path requires
a custom Unpickler that re-routes the class — see _ChunkRoutingUnpickler.
"""
from __future__ import annotations

import pickle
from pathlib import Path
from typing import Any

from src.core.logger import get_logger
from src.rag.clinical.types import ClinicalChunk

logger = get_logger(__name__)


class _ChunkRoutingUnpickler(pickle.Unpickler):
    """Pickle's `find_class` hook lets us swap `__main__.Chunk` (and
    `rag.Chunk` if the user later runs the builder as a module) for our
    `ClinicalChunk` — both are frozen dataclasses with the same fields,
    so the swap is structurally safe.
    """

    def find_class(self, module: str, name: str):
        if name == "Chunk" and module in {"__main__", "rag", "rag.rag"}:
            return ClinicalChunk
        return super().find_class(module, name)


def load_index(path: Path) -> dict[str, Any]:
    """Unpickle a TF-IDF index produced by the user's rag.py."""
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"clinical RAG index not found: {path}")
    with path.open("rb") as f:
        payload = _ChunkRoutingUnpickler(f).load()
    expected = {"chunks", "vectorizer", "matrix"}
    if not expected <= set(payload):
        raise ValueError(
            f"clinical RAG index missing expected keys: have {sorted(payload)}, need {sorted(expected)}"
        )
    logger.info("loaded clinical RAG index: %d chunks from %s", len(payload["chunks"]), path)
    return payload