scripts/query_proof.py

"""Query the proof-tier index with the SAME hybrid search stack as the demo.

Points search.py at the proof-tier Qdrant (populated by the Spark job) and, for
a genuine BM25 + dense fusion, hydrates the lexical index from the documents the
pipeline wrote to Cassandra. This shows the identical hybrid search working over
the live streamed index — the Phase 4 unification.

Usage:
    python scripts/query_proof.py "support vector machine kernels"
    python scripts/query_proof.py "graph neural networks" --top-k 5
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from streamsearch import config
from streamsearch.schema import Document, CASSANDRA_KEYSPACE, CASSANDRA_TABLE
from streamsearch.search import HybridSearcher, tokenize


def hydrate_from_cassandra(searcher: HybridSearcher, host: str, port: int) -> int:
    """Load streamed docs' metadata from Cassandra to power the BM25 stage."""
    from cassandra.cluster import Cluster

    cluster = Cluster([host], port=port)
    session = cluster.connect(CASSANDRA_KEYSPACE)
    rows = session.execute(f"SELECT id, title, text, source FROM {CASSANDRA_TABLE}")
    docs = [Document(id=r.id, title=r.title or "", text=r.text or "",
                     source=r.source or "") for r in rows]
    cluster.shutdown()

    ids = [d.id for d in docs]
    tokens = [tokenize(d.embed_input) for d in docs]
    meta = {d.id: {"title": d.title, "text": d.text, "source": d.source} for d in docs}
    searcher.hydrate(ids, tokens, meta)
    return len(docs)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("query", help="search query")
    ap.add_argument("--top-k", type=int, default=5)
    ap.add_argument("--qdrant-url", default="http://localhost:6533")
    ap.add_argument("--cassandra-host", default=config.CASSANDRA_HOST)
    ap.add_argument("--cassandra-port", type=int, default=config.CASSANDRA_PORT)
    ap.add_argument("--dense-only", action="store_true",
                    help="skip Cassandra/BM25 hydration (dense stage only)")
    args = ap.parse_args()

    from qdrant_client import QdrantClient

    searcher = HybridSearcher(QdrantClient(url=args.qdrant_url))
    if not args.dense_only:
        n = hydrate_from_cassandra(searcher, args.cassandra_host, args.cassandra_port)
        print(f"Hydrated BM25 from Cassandra: {n} streamed docs\n")

    print(f"Query: {args.query!r}\n" + "-" * 60)
    for i, r in enumerate(searcher.search(args.query, top_k=args.top_k), 1):
        print(f"{i}. {r.title or '(untitled)'}")
        print(f"   RRF {r.score:.4f}  dense_rank={r.dense_rank}  bm25_rank={r.bm25_rank}  id={r.id}")


if __name__ == "__main__":
    main()