vantage / scripts /ingest.py
sourabh gupta
feat(pipeline): RAG quality gate, live coverage, streaming API, recursion limit
c24f0b6
Raw
History Blame Contribute Delete
2.65 kB
# File Objective : Orchestrate Phase 2 end-to-end ingestion into Qdrant.
# Scope : Dev script β€” run once (or re-run idempotently) to populate the collection.
# What it does : 1. Streams Documents β†’ chunks β†’ embeds in batches β†’ upserts dense+sparse.
# 2. Respects INGEST_LIMIT env var (default 2000) for fast sample runs.
# 3. Prints a single progress line and a final summary.
# What it does not: 1. Build graph (Phase 7) or compute gaps (Phase 6).
from __future__ import annotations
import os
from pathlib import Path
from vantage_core.adapters.embedder_st import STEmbedder
from vantage_core.adapters.vector_qdrant import QdrantVectorRepo
from vantage_core.chunker import chunk_document
from vantage_core.loader import load_huffpost
DATA_PATH = Path("data/news_category.json")
BATCH_SIZE = 64 # chunks per upsert β€” balances memory vs network round-trips
def main() -> None:
limit = int(os.getenv("INGEST_LIMIT", "2000"))
print(f"\n── Vantage Ingest ──────────────────────────────────────────")
print(f" dataset : {DATA_PATH} | limit : {limit} docs | batch : {BATCH_SIZE}")
repo = QdrantVectorRepo()
embedder = STEmbedder()
repo.ensure_collection()
buf: list = []
total_docs = total_chunks = 0
for doc in load_huffpost(DATA_PATH, limit=limit):
buf.extend(chunk_document(doc))
total_docs += 1
if len(buf) >= BATCH_SIZE:
_flush(buf, embedder, repo)
total_chunks += len(buf)
print(f" docs={total_docs:>5} chunks={total_chunks:>5}", end="\r")
buf = []
if buf: # flush remainder
_flush(buf, embedder, repo)
total_chunks += len(buf)
final = repo.count()
print(f"\n── Done ────────────────────────────────────────────────────")
print(f" docs={total_docs} chunks_sent={total_chunks} qdrant_total={final}\n")
def _flush(chunks: list, embedder: STEmbedder, repo: QdrantVectorRepo) -> None:
import time
texts = [c.chunk_text for c in chunks]
vectors = embedder.embed(texts)
for attempt in range(3):
try:
repo.upsert(chunks, vectors)
return
except Exception:
if attempt == 2:
raise
time.sleep(8 * (attempt + 1)) # 8s, 16s β€” survive Qdrant Cloud free-tier disconnects
if __name__ == "__main__":
main()