backfill_embeddings.py

r"""
One-off backfill: compute embeddings for rows in the DuckDB that are missing them.

Run once (with the app AND DBeaver closed, so the file isn't locked):
  d:\2026\July\Agentic-RAG-with-LangGraph-and-Ollama\.venv\Scripts\python.exe backfill_embeddings.py
"""

import os
import time
import duckdb
import pandas as pd
from dotenv import load_dotenv

load_dotenv()


def env(key):
    val = os.getenv(key)
    if val is None or val == "":
        raise RuntimeError(f"Missing required config '{key}' in .env")
    return val


# All config comes from .env (no hardcoded values here).
DUCKDB_PATH = env("DUCKDB_PATH")
TABLE = env("TABLE")
HNSW_INDEX = f"{TABLE}_embedding_hnsw"

# Subset + throttling for the Gemini free tier (100 embed requests/minute).
MAX_ROWS = int(env("EMBED_MAX_ROWS"))   # cap rows per run
CHUNK = int(env("EMBED_CHUNK"))         # texts per API request
SLEEP = float(env("EMBED_SLEEP"))       # seconds between chunks

# Must match duckdb_app.py's embedding backend (set in .env).
EMBED_PROVIDER = env("EMBED_PROVIDER")
EMBED_MODEL = env("EMBED_MODEL")
EMBED_DIM = int(env("EMBED_DIM"))


def embed_documents(texts):
    if EMBED_PROVIDER == "google":
        from langchain_google_genai import GoogleGenerativeAIEmbeddings
        emb = GoogleGenerativeAIEmbeddings(model=EMBED_MODEL,
                                           output_dimensionality=EMBED_DIM)
        return emb.embed_documents(list(texts))
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer(EMBED_MODEL)
    return [v.tolist() for v in model.encode(list(texts), normalize_embeddings=True)]


def build_embed_text(r):
    return (
        f"{r['make']} {r['model']} {r['modelYear']} | "
        f"manufacturer: {r['manufacturer']} | component: {r['component']} | "
        f"vin: {r['vin']} | crash: {r['crash']} | fire: {r['fire']} | "
        f"injuries: {r['numberOfInjuries']} | deaths: {r['numberOfDeaths']} | "
        f"filed: {r['dateComplaintFiled']} | summary: {r['summary']}"
    )


def main():
    con = duckdb.connect(DUCKDB_PATH)
    con.execute("INSTALL vss; LOAD vss;")
    con.execute("SET hnsw_enable_experimental_persistence = true;")

    # Ensure columns exist (in case this DB predates them).
    con.execute(f"ALTER TABLE {TABLE} ADD COLUMN IF NOT EXISTS embed_text VARCHAR")
    con.execute(f"ALTER TABLE {TABLE} ADD COLUMN IF NOT EXISTS embedding FLOAT[{EMBED_DIM}]")

    # If the embedding dimension changed, recreate the column to match.
    row = con.execute(
        "SELECT data_type FROM information_schema.columns "
        "WHERE table_name = ? AND column_name = 'embedding'", [TABLE]).fetchone()
    want = f"FLOAT[{EMBED_DIM}]"
    if row and row[0] != want:
        print(f"Embedding dim changed {row[0]} -> {want}; recreating column…")
        con.execute(f"DROP INDEX IF EXISTS {HNSW_INDEX}")
        con.execute(f"ALTER TABLE {TABLE} DROP COLUMN embedding")
        con.execute(f"ALTER TABLE {TABLE} ADD COLUMN embedding {want}")

    # Ensure the HNSW vector index exists (UPDATEs below maintain it).
    con.execute(f"""
        CREATE INDEX IF NOT EXISTS {HNSW_INDEX}
        ON {TABLE} USING HNSW (embedding) WITH (metric = 'cosine')
    """)

    # Subset: only embed up to MAX_ROWS per run (free-tier friendly).
    todo = con.execute(f"""
        SELECT odiNumber, make, model, modelYear, manufacturer, component, summary,
               vin, crash, fire, numberOfInjuries, numberOfDeaths, dateComplaintFiled
        FROM {TABLE}
        WHERE embedding IS NULL
        LIMIT {MAX_ROWS}
    """).fetchdf()

    if todo.empty:
        print("Nothing to backfill — all selected rows already have embeddings.")
        return

    todo["embed_text"] = todo.apply(build_embed_text, axis=1)
    total = len(todo)
    print(f"Backfilling {total} rows (subset, MAX_ROWS={MAX_ROWS}) with "
          f"{EMBED_PROVIDER} ({EMBED_MODEL}), chunk={CHUNK}…")

    done = 0
    # Process in small chunks, committing each so progress survives a quota error.
    for start in range(0, total, CHUNK):
        batch = todo.iloc[start:start + CHUNK]
        try:
            vecs = embed_documents(batch["embed_text"].tolist())
        except Exception as e:
            print(f"  Stopped at {done}/{total} — {type(e).__name__}: {e}")
            print("  Re-run later to continue where it left off (resumable).")
            break

        updates = pd.DataFrame({
            "odiNumber": batch["odiNumber"].values,
            "embed_text": batch["embed_text"].values,
            "embedding": list(vecs),
        })
        con.register("upd", updates)
        con.execute(f"""
            UPDATE {TABLE} AS t
            SET embed_text = u.embed_text, embedding = u.embedding
            FROM upd AS u WHERE t.odiNumber = u.odiNumber
        """)
        con.unregister("upd")
        done += len(batch)
        print(f"  embedded {done}/{total}")
        time.sleep(SLEEP)   # throttle to respect the rate limit

    embedded = con.execute(
        f"SELECT COUNT(*) FROM {TABLE} WHERE embedding IS NOT NULL").fetchone()[0]
    remaining = con.execute(
        f"SELECT COUNT(*) FROM {TABLE} WHERE embedding IS NULL").fetchone()[0]
    print(f"Done. Embedded so far: {embedded} · still NULL: {remaining}")


if __name__ == "__main__":
    main()