src/suggest.py

"""Model suggestions submitted through the Space.

Appended to `suggestions/{date}.jsonl` in the dataset via `CommitScheduler`,
which batches writes instead of opening a commit per click. The weekly indexer
reads them; nothing here changes the index directly.

The suggestion box is an unauthenticated text input on a public page, so
everything arriving through it is treated as hostile until validated:

  * the value must look like `author/model-name` -- checked against a strict
    pattern, not merely for a slash
  * length is capped
  * the parsed id is what gets written, never the raw input
  * writes are rate-limited per session

`validate_model_id` is deliberately stricter than the Hub's own rules. A
suggestion that fails it is a suggestion the user can retype; a suggestion that
smuggles a path traversal or a newline into a JSONL file is a bug in us.
"""

from __future__ import annotations

import json
import logging
import os
import re
import threading
from datetime import date, datetime, timezone
from pathlib import Path

log = logging.getLogger("atlas.suggest")

DATASET_REPO = os.environ.get("ATLAS_DATASET", "Bit-Trading-Company/bit-finance-atlas")
SUGGESTIONS_DIR = "suggestions"

MAX_ID_LENGTH = 96

# Hub namespaces and model names allow letters, digits, dot, dash, underscore.
# Anchored, with exactly one slash, so `../`, a newline, or a second path
# segment cannot get through.
MODEL_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,45}/[A-Za-z0-9][A-Za-z0-9._-]{0,45}$")

# Per-process throttle. Not a substitute for real abuse controls, but it stops
# one impatient click-through from opening a hundred commits.
MIN_SECONDS_BETWEEN = 2.0

MAX_MESSAGE_CHARS = 4000
MAX_EMAIL_LENGTH = 160

# Deliberately loose: this is a sanity check so the file stays useful, not an
# attempt to decide what a valid address is. Anything shaped roughly like an
# address passes; anything with a newline or a control character does not.
EMAIL_PATTERN = re.compile(r"^[^@\s;,<>\"\']{1,64}@[^@\s;,<>\"\']{1,90}\.[A-Za-z]{2,24}$")


def validate_email(raw) -> str:
    """Return a clean address, or raise SuggestionError."""
    if not raw or not isinstance(raw, str):
        raise SuggestionError("Enter an email address")
    value = raw.strip()
    if len(value) > MAX_EMAIL_LENGTH or not EMAIL_PATTERN.match(value):
        raise SuggestionError("That does not look like an email address")
    return value


class SuggestionError(ValueError):
    """The submitted value is not something we are willing to write."""


def validate_model_id(raw) -> str:
    """Return a clean `author/name`, or raise SuggestionError."""
    if not raw or not isinstance(raw, str):
        raise SuggestionError("Enter a model id")

    value = raw.strip()
    # A pasted URL is the most common non-id input; accept it by extracting.
    for prefix in ("https://huggingface.co/", "http://huggingface.co/",
                   "huggingface.co/"):
        if value.lower().startswith(prefix):
            value = value[len(prefix):]
            break
    value = value.strip("/").split("?")[0].split("#")[0]

    if len(value) > MAX_ID_LENGTH:
        raise SuggestionError("That id is too long")
    if not MODEL_ID_PATTERN.match(value):
        raise SuggestionError("Use the form author/model-name")
    return value


