sway Public

Watch 0 Fork 0 Star 0
Python · 5664 bytes Raw Blame History
1 """Python client for the ``sway serve`` daemon.
2
3 Surface designed for two use cases:
4
5 1. **Notebook / REPL** — ``ServeClient(url).run(spec)`` returns the
6 parsed report dict; users get the warm-backend speedup with a
7 one-liner.
8 2. **CLI delegation** — ``sway watch --serve-url`` (S34) and any
9 future "client mode" CLI can build a client and forward requests.
10
11 Lives in the ``[serve]`` extra (httpx is the dep). Returns the same
12 JSON shape as ``sway run --json-out`` so notebook code can swap
13 ``sway.run(...)`` for ``ServeClient(url).run(...)`` with no other
14 changes.
15 """
16
17 from __future__ import annotations
18
19 from typing import Any
20
21 from dlm_sway.core.errors import SwayError
22 from dlm_sway.suite.spec import SwaySpec
23
24
25 class ServeClientError(SwayError):
26 """Raised when the daemon returned an error or was unreachable.
27
28 Distinct from generic httpx exceptions so callers can catch a
29 sway-shaped exception family. Wraps the underlying httpx error
30 via ``__cause__``.
31 """
32
33
34 class ServeClient:
35 """HTTP client for ``sway serve``.
36
37 Parameters
38 ----------
39 url:
40 Base URL of the daemon, e.g. ``http://localhost:8787``.
41 timeout:
42 Per-request timeout in seconds. Default 120s — model loads
43 on a cold cache can take 30-60s; the floor is generous.
44 api_key:
45 Bearer token, if the daemon was launched with ``--api-key``.
46
47 The client is **stateless** — no persistent connection pool is
48 held across calls. For high-throughput callers (notebooks doing
49 many quick scores) the per-call connection cost is dwarfed by
50 the inference cost; for low-throughput callers, the simpler
51 contract is worth more than the saved milliseconds.
52 """
53
54 def __init__(self, url: str, *, timeout: float = 120.0, api_key: str | None = None) -> None:
55 self._url = url.rstrip("/")
56 self._timeout = float(timeout)
57 self._headers: dict[str, str] = {"Content-Type": "application/json"}
58 if api_key is not None:
59 self._headers["Authorization"] = f"Bearer {api_key}"
60
61 @property
62 def url(self) -> str:
63 return self._url
64
65 def health(self) -> dict[str, Any]:
66 """Hit ``GET /health``. Returns the daemon's health payload."""
67 return self._get("/health")
68
69 def stats(self) -> dict[str, Any]:
70 """Hit ``GET /stats``. Returns request count / mean latency."""
71 return self._get("/stats")
72
73 def run(self, spec: SwaySpec, *, spec_path: str = "<client>") -> dict[str, Any]:
74 """Hit ``POST /run`` with the spec.
75
76 Returns the parsed JSON response — same shape as the on-disk
77 report ``sway run --json-out`` would write, plus a
78 ``request_seconds`` field with the daemon's measured execution
79 time (excluding HTTP round-trip).
80 """
81 body = {"spec": spec.model_dump(mode="json"), "spec_path": spec_path}
82 return self._post("/run", body)
83
84 def score(
85 self,
86 spec: SwaySpec,
87 *,
88 probe_names: list[str] | None = None,
89 ) -> dict[str, Any]:
90 """Hit ``POST /score``. Same as ``run`` but returns only the
91 probe-level results (no folded SwayScore) and supports a
92 ``probe_names`` filter for partial runs."""
93 body: dict[str, Any] = {"spec": spec.model_dump(mode="json")}
94 if probe_names is not None:
95 body["probe_names"] = list(probe_names)
96 return self._post("/score", body)
97
98 # -- internals -------------------------------------------------------
99
100 def _get(self, path: str) -> dict[str, Any]:
101 try:
102 import httpx
103 except ImportError as exc:
104 raise ServeClientError(
105 "ServeClient requires the [serve] extra: pip install 'dlm-sway[serve]'"
106 ) from exc
107 try:
108 with httpx.Client(timeout=self._timeout) as client:
109 resp = client.get(self._url + path, headers=self._headers)
110 except httpx.HTTPError as exc:
111 raise ServeClientError(f"GET {path} failed: {exc}") from exc
112 return self._parse_response(resp, path=path)
113
114 def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
115 try:
116 import httpx
117 except ImportError as exc:
118 raise ServeClientError(
119 "ServeClient requires the [serve] extra: pip install 'dlm-sway[serve]'"
120 ) from exc
121 try:
122 with httpx.Client(timeout=self._timeout) as client:
123 resp = client.post(self._url + path, headers=self._headers, json=body)
124 except httpx.HTTPError as exc:
125 raise ServeClientError(f"POST {path} failed: {exc}") from exc
126 return self._parse_response(resp, path=path)
127
128 @staticmethod
129 def _parse_response(resp: Any, *, path: str) -> dict[str, Any]:
130 # Trust httpx's status code; reach into ``.detail`` first since
131 # FastAPI's HTTPException uses that key by default.
132 if resp.status_code >= 400:
133 try:
134 payload = resp.json()
135 except Exception: # noqa: BLE001
136 payload = {"detail": resp.text}
137 detail = payload.get("detail") if isinstance(payload, dict) else str(payload)
138 raise ServeClientError(f"{path} returned {resp.status_code}: {detail}")
139 try:
140 data = resp.json()
141 except Exception as exc: # noqa: BLE001
142 raise ServeClientError(f"{path} returned non-JSON body: {resp.text[:200]}") from exc
143 if not isinstance(data, dict):
144 raise ServeClientError(f"{path} returned non-object JSON: {type(data).__name__}")
145 return data