sdm-learn 0.1.0.dev0__py3-none-any.whl

sdm/gateway.py

@@ -0,0 +1,169 @@
+"""The gateway to an LLM, and the environment checks that come with it.
+
+Deliberately not called ``model``: in this project the skill document *is* the
+model. This file is only the wire to a language model that reads it.
+
+The bottom of the stack, so nothing here knows about skills, datasets, or the
+learning loop. It turns a list of content blocks into one completion, retries
+what is worth retrying, and refuses to paper over a truncated response. It is
+also the one seam the offline checks replace, which is what keeps them honest:
+there is no second path to the network.
+
+``SetupNeeded`` lives here because this is the module every other one can
+reach. It marks the failures a traceback cannot fix: no API key, no SVG
+rasterizer, no numpy.
+"""
+
+from __future__ import annotations
+
+import base64
+import contextlib
+import io
+import json
+import os
+import time
+from pathlib import Path
+
+import httpx
+
+
+class SetupNeeded(Exception):
+    """Something about the environment needs a human, not a traceback."""
+
+
+GATEWAY_URL = os.environ.get(
+    "SDM_GATEWAY_URL", "https://inference-api.nvidia.com/v1/chat/completions"
+)
+MODEL = os.environ.get("SDM_MODEL", "azure/anthropic/claude-opus-5")
+EFFORT = os.environ.get("SDM_EFFORT", "low")
+KEY_FILE = Path.home() / ".sdm_key"
+
+
+def _api_key() -> str:
+    key = os.environ.get("SDM_API_KEY")
+    if key:
+        return key.strip()
+    path = Path(os.environ.get("SDM_API_KEY_FILE", KEY_FILE))
+    if path.is_file():
+        return path.read_text().strip()
+    raise SetupNeeded(
+        "no model to talk to yet. Set SDM_API_KEY, or write the key to "
+        f"{KEY_FILE}."
+    )
+
+
+def client(read_timeout: float = 600.0, connect_timeout: float = 30.0):
+    """An authenticated session for the gateway."""
+    return httpx.Client(
+        headers={"Authorization": f"Bearer {_api_key()}",
+                 "Content-Type": "application/json"},
+        timeout=httpx.Timeout(read_timeout, connect=connect_timeout),
+    )
+
+
+@contextlib.contextmanager
+def _borrowed(session):
+    """Yield a session, closing it only if we were the one who opened it."""
+    owned = session is None
+    session = session or client()
+    try:
+        yield session
+    finally:
+        if owned:
+            session.close()
+
+
+def _stream(session, payload: dict, deadline: float = 1800.0):
+    chunks, finish = [], None
+    started = time.monotonic()
+    with session.stream("POST", GATEWAY_URL,
+                        json={**payload, "stream": True}) as response:
+        if response.status_code != 200:
+            response.read()
+            raise httpx.HTTPStatusError(
+                f"{response.status_code}: {response.text[:200]}",
+                request=response.request, response=response,
+            )
+        for line in response.iter_lines():
+            if time.monotonic() - started > deadline:
+                raise httpx.TransportError(f"deadline {deadline}s hit")
+            if not line.startswith("data: "):
+                continue
+            body = line[6:]
+            if body.strip() == "[DONE]":
+                break
+            event = json.loads(body)
+            if not event.get("choices"):
+                continue
+            choice = event["choices"][0]
+            if (choice.get("delta") or {}).get("content"):
+                chunks.append(choice["delta"]["content"])
+            if choice.get("finish_reason"):
+                finish = choice["finish_reason"]
+    return "".join(chunks), finish
+
+
+def complete(session, messages: list, max_tokens: int,
+             effort: str | None = None) -> str:
+    """One chat completion, retrying through the failures a long run hits."""
+    payload = {
+        "model": MODEL,
+        "max_tokens": max_tokens,
+        "output_config": {"effort": effort or EFFORT},
+        "messages": messages,
+    }
+    for attempt in range(30):
+        try:
+            content, finish = _stream(session, payload)
+            if finish == "content_filter":
+                raise RuntimeError("model refused to answer")
+            if not content or finish == "length":
+                payload["max_tokens"] = min(payload["max_tokens"] * 2, 32000)
+                print(f"  [truncated ({finish}), retrying with "
+                      f"max_tokens={payload['max_tokens']}]", flush=True)
+                continue
+            if finish != "stop":
+                print(f"  [incomplete stream (finish={finish!r}), retrying]",
+                      flush=True)
+                continue
+            return content
+        except (httpx.HTTPStatusError, httpx.TransportError,
+                json.JSONDecodeError) as error:
+            wait = min(15 * (attempt + 1), 600)
+            print(f"  [retry in {wait}s: {type(error).__name__}: "
+                  f"{str(error)[:120]}]", flush=True)
+            time.sleep(wait)
+    raise RuntimeError("too many retries")
+
+
+def _png_block(png: bytes) -> dict:
+    """One content block holding already-encoded PNG bytes."""
+    encoded = base64.b64encode(png).decode("ascii")
+    return {"type": "image_url",
+            "image_url": {"url": f"data:image/png;base64,{encoded}"}}
+
+
+def image_block(image, size: int = 224) -> dict:
+    """Encode an RGB array, PIL image, or image path as one content block."""
+    from PIL import Image
+
+    if isinstance(image, Image.Image):
+        pil = image.convert("RGB")
+    elif isinstance(image, (str, Path)):
+        with Image.open(image) as source:
+            pil = source.convert("RGB")
+    elif hasattr(image, "shape"):
+        import numpy as np
+
+        if image.ndim != 3 or image.shape[2] != 3:
+            raise ValueError("expected one HxWx3 RGB image")
+        pil = Image.fromarray(np.asarray(image).astype("uint8"), mode="RGB")
+    else:
+        raise TypeError("expected an RGB array, PIL image, or image path")
+    pil = pil.resize((size, size), Image.Resampling.NEAREST)
+    buffer = io.BytesIO()
+    pil.save(buffer, format="PNG", optimize=False)
+    return _png_block(buffer.getvalue())
+
+
+# --------------------------------------------------------------------------