build_vectorstore.py

"""
Build the Chroma vector database for Week 8 (Day 2 notebook — no Jupyter needed).

Usage:
    cd week8
    uv run python build_vectorstore.py

Requires HF_TOKEN in .env (repo root or week8).
Uses the lite dataset by default (~faster than full 400k products).
"""

import os
import sys
from pathlib import Path

from huggingface_hub import login
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
import chromadb

from agents.items import Item
from env_utils import REPO_ROOT, load_project_env

load_project_env()

DB = "products_vectorstore"
COLLECTION_NAME = "products"
BATCH_SIZE = 1000
LITE_MODE = os.getenv("LITE_MODE", "true").lower() in ("true", "1", "yes")
HF_USER = "ed-donner"


def main():
    if not os.getenv("HF_TOKEN"):
        print(f"ERROR: HF_TOKEN is missing.")
        print(f"  Add HF_TOKEN=hf_... to: {REPO_ROOT / '.env'}")
        sys.exit(1)

    login(token=os.environ["HF_TOKEN"], add_to_git_credential=False)

    dataset = f"{HF_USER}/items_lite" if LITE_MODE else f"{HF_USER}/items_full"
    print(f"Loading dataset: {dataset} (LITE_MODE={LITE_MODE})")
    train, val, test = Item.from_hub(dataset)
    print(f"Loaded {len(train):,} training items, {len(val):,} val, {len(test):,} test")

    client = chromadb.PersistentClient(path=DB)
    existing = [c.name for c in client.list_collections()]

    if COLLECTION_NAME in existing:
        print(f"Collection '{COLLECTION_NAME}' already exists in {DB}/")
        count = client.get_collection(COLLECTION_NAME).count()
        print(f"  → {count:,} vectors. Delete products_vectorstore/ to rebuild.")
        return

    print("Loading embedding model (first run downloads ~90MB)...")
    encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    collection = client.create_collection(COLLECTION_NAME)
    print(f"Encoding {len(train):,} products in batches of {BATCH_SIZE}...")

    for i in tqdm(range(0, len(train), BATCH_SIZE), desc="Building vectorstore"):
        batch = train[i : i + BATCH_SIZE]
        documents = [item.summary for item in batch]
        vectors = encoder.encode(documents).astype(float).tolist()
        metadatas = [{"category": item.category, "price": item.price} for item in batch]
        ids = [f"doc_{j}" for j in range(i, i + len(batch))]
        collection.add(
            ids=ids,
            documents=documents,
            embeddings=vectors,
            metadatas=metadatas,
        )

    count = collection.count()
    print(f"\nDone! {count:,} vectors saved to {Path(DB).resolve()}")


if __name__ == "__main__":
    main()