poolprobe 1.0.0__py3-none-any.whl

poolprobe/safewrite.py

@@ -0,0 +1,66 @@
+"""poolprobe.safewrite — one atomic write path for every file this tool emits.
+
+A report written halfway is worse than no report: the reader sees a truncated
+HTML page or a JSON document that stops mid-object and concludes the run broke.
+Every write in the package goes through here.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+from pathlib import Path
+
+__all__ = ["write_atomic"]
+
+
+def write_atomic(path: str | Path, text: str, *, backup: bool = True) -> Path:
+    """Write ``text`` to ``path`` with no window where the file is half-there.
+
+    The temp file lives in the same directory so ``os.replace`` stays atomic
+    (same filesystem). The rotation copies rather than renames: a reader that
+    opens the path mid-write must never find it missing.
+    """
+    target = Path(path).expanduser()
+    target.parent.mkdir(parents=True, exist_ok=True)
+    temporary = target.with_name(target.name + ".tmp")
+    try:
+        temporary.write_text(text, encoding="utf-8")
+        if backup and target.exists():
+            try:
+                shutil.copy2(target, target.with_name(target.name + ".bak"))
+            except OSError:
+                pass  # a missing backup is never worth failing the write
+        os.replace(temporary, target)
+    finally:
+        temporary.unlink(missing_ok=True)
+    return target
+
+
+def selfcheck() -> None:
+    """Offline asserts. Run: ``python -m poolprobe.safewrite``."""
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as workdir:
+        target = Path(workdir) / "nested" / "report.json"
+        write_atomic(target, '{"first": true}')
+        assert target.read_text(encoding="utf-8") == '{"first": true}'
+        assert not target.with_name(target.name + ".bak").exists()
+
+        write_atomic(target, '{"second": true}')
+        assert target.read_text(encoding="utf-8") == '{"second": true}'
+        backup = target.with_name(target.name + ".bak")
+        assert backup.read_text(encoding="utf-8") == '{"first": true}'
+
+        # no leftovers, and backup rotation is opt-out
+        write_atomic(target, '{"third": true}', backup=False)
+        assert backup.read_text(encoding="utf-8") == '{"first": true}'
+        assert sorted(p.name for p in target.parent.iterdir()) == [
+            "report.json",
+            "report.json.bak",
+        ]
+
+
+if __name__ == "__main__":
+    selfcheck()
+    print("safewrite selfcheck ok")

poolprobe/state.py

@@ -0,0 +1,272 @@
+"""Streak memory: the difference between "failed once" and "failed for a week".
+
+No run in isolation can tell a flaky upstream from a dead one, and nothing in
+this package is allowed to delete a model on a single observation. So verdicts
+are persisted between runs and the prune gate reads the *streak*, not the last
+result.
+
+Two properties are load-bearing:
+
+* **Fail-open on load.** A missing, unreadable or corrupt state file starts
+  from zero streaks. That is safe in exactly one direction: a zeroed streak sits
+  below ``streak_min``, so losing the state can never cause a prune. Failing
+  closed instead would mean refusing to run — which hides a real outage behind a
+  bookkeeping error.
+* **Ghost-drop only on a full sweep.** Models missing from the results are
+  forgotten only when the run covered the whole pool. A two-model diagnostic run
+  once wiped 52 of 54 live rows from the state file; every streak reset to zero
+  and the prune gate went blind for a week.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import shutil
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from . import taxonomy
+from .safewrite import write_atomic
+
+if TYPE_CHECKING:  # runtime duck-typed: any object with these attributes works
+    from .probe import ProbeResult
+
+#: Bumped when the on-disk shape changes. An older file is still read — losing
+#: history is worse than reading a field that is no longer written.
+VERSION = 1
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+
+
+def _as_str(value: object, default: str = "") -> str:
+    return value if isinstance(value, str) else default
+
+
+def _as_int(value: object, default: int = 0) -> int:
+    try:
+        return int(value)  # type: ignore[arg-type]
+    except (TypeError, ValueError):
+        return default
+
+
+def _as_float(value: object, default: float = 0.0) -> float:
+    try:
+        return float(value)  # type: ignore[arg-type]
+    except (TypeError, ValueError):
+        return default
+
+
+@dataclass
+class ModelState:
+    """Everything remembered about one model across runs."""
+
+    verdict: str = ""
+    kind: str = ""
+    note: str = ""
+    provider: str = ""
+    fail_streak: int = 0
+    ok_streak: int = 0
+    last_seen: str = ""
+    last_ok: str = ""
+    latency_s: float = 0.0
+    status: int = 0
+    attempts: int = 1
+
+    @property
+    def dead(self) -> bool:
+        return self.verdict in taxonomy.DEAD
+
+    @classmethod
+    def from_dict(cls, raw: object) -> "ModelState":
+        """Tolerant reader: a state file is a cache, not a contract.
+
+        Every field is coerced, so a half-written or hand-edited file degrades to
+        defaults instead of raising at load time.
+        """
+        data = raw if isinstance(raw, dict) else {}
+        return cls(
+            verdict=_as_str(data.get("verdict")),
+            kind=_as_str(data.get("kind")),
+            note=_as_str(data.get("note")),
+            provider=_as_str(data.get("provider")),
+            fail_streak=max(0, _as_int(data.get("fail_streak"))),
+            ok_streak=max(0, _as_int(data.get("ok_streak"))),
+            last_seen=_as_str(data.get("last_seen")),
+            last_ok=_as_str(data.get("last_ok")),
+            latency_s=max(0.0, _as_float(data.get("latency_s"))),
+            status=_as_int(data.get("status")),
+            attempts=max(1, _as_int(data.get("attempts"), 1)),
+        )
+
+
+@dataclass
+class RunStats:
+    """What one :meth:`State.record` did, for the report footer."""
+
+    updated: int = 0
+    dropped: list[str] = field(default_factory=list)
+    run_count: int = 0
+
+
+class State:
+    """JSON file of per-model verdict history.
+
+    Written atomically (temp file in the same directory + ``os.replace``) with a
+    single ``.bak`` rotation, because this file is read by cron runs that can
+    overlap a manual sweep. A torn write here would look exactly like a lost
+    history.
+    """
+
+    def __init__(self, path: str | Path):
+        self.path = Path(path).expanduser()
+        self.models_: dict[str, ModelState] = {}
+        self.run_count = 0
+        self.updated = ""
+        self.warnings: list[str] = []
+
+    # -- reading -----------------------------------------------------------
+
+    def load(self) -> "State":
+        """Read the file if it is usable; otherwise start empty and say so.
+
+        Corrupt input is preserved as ``.corrupt.bak`` before being discarded —
+        the operator may want to know what the tool forgot, and silently throwing
+        away weeks of history is how a gate goes blind without a trace.
+        """
+        self.warnings = []
+        self.models_ = {}
+        self.run_count = 0
+        self.updated = ""
+        if not self.path.exists():
+            return self
+
+        try:
+            raw = json.loads(self.path.read_text(encoding="utf-8"))
+        except Exception as exc:
+            self._quarantine(f"state file is not readable JSON ({type(exc).__name__}: {exc})")
+            return self
+
+        if not isinstance(raw, dict):
+            self._quarantine("state file is not a JSON object")
+            return self
+
+        version = _as_int(raw.get("version"), VERSION)
+        if version != VERSION:
+            # Still read it: an unknown version means someone upgraded or
+            # downgraded the tool, and the streaks are the valuable part.
+            self.warnings.append(
+                f"state version {version} != {VERSION}, reading it anyway"
+            )
+
+        table = raw.get("models")
+        if not isinstance(table, dict):
+            self._quarantine("state file has no models{} object")
+            return self
+
+        for name, entry in table.items():
+            if not isinstance(name, str) or not name:
+                continue
+            self.models_[name] = ModelState.from_dict(entry)
+        self.run_count = max(0, _as_int(raw.get("run_count")))
+        self.updated = _as_str(raw.get("updated"))
+        return self
+
+    def _quarantine(self, reason: str) -> None:
+        self.warnings.append(f"{self.path}: {reason} — starting from empty state")
+        try:
+            shutil.copy2(self.path, self.path.with_name(self.path.name + ".corrupt.bak"))
+            self.warnings.append(f"previous state preserved at {self.path.name}.corrupt.bak")
+        except Exception:
+            pass
+
+    # -- writing -----------------------------------------------------------
+
+    def record(
+        self,
+        results: "list[ProbeResult]",
+        *,
+        full_sweep: bool = False,
+        now: str | None = None,
+    ) -> RunStats:
+        """Fold one run's results into the history.
+
+        ``full_sweep=True`` means "these results cover the entire pool", which is
+        the only condition under which unseen models are dropped. Streaks move in
+        both directions: a dead verdict increments ``fail_streak`` and zeroes
+        ``ok_streak``, a live one does the opposite. ``EMPTY-LENGTH`` counts as
+        live, because it is — the route answered, it just ran out of budget.
+        """
+        stamp = now or _now()
+        seen: set[str] = set()
+        for result in results:
+            name = getattr(result, "model", "") or ""
+            if not name:
+                continue
+            seen.add(name)
+            previous = self.models_.get(name) or ModelState()
+            dead = getattr(result, "verdict", "") in taxonomy.DEAD
+            provider = getattr(result, "provider", "") or previous.provider
+            self.models_[name] = ModelState(
+                verdict=getattr(result, "verdict", ""),
+                kind=getattr(result, "kind", ""),
+                note=getattr(result, "note", ""),
+                provider=provider,
+                fail_streak=(previous.fail_streak + 1) if dead else 0,
+                ok_streak=0 if dead else (previous.ok_streak + 1),
+                last_seen=stamp,
+                last_ok=stamp if not dead else previous.last_ok,
+                latency_s=_as_float(getattr(result, "latency_s", 0.0)),
+                status=_as_int(getattr(result, "status", 0)),
+                attempts=max(1, _as_int(getattr(result, "attempts", 1), 1)),
+            )
+
+        dropped: list[str] = []
+        if full_sweep:
+            dropped = sorted(set(self.models_) - seen)
+            for name in dropped:
+                del self.models_[name]
+
+        self.run_count += 1
+        self.updated = stamp
+        return RunStats(updated=len(seen), dropped=dropped, run_count=self.run_count)
+
+    def save(self) -> Path:
+        """Atomic write with one ``.bak`` rotation."""
+        payload = {
+            "version": VERSION,
+            "updated": self.updated or _now(),
+            "run_count": self.run_count,
+            "models": {name: asdict(state) for name, state in sorted(self.models_.items())},
+        }
+        return write_atomic(
+            self.path,
+            json.dumps(payload, indent=2, sort_keys=False, ensure_ascii=False) + "\n",
+        )
+
+    # -- queries -----------------------------------------------------------
+
+    def get(self, model: str) -> ModelState | None:
+        return self.models_.get(model)
+
+    def streak(self, model: str) -> int:
+        """Consecutive failing runs for ``model`` (0 when unknown)."""
+        state = self.models_.get(model)
+        return state.fail_streak if state else 0
+
+    def models(self) -> list[str]:
+        return sorted(self.models_)
+
+    def fingerprint(self) -> str:
+        """Stable hash of the current dead/alive picture.
+
+        Used to dedupe alerts: the same pool state should notify once, not once
+        per cron tick. Sorting before hashing keeps it independent of discovery
+        order, which is not stable across runs.
+        """
+        lines = [f"{name}:{state.kind or state.verdict}" for name, state in sorted(self.models_.items())]
+        return hashlib.sha256("\n".join(lines).encode("utf-8")).hexdigest()[:16]

