poolprobe 1.0.0__py3-none-any.whl

poolprobe/__init__.py

@@ -0,0 +1,137 @@
+"""poolprobe — a doctor for OpenAI-compatible LLM gateways.
+
+A pool of routed models rots quietly. An upstream drops a model, a key hits an
+account tier, a reasoning model starts eating its whole token budget on
+thinking, a gateway route begins answering SSE to a non-streaming request. None
+of that raises an exception anywhere: the agent just gets an empty reply, or a
+slower one, or falls through to a worse model and nobody notices for a week.
+
+poolprobe probes the pool and refuses to guess. Three rules drive everything
+here:
+
+* **A verdict is not a death sentence.** Every failure is classified into a
+  kind — ``ALIVE``, ``ALIVE-STARVED``, ``PERMANENT``, ``TRANSIENT``,
+  ``CREDENTIAL``, ``MODALITY``, ``UNKNOWN`` — and only ``PERMANENT`` (the
+  upstream does not know this model) is ever a prune candidate.
+* **One observation is not a verdict.** A measured pool produced 544 scheduled
+  liveness runs in six weeks; 39% went red and every single red was transient.
+  Burst-confirm before believing anything.
+* **A number without evidence is not reported.** Throughput comes from the
+  gateway's own ``usage`` block or it is marked as a lower bound or it is
+  ``0.0``. Never estimated from chunk counts and presented as a measurement.
+
+Library use::
+
+    from poolprobe import Gateway, sweep, taxonomy
+
+    gw = Gateway("http://127.0.0.1:8317/v1", api_key="local")
+    for result in sweep(gw, ["gpt-5.4", "glm-5.3-flash"], attempts=5):
+        print(result.model, result.verdict, result.kind, result.advice)
+
+CLI use::
+
+    poolprobe sweep --base-url http://127.0.0.1:8317/v1 --attempts 5 --html out.html
+    poolprobe prune --base-url ... --adapter cliproxy --config ~/.config/gateway.yaml
+    poolprobe selfcheck
+"""
+
+from __future__ import annotations
+
+__version__ = "1.0.0"
+
+# Submodules first: probe, bench and adapter all resolve ``taxonomy`` through
+# the package attribute at import time, and a partially initialised package must
+# not leave that lookup to CPython's submodule fallback.
+from . import chains, taxonomy
+from .adapter import (
+    Adapter,
+    Attribution,
+    CliproxyAdapter,
+    ConfigError,
+    EditPlan,
+    LiveAdapter,
+    build_adapter,
+)
+from .bench import (
+    BenchResult,
+    bench_many,
+    bench_one,
+    summarize,
+)
+from .chains import (
+    load_config,
+    validate as validate_chains,
+    validate_file as validate_chains_file,
+)
+from .probe import (
+    DEFAULT_MAX_TOKENS,
+    DEFAULT_PROMPT,
+    DEFAULT_TIMEOUT,
+    Gateway,
+    GatewayError,
+    ProbeResult,
+    burst,
+    list_models,
+    parallel_map,
+    ping,
+    probe,
+    sweep,
+)
+from .report import (
+    Report,
+    RunMeta,
+    render_html,
+    render_json,
+    render_terminal,
+)
+from .safewrite import write_atomic
+from .state import ModelState, RunStats, State
+
+__all__ = [
+    "__version__",
+    # taxonomy
+    "taxonomy",
+    # transport
+    "Gateway",
+    "GatewayError",
+    "ProbeResult",
+    "probe",
+    "burst",
+    "sweep",
+    "list_models",
+    "ping",
+    "parallel_map",
+    "DEFAULT_PROMPT",
+    "DEFAULT_MAX_TOKENS",
+    "DEFAULT_TIMEOUT",
+    # memory
+    "State",
+    "ModelState",
+    "RunStats",
+    # bench
+    "BenchResult",
+    "bench_one",
+    "bench_many",
+    "summarize",
+    # adapters
+    "Adapter",
+    "LiveAdapter",
+    "CliproxyAdapter",
+    "Attribution",
+    "EditPlan",
+    "ConfigError",
+    "build_adapter",
+    # chains
+    "chains",
+    "validate_chains",
+    "validate_chains_file",
+    "load_config",
+    # io
+    "write_atomic",
+    # reporting
+    "Report",
+    "RunMeta",
+    "render_terminal",
+    "render_json",
+    "render_html",
+]

poolprobe/__main__.py

