from __future__ import annotations import json import re from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Any import httpx from .models import MediaArtifact, MediaJobRequest from .settings import MediaSettings class DownstreamError(RuntimeError): pass _SAFE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,239}$") @dataclass(slots=True) class SubmitResult: downstream_id: str | None artifacts: list[MediaArtifact] | None = None summary: str | None = None class LocalMediaClient: def __init__( self, settings: MediaSettings, *, client: httpx.AsyncClient | None = None, ) -> None: self.settings = settings self._owns_client = client is None self._request_timeout = httpx.Timeout( settings.request_timeout_s, connect=settings.connect_timeout_s, read=settings.request_timeout_s, write=settings.request_timeout_s, pool=settings.connect_timeout_s, ) self._stream_timeout = httpx.Timeout( settings.stream_timeout_s, connect=settings.connect_timeout_s, read=settings.stream_timeout_s, write=settings.request_timeout_s, pool=settings.connect_timeout_s, ) self.client = client or httpx.AsyncClient( follow_redirects=False, timeout=self._request_timeout, ) async def close(self) -> None: if self._owns_client: await self.client.aclose() async def _json_response( self, method: str, url: str, *, headers: dict[str, str] | None = None, json_body: dict[str, Any] | None = None, ) -> tuple[int, dict[str, Any]]: try: async with self.client.stream( method, url, headers=headers, json=json_body, timeout=self._request_timeout, ) as response: body = bytearray() async for chunk in response.aiter_bytes(): if len(body) + len(chunk) > self.settings.max_response_bytes: raise DownstreamError("downstream response exceeded its configured limit") body.extend(chunk) if not 200 <= response.status_code < 300: raise DownstreamError(f"downstream returned HTTP {response.status_code}") if not body: return response.status_code, {} try: payload = json.loads(body) except (json.JSONDecodeError, UnicodeDecodeError): raise DownstreamError("downstream returned invalid JSON") from None if not isinstance(payload, dict): raise DownstreamError("downstream JSON must be an object") return response.status_code, payload except DownstreamError: raise except httpx.HTTPError as exc: raise DownstreamError("local downstream endpoint was unavailable") from exc async def verify_lease(self, lease_id: str, kind: str) -> None: _, payload = await self._json_response( "GET", f"{self.settings.manager_url}/v1/leases/{lease_id}" ) if payload.get("active") is not True or payload.get("capability") != kind: raise DownstreamError("model lease is inactive or has the wrong capability") def artifacts(self, payload: Any) -> list[MediaArtifact]: if not isinstance(payload, list): raise DownstreamError("completed media response has no artifact list") if len(payload) > self.settings.max_artifacts: raise DownstreamError("completed media response has too many artifacts") try: return [MediaArtifact.model_validate(item) for item in payload] except Exception: raise DownstreamError("completed media response has an invalid artifact reference") from None async def submit( self, request: MediaJobRequest, *, endpoint: str, lease_id: str, idempotency_key: str, ) -> SubmitResult: _, payload = await self._json_response( "POST", f"{endpoint}/v1/jobs", headers={ "Idempotency-Key": idempotency_key, "X-Model-Lease": lease_id, }, json_body=request.model_dump(mode="json"), ) status = str(payload.get("status") or "").casefold() if "artifacts" in payload or status == "completed": artifacts = self.artifacts(payload.get("artifacts")) summary = payload.get("summary") immediate_id = payload.get("job_id") or payload.get("id") if immediate_id is not None and ( not isinstance(immediate_id, str) or not _SAFE_ID.fullmatch(immediate_id) ): raise DownstreamError("downstream response has an invalid job_id") return SubmitResult( downstream_id=immediate_id, artifacts=artifacts, summary=summary[:4000] if isinstance(summary, str) else None, ) downstream_id = payload.get("job_id") or payload.get("id") if not isinstance(downstream_id, str) or not _SAFE_ID.fullmatch(downstream_id): raise DownstreamError("downstream create-job response has no bounded job_id") return SubmitResult(downstream_id=downstream_id) async def events( self, *, endpoint: str, downstream_id: str ) -> AsyncIterator[tuple[str, dict[str, Any]]]: url = f"{endpoint}/v1/jobs/{downstream_id}/events" try: async with self.client.stream( "GET", url, headers={"Accept": "text/event-stream"}, timeout=self._stream_timeout, ) as response: if not 200 <= response.status_code < 300: raise DownstreamError(f"downstream event stream returned HTTP {response.status_code}") event_name = "message" data_lines: list[str] = [] event_bytes = 0 emitted = 0 async for raw_line in response.aiter_lines(): line = raw_line.rstrip("\r\n") event_bytes += len(line.encode("utf-8")) if event_bytes > self.settings.max_sse_event_bytes: raise DownstreamError("downstream SSE event exceeded its configured limit") if not line: if data_lines: try: payload = json.loads("\n".join(data_lines)) except json.JSONDecodeError: raise DownstreamError("downstream SSE contained invalid JSON") from None if not isinstance(payload, dict): raise DownstreamError("downstream SSE JSON must be an object") emitted += 1 if emitted > self.settings.max_downstream_events: raise DownstreamError("downstream emitted too many events") yield str(payload.get("type") or event_name), payload event_name, data_lines, event_bytes = "message", [], 0 continue if line.startswith(":"): continue field, separator, value = line.partition(":") if separator and value.startswith(" "): value = value[1:] if field == "event": event_name = value or "message" elif field == "data": data_lines.append(value) if data_lines: try: payload = json.loads("\n".join(data_lines)) except json.JSONDecodeError: raise DownstreamError("downstream SSE contained invalid JSON") from None if not isinstance(payload, dict): raise DownstreamError("downstream SSE JSON must be an object") yield str(payload.get("type") or event_name), payload except DownstreamError: raise except httpx.HTTPError as exc: raise DownstreamError("local downstream event stream was unavailable") from exc async def cancel(self, *, endpoint: str, downstream_id: str) -> None: try: async with self.client.stream( "POST", f"{endpoint}/v1/jobs/{downstream_id}/cancel", timeout=self._request_timeout, ) as response: if response.status_code not in {200, 202, 204, 404, 409}: raise DownstreamError(f"downstream cancel returned HTTP {response.status_code}") size = 0 async for chunk in response.aiter_bytes(): size += len(chunk) if size > self.settings.max_response_bytes: raise DownstreamError("downstream cancel response exceeded its limit") except DownstreamError: raise except httpx.HTTPError as exc: raise DownstreamError("local downstream cancel endpoint was unavailable") from exc