autotag.py

#!/usr/bin/env python3
"""
autotag.py — Auto-tag GoodLinks untagged links using Claude AI.

Environment:
    GOODLINKS_TOKEN   — GoodLinks API token.
    ANTHROPIC_API_KEY — Anthropic API key.
"""

import concurrent.futures
import html
import json
import os
import re
import sys
import threading
import time
import urllib.request
from html.parser import HTMLParser
from pathlib import Path

import anthropic

# ── Configuration ──────────────────────────────────────────────────────────────
GOODLINKS_BASE    = "http://localhost:9428/api/v1"
UNTAGGED_FILE     = Path("./untagged_links.json")
TAGS_FILE         = Path("./tags.json")
CHECKPOINT_FILE   = Path("./progress.json")
CHECKPOINT_EVERY  = 5
MAX_WORKERS       = 5
MAX_TAGS          = 5
FETCH_TIMEOUT     = 15
MAX_CONTENT_CHARS = 4000
LLM_MODEL         = "claude-haiku-4-5-20251001"
MAX_RETRIES = 3

# ── Globals ────────────────────────────────────────────────────────────────────
current_tags: list[str] = []
processed_ids: set[str] = set()
processed_count: int = 0
new_tags_added: set[str] = set()

# Locks
tags_lock         = threading.RLock()  # Protects current_tags reads/writes
tags_refresh_lock = threading.Lock()   # Serializes GoodLinks tag-list refresh
progress_lock     = threading.RLock()  # Protects checkpoint writes (reentrant: _mark_done → save_checkpoint)

client = anthropic.Anthropic()


# ── Claude rate limiters ───────────────────────────────────────────────────────
class _RequestLimiter:
    """Sliding-window requests-per-minute limiter.

    Tracks timestamps of recent API calls and blocks when the per-minute
    cap is reached, sleeping until the oldest call falls outside the
    60-second window.

    Args:
        max_per_minute: Maximum number of requests allowed within any
            rolling 60-second window.
    """

    def __init__(self, max_per_minute: int):
        self._lock  = threading.Lock()
        self._calls: list[float] = []
        self._max   = max_per_minute

    def acquire(self):
        """Block until a request slot is available, then reserve it.

        Prunes timestamps older than 60 seconds, checks capacity, and
        either records the current time and returns immediately or sleeps
        until the oldest entry expires.
        """
        while True:
            with self._lock:
                now = time.time()
                self._calls = [t for t in self._calls if now - t < 60.0]

                if len(self._calls) < self._max:
                    self._calls.append(now)
                    return

                wait = 60.0 - (now - self._calls[0])

            wait = max(0.1, wait)
            print(f"  [req-limit] sleeping {wait:.1f}s…", flush=True)
            time.sleep(wait)

class _TokenLimiter:
    """Sliding-window input-tokens-per-minute limiter.

    Call ``acquire(estimate)`` before the request, then
    ``record_actual(ts, actual)`` after to correct the running total with
    the real token count from the API response.

    Args:
        max_per_minute: Maximum number of tokens allowed within any
            rolling 60-second window.
    """

    def __init__(self, max_per_minute: int):
        self._lock   = threading.Lock()
        self._window: list[list] = []   # [timestamp, tokens] — mutable so we can update
        self._max    = max_per_minute

    def acquire(self, estimated: int) -> float:
        """Block until there is token capacity, then reserve a slot.

        Prunes entries older than 60 seconds, sums current usage, and
        either appends a new ``[timestamp, estimated]`` entry and returns
        immediately, or sleeps until the oldest entry expires.

        Args:
            estimated: Estimated number of input tokens for the upcoming
                request.

        Returns:
            The slot timestamp (``time.time()`` value) that identifies
            this reservation. Pass it to ``record_actual`` after the
            request completes.
        """
        while True:
            with self._lock:
                now = time.time()
                self._window = [e for e in self._window if now - e[0] < 60.0]
                used = sum(e[1] for e in self._window)

                if used + estimated <= self._max:
                    entry = [now, estimated]
                    self._window.append(entry)
                    return now

                # Wait until the oldest entry falls out of the window
                wait = 60.0 - (now - self._window[0][0])

            wait = max(0.5, wait)
            print(f"  [token-limit] sleeping {wait:.1f}s  ({used:,}/{self._max:,} tokens used)…", flush=True)
            time.sleep(wait)

    def record_actual(self, slot_ts: float, actual: int):
        """Replace the estimated token count for a slot with the real value.

        Should be called after the API response is received so the
        sliding window reflects actual usage rather than the estimate.

        Args:
            slot_ts: The timestamp returned by ``acquire`` that
                identifies the slot to update.
            actual: The real input-token count reported by the API.
        """
        with self._lock:
            for entry in self._window:
                if entry[0] == slot_ts:
                    entry[1] = actual
                    return

    def tokens_used(self) -> int:
        """Return total tokens consumed in the current 60-second window.

        Returns:
            Sum of token counts for all slots whose timestamps fall
            within the last 60 seconds.
        """
        with self._lock:
            now = time.time()
            return sum(e[1] for e in self._window if now - e[0] < 60.0)


