gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/benchmark_hub/dispatch_shard.py

@@ -0,0 +1,34 @@
+"""Explicit hosted-dispatch identity. No identity is inferred from process state."""
+from typing import TypedDict
+from uuid import UUID
+
+
+class DispatchShard(TypedDict):
+    protocol: int
+    dispatchId: str
+    slotId: str
+    index: int
+    total: int
+    jobId: str
+
+
+def validate_dispatch_shard(value: object) -> DispatchShard:
+    keys = {"protocol", "dispatchId", "slotId", "index", "total", "jobId"}
+    if not isinstance(value, dict) or set(value) != keys:
+        raise ValueError("dispatchShard must contain exactly the protocol and dispatch slot identity fields")
+    if type(value["protocol"]) is not int or value["protocol"] != 1:
+        raise ValueError("dispatchShard protocol must be 1")
+    if (type(value["index"]) is not int or type(value["total"]) is not int
+            or not 1 <= value["total"] <= 16 or not 0 <= value["index"] < value["total"]):
+        raise ValueError("dispatchShard must have 0 <= index < total <= 16")
+    for key in ("dispatchId", "slotId", "jobId"):
+        item = value[key]
+        if not isinstance(item, str) or not item or item != item.strip():
+            raise ValueError(f"dispatchShard {key} must be a nonblank identifier without boundary whitespace")
+    try:
+        parsed = UUID(value["slotId"])
+        if str(parsed) != value["slotId"].lower():
+            raise ValueError("noncanonical UUID")
+    except (ValueError, AttributeError):
+        raise ValueError("dispatchShard slotId must be a hyphenated UUID") from None
+    return dict(value)

gatewaysdk/benchmark_hub/eval_config.py

