Files
didactyl/tests/harness/didactyl_client.py

105 lines
4.2 KiB
Python

from __future__ import annotations
import json
import ssl
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from typing import Any, Optional
class DidactylTimeoutError(TimeoutError):
pass
class DidactylConnectionError(ConnectionError):
pass
@dataclass
class APIError(Exception):
status_code: int
body: str
def __str__(self) -> str:
return f"HTTP {self.status_code}: {self.body}"
class DidactylClient:
def __init__(self, base_url: str = "https://127.0.0.1:8485", timeout: float = 60.0, verify_tls: bool = False) -> None:
self.base_url = base_url.rstrip("/")
self.timeout = timeout
self.verify_tls = verify_tls
self.ssl_context = None
if self.base_url.startswith("https://") and not verify_tls:
self.ssl_context = ssl._create_unverified_context()
def _request(self, method: str, path: str, payload: Optional[dict[str, Any]] = None, raw_body: Optional[bytes] = None) -> dict[str, Any]:
url = f"{self.base_url}{path}"
body = raw_body
if payload is not None:
body = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(url=url, method=method.upper(), data=body)
req.add_header("Content-Type", "application/json")
try:
with urllib.request.urlopen(req, timeout=self.timeout, context=self.ssl_context) as resp:
status = resp.getcode()
text = resp.read().decode("utf-8", errors="replace")
data = json.loads(text) if text.strip() else {}
if status < 200 or status >= 300:
raise APIError(status, text)
return data
except urllib.error.HTTPError as e:
text = e.read().decode("utf-8", errors="replace")
raise APIError(e.code, text) from e
except urllib.error.URLError as e:
reason = str(getattr(e, "reason", e))
if "timed out" in reason.lower():
raise DidactylTimeoutError(reason) from e
raise DidactylConnectionError(reason) from e
except TimeoutError as e:
raise DidactylTimeoutError(str(e)) from e
def status(self) -> dict[str, Any]:
return self._request("GET", "/api/status")
def context_current(self) -> dict[str, Any]:
return self._request("GET", "/api/context/current")
def context_parts(self) -> dict[str, Any]:
return self._request("GET", "/api/context/parts")
def prompt(self, message: str, max_turns: int = 4, model: Optional[str] = None) -> dict[str, Any]:
payload: dict[str, Any] = {"message": message, "max_turns": max_turns}
if model:
payload["model"] = model
return self._request("POST", "/api/prompt/agent", payload=payload)
def prompt_raw(self, messages: list[dict[str, str]], max_turns: int = 4, model: Optional[str] = None) -> dict[str, Any]:
payload: dict[str, Any] = {"messages": messages, "max_turns": max_turns}
if model:
payload["model"] = model
return self._request("POST", "/api/prompt/run", payload=payload)
def prompt_simple(self, system: str, user: str, model: Optional[str] = None) -> dict[str, Any]:
payload: dict[str, Any] = {"system": system, "user": user}
if model:
payload["model"] = model
return self._request("POST", "/api/prompt/run-simple", payload=payload)
def fire_webhook(self, d_tag: str, payload: Optional[dict[str, Any]] = None) -> dict[str, Any]:
encoded = urllib.parse.quote(d_tag, safe="")
return self._request("POST", f"/api/trigger/{encoded}", payload=payload or {})
def raw_post(self, path: str, body: bytes) -> tuple[int, str]:
url = f"{self.base_url}{path}"
req = urllib.request.Request(url=url, method="POST", data=body)
req.add_header("Content-Type", "application/json")
try:
with urllib.request.urlopen(req, timeout=self.timeout, context=self.ssl_context) as resp:
return resp.getcode(), resp.read().decode("utf-8", errors="replace")
except urllib.error.HTTPError as e:
return e.code, e.read().decode("utf-8", errors="replace")