_req_limiter   = _RequestLimiter(max_per_minute=45)   # hard limit is 50
_token_limiter = _TokenLimiter(max_per_minute=45_000)  # hard limit is 50k

# ── HTML text extraction ───────────────────────────────────────────────────────
class _TextExtractor(HTMLParser):
    """HTML parser that extracts visible body text, skipping non-content elements.

    Uses a depth-tracking mechanism to handle nested skip tags. When the parser
    enters a tag listed in ``SKIP``, it increments ``_depth``; when it leaves,
    it decrements. Text data is only collected while ``_depth`` is zero,
    ensuring that content inside nested non-content elements (e.g. a
    ``<script>`` inside a ``<nav>``) is also excluded.

    Attributes:
        SKIP: Frozenset of HTML tag names whose content should be ignored.
            Includes ``script``, ``style``, ``noscript``, ``nav``, ``footer``,
            ``header``, ``aside``, ``meta``, and ``link``.
        parts: Accumulated non-empty, stripped text segments.
    """

    SKIP = frozenset({
        "script", "style", "noscript", "nav", "footer",
        "header", "aside", "meta", "link",
    })

    def __init__(self):
        super().__init__()
        self.parts: list[str] = []
        self._depth = 0

    def handle_starttag(self, tag, attrs):
        """Increment the skip-depth counter when entering a non-content tag.

        Args:
            tag: The HTML tag name (e.g. ``"script"``).
            attrs: List of ``(attribute, value)`` pairs (unused).
        """
        if tag.lower() in self.SKIP:
            self._depth += 1

    def handle_endtag(self, tag):
        """Decrement the skip-depth counter when leaving a non-content tag.

        Only decrements if ``_depth`` is already positive, preventing underflow
        from malformed HTML with extra closing tags.

        Args:
            tag: The HTML tag name.
        """
        if tag.lower() in self.SKIP and self._depth:
            self._depth -= 1

    def handle_data(self, data):
        """Collect text data when not inside a skipped element.

        Strips leading/trailing whitespace from the data and appends it to
        ``parts`` only when ``_depth`` is zero (i.e. outside all SKIP tags)
        and the stripped result is non-empty.

        Args:
            data: Raw character data from the HTML document.
        """
        if not self._depth:
            s = data.strip()
            if s:
                self.parts.append(s)

    def text(self) -> str:
        """Return the extracted text as a single whitespace-normalized string.

        Joins all collected ``parts`` with spaces and collapses consecutive
        whitespace characters into a single space.

        Returns:
            The concatenated, whitespace-collapsed visible text.
        """
        return re.sub(r"\s+", " ", " ".join(self.parts)).strip()


