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

arex/__init__.py

@@ -0,0 +1,137 @@
+"""Python SDK for the AREX public API: search, browse, research, usage."""
+
+import importlib.metadata
+
+from arex.async_client import AsyncArex
+from arex.client import Arex
+from arex.errors import (
+    ArexApiError,
+    ArexError,
+    ArexNetworkError,
+    ArexTimeoutError,
+    AuthenticationError,
+    CreditsExhaustedError,
+    InvalidRequestError,
+    NotFoundError,
+    PermissionDeniedError,
+    RateLimitError,
+    ResearchRunCancelledError,
+    ResearchRunFailedError,
+    ServiceUnavailableError,
+    SourceNotAvailableError,
+    api_error_from_response,
+)
+from arex.models import (
+    API_KEY_PREFIX,
+    BROWSE_AUTO_TIMEOUT_MS,
+    BROWSE_FLEX_TIMEOUT_MS,
+    BROWSE_TIERS,
+    DEEP_POLL_TIMEOUT_MS,
+    DEFAULT_BASE_URL,
+    DEFAULT_BROWSE_MAX_CHARACTERS,
+    DEFAULT_BROWSE_TIER,
+    DEFAULT_LIVECRAWL_TIMEOUT_AUTO_MS,
+    DEFAULT_LIVECRAWL_TIMEOUT_FLEX_MS,
+    DEFAULT_MAX_RETRIES,
+    DEFAULT_POLL_INTERVAL_MS,
+    DEFAULT_RESEARCH_MODE,
+    DEFAULT_RETRY_DELAY_MS,
+    DEFAULT_SEARCH_LIMIT,
+    FAST_POLL_TIMEOUT_MS,
+    LIVECRAWL_TIMEOUT_MARGIN_MS,
+    RESEARCH_MODES,
+    RESEARCH_REQUEST_TIMEOUT_MS,
+    SEARCH_TIMEOUT_MS,
+    USAGE_TIMEOUT_MS,
+    BrowseDocument,
+    BrowseTier,
+    Credits,
+    CreditUsage,
+    CreditWeek,
+    CreditWindow,
+    ResearchFailureCode,
+    ResearchMode,
+    ResearchModeRequest,
+    ResearchResult,
+    ResearchRun,
+    ResearchRunError,
+    ResearchSource,
+    ResearchStatus,
+    SearchResponse,
+    SearchResult,
+    is_arex_api_key,
+)
+from arex.research import AsyncResearch, Research
+
+
+def __getattr__(name: str) -> str:
+    # Read lazily from the installed distribution so the version has a single
+    # source (pyproject.toml); an uninstalled tree raises rather than guessing.
+    if name == "__version__":
+        return importlib.metadata.version("arex-python-sdk")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "API_KEY_PREFIX",
+    "BROWSE_AUTO_TIMEOUT_MS",
+    "BROWSE_FLEX_TIMEOUT_MS",
+    "BROWSE_TIERS",
+    "DEEP_POLL_TIMEOUT_MS",
+    "DEFAULT_BASE_URL",
+    "DEFAULT_BROWSE_MAX_CHARACTERS",
+    "DEFAULT_BROWSE_TIER",
+    "DEFAULT_LIVECRAWL_TIMEOUT_AUTO_MS",
+    "DEFAULT_LIVECRAWL_TIMEOUT_FLEX_MS",
+    "DEFAULT_MAX_RETRIES",
+    "DEFAULT_POLL_INTERVAL_MS",
+    "DEFAULT_RESEARCH_MODE",
+    "DEFAULT_RETRY_DELAY_MS",
+    "DEFAULT_SEARCH_LIMIT",
+    "FAST_POLL_TIMEOUT_MS",
+    "LIVECRAWL_TIMEOUT_MARGIN_MS",
+    "RESEARCH_MODES",
+    "RESEARCH_REQUEST_TIMEOUT_MS",
+    "SEARCH_TIMEOUT_MS",
+    "USAGE_TIMEOUT_MS",
+    "Arex",
+    "ArexApiError",
+    # errors
+    "ArexError",
+    "ArexNetworkError",
+    "ArexTimeoutError",
+    "AsyncArex",
+    "AsyncResearch",
+    "AuthenticationError",
+    "BrowseDocument",
+    # models
+    "BrowseTier",
+    "CreditUsage",
+    "CreditWeek",
+    "CreditWindow",
+    "Credits",
+    "CreditsExhaustedError",
+    "InvalidRequestError",
+    "NotFoundError",
+    "PermissionDeniedError",
+    "RateLimitError",
+    "Research",
+    "ResearchFailureCode",
+    "ResearchMode",
+    "ResearchModeRequest",
+    "ResearchResult",
+    "ResearchRun",
+    "ResearchRunCancelledError",
+    "ResearchRunError",
+    "ResearchRunFailedError",
+    "ResearchSource",
+    "ResearchStatus",
+    "SearchResponse",
+    "SearchResult",
+    "ServiceUnavailableError",
+    "SourceNotAvailableError",
+    "__version__",
+    "api_error_from_response",
+    # contract constants
+    "is_arex_api_key",
+]