@@ -0,0 +1,8 @@
+"""``python -m poolprobe`` — the same entry point the console script uses."""
+
+from __future__ import annotations
+
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

poolprobe/adapter.py

@@ -0,0 +1,525 @@
+"""Where the model list comes from, and what a model is allowed to cost you.
+
+Two adapters, one interface:
+
+* :class:`LiveAdapter` — any OpenAI-compatible gateway. Models come from
+  ``GET /models``. Nothing is known about upstreams and nothing can be edited,
+  because there is no config file in reach. This is the default and it works
+  against OpenRouter, one-api, LiteLLM, vLLM, cliproxy or a bare vendor endpoint.
+* :class:`CliproxyAdapter` — adds attribution (which upstream provider serves
+  this alias, is it disabled, what is the real model name) by reading the
+  gateway's own YAML config, and can plan a removal.
+
+The split between *reading* and *removing* is deliberate:
+
+* Attribution is parsed **structurally** with PyYAML. Line-based regex parsing
+  of a config whose key order is not fixed is how a gate silently attributes
+  model A to provider B.
+* Removal is **text surgery** on the original bytes. A YAML round-trip would
+  destroy every comment in the operator's config, and comments in a routed-model
+  config are where the "why is this here" lives.
+
+So: parse to understand, cut to change. Never the other way round.
+"""
+
+from __future__ import annotations
+
+import difflib
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+
+from .probe import Gateway, GatewayError, list_models
+
+#: Top-level YAML keys whose list items are *provider blocks* holding models.
+#: Kept as a tuple so a different gateway layout is a one-line extension.
+PROVIDER_SECTIONS: tuple[str, ...] = ("openai-compatibility", "claude-api-key")
+
+#: A list item is a model entry only when its nearest lower-indented ancestor
+#: line is one of these keys. Provider blocks also carry ``name:`` and would
+#: otherwise be mistaken for models — which would delete a whole provider.
+MODEL_LIST_KEYS: frozenset[str] = frozenset({"models"})
+
+
+class ConfigError(RuntimeError):
+    """The gateway config could not be read or understood."""
+
+
+def _load_yaml(path: Path) -> object:
+    try:
+        import yaml  # PyYAML: hard dependency, imported lazily so the generic
+    except ImportError as exc:  # lane keeps working without it
+        raise ConfigError(
+            "PyYAML is required to read a gateway config: pip install 'PyYAML>=6.0'"
+        ) from exc
+    try:
+        with path.open("r", encoding="utf-8") as handle:
+            return yaml.safe_load(handle)
+    except FileNotFoundError as exc:
+        raise ConfigError(f"config not found: {path}") from exc
+    except Exception as exc:  # includes yaml.YAMLError
+        raise ConfigError(f"cannot parse {path}: {type(exc).__name__}: {exc}") from exc
+
+
+def read_text(path: Path) -> str:
+    """Read without newline translation — removal must not rewrite the file."""
+    with path.open("r", encoding="utf-8", newline="") as handle:
+        return handle.read()
+
+
+def effective_alias(name: str, alias: str = "") -> str:
+    """The id a client actually sends: ``alias`` when set, else ``name``."""
+    return (alias or name or "").strip()
+
+
+# --- attribution ------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class Attribution:
+    """What is known about one published model id."""
+
+    provider: str = ""
+    upstream: str = ""
+    base_url: str = ""
+    disabled: bool = False
+    source: str = "live"  # live | config | both
+
+    @property
+    def known(self) -> bool:
+        return bool(self.provider or self.upstream)
+
+    def as_dict(self) -> dict[str, object]:
+        return {
+            "provider": self.provider,
+            "upstream": self.upstream,
+            "base_url": self.base_url,
+            "disabled": self.disabled,
+            "source": self.source,
+        }
+
+
+@dataclass
+class EditPlan:
+    """A proposed config change. Nothing here writes to disk."""
+
+    path: str
+    before: str
+    after: str
+    removed: list[str] = field(default_factory=list)
+    not_found: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+
+    @property
+    def changed(self) -> bool:
+        return self.before != self.after
+
+    def diff(self, context: int = 3) -> str:
+        name = Path(self.path).name
+        return "".join(
+            difflib.unified_diff(
+                self.before.splitlines(keepends=True),
+                self.after.splitlines(keepends=True),
+                fromfile=f"a/{name}",
+                tofile=f"b/{name}",
+                n=context,
+            )
+        )
+
+
+class Adapter:
+    """Interface both adapters implement."""
+
+    name = "adapter"
+    removable = False
+
+    def models(self) -> list[str]:
+        raise NotImplementedError
+
+    def attribution(self, model: str) -> Attribution:
+        return Attribution()
+
+    def attribution_map(self) -> dict[str, Attribution]:
+        return {}
+
+    def describe(self) -> str:
+        return self.name
+
+    def plan_removal(self, aliases: Sequence[str]) -> EditPlan:
+        raise ConfigError(f"adapter {self.name!r} cannot edit a config")
+
+
+class LiveAdapter(Adapter):
+    """Any OpenAI-compatible gateway, catalog taken from ``GET /models``."""
+
+    name = "generic"
+
+    def __init__(self, gateway: Gateway, catalog: Sequence[str] | None = None):
+        self.gateway = gateway
+        self._catalog: list[str] | None = list(catalog) if catalog is not None else None
+
+    def models(self) -> list[str]:
+        if self._catalog is None:
+            self._catalog = list_models(self.gateway)
+        return list(self._catalog)
+
+    def describe(self) -> str:
+        # The gateway root is already its own field in the report header, so
+        # repeating it here read "URL · URL (? models, read-only)".
+        return "generic adapter (read-only)"
+
+
+class CliproxyAdapter(Adapter):
+    """Live catalog + provider attribution from the gateway's YAML config.
+
+    Models are the **union** of what the gateway publishes and what the config
+    declares. The union matters in both directions: an alias present in config
+    but missing from ``/models`` is exactly the config rot a sweep should show,
+    and an alias published but absent from config means someone edited the
+    running process or another section owns it — either way, not removable here.
+    """
+
+    name = "cliproxy"
+    removable = True
+
+    def __init__(self, gateway: Gateway, config_path: str | Path):
+        self.gateway = gateway
+        self.config_path = Path(config_path).expanduser()
+        self._attr: dict[str, Attribution] | None = None
+        self._live: list[str] | None = None
+        self._providers = 0
+
+    # -- reading ---------------------------------------------------------
+
+    def _parse(self) -> dict[str, Attribution]:
+        doc = _load_yaml(self.config_path)
+        if not isinstance(doc, dict):
+            raise ConfigError(f"{self.config_path}: expected a mapping at top level")
+
+        attr: dict[str, Attribution] = {}
+        providers = 0
+        for section in PROVIDER_SECTIONS:
+            blocks = doc.get(section)
+            if not isinstance(blocks, list):
+                continue
+            for block in blocks:
+                if not isinstance(block, dict):
+                    continue
+                providers += 1
+                provider = str(block.get("name") or section)
+                base_url = str(block.get("base-url") or block.get("base_url") or "")
+                disabled = bool(block.get("disabled"))
+                entries = block.get("models")
+                if not isinstance(entries, list):
+                    continue
+                for entry in entries:
+                    if isinstance(entry, str):
+                        name, alias = entry, ""
+                    elif isinstance(entry, dict):
+                        name = str(entry.get("name") or "")
+                        alias = str(entry.get("alias") or "")
+                    else:
+                        continue
+                    key = effective_alias(name, alias)
+                    if not key:
+                        continue
+                    # First declaration wins: a duplicate alias later in the
+                    # file is a config bug, and silently overwriting attribution
+                    # would hide it from the report.
+                    attr.setdefault(
+                        key,
+                        Attribution(
+                            provider=provider,
+                            upstream=name.strip(),
+                            base_url=base_url,
+                            disabled=disabled,
+                            source="config",
+                        ),
+                    )
+        self._providers = providers
+        return attr
+
+    def attribution_map(self) -> dict[str, Attribution]:
+        if self._attr is None:
+            self._attr = self._parse()
+        return dict(self._attr)
+
+    def attribution(self, model: str) -> Attribution:
+        return self.attribution_map().get(model, Attribution())
+
+    def live_models(self) -> list[str]:
+        if self._live is None:
+            self._live = list_models(self.gateway)
+        return list(self._live)
+
+    def models(self) -> list[str]:
+        known = set(self.attribution_map())
+        try:
+            published = self.live_models()
+        except GatewayError:
+            # The gateway being down is a finding, not a reason to report an
+            # empty pool: fall back to the config so the sweep still has work.
+            published = []
+            self._live = []
+        if self._attr is not None:
+            for key in published:
+                seen = self._attr.get(key)
+                self._attr[key] = (
+                    replace(seen, source="both")
+                    if seen is not None
+                    else Attribution(source="live")
+                )
+        return sorted(known | set(published))
+
+    def describe(self) -> str:
+        total = len(self.attribution_map())
+        head = f"{self.config_path} ({self._providers} provider blocks, {total} declared"
+        # _live stays None until the catalog is fetched *through the adapter*;
+        # ping reads it straight from the gateway, so don't claim a "0 published"
+        # we never measured. An empty list (gateway down) is a real 0 and shows.
+        if self._live is not None:
+            head += f", {len(self._live)} published"
+        return head + ")"
+
+    # -- removal ---------------------------------------------------------
+
+    def plan_removal(self, aliases: Sequence[str]) -> EditPlan:
+        """Cut model entries out of the config text, preserving everything else.
+
+        Removal is driven by the text scan, which locates any ``- name:/alias:``
+        item under a ``models:`` key in the file. Attribution is narrower — it
+        only reads :data:`PROVIDER_SECTIONS` — so an alias from a section this
+        adapter does not know can still be cut, and the diff shown to the
+        operator is the only authority. Aliases that appear nowhere in the text
+        (published by a running process, or already removed) come back in
+        ``not_found`` and are left alone.
+        """
+        wanted = [a for a in dict.fromkeys(aliases) if a]
+        before = read_text(self.config_path)
+        entries = {entry.alias: entry for entry in iter_model_entries(before)}
+
+        removed: list[str] = []
+        not_found: list[str] = []
+        cut_spans: list[tuple[int, int]] = []
+        for alias in wanted:
+            entry = entries.get(alias)
+            if entry is None:
+                not_found.append(alias)
+                continue
+            removed.append(alias)
+            cut_spans.append((entry.start, entry.end))
+
+        lines = before.splitlines(keepends=True)
+        for start, end in sorted(cut_spans, reverse=True):
+            del lines[start:end]
+        after = "".join(lines)
+
+        return EditPlan(
+            path=str(self.config_path),
+            before=before,
+            after=after,
+            removed=removed,
+            not_found=not_found,
+            warnings=_emptied_providers(before, after),
+        )
+
+
+# --- text surgery -----------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class TextEntry:
+    """One ``- name:/alias:`` list item, located by line span."""
+
+    alias: str
+    upstream: str
+    start: int  # inclusive line index of the dash line
+    end: int  # exclusive
+    indent: int  # column of the first character after "- "
+
+
+def _indent(line: str) -> int:
+    return len(line) - len(line.lstrip(" "))
+
+
+def _is_item(line: str) -> bool:
+    stripped = line.lstrip(" ")
+    return stripped.startswith("- ") or stripped.rstrip() == "-"
+
+
+def _scalar(line: str, content_indent: int) -> tuple[str, str] | None:
+    """``key: value`` whose value starts at ``content_indent``, else None.
+
+    The dash line is part of the item: in ``- name: foo`` the key sits two
+    columns right of the indentation, at exactly ``content_indent``. Treating it
+    as "not at this indent" would lose the upstream name of every entry.
+    """
+    indent = _indent(line)
+    stripped = line.lstrip(" ")
+    if stripped.startswith("- "):
+        if indent + 2 != content_indent:
+            return None
+        stripped = stripped[2:]
+    elif indent != content_indent:
+        return None
+
+    stripped = stripped.strip()
+    if not stripped or stripped.startswith("#") or ":" not in stripped:
+        return None
+    key, _, value = stripped.partition(":")
+    if not key.strip():
+        return None
+    return key.strip(), value.split(" #", 1)[0].strip().strip("'\"")
+
+
+def _ancestor_key(lines: list[str], index: int, content_indent: int) -> str:
+    """Key of the mapping line that owns the list item at ``index``.
+
+    Two kinds of line must be stepped over: sibling items (a dash at
+    ``content_indent - 2``, or at the parent indent in the zero-indent list
+    style) and keys belonging to a previous sibling's body (anything at
+    ``content_indent`` or deeper). Stopping at a sibling instead would read
+    ``name`` as the ancestor and drop every entry after the first.
+    """
+    dash_indent = content_indent - 2
+    for back in range(index - 1, -1, -1):
+        line = lines[back]
+        stripped = line.strip()
+        if not stripped or stripped.startswith("#"):
+            continue
+        indent = _indent(line)
+        if stripped.startswith("- ") or stripped.rstrip() == "-":
+            if indent >= dash_indent:
+                continue
+        elif indent >= content_indent:
+            continue
+        return stripped.lstrip("- ").partition(":")[0].strip()
+    return ""
+
+
+def iter_model_entries(text: str) -> Iterator[TextEntry]:
+    """Yield every model list item in a YAML config, with its line span.
+
+    An item counts as a model entry when its nearest lower-indented ancestor is
+    a ``models:`` key. That single rule keeps provider blocks (which also have
+    ``name:`` and sit in a list) out of the result, so a prune can never delete
+    a provider because its name collided with a model alias.
+
+    The span extends over every following line indented at least as deep as the
+    item content, so a multi-key entry (``name`` + ``alias`` + anything else)
+    and its nested lists are removed together. Blank lines and trailing comments
+    at lower indent end the span — they belong to the file, not the entry.
+    """
+    lines = text.splitlines(keepends=True)
+    for index, line in enumerate(lines):
+        if not _is_item(line):
+            continue
+        content_indent = _indent(line) + 2
+        if _ancestor_key(lines, index, content_indent) not in MODEL_LIST_KEYS:
+            continue
+
+        end = index + 1
+        while end < len(lines):
+            candidate = lines[end]
+            if not candidate.strip():
+                break
+            if candidate.lstrip().startswith("#") and _indent(candidate) < content_indent:
+                break
+            if _indent(candidate) < content_indent:
+                break
+            end += 1
+
+        name = alias = ""
+        for offset in range(index, end):
+            pair = _scalar(lines[offset], content_indent)
+            if pair is None:
+                continue
+            key, value = pair
+            if key == "name":
+                name = value
+            elif key == "alias":
+                alias = value
+
+        key = effective_alias(name, alias)
+        if key:
+            yield TextEntry(
+                alias=key, upstream=name, start=index, end=end, indent=content_indent
+            )
+
+
+def _emptied_providers(before: str, after: str) -> list[str]:
+    """Warn when a cut left a provider with no models at all.
+
+    ``models:`` with nothing under it parses as null. Most gateways tolerate it,
+    but the operator should see that a provider block is now dead weight rather
+    than discover it when a route stops resolving.
+    """
+    warnings: list[str] = []
+    before_counts = _models_per_provider(before)
+    after_counts = _models_per_provider(after)
+    for provider, count in before_counts.items():
+        if count and not after_counts.get(provider):
+            warnings.append(f"provider {provider!r} has no models left")
+    return warnings
+
+
+def _models_per_provider(text: str) -> dict[str, int]:
+    """Provider name -> declared model count, from a structural parse."""
+    try:
+        import yaml
+
+        doc = yaml.safe_load(text)
+    except Exception:
+        return {}
+    counts: dict[str, int] = {}
+    if not isinstance(doc, dict):
+        return counts
+    for section in PROVIDER_SECTIONS:
+        blocks = doc.get(section)
+        if not isinstance(blocks, list):
+            continue
+        for block in blocks:
+            if not isinstance(block, dict):
+                continue
+            provider = str(block.get("name") or section)
+            entries = block.get("models")
+            counts[provider] = counts.get(provider, 0) + (
+                len(entries) if isinstance(entries, list) else 0
+            )
+    return counts
+
+
+# --- factory ----------------------------------------------------------------
+
+
+def build_adapter(
+    kind: str,
+    gateway: Gateway,
+    config_path: str | Path | None = None,
+) -> Adapter:
+    """``auto`` picks the config adapter when a usable config is in reach."""
+    kind = (kind or "auto").lower()
+    if kind in ("generic", "live"):
+        return LiveAdapter(gateway)
+    if kind not in ("auto", "cliproxy"):
+        raise ConfigError(f"unknown adapter {kind!r} (auto|cliproxy|generic)")
+
+    if config_path is None:
+        if kind == "cliproxy":
+            raise ConfigError("--adapter cliproxy needs --config PATH")
+        return LiveAdapter(gateway)
+
+    path = Path(config_path).expanduser()
+    if not path.is_file():
+        if kind == "cliproxy":
+            raise ConfigError(f"config not found: {path}")
+        return LiveAdapter(gateway)
+
+    adapter = CliproxyAdapter(gateway, path)
+    try:  # an unreadable config downgrades to read-only, it does not abort
+        adapter.attribution_map()
+    except ConfigError:
+        if kind == "cliproxy":
+            raise
+        return LiveAdapter(gateway)
+    return adapter