controlai_rag/index.py

"""Local Offline Hybrid BM25 Index with metadata preservation and citation support."""

from __future__ import annotations

import json
import pickle
import re
from pathlib import Path
from typing import Any

try:
    from rank_bm25 import BM25Okapi
except ImportError:
    BM25Okapi = None

from controlai_rag.chunker import Chunk

INDEX_DIR = Path("data/rag_index")


def tokenize_corpus(text: str) -> list[str]:
    # Lowercase and extract alphanumeric tokens + mathematical symbols
    return re.findall(r"\b\w+\b|[+\-*/^_]", text.lower())


# Only chunk body text is tokenized into BM25, so a query naming a source
# ("what does Nise say about stability") cannot match the book it refers to.
# These constants drive a post-retrieval boost for chunks whose *filename*
# matches a distinctive query term.
SOURCE_MATCH_BOOST = 1.6
# A filename token is only distinctive enough to boost on if it appears in at
# most this fraction of the indexed files (floored at 2 files, so the rule
# still works on a small corpus). An author surname like "nise" or "ogata"
# appears in exactly one file; generic domain words appear in many.
SOURCE_TOKEN_MAX_FILE_FRACTION = 0.02
# Words that identify the subject rather than a specific source. Even if the
# corpus grows large enough for these to slip under the fraction cutoff, they
# must never trigger a source boost.
GENERIC_SOURCE_TOKENS = {
    "control", "controls", "systems", "system", "engineering", "theory",
    "lecture", "lectures", "notes", "chapter", "solutions", "solution",
    "exercise", "exercises", "book", "textbook", "txtbk", "edition", "vol",
    "part", "final", "exam", "slides", "course", "intro", "introduction",
}


class ControlRAGIndex:
    """Fast, local, offline search index over control engineering documents."""

    def __init__(self, index_dir: Path = INDEX_DIR) -> None:
        self.index_dir = index_dir
        self.chunks: list[dict[str, Any]] = []
        self.bm25: Any | None = None
        self._distinctive_source_tokens: set[str] = set()
        self._load_if_exists()
        self._build_source_token_map()

    def _build_source_token_map(self) -> None:
        """Index which filename tokens are distinctive enough to boost on.

        Author surnames and title words unique to a few files ("nise", "ogata",
        "kharitonov") identify a source; words common across the corpus do not.
        """
        if not self.chunks:
            return
        files_per_token: dict[str, set[str]] = {}
        all_files: set[str] = set()
        for chunk in self.chunks:
            fname = str(chunk.get("metadata", {}).get("filename", ""))
            if not fname:
                continue
            all_files.add(fname)
            for token in set(tokenize_corpus(fname)):
                files_per_token.setdefault(token, set()).add(fname)

        if not all_files:
            return
        cutoff = max(2, int(len(all_files) * SOURCE_TOKEN_MAX_FILE_FRACTION))
        self._distinctive_source_tokens = {
            token
            for token, files in files_per_token.items()
            if len(files) <= cutoff
            and len(token) > 2
            and not token.isdigit()
            and token not in GENERIC_SOURCE_TOKENS
        }

    def build_from_chunks(self, chunks: list[Chunk]) -> None:
        if BM25Okapi is None:
            return
        self.chunks = [c.to_dict() for c in chunks]
        corpus = [tokenize_corpus(c.text) for c in chunks]
        self.bm25 = BM25Okapi(corpus)
        self.save()

    def add_chunks(self, new_chunks: list[Chunk]) -> None:
        """Add new chunks to the existing index and rebuild BM25."""
        new_dict_chunks = [c.to_dict() for c in new_chunks]
        self.chunks.extend(new_dict_chunks)
        corpus = [tokenize_corpus(c["text"]) for c in self.chunks]
        self.bm25 = BM25Okapi(corpus)
        # A newly uploaded document may introduce a new author/title, so the
        # distinctive-source vocabulary has to be recomputed alongside BM25.
        self._build_source_token_map()
        self.save()

    def save(self) -> None:
        self.index_dir.mkdir(parents=True, exist_ok=True)
        with (self.index_dir / "chunks.json").open("w", encoding="utf-8") as f:
            json.dump(self.chunks, f, ensure_ascii=False, indent=2)
        with (self.index_dir / "bm25.pkl").open("wb") as f:
            pickle.dump(self.bm25, f)

    def _load_if_exists(self) -> bool:
        chunks_file = self.index_dir / "chunks.json"
        bm25_file = self.index_dir / "bm25.pkl"
        if chunks_file.exists() and bm25_file.exists():
            try:
                with chunks_file.open("r", encoding="utf-8") as f:
                    self.chunks = json.load(f)
                with bm25_file.open("rb") as f:
                    self.bm25 = pickle.load(f)
                return True
            except Exception as exc:
                print(f"Warning: Failed to load existing index: {exc}")
        return False

    def search(self, query: str, top_k: int = 5, source_filter: str | None = None) -> list[dict[str, Any]]:
        if not self.bm25 or not self.chunks:
            return []

        tokens = tokenize_corpus(query)
        if not tokens:
            return []

        scores = self.bm25.get_scores(tokens)
        ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)

        # Re-rank a widened candidate pool so that a chunk from an explicitly
        # named source can be promoted above a body-text-only BM25 match.
        query_source_tokens = set(tokens) & self._distinctive_source_tokens
        pool = ranked[: max(top_k * 10, 60)]
        rescored: list[tuple[float, int]] = []
        for idx in pool:
            base = float(scores[idx])
            if base <= 0.0:
                continue
            score = base
            if query_source_tokens:
                fname_tokens = set(tokenize_corpus(str(self.chunks[idx].get("metadata", {}).get("filename", ""))))
                if query_source_tokens & fname_tokens:
                    score *= SOURCE_MATCH_BOOST
            rescored.append((score, idx))
        rescored.sort(key=lambda pair: pair[0], reverse=True)

        results = []
        for score, idx in rescored:
            chunk = self.chunks[idx]
            if source_filter and source_filter.lower() not in chunk["source_path"].lower():
                continue
            fname = chunk["metadata"].get("filename", "unknown")
            source_name, is_published = display_source_name(fname)
            results.append({
                "chunk_id": chunk["chunk_id"],
                "score": round(score, 3),
                "source": chunk["source_path"],
                "filename": fname,
                # User-facing label -- never show the raw filename in an answer.
                "source_name": source_name,
                "is_published_work": is_published,
                "page": chunk["metadata"].get("page"),
                "text": chunk["text"],
            })
            if len(results) >= top_k:
                break

        return results