def _parse_html(raw: str) -> tuple[str, str]:
    """Extract the page title and visible body text from raw HTML.

    Searches for a ``<title>`` element to obtain the page title, then feeds
    the full HTML through ``_TextExtractor`` to collect visible text. The
    body text is truncated to ``MAX_CONTENT_CHARS`` characters.

    Args:
        raw: The raw HTML string to parse.

    Returns:
        A ``(title, body_text)`` tuple. *title* is the unescaped content of
        the first ``<title>`` tag, or an empty string if none is found.
        *body_text* is the extracted visible text, truncated to
        ``MAX_CONTENT_CHARS``.
    """
    title_m = re.search(r"<title[^>]*>(.*?)</title>", raw, re.I | re.S)
    title   = html.unescape(title_m.group(1).strip()) if title_m else ""
    ex      = _TextExtractor()
    try:
        ex.feed(raw)
    except Exception:
        pass
    return title, ex.text()[:MAX_CONTENT_CHARS]


def fetch_page(url: str) -> tuple[str, str]:
    """Fetch a web page and extract its title and visible text.

    Sends an HTTP GET request with a browser-like User-Agent header,
    reads up to 200 000 bytes of the response, and delegates HTML
    parsing to ``_parse_html``.

    Args:
        url: The URL to fetch.

    Returns:
        A ``(title, text)`` tuple where *title* is the content of the
        ``<title>`` element (or ``""`` if absent) and *text* is the
        concatenated visible text of the page.

    Raises:
        ValueError: If the response Content-Type is not HTML, XHTML,
            or plain text.
        urllib.error.URLError: If the network request fails or times
            out (``FETCH_TIMEOUT`` seconds).
    """
    req = urllib.request.Request(
        url,
        headers={
            "User-Agent": (
                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/120.0.0.0 Safari/537.36"
            ),
            "Accept": "text/html,application/xhtml+xml,*/*;q=0.9",
            "Accept-Language": "en-US,en;q=0.9",
        },
    )

    with urllib.request.urlopen(req, timeout=FETCH_TIMEOUT) as resp:
        ct = resp.headers.get("Content-Type", "")

        if not any(t in ct for t in ("text/html", "text/plain", "application/xhtml")):
            raise ValueError(f"Unsupported content-type: {ct}")

        raw = resp.read(200_000).decode("utf-8", errors="replace")

    return _parse_html(raw)

# ── GoodLinks API ──────────────────────────────────────────────────────────────
def _gl_request(method: str, path: str, body=None):
    """Send an authenticated request to the GoodLinks API.

    Constructs a JSON request against ``GOODLINKS_BASE`` using the
    bearer token stored in the ``GOODLINKS_TOKEN`` environment
    variable.

    Args:
        method: HTTP method (e.g. ``"GET"``, ``"PATCH"``).
        path: API path appended to ``GOODLINKS_BASE``
            (e.g. ``"/tags"``).
        body: Optional Python object serialised as JSON for the
            request body.  Defaults to ``None``.

    Returns:
        The parsed JSON response body.

    Raises:
        KeyError: If the ``GOODLINKS_TOKEN`` environment variable is
            not set.
        urllib.error.URLError: If the network request fails or times
            out (15-second timeout).
    """
    url  = f"{GOODLINKS_BASE}{path}"
    data = json.dumps(body).encode() if body else None
    req  = urllib.request.Request(
        url,
        data=data,
        method=method,
        headers={
            "Authorization": f"Bearer {os.environ['GOODLINKS_TOKEN']}",
            "Content-Type":  "application/json",
        },
    )
    with urllib.request.urlopen(req, timeout=15) as resp:
        return json.load(resp)

def update_link_tags(link_id: str, tags: list[str]):
    """Replace the tags on a GoodLinks link.

    Sends a ``PATCH`` request to ``/links/{link_id}`` with the
    supplied tag list.

    Args:
        link_id: The GoodLinks identifier of the link to update.
        tags: The full list of tags to set on the link.

    Raises:
        urllib.error.URLError: If the API request fails.
    """
    _gl_request("PATCH", f"/links/{link_id}", {"tags": tags})


