scripts/tag_topics.py

# File Objective  : Tag every indexed chunk with the topics whose string appears in its article.
# Scope           : Dev script — payload-only update to an already-populated Qdrant collection.
# What it does    : 1. Loads the topic vocabulary (Gartner + GDELT) from topics.csv.
#                   2. Loads all HuffPost articles into a DataFrame (doc_id, text).
#                   3. Runs ONE word-boundary regex per topic across the whole column,
#                      collecting which doc_ids mention each topic.
#                   4. Inverts that into doc_id → [topics] and writes it to each chunk's
#                      `topics` payload field in Qdrant — vectors are never touched.
# What it does not: 1. Re-embed, re-chunk, or re-ingest (the slow 2-hour path).
#                   2. Do semantic / fuzzy matching — a topic is "present" only if its
#                      exact string occurs in the article (case-insensitive, word-bounded).
#                   3. Create payload indexes or compute gap scores.
#
# Live logging     : prints a banner, then per-phase progress, then PER TOPIC the number of
#                    docs matched, and for every 50th topic dumps the raw matched documents.
#
# Usage            : PYTHONPATH=packages python scripts/tag_topics.py          # full run
#                    LIMIT=10 PYTHONPATH=packages python scripts/tag_topics.py # 10-doc dry test

from __future__ import annotations

import csv
import json
import os
import re
import time
import uuid
from pathlib import Path

import pandas as pd
from qdrant_client import models as qm

from vantage_core.adapters.vector_qdrant import QdrantVectorRepo
from vantage_core.constants import QDRANT_COLLECTION_NAME

# ── Configuration ────────────────────────────────────────────────────────────
TOPICS_CSV   = Path("packages/vantage_core/seeds/topics.csv")
DATA_PATH    = Path("data/news_category.json")
SCROLL_BATCH = 512    # points fetched per Qdrant scroll page
WRITE_BATCH  = 256    # set-payload ops per batch_update_points request
DUMP_EVERY   = 50     # print raw matched docs on every Nth topic
DUMP_SAMPLE  = 5      # how many raw docs to show in that dump

LIMIT = int(os.environ["LIMIT"]) if os.environ.get("LIMIT") else None


# ── Logging helper ───────────────────────────────────────────────────────────
def log(msg: str = "") -> None:
    """Print immediately (unbuffered) so progress is visible live."""
    print(msg, flush=True)


def banner(title: str) -> None:
    log(f"\n── {title} " + "─" * max(0, 56 - len(title)))


# ── Phase 1: load inputs ─────────────────────────────────────────────────────
def load_topics() -> list[str]:
    """Return the topic vocabulary from topics.csv (the `keyword` column)."""
    with TOPICS_CSV.open() as f:
        return [row["keyword"].strip() for row in csv.DictReader(f) if row["keyword"].strip()]


def load_articles(limit: int | None) -> pd.DataFrame:
    """Stream HuffPost JSONL into a DataFrame of (doc_id, text).

    `text` mirrors exactly what was ingested as chunk_text: headline + short_description.
    `doc_id` is the row's yield-index as a string — identical to the loader/ingest scheme.
    """
    rows: list[dict[str, str]] = []
    doc_id = 0
    with DATA_PATH.open(encoding="utf-8") as fh:
        for line in fh:
            row      = json.loads(line)
            headline = row.get("headline", "").strip()
            desc     = row.get("short_description", "").strip()
            text     = f"{headline}. {desc}" if desc else headline
            if len(text) < 20:                  # skip noise rows (matches loader.py)
                continue
            rows.append({"doc_id": str(doc_id), "text": text})
            doc_id += 1
            if limit is not None and doc_id >= limit:
                break
    return pd.DataFrame(rows)


# ── Phase 2: match topics → documents ────────────────────────────────────────
def match_topics_to_docs(df: pd.DataFrame, topics: list[str]) -> dict[str, list[str]]:
    """For each topic, find the docs whose text contains it; invert to doc_id → [topics].

    One vectorized regex scan per topic across the whole column (pandas str.contains in C).
    Prints the match count for every topic, and the raw matched docs every DUMP_EVERY topics.
    """
    doc_topics: dict[str, list[str]] = {doc_id: [] for doc_id in df["doc_id"]}
    t0 = time.time()

    for i, topic in enumerate(topics, start=1):
        pattern = rf"\b{re.escape(topic)}\b"        # word-bounded so 'ai' ≠ 'maintain'
        hit_mask = df["text"].str.contains(pattern, case=False, regex=True, na=False)
        matched = df.loc[hit_mask]

        for doc_id in matched["doc_id"]:
            doc_topics[doc_id].append(topic)

        log(f"  [{i:>3}/{len(topics)}] {topic[:45]:<45} → {len(matched):>6} docs")

        if i % DUMP_EVERY == 0 and not matched.empty:
            log(f"      ── raw matches for '{topic}' (showing {DUMP_SAMPLE}) ──")
            for _, r in matched.head(DUMP_SAMPLE).iterrows():
                log(f"        doc {r['doc_id']:>6} | {r['text'][:90]}")

    log(f"\n  matched {len(topics)} topics over {len(df)} docs in {time.time()-t0:.1f}s")
    return doc_topics


