OBLITERATUS/obliteratus/watchtower.py

"""Watchtower — Scan HuggingFace for new popular open-weight models.

Continuously monitors the HuggingFace Hub for instruction-tuned,
open-license text-generation models that are gaining traction.  Tracks
which models have been seen, queued, and obliterated in a persistent
JSON state file (~/.obliteratus/watchtower_state.json).

Usage:
    from obliteratus.watchtower import Watchtower

    wt = Watchtower()
    new_models = wt.scan()          # returns list of newly discovered models
    trending  = wt.get_trending()   # returns current hot models sorted by downloads
    wt.start_scheduler(interval=3600)  # background hourly scan
    wt.stop_scheduler()
"""

from __future__ import annotations

import json
import logging
import os
import threading
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# ── Constants ─────────────────────────────────────────────────────────

STATE_DIR = Path.home() / ".obliteratus"
STATE_FILE = STATE_DIR / "watchtower_state.json"

# Minimum downloads in last 7 days to qualify as "popular"
MIN_RECENT_DOWNLOADS = 1000

# Open-weight licenses we care about (lowercase, prefix-matched)
OPEN_LICENSES = {
    "apache-2.0", "mit", "bsd-2-clause", "bsd-3-clause",
    "llama2", "llama3", "llama3.1", "llama3.2", "llama3.3", "llama4",
    "gemma", "qwen", "deepseek",
    "cc-by-4.0", "cc-by-sa-4.0", "cc-by-nc-4.0",
    "openrail", "openrail++", "bigscience-openrail-m",
    "artistic-2.0", "wtfpl", "unlicense", "zlib",
    "other",  # many open models use "other" + a permissive custom license
}

# Keywords that indicate instruction/chat tuning
INSTRUCT_KEYWORDS = {
    "instruct", "chat", "it", "rlhf", "dpo", "sft",
    "aligned", "conversational", "assistant", "dialogue",
}

# ── Organizations to watch ────────────────────────────────────────────

MAJOR_LABS = {
    "google", "meta-llama", "mistralai", "Qwen", "deepseek-ai",
    "microsoft", "nvidia", "01-ai", "internlm", "THUDM",
    "tiiuae", "CohereForAI", "allenai", "openai", "openai-community",
    "moonshotai", "openbmb", "stabilityai", "stepfun-ai", "zai-org",
    "MiniMaxAI",
}

MEDIUM_LABS = {
    "teknium", "NousResearch", "OpenBuddy", "lmsys",
    "HuggingFaceH4", "mosaicml", "bigcode",
    "cognitivecomputations", "mlabonne", "huihui-ai",
    "Orenguteng", "WhiteRabbitNeo",
}

ALL_WATCHED_ORGS = MAJOR_LABS | MEDIUM_LABS


# ── Data classes ──────────────────────────────────────────────────────

@dataclass
class DiscoveredModel:
    """A model discovered by the Watchtower."""
    model_id: str                     # e.g. "meta-llama/Llama-3.1-8B-Instruct"
    name: str                         # short display name
    org: str                          # organization
    downloads_7d: int = 0             # downloads in last 7 days
    total_downloads: int = 0          # all-time downloads
    likes: int = 0                    # HF likes
    size_category: str = ""           # e.g. "7B", "13B", "70B"
    license: str = ""                 # license identifier
    pipeline_tag: str = ""            # should be "text-generation"
    discovered_at: str = ""           # ISO timestamp
    status: str = "new"               # new | queued | obliterating | obliterated | failed
    obliteration_metrics: dict = field(default_factory=dict)
    last_updated: str = ""            # ISO timestamp of last status change

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "DiscoveredModel":
        # Handle extra/missing keys gracefully
        known = {f.name for f in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in d.items() if k in known})


# ── Watchtower class ──────────────────────────────────────────────────

