runtime/jevlike/client.py

"""Tiny client for `python -m jevlike.server` (stdlib only).

    from jevlike import Choice, Score, Noul
    from jevlike.client import SystemOneClient

    c = SystemOneClient("http://127.0.0.1:8077")
    res = c.predict("My invoice is wrong and I'm cancelling!", {
        "dept": Choice("Which team?", {"billing": "invoices, refunds", "technical": "bugs, outages"}),
        "urgency": Score("How urgent?", ["not urgent", "soon", "blocking"]),
        "churn": Noul("The user threatens to cancel"),
    })
    res["answers"]["dept"]["choice"]

Answers have exactly the shapes of `SystemOne.predict` (see SPEC.md). Questions may also be given
as plain wire dicts ({"type": ..., "instructions": ..., "criteria": ...}).
"""
from __future__ import annotations

import json
import urllib.error
import urllib.request
from typing import Any, Optional, Union

from jevlike.types import Choice, Noul, Score

QuestionLike = Union[Choice, Score, Noul, dict]


class SystemOneError(RuntimeError):
    def __init__(self, status: int, detail: Any):
        self.status, self.detail = status, detail
        super().__init__(f"HTTP {status}: {detail}")


def question_to_wire(q: QuestionLike) -> dict:
    if isinstance(q, dict):
        return q
    if isinstance(q, Choice):
        crit = dict(q.criteria) if isinstance(q.criteria, dict) else list(q.criteria)
        return {"type": "choice", "instructions": q.instructions, "criteria": crit}
    if isinstance(q, Score):
        return {"type": "score", "instructions": q.instructions, "criteria": list(q.criteria)}
    if isinstance(q, Noul):
        return {"type": "noul", "instructions": q.instructions}
    raise TypeError(f"not a question: {type(q).__name__}")


def questions_to_wire(questions: dict[str, QuestionLike]) -> dict[str, dict]:
    return {k: question_to_wire(q) for k, q in questions.items()}


class _Response:
    def __init__(self, status_code: int, body: bytes):
        self.status_code, self._body = status_code, body

    def json(self):
        return json.loads(self._body.decode("utf-8")) if self._body else None


class _UrllibHTTP:
    def __init__(self, timeout: float):
        self.timeout = timeout

    def _send(self, req: urllib.request.Request) -> _Response:
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as r:
                return _Response(r.status, r.read())
        except urllib.error.HTTPError as e:
            return _Response(e.code, e.read())

    def get(self, url: str) -> _Response:
        return self._send(urllib.request.Request(url, method="GET"))

    def post(self, url: str, json: Any) -> _Response:  # noqa: A002 - mirrors httpx
        data = _json_dumps(json).encode("utf-8")
        return self._send(urllib.request.Request(url, data=data, method="POST",
                                                 headers={"Content-Type": "application/json"}))


def _json_dumps(x) -> str:
    return json.dumps(x, ensure_ascii=False)


class SystemOneClient:
    """`http` may be any object with httpx-style `get(url)` / `post(url, json=...)` (e.g. an
    httpx.Client or FastAPI's TestClient); by default a urllib-based one is used."""

    def __init__(self, base_url: str = "http://127.0.0.1:8077", timeout: float = 60.0, http: Optional[Any] = None):
        self.base_url = base_url.rstrip("/")
        self.http = http if http is not None else _UrllibHTTP(timeout)

    def _check(self, r) -> Any:
        body = r.json()
        if r.status_code >= 400:
            raise SystemOneError(r.status_code, body.get("detail", body) if isinstance(body, dict) else body)
        return body

    def health(self) -> dict:
        return self._check(self.http.get(f"{self.base_url}/health"))

    def predict(self, state: Union[str, dict, list], questions: dict[str, QuestionLike]) -> dict:
        body = {"state": state, "questions": questions_to_wire(questions)}
        return self._check(self.http.post(f"{self.base_url}/v1/systemone", json=body))

    def predict_batch(self, items: list[tuple[Union[str, dict, list], dict[str, QuestionLike]]]) -> list[dict]:
        body = {"items": [{"state": s, "questions": questions_to_wire(q)} for s, q in items]}
        return self._check(self.http.post(f"{self.base_url}/v1/systemone/batch", json=body))["results"]