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 threading
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
        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."""
        # 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