class Watchtower:
    """Scans HuggingFace Hub for new popular open-weight instruction models."""

    def __init__(self, state_file: Path | str | None = None):
        self.state_file = Path(state_file) if state_file else STATE_FILE
        self._models: dict[str, DiscoveredModel] = {}
        self._last_scan: str | None = None
        self._scan_count: int = 0
        self._lock = threading.Lock()
        self._scheduler_thread: threading.Thread | None = None
        self._scheduler_stop = threading.Event()
        self._on_new_model_callbacks: list = []

        # Load persisted state
        self._load_state()

    # ── Persistence ───────────────────────────────────────────────────

    def _load_state(self):
        """Load watchtower state from disk."""
        try:
            if self.state_file.exists():
                data = json.loads(self.state_file.read_text(encoding="utf-8"))
                self._last_scan = data.get("last_scan")
                self._scan_count = data.get("scan_count", 0)
                for mid, mdata in data.get("models", {}).items():
                    self._models[mid] = DiscoveredModel.from_dict(mdata)
                logger.info(
                    "Watchtower: loaded %d models from %s",
                    len(self._models), self.state_file,
                )
        except Exception as e:
            logger.warning("Watchtower: failed to load state: %s", e)

    def _save_state(self):
        """Persist watchtower state to disk."""
        try:
            self.state_file.parent.mkdir(parents=True, exist_ok=True)
            data = {
                "last_scan": self._last_scan,
                "scan_count": self._scan_count,
                "models": {mid: m.to_dict() for mid, m in self._models.items()},
            }
            # Atomic write via temp file
            tmp = self.state_file.with_suffix(".tmp")
            tmp.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
            tmp.replace(self.state_file)
        except Exception as e:
            logger.warning("Watchtower: failed to save state: %s", e)

    # ── HuggingFace API helpers ───────────────────────────────────────

    @staticmethod
    def _fetch_models_from_hf(
        *,
        orgs: set[str] | None = None,
        limit_per_org: int = 50,
        min_downloads: int = MIN_RECENT_DOWNLOADS,
    ) -> list[dict[str, Any]]:
        """Fetch model metadata from HuggingFace Hub API.

        Returns a list of raw model info dicts. Gracefully returns []
        if the API is unreachable.
        """
        try:
            from huggingface_hub import HfApi, ModelFilter
        except ImportError:
            logger.error("huggingface_hub not installed — cannot scan HF")
            return []

        api = HfApi()
        results = []
        search_orgs = orgs or ALL_WATCHED_ORGS

        for org in search_orgs:
            try:
                models = api.list_models(
                    author=org,
                    pipeline_tag="text-generation",
                    sort="downloads",
                    direction=-1,
                    limit=limit_per_org,
                )
                for m in models:
                    results.append(m)
            except Exception as e:
                logger.debug("Watchtower: error scanning org '%s': %s", org, e)
                continue

        return results

    @staticmethod
    def _is_instruction_tuned(model_id: str, tags: list[str] | None = None) -> bool:
        """Heuristic check if a model is instruction/chat tuned."""
        name_lower = model_id.lower()
        # Check model name
        for kw in INSTRUCT_KEYWORDS:
            if kw in name_lower:
                return True
        # Check tags
        if tags:
            tags_lower = {t.lower() for t in tags}
            for kw in INSTRUCT_KEYWORDS:
                if kw in tags_lower:
                    return True
            # Explicit tag checks
            if "conversational" in tags_lower:
                return True
        return False

    @staticmethod
    def _has_open_license(license_id: str | None) -> bool:
        """Check if the license is considered open-weight."""
        if not license_id:
            return False
        lid = license_id.lower().strip()
        for allowed in OPEN_LICENSES:
            if lid == allowed or lid.startswith(allowed):
                return True
        return False

    @staticmethod
    def _estimate_size(model_id: str, config: dict | None = None) -> str:
        """Estimate model size from the name or config."""
        name = model_id.lower()
        # Try to extract a size like "7b", "70b", "1.5b", "397b"
        import re
        match = re.search(r'(\d+\.?\d*)\s*[bB]', model_id)
        if match:
            size = float(match.group(1))
            if size >= 1:
                return f"{match.group(1)}B"
            else:
                return f"{size * 1000:.0f}M"
        return "unknown"

    # ── Core scan logic ───────────────────────────────────────────────

    def scan(self, on_log=None) -> list[DiscoveredModel]:
        """Scan HuggingFace for new popular open-weight instruction models.

        Returns a list of *newly discovered* models (not previously seen).
        Thread-safe.
        """
        def _log(msg):
            logger.info(msg)
            if on_log:
                on_log(msg)

        _log("Watchtower: starting scan...")
        new_models = []

        try:
            raw_models = self._fetch_models_from_hf()
        except Exception as e:
            _log(f"Watchtower: HF API error — {e}")
            return []

        _log(f"Watchtower: fetched {len(raw_models)} candidate models from HF Hub")

        now = datetime.now(timezone.utc).isoformat()
        processed = 0

        for m in raw_models:
            model_id = getattr(m, "id", None) or getattr(m, "modelId", "")
            if not model_id:
                continue

            # Extract metadata safely
            downloads = getattr(m, "downloads", 0) or 0
            likes = getattr(m, "likes", 0) or 0
            tags = getattr(m, "tags", []) or []
            license_id = getattr(m, "license", None) or ""
            # Some models store license in tags
            if not license_id:
                for t in tags:
                    if t.startswith("license:"):
                        license_id = t.split(":", 1)[1]
                        break

            pipeline_tag = getattr(m, "pipeline_tag", "") or ""

            # Filter: must be text-generation
            if pipeline_tag and pipeline_tag != "text-generation":
                continue

            # Filter: minimum popularity
            if downloads < MIN_RECENT_DOWNLOADS:
                continue

            # Filter: open license
            if not self._has_open_license(license_id):
                continue

            # Filter: instruction-tuned (or from a major lab — they usually are)
            org = model_id.split("/")[0] if "/" in model_id else ""
            is_major = org in MAJOR_LABS
            if not is_major and not self._is_instruction_tuned(model_id, tags):
                continue

            processed += 1

            # Check if we've already seen this model
            with self._lock:
                if model_id in self._models:
                    # Update download counts
                    existing = self._models[model_id]
                    existing.downloads_7d = downloads  # HF "downloads" is ~recent
                    existing.total_downloads = downloads
                    existing.likes = likes
                    existing.last_updated = now
                    continue

                # New model!
                dm = DiscoveredModel(
                    model_id=model_id,
                    name=model_id.split("/")[-1] if "/" in model_id else model_id,
                    org=org,
                    downloads_7d=downloads,
                    total_downloads=downloads,
                    likes=likes,
                    size_category=self._estimate_size(model_id),
                    license=license_id,
                    pipeline_tag=pipeline_tag or "text-generation",
                    discovered_at=now,
                    status="new",
                    last_updated=now,
                )
                self._models[model_id] = dm
                new_models.append(dm)

        with self._lock:
            self._last_scan = now
            self._scan_count += 1

        self._save_state()

        _log(
            f"Watchtower: scan complete — {processed} qualifying models, "
            f"{len(new_models)} new discoveries"
        )

        # Fire callbacks for new models
        for dm in new_models:
            for cb in self._on_new_model_callbacks:
                try:
                    cb(dm)
                except Exception as e:
                    logger.warning("Watchtower callback error: %s", e)

        return new_models

    # ── Query methods ─────────────────────────────────────────────────

    def get_trending(self, limit: int = 50) -> list[DiscoveredModel]:
        """Return current hot models sorted by recent downloads (descending)."""
        with self._lock:
            models = list(self._models.values())
        models.sort(key=lambda m: m.downloads_7d, reverse=True)
        return models[:limit]

    def get_new_models(self) -> list[DiscoveredModel]:
        """Return models with status 'new' (not yet queued or obliterated)."""
        with self._lock:
            return [m for m in self._models.values() if m.status == "new"]

    def get_obliterated(self) -> list[DiscoveredModel]:
        """Return models that have been obliterated."""
        with self._lock:
            return [m for m in self._models.values() if m.status == "obliterated"]

    def get_all_models(self) -> list[DiscoveredModel]:
        """Return all tracked models."""
        with self._lock:
            return list(self._models.values())

    def get_model(self, model_id: str) -> DiscoveredModel | None:
        """Get a specific model by ID."""
        with self._lock:
            return self._models.get(model_id)

    def set_status(self, model_id: str, status: str, metrics: dict | None = None):
        """Update a model's status (new/queued/obliterating/obliterated/failed)."""
        with self._lock:
            if model_id in self._models:
                m = self._models[model_id]
                m.status = status
                m.last_updated = datetime.now(timezone.utc).isoformat()
                if metrics:
                    m.obliteration_metrics = metrics
        self._save_state()

    def get_stats(self) -> dict[str, Any]:
        """Return summary statistics."""
        with self._lock:
            total = len(self._models)
            by_status = {}
            for m in self._models.values():
                by_status[m.status] = by_status.get(m.status, 0) + 1
        return {
            "total_tracked": total,
            "last_scan": self._last_scan,
            "scan_count": self._scan_count,
            "by_status": by_status,
        }

    def get_model_choices(self) -> list[str]:
        """Return model IDs suitable for a dropdown, trending first."""
        trending = self.get_trending(limit=100)
        return [m.model_id for m in trending]

    # ── Callbacks ─────────────────────────────────────────────────────

    def on_new_model(self, callback):
        """Register a callback for when a new model is discovered.

        Signature: callback(model: DiscoveredModel)
        """
        self._on_new_model_callbacks.append(callback)

    # ── Scheduler ─────────────────────────────────────────────────────

    def start_scheduler(self, interval: int = 3600, on_log=None):
        """Start background scanning at the given interval (seconds).

        Safe to call multiple times — restarts with new interval.
        """
        self.stop_scheduler()
        self._scheduler_stop.clear()

        def _run():
            while not self._scheduler_stop.is_set():
                try:
                    self.scan(on_log=on_log)
                except Exception as e:
                    logger.error("Watchtower scheduler error: %s", e)
                self._scheduler_stop.wait(timeout=interval)

        self._scheduler_thread = threading.Thread(
            target=_run, daemon=True, name="watchtower-scheduler"
        )
        self._scheduler_thread.start()
        logger.info("Watchtower: scheduler started (interval=%ds)", interval)

    def stop_scheduler(self):
        """Stop the background scheduler if running."""
        if self._scheduler_thread and self._scheduler_thread.is_alive():
            self._scheduler_stop.set()
            self._scheduler_thread.join(timeout=5)
            self._scheduler_thread = None
            logger.info("Watchtower: scheduler stopped")

    @property
    def is_scanning(self) -> bool:
        """True if the scheduler is actively running."""
        return (
            self._scheduler_thread is not None
            and self._scheduler_thread.is_alive()
        )

    # ── Table formatting ──────────────────────────────────────────────

    def format_table(self, models: list[DiscoveredModel] | None = None) -> list[list[str]]:
        """Format models as a list of rows for a Gradio Dataframe.

        Columns: [Name, Org, Size, Downloads, Likes, License, Discovered, Status]
        """
        if models is None:
            models = self.get_trending()

        rows = []
        for m in models:
            discovered = ""
            if m.discovered_at:
                try:
                    dt = datetime.fromisoformat(m.discovered_at)
                    discovered = dt.strftime("%Y-%m-%d %H:%M")
                except Exception:
                    discovered = m.discovered_at[:16]

            status_emoji = {
                "new": "🆕",
                "queued": "⏳",
                "obliterating": "⚡",
                "obliterated": "✅",
                "failed": "❌",
            }.get(m.status, "❓")

            rows.append([
                m.model_id,
                m.org,
                m.size_category,
                f"{m.downloads_7d:,}",
                str(m.likes),
                m.license,
                discovered,
                f"{status_emoji} {m.status}",
            ])
        return rows

    TABLE_HEADERS = [
        "Model ID", "Org", "Size", "Downloads (7d)",
        "Likes", "License", "Discovered", "Status",
    ]


# ── Module-level singleton ────────────────────────────────────────────
# Lazy-initialized so importing the module doesn't trigger disk I/O.

_watchtower_instance: Watchtower | None = None
_watchtower_lock = threading.Lock()


def get_watchtower() -> Watchtower:
    """Get or create the global Watchtower singleton."""
    global _watchtower_instance
    if _watchtower_instance is None:
        with _watchtower_lock:
            if _watchtower_instance is None:
                _watchtower_instance = Watchtower()
    return _watchtower_instance