def refresh_tags() -> list[str]:
    """Fetch the current tag list from GoodLinks and update local state.

    Retrieves all tags via the GoodLinks API, writes them to
    ``tags.json``, and replaces the in-memory ``current_tags`` list.
    Access is serialised with ``tags_refresh_lock`` so only one thread
    refreshes at a time.

    Returns:
        The full list of tag names returned by the API.

    Raises:
        urllib.error.URLError: If the API request fails.
    """
    with tags_refresh_lock:
        tags = _gl_request("GET", "/tags")
        TAGS_FILE.write_text(json.dumps(tags, indent=2))
        with tags_lock:
            current_tags.clear()
            current_tags.extend(tags)

        return tags

# ── Claude tag suggestion ──────────────────────────────────────────────────────
def suggest_tags(
    link: dict, page_title: str, page_text: str
) -> tuple[list[str], list[str]]:
    """Ask Claude to choose tags for a saved link.

    Builds a prompt containing the link metadata, page text, and the
    current tag snapshot, then sends it to the Claude API.  The
    response is parsed as JSON and split into existing and new tags.

    On ``anthropic.RateLimitError`` the call is retried up to
    ``MAX_RETRIES`` times with exponential back-off (60 s on the
    first retry, then 2 ** *attempt* seconds).  Any other exception
    aborts immediately.

    Args:
        link: A GoodLinks link dict containing at least ``"url"`` and
            optionally ``"title"`` and ``"summary"``.
        page_title: The ``<title>`` extracted from the fetched page.
        page_text: The visible text extracted from the fetched page.

    Returns:
        A ``(existing_tags, new_tags)`` tuple.  *existing_tags* are
        tags already present in the current known-tags snapshot.
        *new_tags* are tags proposed by Claude that do not yet exist.
        The combined length never exceeds ``MAX_TAGS``.

    Raises:
        RuntimeError: If all retry attempts are exhausted or a
            non-rate-limit error occurs.
    """
    with tags_lock:
        snapshot = list(current_tags)

    prompt = f"""You are tagging a saved web link for a personal bookmarks manager.

Link title : {link.get("title") or page_title or "(unknown)"}
Link URL   : {link["url"]}
Summary    : {link.get("summary") or "(none)"}
Page text  : {page_text or "(empty)"}

Available tags (always prefer these when they fit):
{json.dumps(snapshot)}

Instructions:
- Choose up to {MAX_TAGS} tags total that best describe this content.
- Only pick tags that genuinely apply — do not force irrelevant tags.
- If 1-2 important concepts are clearly absent from the available list, you may propose new tags.
- New tags must be a single word or two-word kebab-case lowercase (e.g. "rust", "load-testing").
- If you are not confident any tags fit, return empty lists.

Respond with JSON only (no markdown fences):
{{
  "existing_tags": ["tag1", "tag2"],
  "new_tags": ["new-tag"]
}}
"""

    estimated_tokens = len(prompt) // 4 + 256  # rough estimate; 256 for output headroom

    slot_ts: float = 0.0
    last_err: Exception | None = None

    for attempt in range(MAX_RETRIES):
        try:
            _req_limiter.acquire()
            slot_ts = _token_limiter.acquire(estimated_tokens)
            msg = client.messages.create(
                model=LLM_MODEL,
                max_tokens=256,
                messages=[{"role": "user", "content": prompt}],
            )

            _token_limiter.record_actual(slot_ts, msg.usage.input_tokens)
            used = _token_limiter.tokens_used()
            print(f"  [tokens] {used:,}/50,000 used in last 60s", flush=True)
            raw = msg.content[0].text.strip()

            # Strip markdown fences, then extract the first JSON object
            raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw, flags=re.M).strip()
            m   = re.search(r"\{.*\}", raw, re.S)

            if not m:
                raise ValueError(f"No JSON object in response: {raw!r}")

            result = json.loads(m.group())

            # Validate: only keep existing_tags that are actually in the snapshot
            existing = [t for t in result.get("existing_tags", []) if t in snapshot]
            new_tags = [t for t in result.get("new_tags", [])]

            # Enforce total cap (trim new_tags first)
            if len(existing) + len(new_tags) > MAX_TAGS:
                new_tags = new_tags[: max(0, MAX_TAGS - len(existing))]
                existing = existing[:MAX_TAGS]

            return existing, new_tags

        except anthropic.RateLimitError as e:
            # Correct the token slot to full estimated so the limiter backs off
            _token_limiter.record_actual(slot_ts, estimated_tokens)
            wait = 60 if attempt == 0 else 2 ** attempt
            print(f"  [rate-limit] sleeping {wait}s…", flush=True)
            time.sleep(wait)
            last_err = e
        except Exception as e:
            last_err = e
            break

    raise RuntimeError(f"Claude failed after {MAX_RETRIES} attempts: {last_err}")

