arex-python-sdk 0.2.0__py3-none-any.whl

arex/errors.py

@@ -0,0 +1,334 @@
+"""Error hierarchy raised by the AREX SDK."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from arex.models import ResearchRun
+
+__all__ = [
+    "ArexApiError",
+    "ArexError",
+    "ArexNetworkError",
+    "ArexTimeoutError",
+    "AuthenticationError",
+    "CreditsExhaustedError",
+    "InvalidRequestError",
+    "NotFoundError",
+    "PermissionDeniedError",
+    "RateLimitError",
+    "ResearchRunCancelledError",
+    "ResearchRunFailedError",
+    "ServiceUnavailableError",
+    "SourceNotAvailableError",
+    "api_error_from_response",
+]
+
+NetworkErrorCode = Literal["request_failed", "invalid_response"]
+TimeoutPhase = Literal["request", "poll"]
+
+
+class ArexError(Exception):
+    """Base class for every error raised by this SDK.
+
+    `run_id` names the research run being polled when the error surfaced, so
+    the caller can `wait(run_id)` again or `cancel(run_id)`. It is `None`
+    outside `wait`/`run`.
+    """
+
+    run_id: str | None = None
+
+
+class ArexApiError(ArexError):
+    """An error response returned by the AREX API."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: int,
+        code: str,
+        retryable: bool = False,
+        retry_after_ms: int | None = None,
+        request_id: str | None = None,
+        param: str | None = None,
+        details: Mapping[str, Any] | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.code = code
+        self.retryable = retryable
+        self.retry_after_ms = retry_after_ms
+        self.request_id = request_id
+        self.param = param
+        self.details: Mapping[str, Any] = dict(details or {})
+
+
+class AuthenticationError(ArexApiError):
+    """401 `unauthorized`."""
+
+
+class PermissionDeniedError(ArexApiError):
+    """403 permission or account access denied."""
+
+
+class InvalidRequestError(ArexApiError):
+    """Invalid request, blocked content, reused key, or client-side validation."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: int = 0,
+        code: str = "invalid_request",
+        retryable: bool = False,
+        retry_after_ms: int | None = None,
+        request_id: str | None = None,
+        param: str | None = None,
+        details: Mapping[str, Any] | None = None,
+    ) -> None:
+        super().__init__(
+            message,
+            status=status,
+            code=code,
+            retryable=retryable,
+            retry_after_ms=retry_after_ms,
+            request_id=request_id,
+            param=param,
+            details=details,
+        )
+
+
+class SourceNotAvailableError(ArexApiError):
+    """422 `source_not_available`, raised before the request is charged."""
+
+
+class NotFoundError(ArexApiError):
+    """404 `not_found`."""
+
+
+class RateLimitError(ArexApiError):
+    """Rate limits, outstanding runs, or an idempotent request still in progress."""
+
+
+class CreditsExhaustedError(ArexApiError):
+    """429 `daily_credits_exhausted` / `weekly_credits_exhausted`."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: int,
+        code: str,
+        retryable: bool = False,
+        retry_after_ms: int | None = None,
+        request_id: str | None = None,
+        param: str | None = None,
+        details: Mapping[str, Any] | None = None,
+        resource: str | None = None,
+        used: int | float | None = None,
+        limit: int | float | None = None,
+        requested: int | float | None = None,
+        reset_at: str | None = None,
+    ) -> None:
+        super().__init__(
+            message,
+            status=status,
+            code=code,
+            retryable=retryable,
+            retry_after_ms=retry_after_ms,
+            request_id=request_id,
+            param=param,
+            details=details,
+        )
+        self.resource = resource
+        self.used = used
+        self.limit = limit
+        self.requested = requested
+        self.reset_at = reset_at
+
+
+class ServiceUnavailableError(ArexApiError):
+    """503 `temporarily_unavailable` / 500 `internal_error`."""
+
+
+class ArexNetworkError(ArexError):
+    """The request never produced a usable response."""
+
+    def __init__(self, code: NetworkErrorCode, message: str | None = None) -> None:
+        super().__init__(message or _NETWORK_MESSAGES[code])
+        self.code = code
+
+
+_NETWORK_MESSAGES: dict[str, str] = {
+    "request_failed": "The request to the AREX API failed.",
+    "invalid_response": "The AREX API returned a response the SDK cannot parse.",
+}
+
+
+class ArexTimeoutError(ArexError):
+    """A request or a research poll exceeded its deadline."""
+
+    def __init__(
+        self,
+        phase: TimeoutPhase,
+        timeout_ms: int,
+        run_id: str | None = None,
+    ) -> None:
+        subject = "Research run" if phase == "poll" else "Request"
+        suffix = f" (run {run_id})" if run_id else ""
+        super().__init__(f"{subject} timed out after {timeout_ms} ms{suffix}.")
+        self.phase = phase
+        self.timeout_ms = timeout_ms
+        self.run_id = run_id
+
+
+class ResearchRunFailedError(ArexError):
+    """A research run reached the terminal `failed` status."""
+
+    def __init__(self, run: ResearchRun) -> None:
+        retryable = run.error.retryable if run.error else False
+        super().__init__(f"Research run {run.id} failed (retryable={retryable}).")
+        self.run = run
+        self.run_id = run.id
+
+
+class ResearchRunCancelledError(ArexError):
+    """A research run reached the terminal `cancelled` status."""
+
+    def __init__(self, run: ResearchRun) -> None:
+        super().__init__(f"Research run {run.id} was cancelled.")
+        self.run = run
+        self.run_id = run.id
+
+
+_ERROR_CLASSES: dict[str, type[ArexApiError]] = {
+    "invalid_request": InvalidRequestError,
+    "content_blocked": InvalidRequestError,
+    "idempotency_key_reused": InvalidRequestError,
+    "source_not_available": SourceNotAvailableError,
+    "not_found": NotFoundError,
+    "rate_limit_exceeded": RateLimitError,
+    "idempotency_key_in_progress": RateLimitError,
+    "outstanding_run_limit": RateLimitError,
+    "temporarily_unavailable": ServiceUnavailableError,
+    "internal_error": ServiceUnavailableError,
+}
+
+
+_CREDIT_CODES = frozenset({"daily_credits_exhausted", "weekly_credits_exhausted"})
+
+
+def _header(headers: Mapping[str, str], name: str) -> str | None:
+    target = name.lower()
+    for key, value in headers.items():
+        if key.lower() == target:
+            return value
+    return None
+
+
+# Longest wait the SDK will report. A server asking for more than this means
+# "not soon", and clamping keeps the arithmetic and `sleep` calls in range.
+_MAX_RETRY_AFTER_MS = 2**31 - 1
+_DELTA_SECONDS = re.compile(r"[0-9]+(?:\.[0-9]+)?")
+
+
+def _clamp_ms(ms: float) -> int:
+    return _MAX_RETRY_AFTER_MS if ms >= _MAX_RETRY_AFTER_MS else max(0, int(ms))
+
+
+def retry_after_ms_from_headers(headers: Mapping[str, str]) -> int | None:
+    """`Retry-After` in milliseconds; accepts delta-seconds or an HTTP date.
+
+    Malformed values give `None`; huge ones clamp instead of overflowing.
+    """
+    raw = _header(headers, "Retry-After")
+    if raw is None:
+        return None
+    text = raw.strip()
+    if _DELTA_SECONDS.fullmatch(text):
+        # ASCII digits never fail `float`, but a long run of them is `inf`.
+        return _clamp_ms(float(text) * 1000)
+    try:
+        at = parsedate_to_datetime(text)
+    except (TypeError, ValueError, IndexError, OverflowError):
+        return None
+    if at.tzinfo is None:
+        # A date without a zone (`-0000`) is still GMT per RFC 9110.
+        at = at.replace(tzinfo=timezone.utc)
+    return _clamp_ms((at - datetime.now(timezone.utc)).total_seconds() * 1000)
+
+
+def _optional_number(value: Any) -> int | float | None:
+    return (
+        value
+        if isinstance(value, (int, float)) and not isinstance(value, bool)
+        else None
+    )
+
+
+def api_error_from_response(
+    status: int,
+    body: Any,
+    headers: Mapping[str, str],
+) -> ArexApiError:
+    """Map a wire error response onto the matching SDK error class."""
+    retry_after_ms = retry_after_ms_from_headers(headers)
+
+    if not isinstance(body, dict) or not isinstance(body.get("error"), dict):
+        raise ArexNetworkError("invalid_response")
+    details: dict[str, Any] = dict(body["error"])
+    code = details.get("code")
+    message = details.get("message")
+    retryable = details.get("retryable")
+    request_id = details.get("request_id")
+    param = details.get("param")
+    if (
+        not isinstance(code, str)
+        or not isinstance(message, str)
+        or not isinstance(retryable, bool)
+        or (request_id is not None and not isinstance(request_id, str))
+        or (param is not None and not isinstance(param, str))
+    ):
+        raise ArexNetworkError("invalid_response")
+
+    if code in _CREDIT_CODES:
+        resource = details.get("resource")
+        reset_at = details.get("reset_at")
+        return CreditsExhaustedError(
+            message,
+            status=status,
+            code=code,
+            retryable=retryable,
+            retry_after_ms=retry_after_ms,
+            request_id=request_id,
+            param=param,
+            details=details,
+            resource=resource if isinstance(resource, str) else None,
+            used=_optional_number(details.get("used")),
+            limit=_optional_number(details.get("limit")),
+            requested=_optional_number(details.get("requested")),
+            reset_at=reset_at if isinstance(reset_at, str) else None,
+        )
+
+    if status == 401:
+        cls: type[ArexApiError] = AuthenticationError
+    elif status == 403:
+        cls = PermissionDeniedError
+    else:
+        cls = _ERROR_CLASSES.get(code, ArexApiError)
+    return cls(
+        message,
+        status=status,
+        code=code,
+        retryable=retryable,
+        retry_after_ms=retry_after_ms,
+        request_id=request_id,
+        param=param,
+        details=details,
+    )