scripts/normalize_locations.py

#!/usr/bin/env python3
"""Normalize raw location strings from scene_video_locations into structured geo data.

Reads every unique place_name from the DB, batches them to a cheap text LLM on
OpenRouter, and writes normalized records (city, country, country_iso, region,
continent) to the new location_normalized table.

Run AFTER the main scene indexing batch is complete.

Usage (from repo root, venv active):
    export SEARCH_UI_DATA_ROOT="$PWD/backend"
    python3 scripts/normalize_locations.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 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"
# 15 places/batch keeps the structured JSON output well under max_tokens —
# 40 was overflowing 1500 tokens and truncating mid-object.
BATCH_SIZE = 15

CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS location_normalized (
    place_name    TEXT PRIMARY KEY,
    city          TEXT,
    country       TEXT,
    country_iso   TEXT,
    region        TEXT,
    continent     TEXT,
    place_type    TEXT,
    normalized_at TEXT
);
"""

SYSTEM_PROMPT = """\
You are a geographic data normalizer. Given a list of place name strings \
(which may be misspelled, abbreviated, or in mixed formats), return structured \
geographic data for each one.

Return ONLY a JSON object — no markdown, no explanation — mapping each input \
place_name exactly to its structured data:

{
  "<exact input string>": {
    "city": "<city name or null>",
    "country": "<full English country name or null>",
    "country_iso": "<ISO 3166-1 alpha-2 code or null>",
    "region": "<geographic sub-region e.g. 'Southern Europe', 'Southeast Asia' or null>",
    "continent": "<one of: Africa, Americas, Asia, Europe, Oceania or null>",
    "type": "<one of: city, country, region, other>"
  },
  ...
}

Rules:
- For a place like "Naples, Italy": city=Naples, country=Italy, country_iso=IT, continent=Europe
- For "Myanmar": city=null, country=Myanmar, country_iso=MM, continent=Asia
- For "Southern Europe": city=null, country=null, region=Southern Europe, continent=Europe, type=region
- If genuinely unknown, set all fields to null but still include the key."""

USER_TEMPLATE = "Normalize these place names:\n{places}"


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 normalize_batch(client, model: str, place_names: list[str]) -> dict:
    """Send one batch to LLM and return parsed dict. Raises on failure."""
    places_text = "\n".join(f"- {p}" for p in place_names)
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": USER_TEMPLATE.format(places=places_text)},
        ],
        max_tokens=3000,
        temperature=0.0,
    )
    raw = response.choices[0].message.content or ""

    import re
    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}")
    return json.loads(raw[start:end])


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")

    # Create normalized table if needed
    with scene_db.open_db(db_path) as conn:
        conn.execute(CREATE_TABLE_SQL)

    # Load all unique place names not yet normalized
    with scene_db.open_db(db_path) as conn:
        all_places = [
            r[0] for r in conn.execute(
                "SELECT DISTINCT place_name FROM scene_video_locations WHERE place_name IS NOT NULL"
            ).fetchall()
        ]
        already_done = {
            r[0] for r in conn.execute("SELECT place_name FROM location_normalized").fetchall()
        }

    to_process = [p for p in all_places if p not in already_done]
    log.info("Unique locations in DB : %d", len(all_places))
    log.info("Already normalized     : %d", len(already_done))
    log.info("To process             : %d", len(to_process))

    if not to_process:
        print("All locations already normalized. Done.")
        return 0

    if args.dry_run:
        print(f"DRY RUN — would normalize {len(to_process)} locations in "
              f"{(len(to_process) + BATCH_SIZE - 1) // BATCH_SIZE} batches")
        for p in to_process[:10]:
            print(f"  {p}")
        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 places) ...", i, len(batches), len(batch))
        try:
            results = normalize_batch(client, args.model, batch)
        except Exception as exc:
            log.error("Batch %d failed: %s", i, exc)
            log.error("Places in batch: %s", batch)
            raise

        rows = []
        now = time.strftime("%Y-%m-%dT%H:%M:%S")
        for place_name in batch:
            data = results.get(place_name, {})
            if not isinstance(data, dict):
                log.warning("No data returned for %r — skipping", place_name)
                continue
            rows.append((
                place_name,
                data.get("city"),
                data.get("country"),
                data.get("country_iso"),
                data.get("region"),
                data.get("continent"),
                data.get("type"),
                now,
            ))

        with scene_db.open_db(db_path) as conn:
            conn.executemany(
                """INSERT OR REPLACE INTO location_normalized
                   (place_name, city, country, country_iso, region, continent, place_type, normalized_at)
                   VALUES (?,?,?,?,?,?,?,?)""",
                rows,
            )
        total_written += len(rows)
        log.info("  wrote %d rows (total so far: %d)", len(rows), total_written)

    # Print a summary
    with scene_db.open_db(db_path) as conn:
        continents = conn.execute(
            "SELECT continent, COUNT(*) FROM location_normalized WHERE continent IS NOT NULL "
            "GROUP BY continent ORDER BY COUNT(*) DESC"
        ).fetchall()
    print(f"\nNormalized {total_written} locations. Breakdown by continent:")
    for c, n in continents:
        print(f"  {c:<15} {n} locations")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--api-key", required=True, help="OpenRouter API key")
    parser.add_argument("--base-url", default="https://openrouter.ai/api/v1")
    parser.add_argument("--model", default=DEFAULT_MODEL,
                        help=f"Text model for normalization (default: {DEFAULT_MODEL})")
    parser.add_argument("--db", default=None)
    parser.add_argument("--dry-run", action="store_true", help="Show what would be processed, don't call API")
    return run(parser.parse_args())


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