# ── Checkpoint helpers ─────────────────────────────────────────────────────────
def load_checkpoint() -> set[str]:
    """Load previously processed link IDs from the checkpoint file.

    Reads ``CHECKPOINT_FILE`` (``progress.json``) and returns the set of
    link IDs recorded under the ``"processed_ids"`` key.  If the file
    does not exist, returns an empty set so processing starts from
    scratch.

    Returns:
        A set of link-ID strings that have already been processed.
    """
    if CHECKPOINT_FILE.exists():
        return set(json.loads(CHECKPOINT_FILE.read_text()).get("processed_ids", []))

    return set()

def save_checkpoint():
    """Persist the current set of processed link IDs to disk.

    Writes the global ``processed_ids`` set to ``CHECKPOINT_FILE``
    (``progress.json``) as a JSON object.  The write is protected by
    ``progress_lock`` to prevent concurrent corruption from worker
    threads.

    Side Effects:
        Overwrites ``CHECKPOINT_FILE`` with the latest processed-ID
        snapshot.
    """
    with progress_lock:
        CHECKPOINT_FILE.write_text(
            json.dumps({"processed_ids": list(processed_ids)}, indent=2)
        )

# ── Per-link worker ────────────────────────────────────────────────────────────
def process_link(link: dict) -> str:
    """Fetch, tag, and update a single GoodLinks link.

    Orchestrates the full per-link pipeline:

    1. Fetch the page content via ``fetch_page``.
    2. Ask Claude for tag suggestions via ``suggest_tags``.
    3. Merge any new tags into the global ``current_tags`` list.
    4. Apply the chosen tags (or a fallback label) via the GoodLinks API.
    5. Refresh the remote tag list if new tags were created.
    6. Mark the link as done and periodically checkpoint progress.

    If fetching fails the link is tagged ``"problem"``; if no tags are
    suggested it is tagged ``"notags"``.

    Args:
        link: A GoodLinks link dict containing at least ``"id"`` and
            ``"url"``, and optionally ``"title"`` and ``"summary"``.

    Returns:
        A human-readable status line prefixed with ``OK``, ``PROB``,
        ``NONE``, or ``ERR`` describing the outcome.

    Side Effects:
        Mutates the global ``current_tags`` list when new tags are
        accepted.  Calls ``_mark_done`` which updates ``processed_ids``,
        increments ``processed_count``, and may write a checkpoint.
    """
    global processed_count

    lid  = link["id"]
    url  = link["url"]
    desc = (link.get("title") or url)[:72]

    print(f"FETCH {desc!r}", flush=True)

    # ── 1. Fetch page ────────────────────────────────────────────────────────
    try:
        page_title, page_text = fetch_page(url)
    except Exception as e:
        try:
            update_link_tags(lid, ["__problem"])
        except Exception:
            pass

        _mark_done(lid)

        return f"PROB  {desc!r}  ({type(e).__name__}: {e})"

    # ── 2. Ask Claude for tags ───────────────────────────────────────────────
    try:
        existing, new_tags = suggest_tags(link, page_title, page_text)
    except Exception as e:
        _mark_done(lid)
        return f"ERR   {desc!r}  (Claude: {e})"

    # ── 3. Auto-apply new tags, adding them to the in-memory list ───────────
    if new_tags:
        with tags_lock:
            for t in new_tags:
                if t not in current_tags:
                    current_tags.append(t)
                    new_tags_added.add(t)

    final = (existing + new_tags)[:MAX_TAGS]

    # ── 4. No confident match → "notags" ────────────────────────────────────
    if not final:
        try:
            update_link_tags(lid, ["__notags"])
        except Exception:
            pass

        _mark_done(lid)

        return f"NONE  {desc!r}"

    # ── 5. Apply tags via GoodLinks API ──────────────────────────────────────
    try:
        update_link_tags(lid, final)
    except Exception as e:
        _mark_done(lid)
        return f"ERR   {desc!r}  (API: {e})"

    # ── 6. Refresh tag list if new tags were created ─────────────────────────
    if new_tags:
        try:
            refresh_tags()
        except Exception:
            pass

    _mark_done(lid)
    tag_str = ", ".join(final)
    new_str = f"  [new: {', '.join(new_tags)}]" if new_tags else ""

    return f"OK    {desc!r}  → {tag_str}{new_str}"