arex/_poll.py

@@ -0,0 +1,126 @@
+"""Fixed-interval polling with a deadline.
+
+`get` receives the milliseconds left before the deadline and returns
+`(value, retry_after_ms)`: a non-`None` value ends the loop, and the
+`Retry-After` of the most recent response overrides the fixed interval.
+Inputs are validated by the resource before any charged call.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import time
+from collections.abc import Awaitable, Callable
+from typing import TypeVar
+
+from arex.errors import ArexError, ArexTimeoutError
+
+__all__ = ["async_poll_until", "poll_until"]
+
+T = TypeVar("T")
+
+
+def _delay_seconds(retry_after_ms: int | None, poll_interval_ms: int) -> float:
+    return (retry_after_ms if retry_after_ms is not None else poll_interval_ms) / 1000
+
+
+def _remaining_ms(deadline: float) -> int:
+    return math.ceil((deadline - time.monotonic()) * 1000)
+
+
+def _poll_timeout(
+    exc: ArexError,
+    *,
+    deadline: float,
+    remaining_ms: int,
+    timeout_ms: int,
+    run_id: str,
+) -> ArexTimeoutError | None:
+    """Only elapsed time or a request capped by the poll budget is a poll timeout.
+
+    Retry eligibility belongs to the transport. A server Retry-After header
+    must not relabel a non-retried API error before the poll deadline.
+    """
+    now = time.monotonic()
+    cut = (
+        isinstance(exc, ArexTimeoutError)
+        and exc.phase == "request"
+        and exc.timeout_ms >= remaining_ms
+    )
+    if cut or now >= deadline:
+        return ArexTimeoutError("poll", timeout_ms, run_id)
+    return None
+
+
+def poll_until(
+    get: Callable[[int], tuple[T | None, int | None]],
+    *,
+    poll_interval_ms: int,
+    timeout_ms: int,
+    run_id: str,
+) -> T:
+    """Call `get` until it yields a value, or raise `ArexTimeoutError`.
+
+    An outcome that arrives after the deadline is a poll timeout, whatever it is.
+    """
+    deadline = time.monotonic() + timeout_ms / 1000
+    while True:
+        remaining_ms = _remaining_ms(deadline)
+        if remaining_ms <= 0:
+            raise ArexTimeoutError("poll", timeout_ms, run_id)
+        try:
+            value, retry_after_ms = get(remaining_ms)
+        except ArexError as exc:
+            timeout = _poll_timeout(
+                exc,
+                deadline=deadline,
+                remaining_ms=remaining_ms,
+                timeout_ms=timeout_ms,
+                run_id=run_id,
+            )
+            if timeout is None:
+                raise
+            raise timeout from exc
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise ArexTimeoutError("poll", timeout_ms, run_id)
+        if value is not None:
+            return value
+        time.sleep(min(_delay_seconds(retry_after_ms, poll_interval_ms), remaining))
+
+
+async def async_poll_until(
+    get: Callable[[int], Awaitable[tuple[T | None, int | None]]],
+    *,
+    poll_interval_ms: int,
+    timeout_ms: int,
+    run_id: str,
+) -> T:
+    """Await `get` until it yields a value, or raise `ArexTimeoutError`."""
+    deadline = time.monotonic() + timeout_ms / 1000
+    while True:
+        remaining_ms = _remaining_ms(deadline)
+        if remaining_ms <= 0:
+            raise ArexTimeoutError("poll", timeout_ms, run_id)
+        try:
+            value, retry_after_ms = await get(remaining_ms)
+        except ArexError as exc:
+            timeout = _poll_timeout(
+                exc,
+                deadline=deadline,
+                remaining_ms=remaining_ms,
+                timeout_ms=timeout_ms,
+                run_id=run_id,
+            )
+            if timeout is None:
+                raise
+            raise timeout from exc
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise ArexTimeoutError("poll", timeout_ms, run_id)
+        if value is not None:
+            return value
+        await asyncio.sleep(
+            min(_delay_seconds(retry_after_ms, poll_interval_ms), remaining)
+        )

arex/_transport.py

@@ -0,0 +1,303 @@
+"""HTTP plumbing: auth, per-call timeouts, credit headers, errors, retries."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, NamedTuple
+
+import httpx
+
+from arex._validate import is_timeout_ms, validate_timeout_ms
+from arex.errors import (
+    ArexNetworkError,
+    ArexTimeoutError,
+    RateLimitError,
+    ServiceUnavailableError,
+    api_error_from_response,
+    retry_after_ms_from_headers,
+)
+from arex.models import (
+    DEFAULT_MAX_RETRIES,
+    DEFAULT_RETRY_DELAY_MS,
+    Credits,
+)
+
+__all__ = ["ApiResult", "AsyncTransport", "Transport"]
+
+
+class ApiResult(NamedTuple):
+    """One decoded successful response."""
+
+    data: Any
+    credits: Credits | None
+    retry_after_ms: int | None
+
+
+def _prepare(
+    base_url: str,
+    api_key: str,
+    path: str,
+    idempotency_key: str | None,
+) -> tuple[str, dict[str, str]]:
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "application/json",
+    }
+    if idempotency_key is not None:
+        headers["Idempotency-Key"] = idempotency_key
+    return base_url.rstrip("/") + path, headers
+
+
+def _decode(response: httpx.Response) -> ApiResult:
+    retry_after_ms = retry_after_ms_from_headers(response.headers)
+
+    if not 200 <= response.status_code < 300:
+        try:
+            body: Any = response.json()
+        except ValueError as exc:
+            raise ArexNetworkError("invalid_response") from exc
+        raise api_error_from_response(response.status_code, body, response.headers)
+
+    if response.status_code == 204 or not response.content:
+        return ApiResult(None, None, retry_after_ms)
+
+    try:
+        data = response.json()
+    except ValueError as exc:
+        raise ArexNetworkError("invalid_response") from exc
+
+    return ApiResult(data, Credits.from_headers(response.headers), retry_after_ms)
+
+
+class _Base:
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        base_url: str,
+        timeout_ms: int | None,
+        max_retries: int,
+        retry_unavailable: bool,
+    ) -> None:
+        if timeout_ms is not None and not is_timeout_ms(timeout_ms):
+            raise ValueError("timeout_ms must be a positive integer")
+        if (
+            isinstance(max_retries, bool)
+            or not isinstance(max_retries, int)
+            or max_retries < 0
+        ):
+            raise ValueError("max_retries must be a non-negative integer")
+        self.api_key = api_key
+        self.base_url = base_url
+        self.timeout_ms = timeout_ms
+        self.max_retries = max_retries
+        self.retry_unavailable = retry_unavailable
+
+    def resolve_timeout(self, default_ms: int, override_ms: int | None) -> int:
+        """Per-call argument beats the client-level override beats the default."""
+        if override_ms is not None:
+            validate_timeout_ms(override_ms)
+            return override_ms
+        if self.timeout_ms is not None:
+            return self.timeout_ms
+        return default_ms
+
+    def _retry_delay_ms(
+        self,
+        error: RateLimitError | ServiceUnavailableError,
+        attempt: int,
+    ) -> int | None:
+        """Milliseconds to wait before retrying, or `None` to give up."""
+        if attempt >= self.max_retries:
+            return None
+        if isinstance(error, RateLimitError):
+            # `outstanding_run_limit` is a state, not a burst; retrying never helps.
+            if error.code not in {"rate_limit_exceeded", "idempotency_key_in_progress"}:
+                return None
+        elif not self.retry_unavailable or error.code != "temporarily_unavailable":
+            return None
+        if error.retry_after_ms is not None:
+            return error.retry_after_ms
+        return DEFAULT_RETRY_DELAY_MS
+
+    def _retry_sleep(
+        self,
+        error: RateLimitError | ServiceUnavailableError,
+        attempt: int,
+        deadline: float,
+    ) -> float | None:
+        """Seconds to sleep, capped to the deadline, or `None` if not retryable."""
+        delay_ms = self._retry_delay_ms(error, attempt)
+        if delay_ms is None:
+            return None
+        return min(delay_ms / 1000, max(0.0, deadline - time.monotonic()))
+
+
+def _attempt_timeout(attempt: int, timeout_ms: int, deadline: float) -> httpx.Timeout:
+    """httpx timeouts are per phase, so each attempt gets what is left of the
+    one budget (the first gets all of it)."""
+    seconds = timeout_ms / 1000 if attempt == 0 else deadline - time.monotonic()
+    if seconds <= 0:
+        raise ArexTimeoutError("request", timeout_ms)
+    return httpx.Timeout(seconds)
+
+
+class Transport(_Base):
+    """Synchronous transport over `httpx.Client`.
+
+    Each attempt gives its I/O phases the remaining budget. The SDK checks
+    elapsed time after a response, but cannot interrupt a blocking call, so
+    a server that keeps trickling bytes can outlast the budget.
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        base_url: str,
+        timeout_ms: int | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_unavailable: bool = False,
+        client: httpx.Client | None = None,
+    ) -> None:
+        super().__init__(
+            api_key=api_key,
+            base_url=base_url,
+            timeout_ms=timeout_ms,
+            max_retries=max_retries,
+            retry_unavailable=retry_unavailable,
+        )
+        self.owns_client = client is None
+        self.client = client if client is not None else httpx.Client(base_url=base_url)
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json: Any | None = None,
+        timeout_ms: int,
+        idempotency_key: str | None = None,
+    ) -> ApiResult:
+        url, headers = _prepare(self.base_url, self.api_key, path, idempotency_key)
+        deadline = time.monotonic() + timeout_ms / 1000
+        attempt = 0
+        while True:
+            timeout = _attempt_timeout(attempt, timeout_ms, deadline)
+            try:
+                response = self.client.request(
+                    method, url, headers=headers, json=json, timeout=timeout
+                )
+            except httpx.TimeoutException as exc:
+                raise ArexTimeoutError("request", timeout_ms) from exc
+            except httpx.HTTPError as exc:
+                raise ArexNetworkError("request_failed") from exc
+            if time.monotonic() >= deadline:
+                raise ArexTimeoutError("request", timeout_ms)
+            try:
+                return _decode(response)
+            except (RateLimitError, ServiceUnavailableError) as error:
+                delay = self._retry_sleep(error, attempt, deadline)
+                if delay is None:
+                    raise
+                time.sleep(delay)
+                attempt += 1
+
+    def close(self) -> None:
+        if self.owns_client:
+            self.client.close()
+
+
+class AsyncTransport(_Base):
+    """Asynchronous transport over `httpx.AsyncClient`.
+
+    The budget is a wall-clock deadline: when it passes, the in-flight request
+    or retry sleep is cancelled.
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        base_url: str,
+        timeout_ms: int | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_unavailable: bool = False,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        super().__init__(
+            api_key=api_key,
+            base_url=base_url,
+            timeout_ms=timeout_ms,
+            max_retries=max_retries,
+            retry_unavailable=retry_unavailable,
+        )
+        self.owns_client = client is None
+        self.client = (
+            client if client is not None else httpx.AsyncClient(base_url=base_url)
+        )
+
+    async def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json: Any | None = None,
+        timeout_ms: int,
+        idempotency_key: str | None = None,
+    ) -> ApiResult:
+        url, headers = _prepare(self.base_url, self.api_key, path, idempotency_key)
+        deadline = time.monotonic() + timeout_ms / 1000
+        work = asyncio.ensure_future(
+            self._send(method, url, headers, json, timeout_ms, deadline)
+        )
+        # Not `asyncio.wait_for`: before Python 3.12 it can swallow a caller's
+        # cancellation that lands as the request finishes, and `asyncio.timeout`
+        # needs 3.11. Waiting on the task keeps both cases explicit.
+        try:
+            done, _ = await asyncio.wait({work}, timeout=timeout_ms / 1000)
+        except asyncio.CancelledError:
+            work.cancel()
+            await asyncio.gather(work, return_exceptions=True)
+            raise
+        if work in done:
+            return work.result()
+        work.cancel()
+        await asyncio.gather(work, return_exceptions=True)
+        raise ArexTimeoutError("request", timeout_ms)
+
+    async def _send(
+        self,
+        method: str,
+        url: str,
+        headers: dict[str, str],
+        json: Any | None,
+        timeout_ms: int,
+        deadline: float,
+    ) -> ApiResult:
+        attempt = 0
+        while True:
+            timeout = _attempt_timeout(attempt, timeout_ms, deadline)
+            try:
+                response = await self.client.request(
+                    method, url, headers=headers, json=json, timeout=timeout
+                )
+            except httpx.TimeoutException as exc:
+                raise ArexTimeoutError("request", timeout_ms) from exc
+            except httpx.HTTPError as exc:
+                raise ArexNetworkError("request_failed") from exc
+            if time.monotonic() >= deadline:
+                raise ArexTimeoutError("request", timeout_ms)
+            try:
+                return _decode(response)
+            except (RateLimitError, ServiceUnavailableError) as error:
+                delay = self._retry_sleep(error, attempt, deadline)
+                if delay is None:
+                    raise
+                await asyncio.sleep(delay)
+                attempt += 1
+
+    async def aclose(self) -> None:
+        if self.owns_client:
+            await self.client.aclose()