data_indexer.py

"""Build the Qdrant vector index from the curated Markdown policy docs.

Reads every ``Contextmd/*.md`` file (the hand-verified policy Markdown) and
applies a two-stage chunking strategy tuned to keep *named rules* intact:

  Stage 1 — MarkdownHeaderTextSplitter splits on Markdown headers (#, ##, ###).
            Because the policy gives each rule its own header (e.g.
            "### 4.1 Basic Offer Rule (1× Offer)"), this alone keeps most named
            rules together as a single section, and preserves the header text so
            the rule/section name travels with the chunk for citation.

  Stage 2 — RecursiveCharacterTextSplitter only sub-splits sections that exceed
            CHUNK_SIZE (with CHUNK_OVERLAP). Sections shorter than CHUNK_SIZE
            pass through untouched, so short named rules remain single chunks.

Chunks are embedded with FastEmbed (ONNX, local CPU) and upserted into a Qdrant
local collection using cosine distance.
"""

from __future__ import annotations

import logging
import os
import sys
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | data_indexer | %(message)s",
)
log = logging.getLogger("data_indexer")

# --- Config (non-secret operational values; env with build-safe defaults) -----
EMBEDDING_MODEL_ID = os.environ.get("EMBEDDING_MODEL_ID", "BAAI/bge-small-en-v1.5")
QDRANT_STORAGE_PATH = os.environ.get("QDRANT_STORAGE_PATH", "./qdrant_storage")
QDRANT_COLLECTION_NAME = os.environ.get("QDRANT_COLLECTION_NAME", "placement_policy")
CHUNK_SIZE = int(os.environ.get("CHUNK_SIZE", "1000"))
CHUNK_OVERLAP = int(os.environ.get("CHUNK_OVERLAP", "150"))
FASTEMBED_CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", ".fastembed_cache")

# Curated, pre-verified Markdown is the indexed source of truth. (The main
# policy PDF is scanned/image-only, so this hand-checked Markdown is more
# accurate than any automated PDF extraction.) convert_docs.py validates these
# files against the source PDFs before indexing.
SOURCE_DIR = Path("Contextmd")

# Header levels to split on. Metadata keys are ordered from broadest to
# narrowest so we can reconstruct a readable "section path" for citation.
HEADERS_TO_SPLIT_ON = [
    ("#", "doc_title"),
    ("##", "section"),
    ("###", "subsection"),
]


def load_chunks():
    """Run the two-stage split over every Contextmd/*.md file. Returns a list of
    LangChain Documents (page_content + metadata)."""
    from langchain_text_splitters import (
        MarkdownHeaderTextSplitter,
        RecursiveCharacterTextSplitter,
    )

    md_files = sorted(SOURCE_DIR.glob("*.md"))
    if not md_files:
        raise FileNotFoundError(
            f"No Markdown files in {SOURCE_DIR.resolve()}. "
            f"Expected the curated policy Markdown there."
        )

    # strip_headers=False keeps the header line inside the chunk so the rule /
    # section name is embedded and retrievable.
    header_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=HEADERS_TO_SPLIT_ON,
        strip_headers=False,
    )
    char_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
    )

    header_docs = []
    for md_file in md_files:
        text = md_file.read_text(encoding="utf-8")
        sections = header_splitter.split_text(text)
        for doc in sections:
            doc.metadata["source"] = md_file.name
            doc.metadata["section_path"] = _section_path(doc.metadata)
        header_docs.extend(sections)
        log.info("%s -> %d header section(s)", md_file.name, len(sections))

    # Stage 2: only oversized sections get sub-split; short named rules stay whole.
    chunks = char_splitter.split_documents(header_docs)
    log.info(
        "Chunking complete: %d header section(s) -> %d final chunk(s) "
        "(chunk_size=%d, overlap=%d)",
        len(header_docs),
        len(chunks),
        CHUNK_SIZE,
        CHUNK_OVERLAP,
    )
    return chunks


def _section_path(metadata: dict) -> str:
    """Human-readable breadcrumb from header metadata, e.g.
    'RV UNIVERSITY – PLACEMENT POLICY > SECTION 4 — OFFER RULES > 4.1 Basic Offer Rule'."""
    parts = [
        metadata.get("doc_title"),
        metadata.get("section"),
        metadata.get("subsection"),
    ]
    return " > ".join(p for p in parts if p)


def embed_texts(texts):
    from fastembed import TextEmbedding

    log.info("Loading FastEmbed model '%s' ...", EMBEDDING_MODEL_ID)
    model = TextEmbedding(model_name=EMBEDDING_MODEL_ID, cache_dir=FASTEMBED_CACHE_DIR)
    log.info("Embedding %d chunk(s) ...", len(texts))
    vectors = [v.tolist() for v in model.embed(texts)]
    dim = len(vectors[0])
    log.info("Embedded %d chunk(s); vector dim = %d", len(vectors), dim)
    return vectors, dim


def build_index(chunks) -> None:
    from qdrant_client import QdrantClient
    from qdrant_client.models import Distance, PointStruct, VectorParams

    texts = [c.page_content for c in chunks]
    vectors, dim = embed_texts(texts)

    client = QdrantClient(path=QDRANT_STORAGE_PATH)
    try:
        if client.collection_exists(QDRANT_COLLECTION_NAME):
            log.info("Dropping existing collection '%s'", QDRANT_COLLECTION_NAME)
            client.delete_collection(QDRANT_COLLECTION_NAME)

        log.info(
            "Creating collection '%s' (size=%d, distance=COSINE)",
            QDRANT_COLLECTION_NAME,
            dim,
        )
        client.create_collection(
            collection_name=QDRANT_COLLECTION_NAME,
            vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
        )

        points = [
            PointStruct(
                id=i,
                vector=vectors[i],
                payload={
                    "text": chunks[i].page_content,
                    "source": chunks[i].metadata.get("source", ""),
                    "section_path": chunks[i].metadata.get("section_path", ""),
                },
            )
            for i in range(len(chunks))
        ]
        client.upsert(collection_name=QDRANT_COLLECTION_NAME, points=points)

        count = client.count(QDRANT_COLLECTION_NAME).count
        log.info("Upserted %d point(s). Collection now holds %d point(s).", len(points), count)
    finally:
        client.close()


def main() -> int:
    try:
        chunks = load_chunks()
        if not chunks:
            log.error("No chunks produced; nothing to index.")
            return 1
        build_index(chunks)
    except Exception:  # noqa: BLE001 — fail the build loudly on any indexing error
        log.exception("Indexing FAILED")
        return 1
    log.info("Indexing complete.")
    return 0


if __name__ == "__main__":
    sys.exit(main())