scripts/generate_taxonomy.py

#!/usr/bin/env python3
"""Assign consistent topic categories to every indexed video.

Reads tldr + themes from video_summaries, batches them to a text LLM on
OpenRouter, and writes 1–3 categories per video to the new video_categories
table. Enables topic browsing in the Search UI.

Run AFTER the main scene indexing batch is complete.

Usage (from repo root, venv active):
    export SEARCH_UI_DATA_ROOT="$PWD/backend"
    python3 scripts/generate_taxonomy.py \\
        --api-key sk-or-v1-... \\
        [--model qwen/qwen3-30b-a3b-instruct-2507] \\
        [--dry-run]
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import re
import sys
import time

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.dirname(SCRIPT_DIR)
BACKEND_DIR = os.path.join(REPO_ROOT, "backend")
if BACKEND_DIR not in sys.path:
    sys.path.insert(0, BACKEND_DIR)

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
log = logging.getLogger(__name__)

DEFAULT_MODEL = "qwen/qwen3-30b-a3b-instruct-2507"
BATCH_SIZE = 15

# Fixed taxonomy — LLM must pick from this list only.
TAXONOMY = [
    "Preaching & Witnessing",
    "Construction & Volunteer Work",
    "Personal Experience & Life Story",
    "Legal & Religious Freedom",
    "Medical & Healthcare",
    "Assembly & Convention",
    "Historical Archive",
    "Faith Under Trial",
    "Branch & Country Report",
    "Family & Youth",
    "Natural Disaster & Relief",
    "JW Broadcasting Feature",
    "Interview & Profile",
    "Education & Training",
    "Music & Song",
]

CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS video_categories (
    natural_key  TEXT NOT NULL,
    language     TEXT NOT NULL,
    category     TEXT NOT NULL,
    assigned_at  TEXT NOT NULL,
    PRIMARY KEY (natural_key, language, category)
);
"""

SYSTEM_PROMPT = f"""\
You are a video content categorizer for a JW.org video library. \
Assign each video 1–3 categories from the EXACT list below — use the \
exact strings, no variations.

VALID CATEGORIES:
{chr(10).join(f"- {c}" for c in TAXONOMY)}

Return ONLY a JSON object mapping each video's natural_key to a list of \
1–3 categories (no markdown, no explanation):

{{
  "<natural_key>": ["<category1>", "<category2>"],
  ...
}}

Rules:
- Use ONLY categories from the list above.
- Assign 1 category minimum, 3 maximum.
- Choose the most specific applicable category.
- "JW Broadcasting Feature" = monthly broadcast episodes or studio-produced features.
- "Historical Archive" = footage from before ~2000, vintage assemblies, old newsreel style.
- "Branch & Country Report" = news/report about a specific country branch office or country situation.
- "Interview & Profile" = structured interview with one or a few named individuals as primary format."""

USER_TEMPLATE = """\
Categorize these videos:

{videos}"""


def load_client(api_key: str, base_url: str) -> object:
    try:
        from openai import OpenAI
        return OpenAI(base_url=base_url, api_key=api_key)
    except ImportError as exc:
        raise RuntimeError("pip install openai") from exc


def categorize_batch(client, model: str, videos: list[dict]) -> dict:
    """Send one batch to LLM and return {natural_key: [categories]}. Raises on failure."""
    lines = []
    for v in videos:
        themes_str = "; ".join((v.get("themes") or [])[:5])
        lines.append(
            f"key: {v['natural_key']}\n"
            f"tldr: {(v.get('tldr') or '')[:300]}\n"
            f"themes: {themes_str}"
        )
    videos_text = "\n\n".join(lines)

    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": USER_TEMPLATE.format(videos=videos_text)},
        ],
        max_tokens=1000,
        temperature=0.0,
    )
    raw = response.choices[0].message.content or ""
    raw = re.sub(r"```(?:json)?\s*", "", raw).strip()
    start = raw.find("{")
    end = raw.rfind("}") + 1
    if start == -1 or end == 0:
        raise ValueError(f"No JSON in response. Raw: {raw[:400]!r}")
    data = json.loads(raw[start:end])

    # Validate: only allow categories from the taxonomy
    valid = set(TAXONOMY)
    cleaned = {}
    for key, cats in data.items():
        if not isinstance(cats, list):
            log.warning("Non-list categories for %s: %r — skipping", key, cats)
            continue
        good = [c for c in cats if c in valid]
        bad = [c for c in cats if c not in valid]
        if bad:
            log.warning("Invalid categories for %s (ignored): %s", key, bad)
        if good:
            cleaned[key] = good[:3]
        else:
            log.warning("No valid categories for %s", key)
    return cleaned