def _mark_done(lid: str):
    """Record a link as processed and checkpoint periodically.

    Adds the link ID to the global ``processed_ids`` set, increments
    ``processed_count``, and writes a checkpoint every
    ``CHECKPOINT_EVERY`` links.  All mutations are protected by
    ``progress_lock``.

    Args:
        lid: The GoodLinks identifier of the link that was processed.

    Side Effects:
        Mutates ``processed_ids`` and ``processed_count``.  Calls
        ``save_checkpoint`` every ``CHECKPOINT_EVERY`` processed links,
        writing ``progress.json`` to disk.
    """
    global processed_count
    with progress_lock:
        processed_ids.add(lid)
        processed_count += 1
        if processed_count % CHECKPOINT_EVERY == 0:
            save_checkpoint()
            print(f"  ── checkpoint: {processed_count} processed ──", flush=True)

# ── Entry point ────────────────────────────────────────────────────────────────
def main():
    """Entry point: auto-tag all untagged GoodLinks links.

    Validates required environment variables, loads the untagged-links
    file and tag list, restores checkpoint state, then processes
    remaining links concurrently using a ``ThreadPoolExecutor`` with
    ``MAX_WORKERS`` threads.  A final checkpoint is written when all
    links have been processed.

    Side Effects:
        Populates the global ``current_tags``, ``processed_ids``, and
        ``processed_count`` from disk.  Writes ``progress.json`` on
        completion.  Prints progress and per-link status lines to
        stdout.

    Raises:
        SystemExit: If ``GOODLINKS_TOKEN`` or ``ANTHROPIC_API_KEY``
            environment variables are not set.
    """
    global processed_ids, processed_count

    for var in ("GOODLINKS_TOKEN", "ANTHROPIC_API_KEY"):
        if not os.environ.get(var):
            sys.exit(f"Error: environment variable {var} is not set.")

    links           = json.loads(UNTAGGED_FILE.read_text())
    current_tags.extend(json.loads(TAGS_FILE.read_text()))
    processed_ids   = load_checkpoint()
    processed_count = len(processed_ids)

    pending = [lnk for lnk in links if lnk["id"] not in processed_ids]

    print(f"Links total    : {len(links)}")
    print(f"Already done   : {processed_count}")
    print(f"To process     : {len(pending)}")
    print(f"Tags available : {len(current_tags)}")
    print(f"Workers        : {MAX_WORKERS}")
    print()

    if not pending:
        print("Nothing to do — all links already processed.")
        return

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
        futures = {pool.submit(process_link, lnk): lnk for lnk in pending}
        for fut in concurrent.futures.as_completed(futures):
            try:
                print(fut.result(), flush=True)
            except Exception as exc:
                lnk = futures[fut]
                print(f"FATAL {lnk['url']!r}: {exc}", flush=True)

    save_checkpoint()
    print(f"\nFinished. {processed_count}/{len(links)} links processed.")

    if new_tags_added:
        print(f"\nNew tags added ({len(new_tags_added)}):")
        for tag in sorted(new_tags_added):
            print(f"  {tag}")

if __name__ == "__main__":
    main()