# A single process-wide index instance. The agent, the retrieval tool, and the
# document-upload endpoint must all read and mutate the SAME in-memory object:
# with separate instances, a document uploaded through the web UI is written to
# disk but stays invisible to the running agent until the server restarts.
_shared_index: ControlRAGIndex | None = None


def get_shared_index(index_dir: Path = INDEX_DIR) -> ControlRAGIndex:
    """Return the process-wide shared RAG index, loading it on first use."""
    global _shared_index
    if _shared_index is None:
        _shared_index = ControlRAGIndex(index_dir)
    return _shared_index


# --- Human-readable source names -------------------------------------------
# Indexed filenames carry private organisational cruft -- owner initials
# ("JD_", "B&B_"), course codes ("AMC", "PLMMR"), and scan artefacts
# ("txtbk", "DEFINITIVO", a doubled ".pdf.pdf"). Those must never reach a user
# as a citation, so every hit also carries a cleaned display name plus whether
# it is a published work (citable by author/title) or personal course notes
# (referred to generically).

_PUBLISHED_SOURCES: dict[str, str] = {
    "norman s. nise - control systems engineering": "Nise, *Control Systems Engineering*",
    "ogata modern control engineering 5th txtbk": "Ogata, *Modern Control Engineering* (5th ed.)",
}

# Course code -> the subject it stands for, so a generic remainder such as
# "CAM Course Notes Part 2" still resolves to something meaningful.
_COURSE_CODES: dict[str, str] = {
    "SAS": "Safety in Automation Systems",
    "AMC": "Advanced and Multivariable Control",
    "CIR": "Control of Industrial Robots",
    "PLMMR": "Perception, Localization and Mapping for Mobile Robots",
    "NC": "Networked Control",
    "CAM": "Computer-Aided Manufacturing",
    "PSC": "Production Systems Control",
    "MIDA": "Model Identification and Data Analysis",
    "MIDA1": "Model Identification and Data Analysis",
    "ACEHV": "Autonomous and Connected Electric/Hybrid Vehicles",
    "ACHEV": "Autonomous and Connected Electric/Hybrid Vehicles",
    "ACAV": "Autonomous and Connected Vehicles",
    "DDCSD": "Data-Driven Control System Design",
    "SACI": "Industrial Automation and Communication Systems",
    "ICT": "Information and Communication Technology",
}

_OWNER_PREFIX_RE = re.compile(r"^(?:B&B|BBB|JD|LP|EC|AC|FG|RB|XX)[_\-\s]+", re.IGNORECASE)
_NOISE_RE = re.compile(
    r"\b(?:txtbk|definitivo|margini\s+larghi|theory\s+notes|practice\s+notes|final)\b|\(.*?\)",
    re.IGNORECASE,
)


def display_source_name(filename: str) -> tuple[str, bool]:
    """Map an indexed filename to (display name, is_published_work)."""
    if not filename:
        return ("local reference", False)

    stem = str(filename)
    while True:
        lowered = stem.lower()
        for ext in (".pdf", ".md", ".txt", ".json", ".jsonl"):
            if lowered.endswith(ext):
                stem = stem[: -len(ext)]
                break
        else:
            break

    key = " ".join(stem.split()).lower()
    if key in _PUBLISHED_SOURCES:
        return (_PUBLISHED_SOURCES[key], True)

    name = _OWNER_PREFIX_RE.sub("", stem)

    # A leading all-caps token is a course code; swap it for its subject.
    subject = ""
    parts = name.replace("_", " ").split()
    if parts:
        head = parts[0].strip(":-").upper()
        if head in _COURSE_CODES:
            subject = _COURSE_CODES[head]
            parts = parts[1:]
        elif len(parts) > 1 and 2 <= len(head) <= 6 and head.isalpha() and parts[0].isupper():
            parts = parts[1:]

    remainder = _NOISE_RE.sub("", " ".join(parts))
    remainder = " ".join(remainder.replace("_", " ").split()).strip(" -–—")

    generic = remainder.lower() in {
        "", "course notes", "lecture notes", "notes", "exercise sessions",
        "summary", "course notes part 1", "course notes part 2", "lectures",
    }
    if subject and generic:
        label = subject
    elif subject and remainder:
        label = subject if remainder.lower() in subject.lower() else f"{subject} - {remainder}"
    else:
        label = remainder or subject or "local reference"

    return (f"{label} (course notes)", False)