def run(args: argparse.Namespace) -> int:
    from runtime_paths import get_data_root
    from scene_processing import scene_db

    db_path = args.db or os.path.join(get_data_root(), "scene_index.db")
    language = args.language

    with scene_db.open_db(db_path) as conn:
        conn.execute(CREATE_TABLE_SQL)

    # Load videos not yet categorized
    with scene_db.open_db(db_path) as conn:
        all_videos = conn.execute(
            "SELECT natural_key, tldr, themes_json FROM video_summaries WHERE language=?",
            (language,),
        ).fetchall()
        already_done = {
            r[0] for r in conn.execute(
                "SELECT DISTINCT natural_key FROM video_categories WHERE language=?",
                (language,),
            ).fetchall()
        }

    to_process = [
        {
            "natural_key": r[0],
            "tldr": r[1],
            "themes": json.loads(r[2] or "[]"),
        }
        for r in all_videos
        if r[0] not in already_done
    ]

    log.info("Videos in DB          : %d", len(all_videos))
    log.info("Already categorized   : %d", len(already_done))
    log.info("To process            : %d", len(to_process))

    if not to_process:
        print("All videos already categorized. Done.")
        _print_summary(db_path, language)
        return 0

    if args.dry_run:
        print(f"DRY RUN — would categorize {len(to_process)} videos in "
              f"{(len(to_process) + BATCH_SIZE - 1) // BATCH_SIZE} batches")
        for v in to_process[:3]:
            print(f"  {v['natural_key']}: {v['tldr'][:80]}...")
        return 0

    client = load_client(args.api_key, args.base_url)
    batches = [to_process[i:i + BATCH_SIZE] for i in range(0, len(to_process), BATCH_SIZE)]
    log.info("Sending %d batch(es) to %s ...", len(batches), args.model)

    total_written = 0
    for i, batch in enumerate(batches, 1):
        log.info("Batch %d/%d (%d videos) ...", i, len(batches), len(batch))
        try:
            results = categorize_batch(client, args.model, batch)
        except Exception as exc:
            log.error("Batch %d failed: %s", i, exc)
            raise

        now = time.strftime("%Y-%m-%dT%H:%M:%S")
        rows = [
            (key, language, cat, now)
            for key, cats in results.items()
            for cat in cats
        ]
        with scene_db.open_db(db_path) as conn:
            conn.executemany(
                "INSERT OR REPLACE INTO video_categories "
                "(natural_key, language, category, assigned_at) VALUES (?,?,?,?)",
                rows,
            )
        total_written += len(rows)
        log.info("  wrote %d category rows", len(rows))

    print(f"\nCategorized {len(to_process)} videos → {total_written} category assignments")
    _print_summary(db_path, language)
    return 0


def _print_summary(db_path: str, language: str) -> None:
    from scene_processing import scene_db
    with scene_db.open_db(db_path) as conn:
        rows = conn.execute(
            "SELECT category, COUNT(*) as n FROM video_categories WHERE language=? "
            "GROUP BY category ORDER BY n DESC",
            (language,),
        ).fetchall()
    if rows:
        print("\nVideos per category:")
        for cat, n in rows:
            print(f"  {n:4d}  {cat}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--api-key", required=True)
    parser.add_argument("--base-url", default="https://openrouter.ai/api/v1")
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--language", default="E")
    parser.add_argument("--db", default=None)
    parser.add_argument("--dry-run", action="store_true")
    return run(parser.parse_args())


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