File size: 3,363 Bytes
f8b48da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Initialize the AskMyDocs platform.
Creates all PostgreSQL tables (including keyword postings), bootstraps the
Qdrant collection and the keyword search objects, and ensures the
evaluation output directories exist.
Usage:
python scripts/init_db.py # full bootstrap (PG + Qdrant + keyword)
python scripts/init_db.py --no-vectors # PG only (offline / unit-test setup)
The default .env connects to the services started by
`docker compose up -d`. Each step is independent — a failure in vector or
keyword bootstrap does not roll back the database schema.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from app.config import get_settings # noqa: E402
from app.db.session import init_db # noqa: E402
def ensure_eval_dirs() -> None:
for rel in ("evals/reports", "evals/results", "evals/baselines"):
(ROOT / rel).mkdir(parents=True, exist_ok=True)
def bootstrap_vectors() -> tuple[bool, bool]:
"""Best-effort bootstrap of Qdrant + keyword search.
Returns ``(qdrant_ok, keyword_ok)``. Failures are logged but do not
raise, so this script can run in environments where the vector stores
are not yet up.
"""
qdrant_ok = keyword_ok = False
try:
from app.retrieval.vector import VectorStore
VectorStore().ensure_collection()
qdrant_ok = True
print("Qdrant collection ready.")
except Exception as exc: # noqa: BLE001
print(f"[warn] Qdrant bootstrap skipped: {exc}")
try:
from app.retrieval.bm25 import BM25Indexer
BM25Indexer().ensure_index()
keyword_ok = True
print("Keyword (tsvector) search ready.")
except Exception as exc: # noqa: BLE001
print(f"[warn] Keyword bootstrap skipped: {exc}")
return qdrant_ok, keyword_ok
def main() -> int:
parser = argparse.ArgumentParser(description="Initialize AskMyDocs storage.")
parser.add_argument(
"--no-vectors",
action="store_true",
help="Only initialize PostgreSQL (skip Qdrant + keyword bootstrap).",
)
args = parser.parse_args()
settings = get_settings()
print(f"Connecting to: {settings.database_url}")
try:
init_db()
except Exception as exc: # noqa: BLE001
print(f"Failed to initialize database: {exc}")
print("Is PostgreSQL running? Try `make docker-up` first.")
return 1
print("Database schema initialized.")
# Keyword tsvector objects live outside the ORM mapping (trigger-owned),
# so create_all doesn't cover them — ensure here for fresh and old DBs.
try:
from app.retrieval.bm25 import BM25Indexer
BM25Indexer().ensure_index()
print("Keyword (tsvector) objects ensured.")
except Exception as exc: # noqa: BLE001
print(f"[warn] Keyword DDL skipped: {exc}")
if not args.no_vectors:
bootstrap_vectors()
else:
print("Skipping Qdrant + keyword bootstrap (--no-vectors).")
ensure_eval_dirs()
print("Evaluation directories ensured.")
return 0
if __name__ == "__main__":
raise SystemExit(main()) |