class Suggestions:
    """Batched appender for the dataset's suggestions folder.

    Falls back to writing locally when no write token is present, so the Space
    still runs (and the suite still passes) without credentials. A dropped
    suggestion is logged rather than surfaced as a crash -- the index is not
    damaged by losing one, and a stack trace helps nobody looking at a form.
    """

    def __init__(self, repo_id: str = None, token: str = None,
                 local_dir: str | Path = None, every_minutes: float = 5.0):
        self.repo_id = repo_id or DATASET_REPO
        self.token = token or os.environ.get("HF_WRITE_TOKEN") or ""
        self.folder = Path(local_dir or "suggestions_out")
        self.folder.mkdir(parents=True, exist_ok=True)
        self.path = self.folder / f"{date.today().isoformat()}.jsonl"
        self._lock = threading.Lock()
        self._last_write = 0.0
        self._scheduler = None

        if self.token:
            try:
                from huggingface_hub import CommitScheduler

                self._scheduler = CommitScheduler(
                    repo_id=self.repo_id,
                    repo_type="dataset",
                    folder_path=str(self.folder),
                    path_in_repo=SUGGESTIONS_DIR,
                    every=every_minutes,
                    token=self.token,
                    squash_history=False,
                )
                log.info("suggestions -> %s/%s every %sm",
                         self.repo_id, SUGGESTIONS_DIR, every_minutes)
            except Exception as exc:  # noqa: BLE001 - never block boot on this
                log.warning("CommitScheduler unavailable (%s) -- "
                            "suggestions stay local", exc)
        else:
            log.info("no HF_WRITE_TOKEN -- suggestions stay local")

    @property
    def syncing(self) -> bool:
        return self._scheduler is not None

    def submit(self, raw, source: str = "space") -> str:
        """Validate and append one suggestion. Returns the accepted model id."""
        import time

        model_id = validate_model_id(raw)

        with self._lock:
            now = time.monotonic()
            if now - self._last_write < MIN_SECONDS_BETWEEN:
                raise SuggestionError("Slow down a moment")
            self._last_write = now

            record = {
                "id": model_id,
                "submitted_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
                "source": source,
            }
            try:
                # The scheduler holds a lock on the folder while it uploads;
                # appending under it keeps a write from racing a commit.
                context = self._scheduler.lock if self._scheduler else _NullLock()
                with context:
                    with self.path.open("a", encoding="utf-8") as handle:
                        handle.write(json.dumps(record) + "\n")
            except OSError as exc:
                log.error("could not append suggestion: %s", exc)
                raise SuggestionError("Could not save that — try again later") from exc

        return model_id


    # ----------------------------------------------------------------
    # The other two things the chrome can submit
    # ----------------------------------------------------------------
    #
    # Same file, same validation posture: everything arriving from a public
    # form is hostile until proven otherwise, is length-capped, and is written
    # as one JSON object per line so a newline cannot forge a second record.

    def _append(self, record: dict, filename: str) -> None:
        import time

        with self._lock:
            now = time.monotonic()
            if now - self._last_write < MIN_SECONDS_BETWEEN:
                raise SuggestionError("Slow down a moment")
            self._last_write = now
            path = self.folder / filename
            record["submitted_at"] = datetime.now(timezone.utc).isoformat(
                timespec="seconds")
            try:
                context = self._scheduler.lock if self._scheduler else _NullLock()
                with context:
                    with path.open("a", encoding="utf-8") as handle:
                        handle.write(json.dumps(record) + "\n")
            except OSError as exc:
                log.error("could not append to %s: %s", filename, exc)
                raise SuggestionError("Could not save that — try again later") from exc

    def submit_interest(self, email: str, module: str) -> None:
        """"Notify me when <module> launches"."""
        address = validate_email(email)
        self._append({"module": str(module)[:80], "email": address,
                      "kind": "interest"},
                     f"interest-{date.today().isoformat()}.jsonl")

    def submit_contact(self, name: str, email: str, topic: str,
                       message: str) -> None:
        address = validate_email(email)
        body = (message or "").strip()
        if len(body) < 10:
            raise SuggestionError("Add a little more detail")
        self._append({
            "kind": "contact",
            "name": (name or "").strip()[:120],
            "email": address,
            "topic": (topic or "")[:60],
            "message": body[:MAX_MESSAGE_CHARS],
        }, f"contact-{date.today().isoformat()}.jsonl")


class _NullLock:
    def __enter__(self):
        return self

    def __exit__(self, *exc):
        return False