poolprobe/taxonomy.py

@@ -0,0 +1,312 @@
+"""Verdict and death-kind classification for OpenAI-compatible gateway probes.
+
+This module is the reason poolprobe exists. Every rule in it was paid for on a
+live pool of ~60 routed models behind a self-hosted gateway, over six weeks of
+scheduled sweeps. The rules read as paranoid because each one is a scar:
+
+* ``EMPTY-LENGTH`` is alive, not dead — a reasoning model burned ``max_tokens``
+  on thinking and never reached content. Pruning it kills a working route.
+* ``only supported on /v1/images/generations`` is a wrong-modality error, not a
+  dead model. Image, video, TTS and music models legitimately fail a chat probe.
+* A single probe lies. Rate limits, cold starts and one hung upstream produced
+  red verdicts for models that answered 5/5 two minutes later. Nothing in this
+  package is allowed to delete a model on one observation.
+* Only a hard upstream "this model does not exist" is permanent. Empty bodies,
+  429s, quota errors and transport timeouts are transient or credential-level
+  and must never drive an automatic prune.
+
+Nothing here performs I/O. Pure functions over already-collected evidence, so
+the whole taxonomy is testable offline (see ``selfcheck.py``).
+"""
+
+from __future__ import annotations
+
+# --- probe verdicts ---------------------------------------------------------
+
+OK = "OK"
+EMPTY = "EMPTY"
+EMPTY_LENGTH = "EMPTY-LENGTH"
+ERR = "ERR"
+TIMEOUT = "TIMEOUT"
+
+VERDICTS: tuple[str, ...] = (OK, EMPTY_LENGTH, EMPTY, ERR, TIMEOUT)
+
+#: Verdicts that mean "this probe got nothing usable back".
+#: Dead is NOT the same as prunable — see :func:`prunable`.
+DEAD: frozenset[str] = frozenset({EMPTY, ERR, TIMEOUT})
+
+# --- death kinds ------------------------------------------------------------
+
+ALIVE = "ALIVE"
+ALIVE_STARVED = "ALIVE-STARVED"
+PERMANENT = "PERMANENT"
+TRANSIENT = "TRANSIENT"
+CREDENTIAL = "CREDENTIAL"
+MODALITY = "MODALITY"
+UNKNOWN = "UNKNOWN"
+
+KINDS: tuple[str, ...] = (
+    ALIVE,
+    ALIVE_STARVED,
+    PERMANENT,
+    TRANSIENT,
+    CREDENTIAL,
+    MODALITY,
+    UNKNOWN,
+)
+
+#: The only kind that may ever justify removing a model from a config.
+PRUNABLE_KINDS: frozenset[str] = frozenset({PERMANENT})
+
+# --- signature tables -------------------------------------------------------
+# Matched case-insensitively as substrings of the note. Order of the *tables*
+# matters; see classify_death for the precedence and why it is that way.
+
+MODALITY_SIGNS: tuple[str, ...] = (
+    "only supported",
+    "not supported on",
+    "use the /v1/images",
+    "use the /v1/audio",
+    "images/generations",
+    "audio/speech",
+    "audio/transcriptions",
+    "videos/generations",
+    "wrong endpoint",
+    "invalid endpoint",
+)
+
+PERMANENT_SIGNS: tuple[str, ...] = (
+    "not found",
+    "does not exist",
+    "unknown model",
+    "no such model",
+    "model_not_found",
+    "invalid model",
+    "unsupported model",
+    "is not supported",
+    "model is not available",
+    "not a valid model",
+)
+
+CREDENTIAL_SIGNS: tuple[str, ...] = (
+    "unauthorized",
+    "forbidden",
+    "invalid api key",
+    "invalid_api_key",
+    "incorrect api key",
+    "authentication",
+    "auth_unavailable",
+    "permission denied",
+    "quota",
+    "credit",
+    "insufficient",
+    "billing",
+    "account tier",
+    "x-ratelimit-limit-req-minute: 0",
+    "用户额度不足",  # "user quota insufficient" — the literal upstream text
+    "余额不足",
+)
+
+TRANSIENT_SIGNS: tuple[str, ...] = (
+    "rate limit",
+    "rate_limit",
+    "too many requests",
+    "overloaded",
+    "upstream",
+    "timed out",
+    "timeout",
+    "temporarily",
+    "try again",
+    "bad gateway",
+    "service unavailable",
+    "gateway time-out",
+    "connection reset",
+    "connection refused",
+    "remote end closed",
+    "expecting value",  # non-JSON body from a proxy that was mid-failure
+    "ledger_unavailable",
+)
+
+#: HTTP statuses that mean "the credential or the account, not the model".
+CREDENTIAL_STATUSES: frozenset[int] = frozenset({401, 403})
+
+#: HTTP statuses that mean "the upstream is having a moment".
+TRANSIENT_STATUSES: frozenset[int] = frozenset({408, 425, 429, 500, 502, 503, 504, 529})
+
+# --- generation (non-chat) model detection ----------------------------------
+# A chat probe against an image generator is a category error, not a liveness
+# signal. These substrings are deliberately loose: pool aliases are named by
+# whoever wired the route ("banana", "flux", "suno"), not by a schema.
+
+GENERATION_TOKENS: tuple[str, ...] = (
+    "image",
+    "imagen",
+    "img",
+    "video",
+    "tts",
+    "speech",
+    "audio",
+    "voice",
+    "music",
+    "song",
+    "suno",
+    "flux",
+    "banana",
+    "dall",
+    "midjourney",
+    "kling",
+    "wan-2",
+    "imagine",
+    "embed",
+    "rerank",
+)
+
+#: One line of human-facing advice per death kind. Printed in every report so
+#: the reader gets an action, not just a colour.
+ADVICE: dict[str, str] = {
+    ALIVE: "",
+    ALIVE_STARVED: "alive but starved: raise --max-tokens (>=1024) for reasoning models",
+    MODALITY: "not a chat model: probe it on its own endpoint (images/audio/video)",
+    CREDENTIAL: "credential or account level: rotate/top up the key, never prune the model",
+    TRANSIENT: "upstream hiccup: burst-confirm before believing it, then retry later",
+    PERMANENT: "upstream does not know this model: prunable after burst confirmation",
+    UNKNOWN: "unrecognised failure: read the note by hand, do not automate on it",
+}
+
+
+def base_name(model: str) -> str:
+    """Strip a namespace prefix: ``Lite/glm-5.3-flash`` -> ``glm-5.3-flash``.
+
+    Gateways that expose several upstreams publish namespaced aliases. The
+    namespace is a routing hint, not a different model, so protection rules
+    keyed on the real name have to be checked against both forms.
+    """
+    return model.split("/", 1)[-1].strip()
+
+
+def is_generation(model: str) -> bool:
+    """True if the model name advertises a non-chat modality."""
+    lowered = model.lower()
+    return any(token in lowered for token in GENERATION_TOKENS)
+
+
+def classify(status: int, body: object, error: str | None = None) -> tuple[str, str]:
+    """Turn one HTTP exchange into ``(verdict, note)``.
+
+    ``body`` is the parsed JSON reply (may be any shape — upstreams are not
+    trustworthy), ``error`` is the exception text when the exchange itself
+    failed. Parsing lives inside a guard: a model that once replied with a
+    top-level JSON list killed an entire sweep run because the extraction was
+    unguarded.
+    """
+    if error is not None:
+        return (TIMEOUT, "no answer within timeout") if _is_timeout(error) else (
+            ERR,
+            error[:160],
+        )
+
+    try:
+        choices = (body or {}).get("choices") or [{}]  # type: ignore[union-attr]
+        message = choices[0].get("message") or {}
+        content = (message.get("content") or "").strip()
+        reasoning = (message.get("reasoning_content") or message.get("reasoning") or "")
+        finish = choices[0].get("finish_reason") or "?"
+    except Exception as exc:  # reply shape is untrusted upstream JSON
+        return ERR, f"unparsable reply: {type(exc).__name__} {str(body)[:80]}"
+
+    if content or (isinstance(reasoning, str) and reasoning.strip()):
+        return OK, (content or reasoning.strip())[:48].replace("\n", " ")
+    if finish == "length":
+        return EMPTY_LENGTH, "thinking ate max_tokens"
+    if status and status != 200:
+        return EMPTY, f"http {status} finish={finish}"
+    return EMPTY, f"finish={finish} reasoning={bool(reasoning)}"
+
+
+def classify_death(verdict: str, note: str, status: int = 0) -> str:
+    """Map a verdict + note onto a :data:`KINDS` member.
+
+    Precedence is the whole design:
+
+    1. ``TIMEOUT`` is always transient. A 30s budget once timed out on a route
+       that answered ``ok`` in 66.8s — a timeout proves nothing about existence.
+    2. 401/403 is the credential, not the model. Pruning here deletes a good
+       alias because a key ran dry.
+    3. Modality beats permanence. ``model X only supported on /v1/images/...``
+       contains "supported" and would otherwise match a permanent signature.
+    4. Only then does an explicit "not found" family count as permanent.
+    5. A bare 404 with no readable body is a permanent *candidate* — the caller
+       still has to burst-confirm and check the streak before pruning.
+    """
+    if verdict == OK:
+        return ALIVE
+    if verdict == EMPTY_LENGTH:
+        return ALIVE_STARVED
+    if verdict == TIMEOUT:
+        return TRANSIENT
+
+    lowered = note.lower()
+    if status in CREDENTIAL_STATUSES:
+        return CREDENTIAL
+    if any(sign in lowered for sign in MODALITY_SIGNS):
+        return MODALITY
+    if any(sign in lowered for sign in PERMANENT_SIGNS):
+        return PERMANENT
+    if status == 404:
+        return PERMANENT
+    if any(sign in lowered for sign in CREDENTIAL_SIGNS):
+        return CREDENTIAL
+    if status in TRANSIENT_STATUSES or any(sign in lowered for sign in TRANSIENT_SIGNS):
+        return TRANSIENT
+    return UNKNOWN
+
+
+def advice(kind: str) -> str:
+    return ADVICE.get(kind, "")
+
+
+def protected(model: str, protect: frozenset[str] | set[str]) -> bool:
+    """True if the model (or its base name) is on the protect list."""
+    if not protect:
+        return False
+    return model in protect or base_name(model) in protect
+
+
+def prunable(
+    model: str,
+    verdict: str,
+    kind: str,
+    fail_streak: int,
+    streak_min: int,
+    protect: frozenset[str] | set[str] = frozenset(),
+) -> bool:
+    """The single prune predicate. Report and prune loop both call this.
+
+    One function on purpose: a report that says "prunable" and a prune loop
+    that deletes something else is the worst possible failure mode for a tool
+    whose job is to not delete working routes.
+
+    All five conditions must hold:
+
+    * the verdict is dead (empty body, error, or timeout),
+    * the death kind is permanent — upstream does not know the model,
+    * the streak reached ``streak_min`` consecutive failing sweeps,
+    * the model is not on the protect list,
+    * the model is not a non-chat (generation/embedding) modality.
+    """
+    if verdict not in DEAD:
+        return False
+    if kind not in PRUNABLE_KINDS:
+        return False
+    if fail_streak < streak_min:
+        return False
+    if protected(model, protect):
+        return False
+    if is_generation(model) or is_generation(base_name(model)):
+        return False
+    return True
+
+
+def _is_timeout(error: str) -> bool:
+    lowered = error.lower()
+    return "timed out" in lowered or "timeout" in lowered