@@ -0,0 +1,20 @@
+"""Config for a local ``bench eval`` run (the in-scope subset of verifiers'
+EvalConfig). Drives the resumable results artifact + the Gateway harvest."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class EvalConfig:
+    env_id: str  # container slug/name (the verifiers env id)
+    model: str  # model id for the OpenAI-compatible client
+    num_examples: int = -1
+    rollouts_per_example: int = 1
+    shuffle: bool = False
+    shuffle_seed: Optional[int] = None
+    output_dir: str = "./outputs"
+    save_results: bool = False
+    resume_path: Optional[str] = None

gatewaysdk/benchmark_hub/evals.py

@@ -0,0 +1,351 @@
+"""Gateway eval-harvest client.
+
+Mirrors the prime-evals ``EvalsClient`` shape (create_evaluation ->
+push_samples[batched] -> finalize_evaluation) but harvests INTO Gateway — pointed
+at our version-pinned eval routes, never at Prime's platform. An evaluation is
+anchored to an immutable container version, so scores stay reproducible across
+re-pushes; each sample carries the multi-component reward breakdown.
+
+Stdlib only. Writes require a secret key (Basic auth).
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import logging
+import os
+import sys
+import urllib.error
+import urllib.request
+import uuid
+from collections import OrderedDict
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+from .dispatch_shard import DispatchShard, validate_dispatch_shard
+
+_BASE = "/api/public/benchmark-containers/evaluations"
+
+_log = logging.getLogger("gatewaysdk.benchmark_hub.evals")
+
+# Each rollout is emitted as a Gateway SESSION (a session-tagged trace) so it
+# opens in the platform's session view. Set GATEWAY_EMIT_ROLLOUT_SESSIONS=0 to
+# disable. The resulting {sessionId, traceId} are stamped into sample.info so
+# the RolloutViewer links straight to the session/trace.
+
+#: The single environment every rollout trace is stamped with. Rollouts land
+#: in the SAME project as the agent's live traffic (agent projects own their
+#: environment containers), and this well-known name is what lets the
+#: platform hide them by default while keeping them selectable. Keep in sync
+#: with GATEWAY_ROLLOUTS_ENVIRONMENT in gateway/packages/shared/src/constants.ts
+#: and ROLLOUTS_ENVIRONMENT in deploy/benchmark-runner/engines/.
+ROLLOUTS_ENVIRONMENT = "rollouts"
+
+
+def _emit_rollout_sessions_enabled() -> bool:
+    return os.environ.get("GATEWAY_EMIT_ROLLOUT_SESSIONS", "1") not in (
+        "0",
+        "false",
+        "no",
+    )
+
+
+class EvalHarvestError(RuntimeError):
+    """Raised when a harvest call fails."""
+
+
+def _chunk(items: List[Any], size: int) -> List[List[Any]]:
+    if size <= 0:
+        raise ValueError("size must be > 0")
+    return [items[i : i + size] for i in range(0, len(items), size)]
+
+
+def normalize_sample(sample: Dict[str, Any]) -> Dict[str, Any]:
+    """Map a sample (snake_case prime-evals style OR camelCase) to our API shape.
+    Named sub-rewards belong in `rewards`; the caller (e.g. the verifiers worker)
+    builds that dict from the rubric's components."""
+
+    def pick(*keys: str) -> Any:
+        for k in keys:
+            if k in sample and sample[k] is not None:
+                return sample[k]
+        return None
+
+    mapping = [
+        ("exampleId", ("exampleId", "example_id")),
+        ("task", ("task",)),
+        ("rolloutIndex", ("rolloutIndex", "rollout_index")),
+        ("prompt", ("prompt",)),
+        ("completion", ("completion",)),
+        ("answer", ("answer",)),
+        ("reward", ("reward",)),
+        ("score", ("score",)),
+        ("correct", ("correct",)),
+        ("rewards", ("rewards",)),
+        ("info", ("info",)),
+    ]
+    out: Dict[str, Any] = {}
+    for our_key, in_keys in mapping:
+        v = pick(*in_keys)
+        if v is not None:
+            out[our_key] = v
+    return out
+
+
+class GatewayEvalsClient:
+    """Create/push/finalize evaluations against a Gateway host."""
+
+    def __init__(self, host: str, public_key: str, secret_key: str) -> None:
+        if not host:
+            raise ValueError("host is required (e.g. https://withgateway.ai)")
+        if not public_key or not secret_key:
+            raise ValueError("both public_key and secret_key are required")
+        self.host = host.rstrip("/")
+        self._auth = base64.b64encode(
+            f"{public_key}:{secret_key}".encode("utf-8")
+        ).decode("ascii")
+
+    @classmethod
+    def from_env(cls, host: Optional[str] = None) -> "GatewayEvalsClient":
+        return cls(
+            host=host or os.environ.get("GATEWAY_HOST", ""),
+            public_key=os.environ.get("GATEWAY_PUBLIC_KEY", ""),
+            secret_key=os.environ.get("GATEWAY_SECRET_KEY", ""),
+        )
+
+    def _post(self, path: str, body: Dict[str, Any]) -> Dict[str, Any]:
+        req = urllib.request.Request(
+            f"{self.host}{path}",
+            data=json.dumps(body).encode("utf-8"),
+            method="POST",
+            headers={
+                "Authorization": f"Basic {self._auth}",
+                "Content-Type": "application/json",
+            },
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=180) as resp:
+                return json.loads(resp.read().decode("utf-8"))
+        except urllib.error.HTTPError as e:
+            detail = e.read().decode("utf-8", "replace")[:500]
+            raise EvalHarvestError(f"POST {path} -> {e.code}: {detail}") from None
+
+    def create_evaluation(
+        self,
+        version_id: str,
+        *,
+        model: Optional[str] = None,
+        run_id: Optional[str] = None,
+        dispatch_shard: Optional[DispatchShard] = None,
+        # The world condition this evaluation executed under (identity only —
+        # the bindings that changed behavior travel in `metadata`). Recorded so
+        # results can be compared BY condition instead of by whichever label
+        # someone remembers attaching to the run.
+        condition_id: Optional[str] = None,
+        # The agent build this evaluation tested. Identity only, like
+        # condition_id: nothing in the run branches on it, but without it a run
+        # records only which MODEL answered, and a release citing that run
+        # cannot show the thing shipped is the thing tested.
+        agent_build_id: Optional[str] = None,
+        # WHAT THE RUN ACTUALLY DID with the agent variation it was sent —
+        # `AppliedAgentVariation` from
+        # packages/shared/src/agent-variation/schemas.ts.
+        #
+        # Unlike agent_build_id this is NOT identity: it is the only thing that
+        # ever writes `benchmark_container_evaluations
+        # .agent_variation_revision_id`. The platform can dispatch a variation
+        # and cannot observe that the engine honoured it — a stale image, a
+        # dropped env var, or an engine that never learnt to read the variable
+        # each produce a run that used the world bundle's own prompt. So the
+        # pin comes from the runner's echo and from nothing else. Omitted = not
+        # recorded, which is the honest state for every runner that predates
+        # this field.
+        applied_agent_variation: Optional[Dict[str, Any]] = None,
+        num_examples: Optional[int] = None,
+        rollouts_per_example: Optional[int] = None,
+        environment_ids: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        body: Dict[str, Any] = {"versionId": version_id}
+        if dispatch_shard is not None:
+            if not isinstance(run_id, str) or not run_id or run_id != run_id.strip():
+                raise ValueError("dispatchShard requires a nonblank run_id without boundary whitespace")
+            body["dispatchShard"] = validate_dispatch_shard(dispatch_shard)
+        if model is not None:
+            body["model"] = model
+        if run_id is not None:
+            body["runId"] = run_id
+        if condition_id is not None:
+            body["conditionId"] = condition_id
+        if agent_build_id is not None:
+            body["agentBuildId"] = agent_build_id
+        if applied_agent_variation is not None:
+            body["appliedAgentVariation"] = applied_agent_variation
+        if num_examples is not None:
+            body["numExamples"] = num_examples
+        if rollouts_per_example is not None:
+            body["rolloutsPerExample"] = rollouts_per_example
+        if environment_ids is not None:
+            body["environmentIds"] = environment_ids
+        if metadata is not None:
+            body["metadata"] = metadata
+        evaluation_id = self._post(_BASE, body)["evaluationId"]
+        _OPEN_EVALUATIONS[evaluation_id] = self
+        return evaluation_id
+
+    def _stringify(self, value: Any) -> str:
+        if value is None:
+            return ""
+        if isinstance(value, str):
+            return value
+        try:
+            return json.dumps(value, ensure_ascii=False)
+        except (TypeError, ValueError):
+            return str(value)
+
+    def _emit_rollout_session(
+        self,
+        session_id: str,
+        *,
+        name: str,
+        prompt: Any,
+        completion: Any,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, str]:
+        """Emit one rollout as a session-tagged trace via the ingestion API.
+
+        Mirrors the runner's emit_rollout_session (and voice sessions): a single
+        `trace-create` event whose `sessionId` ties the trace into a session,
+        with input=prompt / output=completion. Fail-soft — returns {} on any
+        error and never raises (harvest must not break on tracing).
+        """
+        trace_id = uuid.uuid4().hex
+        now = datetime.now(timezone.utc).isoformat()
+        try:
+            self._post(
+                "/api/public/ingestion",
+                {
+                    "batch": [
+                        {
+                            "id": uuid.uuid4().hex,
+                            "type": "trace-create",
+                            "timestamp": now,
+                            "body": {
+                                "id": trace_id,
+                                "sessionId": session_id,
+                                "name": name or "rollout",
+                                "environment": ROLLOUTS_ENVIRONMENT,
+                                "input": self._stringify(prompt),
+                                "output": self._stringify(completion),
+                                "metadata": metadata or {},
+                                "timestamp": now,
+                            },
+                        }
+                    ]
+                },
+            )
+            return {"sessionId": session_id, "traceId": trace_id}
+        except Exception as exc:  # noqa: BLE001 — tracing is best-effort
+            _log.warning("rollout session emit failed for %s: %s", session_id, exc)
+            return {}
+
+    def _tag_sample_session(self, evaluation_id: str, sample: Dict[str, Any]) -> None:
+        """Emit a session for `sample` and stamp {sessionId, traceId} into info."""
+        info = sample.get("info")
+        if not isinstance(info, dict):
+            info = {}
+            sample["info"] = info
+        # Respect an id the caller already set (e.g. a live-traced rollout).
+        if info.get("sessionId") or info.get("session_id"):
+            return
+        example_id = sample.get("exampleId") or sample.get("example_id") or "ex"
+        rollout_index = sample.get("rolloutIndex") or sample.get("rollout_index") or 0
+        session_id = f"{evaluation_id}:{example_id}:{rollout_index}"
+        ids = self._emit_rollout_session(
+            session_id,
+            name=str(sample.get("task") or example_id),
+            prompt=sample.get("prompt"),
+            completion=sample.get("completion") or sample.get("answer"),
+            metadata={
+                "evaluationId": evaluation_id,
+                "task": sample.get("task"),
+                "reward": sample.get("reward"),
+            },
+        )
+        info.update(ids)
+
+    def push_samples(
+        self,
+        evaluation_id: str,
+        samples: List[Dict[str, Any]],
+        *,
+        batch_size: int = 100,
+    ) -> int:
+        emit_sessions = _emit_rollout_sessions_enabled()
+        total = 0
+        for batch in _chunk(samples, batch_size):
+            if emit_sessions:
+                # Each rollout becomes a session-tagged trace; sample.info then
+                # carries {sessionId, traceId} for the RolloutViewer links.
+                for s in batch:
+                    self._tag_sample_session(evaluation_id, s)
+            resp = self._post(
+                f"{_BASE}/{evaluation_id}/samples",
+                {"samples": [normalize_sample(s) for s in batch]},
+            )
+            total += int(resp.get("inserted", len(batch)))
+        return total
+
+    def finalize_evaluation(
+        self,
+        evaluation_id: str,
+        *,
+        status: str = "COMPLETED",
+        metrics: Optional[Dict[str, Any]] = None,
+        cost: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        body: Dict[str, Any] = {"status": status}
+        if metrics is not None:
+            body["metrics"] = metrics
+        if cost is not None:
+            body["cost"] = cost
+        result = self._post(f"{_BASE}/{evaluation_id}/finalize", body)["status"]
+        _OPEN_EVALUATIONS.pop(evaluation_id, None)
+        return result
+
+
+# Evaluations created but not yet finalized. A runner that dies mid-rollout
+# would otherwise leave its evaluation RUNNING forever, so the process-level
+# error handler finalizes whatever is still open as FAILED.
+_OPEN_EVALUATIONS: "OrderedDict[str, GatewayEvalsClient]" = OrderedDict()
+
+
+def open_evaluations() -> List[str]:
+    """Ids of evaluations created in this process that were never finalized."""
+    return list(_OPEN_EVALUATIONS)
+
+
+def fail_open_evaluations(reason: str) -> List[str]:
+    """Mark every still-open evaluation FAILED, recording `reason`.
+
+    Best-effort by design: it runs on a path where something has ALREADY gone
+    wrong, so a failure to report the failure must not mask the original error.
+    Returns the ids it managed to finalize.
+    """
+    failed: List[str] = []
+    for evaluation_id, client in list(_OPEN_EVALUATIONS.items()):
+        # Deregister FIRST: whatever happens next, this evaluation has had its
+        # one attempt. Leaving it registered risks reporting it twice or looping.
+        _OPEN_EVALUATIONS.pop(evaluation_id, None)
+        try:
+            client.finalize_evaluation(
+                evaluation_id, status="FAILED", metrics={"error": reason[:2000]}
+            )
+            failed.append(evaluation_id)
+        except Exception as exc:  # noqa: BLE001 - never mask the original error
+            print(
+                f"[evals] could not mark evaluation {evaluation_id} FAILED: {exc}",
+                file=sys.stderr,
+            )
+    return failed