scripts/restore_snapshot.py

#!/usr/bin/env python3
"""
Restore the committed CDMS Qdrant snapshot into a running Docker Qdrant.

Run once per deploy, after starting the Qdrant container and before serving.
Both the Streamlit UI and the FastAPI server then read from this Docker Qdrant
(localhost:6333) — which is why the deploy uses Docker rather than the
single-writer on-disk store.

Usage:
    docker run -d -p 6333:6333 qdrant/qdrant
    python scripts/restore_snapshot.py
    # then serve with: CDMS_OFFLINE_INDEX=1 QDRANT_REQUIRE_DOCKER=1 streamlit run ...
"""

import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

import os

import requests

from src.config.paths import SNAPSHOT_DIR

HOST = os.environ.get("QDRANT_HOST", "localhost")
PORT = int(os.environ.get("QDRANT_PORT", "6333"))
COLLECTION = "cdms_documents"


def main():
    snapshot = SNAPSHOT_DIR / "cdms_documents.snapshot"
    if not snapshot.exists():
        sys.exit(f"❌ Snapshot not found: {snapshot}\n"
                 f"   Did you run `git lfs pull`? Build it with scripts/build_index.py.")

    upload_url = (f"http://{HOST}:{PORT}/collections/{COLLECTION}"
                  f"/snapshots/upload?priority=snapshot")
    print(f"⬆️  Uploading {snapshot.name} ({snapshot.stat().st_size / 1e6:.1f} MB) "
          f"-> {HOST}:{PORT}/{COLLECTION}")
    try:
        with open(snapshot, "rb") as f:
            resp = requests.post(upload_url, files={"snapshot": f}, timeout=300)
        resp.raise_for_status()
    except Exception as e:
        sys.exit(f"❌ Restore failed: {e}\n"
                 f"   Is the Docker Qdrant running? `docker run -d -p 6333:6333 qdrant/qdrant`")

    # Confirm the collection is populated.
    try:
        info = requests.get(f"http://{HOST}:{PORT}/collections/{COLLECTION}", timeout=30).json()
        count = info.get("result", {}).get("points_count")
        print(f"✅ Restored. {COLLECTION} points_count = {count}")
    except Exception:
        print("✅ Restore request accepted (could not read point count).")


if __name__ == "__main__":
    main()