# ── Phase 3a: write topic→doc mapping to CSV ─────────────────────────────────
COVERAGE_CSV = Path("data/topic_coverage.csv")


def write_coverage_csv(doc_topics: dict[str, list[str]]) -> None:
    """Write (topic, doc_id) pairs to CSV — one row per match."""
    rows = [
        {"topic": topic, "doc_id": doc_id}
        for doc_id, topics in doc_topics.items()
        for topic in topics
    ]
    COVERAGE_CSV.parent.mkdir(parents=True, exist_ok=True)
    with COVERAGE_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["topic", "doc_id"])
        writer.writeheader()
        writer.writerows(rows)
    log(f"  [csv] wrote {len(rows)} rows → {COVERAGE_CSV}")


# ── Phase 3b: write topics into Qdrant payloads (commented out) ───────────────
def point_id(chunk_id: str) -> str:
    """Stable Qdrant point UUID from chunk_id — mirrors Chunk.qdrant_uuid()."""
    return str(uuid.uuid5(uuid.NAMESPACE_DNS, chunk_id))


def write_matched_only(doc_topics: dict[str, list[str]]) -> None:
    """Full run: only write to docs that actually matched ≥1 topic.

    Computes point IDs directly from doc_id (no collection scroll needed).
    Docs with zero topic matches are left untouched — no empty writes.
    Note: uses {doc_id}-0 (covers 99.87% of docs; multi-chunk docs get chunk-0 updated).
    """
    client = QdrantVectorRepo()._client
    matched = {doc_id: topics for doc_id, topics in doc_topics.items() if topics}
    ops: list[qm.SetPayloadOperation] = []
    written = 0
    t0 = time.time()

    def flush() -> None:
        nonlocal written
        if ops:
            client.batch_update_points(collection_name=QDRANT_COLLECTION_NAME, update_operations=ops)
            written += len(ops)
            ops.clear()

    for doc_id, topics in matched.items():
        ops.append(qm.SetPayloadOperation(set_payload=qm.SetPayload(
            payload={"topics": topics}, points=[point_id(f"{doc_id}-0")],
        )))
        if len(ops) >= WRITE_BATCH:
            flush()
            log(f"  [qdrant] written={written:>6}/{len(matched)}  elapsed={time.time()-t0:.0f}s")

    flush()
    log(f"  [qdrant] done — wrote to {written} matched docs (skipped {len(doc_topics)-written} with no topics)  elapsed={time.time()-t0:.1f}s")


def write_targeted(doc_topics: dict[str, list[str]]) -> list[str]:
    """LIMIT/dry-run mode: write directly to computed point IDs ({doc_id}-0).

    No filter, no index, no full scroll — touches only the docs under test.
    """
    client = QdrantVectorRepo()._client
    ops = [
        qm.SetPayloadOperation(set_payload=qm.SetPayload(
            payload={"topics": topics}, points=[point_id(f"{doc_id}-0")],
        ))
        for doc_id, topics in doc_topics.items()
    ]
    client.batch_update_points(collection_name=QDRANT_COLLECTION_NAME, update_operations=ops)
    log(f"  [qdrant] wrote topics to {len(ops)} points (doc_id-0)")
    return [point_id(f"{doc_id}-0") for doc_id in doc_topics]


def read_back(ids: list[str]) -> None:
    """Re-fetch the written points from Qdrant to confirm the payload persisted."""
    client = QdrantVectorRepo()._client
    points = client.retrieve(
        collection_name=QDRANT_COLLECTION_NAME, ids=ids,
        with_payload=["doc_id", "title", "topics"], with_vectors=False,
    )
    banner("Read-back from Qdrant")
    for p in sorted(points, key=lambda x: int(x.payload["doc_id"])):
        pl = p.payload
        log(f"  doc {pl['doc_id']:>3} | topics={pl.get('topics')}")
        log(f"         {pl['title'][:80]}")


# ── Orchestration ────────────────────────────────────────────────────────────
def main() -> None:
    topics = load_topics()
    dry_run = LIMIT is not None

    banner(f"Tag Topics{'  [LIMIT=' + str(LIMIT) + ']' if dry_run else ''}")
    log(f"   topics={len(topics)}  collection={QDRANT_COLLECTION_NAME}  source={DATA_PATH}")

    banner("Phase 1 — load articles")
    df = load_articles(LIMIT)
    log(f"   loaded {len(df)} articles")

    banner("Phase 2 — match topics → documents")
    doc_topics = match_topics_to_docs(df, topics)
    matched_docs = sum(1 for tps in doc_topics.values() if tps)
    log(f"   docs with ≥1 topic: {matched_docs}/{len(doc_topics)}")

    banner("Phase 3a — write coverage CSV")
    write_coverage_csv(doc_topics)

    # banner("Phase 3b — write topics to Qdrant")
    # if dry_run:
    #     ids = write_targeted(doc_topics)
    #     read_back(ids)
    # else:
    #     write_matched_only(doc_topics)

    banner("Done")
    log("")


if __name__ == "__main__":
    main()