ingest.py

"""ingest.py - the data connector run: sources -> chunks -> embeddings -> store
                AND -> entity graph, in one pass.

  cd ask-london
  EMBED_MODE=hash STORE_MODE=local python ingest.py            # offline, today
  EMBED_MODE=hash STORE_MODE=local python ingest.py --verify   # + sample query
  python ingest.py                                             # uses .env modes (NIM/Milvus)

Writes:  index/store.npz + index/store.json   (vectors + chunks)
         index/graph.json                     (borough/indicator/source graph)
"""
from __future__ import annotations

import argparse
import time

import config
import connector
import embeddings
import store as store_mod
from graphstore import GraphStore


def run(limit: int | None, batch: int = 64) -> tuple[int, int]:
    print(config.banner())
    facts = connector.load_fixtures(config.DATA_DIR, limit, include_economy=True)
    if not facts:
        raise SystemExit(f"no facts found in {config.DATA_DIR} "
                         "(expected fingertips_ind90366.csv / nomis_nm_1_1.csv)")
    print(f"connector: {len(facts)} fact chunks from "
          f"{len({f['source'] for f in facts})} source(s)")

    store = store_mod.get_store(config.EMBED_DIM)
    graph = GraphStore()
    t0 = time.time()
    for i in range(0, len(facts), batch):
        chunk = facts[i:i + batch]
        vecs = embeddings.embed([f["text"] for f in chunk], "passage")
        store.insert(
            ids=[f["id"] for f in chunk], vecs=vecs,
            contents=[f["text"] for f in chunk],
            metas=[{"gss": f["gss"], "borough": f["borough"],
                    "indicator": f["indicator"], "period": f["period"],
                    "value": f["value"], "source": f["source"],
                    "source_title": f["source_title"], "source_url": f["source_url"]}
                   for f in chunk],
        )
        for f in chunk:
            graph.add_fact(f)
        print(f"  embedded+stored {min(i + batch, len(facts))}/{len(facts)}")

    store.save(config.STORE_PATH)
    graph.save(config.GRAPH_PATH)
    # SQL lane: a DuckDB facts table for exact aggregates (best-effort; needs duckdb)
    n_sql = 0
    try:
        import sql_db
        n_sql = sql_db.build_facts_db(facts, config.FACTS_DB)
    except Exception as e:  # noqa: BLE001
        print(f"  [sql] facts DB skipped: {repr(e)[:100]}")
    dt = time.time() - t0
    gs = graph.stats()
    print(f"\nDONE in {dt:.1f}s")
    print(f"  store : {len(store)} chunks ({config.EMBED_DIM}D) -> {config.STORE_PATH}.*")
    print(f"  graph : {gs['nodes']} nodes {gs['edges']} edges "
          f"(nodes {gs['node_types']}, edges {gs['edge_types']}) -> {config.GRAPH_PATH}")
    if n_sql:
        print(f"  sql   : {n_sql} rows -> {config.FACTS_DB} (facts table for the SQL lane)")
    return len(store), gs["nodes"]


def verify() -> None:
    """Self-contained smoke: a vector query + a graph cross-source join."""
    store = store_mod.LocalStore.load(config.STORE_PATH) if config.STORE_MODE == "local" \
        else store_mod.get_store(config.EMBED_DIM)
    graph = GraphStore.load(config.GRAPH_PATH)

    print("\n=== VERIFY 1: vector search 'life expectancy in Camden' ===")
    q = embeddings.embed_one("life expectancy in Camden", "query")
    for i, h in enumerate(store.search(q, 3), 1):
        print(f"  {i}. (score={h['score']:.3f}) {h['content'][:130]}")

    print("\n=== VERIFY 2: graph cross-source join for Camden (all indicators) ===")
    rows = graph.connected_facts(boroughs=["Camden"], limit=6)
    for r in rows:
        print(f"  - [{r['source']}] {r['indicator']}: {r['value']} {r['unit']} "
              f"({r['period'] or 'latest'})")
    if not rows:
        print("  (no Camden facts - check the corpus)")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None, help="max facts per source")
    ap.add_argument("--fetch", action="store_true",
                    help="download sources first (the Spark deploy ingests from source)")
    ap.add_argument("--verify", action="store_true")
    args = ap.parse_args()
    if args.fetch:
        import fetch
        fetch.ensure(do_fetch=True)
    run(args.limit)
    if args.verify:
        verify()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())