scripts/build_index.py

"""Build the baked-in demo index from data/corpus_sample.jsonl.

Produces, under app/prebuilt/:
  - qdrant/        on-disk Qdrant collection with the dense vectors
  - lexical.pkl    exported BM25 state (ids, tokens, metadata)

Baking both into the image means the demo is never empty, even after the host
restarts its ephemeral storage (spec, Phase 1 acceptance).

Usage:
    python scripts/build_index.py
"""

from __future__ import annotations

import json
import logging
import pickle
import sys
from pathlib import Path

# Make src/ importable when run as a script.
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from streamsearch import config
from streamsearch.schema import Document
from streamsearch.search import HybridSearcher

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
log = logging.getLogger("build_index")


def load_corpus(path: Path) -> list[Document]:
    docs = []
    with path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                docs.append(Document.from_json(json.loads(line)))
    return docs


def main() -> None:
    if config.QDRANT_URL:
        log.info("QDRANT_URL set -> building into cloud Qdrant at %s", config.QDRANT_URL)
    else:
        # Fresh local build: wipe any stale on-disk index first.
        import shutil

        if config.QDRANT_LOCAL_PATH.exists():
            shutil.rmtree(config.QDRANT_LOCAL_PATH)
        log.info("Building embedded on-disk index at %s", config.QDRANT_LOCAL_PATH)

    docs = load_corpus(config.CORPUS_PATH)
    log.info("Loaded %d documents from %s", len(docs), config.CORPUS_PATH)

    client = config.make_qdrant_client()
    searcher = HybridSearcher(client)
    searcher.ensure_collection()
    searcher.add_documents(docs)

    config.PREBUILT_DIR.mkdir(parents=True, exist_ok=True)
    with config.LEXICAL_PATH.open("wb") as f:
        pickle.dump(searcher.export_lexical(), f)
    log.info("Wrote lexical export to %s", config.LEXICAL_PATH)

    # Release the embedded client's file lock so the app can open it.
    del client
    log.info("Done. Indexed %d docs.", len(docs))


if __name__ == "__main__":
    main()