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

arex/_validate.py

@@ -0,0 +1,156 @@
+"""Validate requests before sending them using generated wire constraints."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, get_args
+from urllib.parse import urlsplit
+from uuid import uuid4
+
+from arex._wire import WIRE_LIMITS, BrowseTier, ResearchModeRequest
+from arex.errors import InvalidRequestError
+
+__all__ = [
+    "MAX_TIMEOUT_MS",
+    "is_timeout_ms",
+    "resolve_idempotency_key",
+    "validate_browse",
+    "validate_research_create",
+    "validate_research_wait",
+    "validate_run_id",
+    "validate_search",
+    "validate_timeout_ms",
+]
+
+_RUN_ID_PATTERN = re.compile(
+    r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-8][0-9a-fA-F]{3}"
+    r"-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$"
+)
+
+
+def _integer(field: str, value: object, bounds: dict[str, Any]) -> None:
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise InvalidRequestError(f"{field} must be an integer.", param=field)
+    if not bounds["minimum"] <= value <= bounds["maximum"]:
+        raise InvalidRequestError(
+            f"{field} must be between {bounds['minimum']} and {bounds['maximum']}.",
+            param=field,
+        )
+
+
+def _query(value: str, bounds: dict[str, Any]) -> None:
+    if not isinstance(value, str) or not value.strip():
+        raise InvalidRequestError("query must be a non-empty string.", param="query")
+    # Python string lengths count Unicode code points rather than UTF-16 units.
+    if not bounds["minLength"] <= len(value) <= bounds["maxLength"]:
+        raise InvalidRequestError(
+            f"query must be at most {bounds['maxLength']} characters.", param="query"
+        )
+
+
+def validate_search(query: str, limit: int | None) -> None:
+    bounds = WIRE_LIMITS["SearchRequest"]
+    _query(query, bounds["query"])
+    if limit is not None:
+        _integer("limit", limit, bounds["limit"])
+
+
+def validate_browse(
+    url: str,
+    *,
+    tier: str | None,
+    livecrawl: bool | None,
+    livecrawl_timeout_ms: int | None,
+    max_characters: int | None,
+) -> None:
+    bounds = WIRE_LIMITS["BrowseRequest"]
+    if not isinstance(url, str) or not url.strip():
+        raise InvalidRequestError("url must be a non-empty string.", param="url")
+    if len(url) > bounds["url"]["maxLength"]:
+        raise InvalidRequestError(
+            f"url must be at most {bounds['url']['maxLength']} characters.", param="url"
+        )
+    try:
+        parts = urlsplit(url)
+        valid_host = parts.hostname
+    except ValueError as exc:
+        raise InvalidRequestError(
+            "url must be an absolute HTTP(S) URL.", param="url"
+        ) from exc
+    if parts.scheme not in ("http", "https") or not valid_host:
+        raise InvalidRequestError("url must be an absolute HTTP(S) URL.", param="url")
+    if parts.username is not None or parts.password is not None:
+        raise InvalidRequestError("url must not contain credentials.", param="url")
+    if tier is not None and tier not in get_args(BrowseTier):
+        raise InvalidRequestError("tier must be auto or flex.", param="tier")
+    if livecrawl is not None and not isinstance(livecrawl, bool):
+        raise InvalidRequestError("livecrawl must be a boolean.", param="livecrawl")
+    if livecrawl_timeout_ms is not None:
+        if livecrawl is not True:
+            raise InvalidRequestError(
+                "livecrawl_timeout_ms is only allowed with livecrawl=True.",
+                param="livecrawl_timeout_ms",
+            )
+        _integer(
+            "livecrawl_timeout_ms", livecrawl_timeout_ms, bounds["livecrawl_timeout_ms"]
+        )
+    if max_characters is not None:
+        _integer("max_characters", max_characters, bounds["max_characters"])
+
+
+def validate_research_create(query: str, mode: str | None) -> None:
+    _query(query, WIRE_LIMITS["ResearchCreateRequest"]["query"])
+    if mode is not None and mode not in get_args(ResearchModeRequest):
+        raise InvalidRequestError("mode must be fast, deep or auto.", param="mode")
+
+
+def validate_run_id(run_id: str) -> None:
+    if not isinstance(run_id, str) or not _RUN_ID_PATTERN.fullmatch(run_id):
+        raise InvalidRequestError("run_id must be a UUID.", param="run_id")
+
+
+def validate_research_wait(poll_interval_ms: int, timeout_ms: int | None) -> None:
+    """Runs before `create` in `run`, so a bad poll setting never pays for a run
+    it cannot wait for."""
+    if not is_timeout_ms(poll_interval_ms):
+        raise InvalidRequestError(
+            f"poll_interval_ms must be an integer between 1 and {MAX_TIMEOUT_MS}.",
+            param="poll_interval_ms",
+        )
+    if timeout_ms is not None:
+        validate_timeout_ms(timeout_ms)
+
+
+# The largest delay a JavaScript timer accepts; also keeps the millisecond to
+# second conversion and `socket` timeouts far from float overflow.
+MAX_TIMEOUT_MS = 2**31 - 1
+
+
+def is_timeout_ms(value: object) -> bool:
+    return (
+        isinstance(value, int)
+        and not isinstance(value, bool)
+        and 0 < value <= MAX_TIMEOUT_MS
+    )
+
+
+def validate_timeout_ms(value: object) -> None:
+    if not is_timeout_ms(value):
+        raise InvalidRequestError(
+            f"timeout_ms must be a positive integer of at most {MAX_TIMEOUT_MS}.",
+            param="timeout_ms",
+        )
+
+
+def resolve_idempotency_key(key: str | None) -> str:
+    """The caller's key once validated, else a fresh UUID4 for this one call."""
+    if key is None:
+        return str(uuid4())
+    pattern = WIRE_LIMITS["IdempotencyKey"]["pattern"]
+    # `fullmatch` so the pattern's `$` cannot accept a trailing newline.
+    if not isinstance(key, str) or not re.fullmatch(pattern, key):
+        raise InvalidRequestError(
+            "idempotency_key must be 1 to 255 characters from A-Z a-z 0-9 . _ -",
+            param="idempotency_key",
+        )
+    return key

arex/_wire.py

@@ -0,0 +1,126 @@
+# Generated from openapi/arex-v1.json by generate_wire.py; do not edit.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
+
+if TYPE_CHECKING:
+    from typing_extensions import Required
+
+ErrorCode = Literal['invalid_request', 'content_blocked', 'unauthorized', 'registration_required', 'account_deletion_pending', 'feature_not_enabled', 'not_found', 'method_not_allowed', 'run_not_cancellable', 'idempotency_key_in_progress', 'source_not_available', 'idempotency_key_reused', 'rate_limit_exceeded', 'outstanding_run_limit', 'daily_credits_exhausted', 'weekly_credits_exhausted', 'internal_error', 'temporarily_unavailable']
+ResearchFailureCode = Literal['timeout', 'execution_failed', 'worker_unavailable', 'interrupted']
+BrowseTier = Literal['auto', 'flex']
+ResearchMode = Literal['fast', 'deep']
+ResearchModeRequest = Literal['fast', 'deep', 'auto']
+ResearchStatus = Literal['queued', 'running', 'completed', 'failed', 'cancelled']
+
+class BrowseRequest(TypedDict, total=False):
+    url: Required[str]
+    tier: BrowseTier
+    livecrawl: bool
+    livecrawl_timeout_ms: int
+    max_characters: int
+
+class BrowseResponse(TypedDict, total=False):
+    url: Required[str]
+    title: str
+    published_at: str
+    image_url: str
+    image_urls: list[str]
+    content: Required[str]
+    truncated: Required[bool]
+
+class ErrorError(TypedDict, total=False):
+    code: Required[ErrorCode]
+    message: Required[str]
+    retryable: Required[bool]
+    request_id: Required[str]
+    param: str
+    resource: str
+    used: int
+    limit: int
+    requested: int
+    reset_at: str
+
+class Error(TypedDict, total=False):
+    error: Required[ErrorError]
+
+class ResearchCreateRequest(TypedDict, total=False):
+    query: Required[str]
+    mode: ResearchModeRequest
+
+class ResearchRunResultSourcesItem(TypedDict, total=False):
+    title: Required[str]
+    url: Required[str]
+
+class ResearchRunResult(TypedDict, total=False):
+    answer: Required[str]
+    sources: Required[list[ResearchRunResultSourcesItem]]
+
+class ResearchRunError(TypedDict, total=False):
+    code: Required[ResearchFailureCode]
+    retryable: Required[bool]
+
+class ResearchRun(TypedDict, total=False):
+    id: Required[str]
+    status: Required[ResearchStatus]
+    mode: Required[ResearchMode]
+    requested_mode: Required[ResearchModeRequest]
+    created_at: Required[str]
+    started_at: str
+    completed_at: str
+    result: ResearchRunResult
+    error: ResearchRunError
+
+class SearchRequest(TypedDict, total=False):
+    query: Required[str]
+    limit: int
+
+class SearchResponseResultsItem(TypedDict, total=False):
+    title: Required[str]
+    url: Required[str]
+    snippet: str
+
+class SearchResponse(TypedDict, total=False):
+    results: Required[list[SearchResponseResultsItem]]
+
+class UsageResponseWeek(TypedDict, total=False):
+    start: Required[str]
+    used: Required[int]
+    limit: int
+    remaining: int
+    reset_at: Required[str]
+
+class UsageResponse(TypedDict, total=False):
+    date: Required[str]
+    used: Required[int]
+    limit: int
+    remaining: int
+    reset_at: Required[str]
+    week: UsageResponseWeek
+
+
+WIRE_LIMITS: dict[str, Any] = {'BrowseRequest': {'livecrawl': {'default': False},
+                   'livecrawl_timeout_ms': {'maximum': 90000, 'minimum': 1000},
+                   'max_characters': {'default': 12000,
+                                      'maximum': 131072,
+                                      'minimum': 1},
+                   'tier': {'default': 'auto', 'enum': ['auto', 'flex']},
+                   'url': {'maxLength': 8192,
+                           'pattern': '^[Hh][Tt][Tt][Pp][Ss]?://[^/@?#]+(?:[/?#]|$)'}},
+ 'BrowseResponse': {'image_urls': {'maxItems': 50}},
+ 'IdempotencyKey': {'pattern': '^[A-Za-z0-9._-]{1,255}$', 'type': 'string'},
+ 'ResearchCreateRequest': {'mode': {'default': 'auto',
+                                    'enum': ['fast', 'deep', 'auto']},
+                           'query': {'maxLength': 2000, 'minLength': 1}},
+ 'ResearchRun': {'mode': {'enum': ['fast', 'deep']},
+                 'requested_mode': {'enum': ['fast', 'deep', 'auto']},
+                 'status': {'enum': ['queued',
+                                     'running',
+                                     'completed',
+                                     'failed',
+                                     'cancelled']}},
+ 'SearchRequest': {'limit': {'default': 10, 'maximum': 10, 'minimum': 1},
+                   'query': {'maxLength': 2000, 'minLength': 1}},
+ 'UsageResponse': {'limit': {'minimum': 0},
+                   'remaining': {'minimum': 0},
+                   'used': {'minimum': 0}}}

arex/async_client.py

@@ -0,0 +1,133 @@
+"""The asynchronous `AsyncArex` client."""
+
+from __future__ import annotations
+
+from types import TracebackType
+from typing import Any
+
+import httpx
+
+from arex._transport import AsyncTransport
+from arex._validate import resolve_idempotency_key, validate_browse, validate_search
+from arex.client import (
+    browse_body,
+    browse_timeout_ms,
+    require_credits,
+    resolve_api_key,
+)
+from arex.models import (
+    DEFAULT_BASE_URL,
+    DEFAULT_MAX_RETRIES,
+    SEARCH_TIMEOUT_MS,
+    USAGE_TIMEOUT_MS,
+    BrowseDocument,
+    BrowseTier,
+    CreditUsage,
+    SearchResponse,
+)
+from arex.research import AsyncResearch
+
+__all__ = ["AsyncArex"]
+
+
+class AsyncArex:
+    """Asynchronous client for the AREX public API."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        *,
+        base_url: str | None = None,
+        timeout_ms: int | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_unavailable: bool = False,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self._transport = AsyncTransport(
+            api_key=resolve_api_key(api_key),
+            base_url=base_url or DEFAULT_BASE_URL,
+            timeout_ms=timeout_ms,
+            max_retries=max_retries,
+            retry_unavailable=retry_unavailable,
+            client=http_client,
+        )
+        self.research = AsyncResearch(self._transport)
+
+    async def search(
+        self,
+        query: str,
+        *,
+        limit: int | None = None,
+        idempotency_key: str | None = None,
+        timeout_ms: int | None = None,
+    ) -> SearchResponse:
+        """Run a web search."""
+        validate_search(query, limit)
+        key = resolve_idempotency_key(idempotency_key)
+        body: dict[str, Any] = {"query": query}
+        if limit is not None:
+            body["limit"] = limit
+        result = await self._transport.request(
+            "POST",
+            "/v1/search",
+            json=body,
+            idempotency_key=key,
+            timeout_ms=self._transport.resolve_timeout(SEARCH_TIMEOUT_MS, timeout_ms),
+        )
+        return SearchResponse.from_wire(result.data, require_credits(result.credits))
+
+    async def browse(
+        self,
+        url: str,
+        *,
+        tier: BrowseTier | None = None,
+        livecrawl: bool | None = None,
+        livecrawl_timeout_ms: int | None = None,
+        max_characters: int | None = None,
+        idempotency_key: str | None = None,
+        timeout_ms: int | None = None,
+    ) -> BrowseDocument:
+        """Fetch and extract one page."""
+        validate_browse(
+            url,
+            tier=tier,
+            livecrawl=livecrawl,
+            livecrawl_timeout_ms=livecrawl_timeout_ms,
+            max_characters=max_characters,
+        )
+        key = resolve_idempotency_key(idempotency_key)
+        default_ms = browse_timeout_ms(tier, livecrawl, livecrawl_timeout_ms)
+        result = await self._transport.request(
+            "POST",
+            "/v1/browse",
+            idempotency_key=key,
+            json=browse_body(
+                url, tier, livecrawl, livecrawl_timeout_ms, max_characters
+            ),
+            timeout_ms=self._transport.resolve_timeout(default_ms, timeout_ms),
+        )
+        return BrowseDocument.from_wire(result.data, require_credits(result.credits))
+
+    async def usage(self, *, timeout_ms: int | None = None) -> CreditUsage:
+        """Read the current credit windows."""
+        result = await self._transport.request(
+            "GET",
+            "/v1/usage",
+            timeout_ms=self._transport.resolve_timeout(USAGE_TIMEOUT_MS, timeout_ms),
+        )
+        return CreditUsage.from_wire(result.data)
+
+    async def aclose(self) -> None:
+        """Close the HTTP client, unless it was injected."""
+        await self._transport.aclose()
+
+    async def __aenter__(self) -> AsyncArex:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        await self.aclose()

arex/client.py

@@ -0,0 +1,197 @@
+"""The synchronous `Arex` client."""
+
+from __future__ import annotations
+
+import os
+from types import TracebackType
+from typing import Any
+
+import httpx
+
+from arex._transport import Transport
+from arex._validate import resolve_idempotency_key, validate_browse, validate_search
+from arex.errors import ArexNetworkError
+from arex.models import (
+    BROWSE_AUTO_TIMEOUT_MS,
+    BROWSE_FLEX_TIMEOUT_MS,
+    DEFAULT_BASE_URL,
+    DEFAULT_BROWSE_TIER,
+    DEFAULT_LIVECRAWL_TIMEOUT_AUTO_MS,
+    DEFAULT_LIVECRAWL_TIMEOUT_FLEX_MS,
+    DEFAULT_MAX_RETRIES,
+    LIVECRAWL_TIMEOUT_MARGIN_MS,
+    SEARCH_TIMEOUT_MS,
+    USAGE_TIMEOUT_MS,
+    BrowseDocument,
+    BrowseTier,
+    Credits,
+    CreditUsage,
+    SearchResponse,
+)
+from arex.research import Research
+
+__all__ = ["Arex"]
+
+
+def resolve_api_key(api_key: str | None) -> str:
+    """The explicit key, else `AREX_API_KEY`."""
+    key = api_key if api_key is not None else os.environ.get("AREX_API_KEY")
+    if not key:
+        raise ValueError(
+            "An AREX API key is required: pass api_key= or set AREX_API_KEY."
+        )
+    return key
+
+
+def browse_body(
+    url: str,
+    tier: BrowseTier | None,
+    livecrawl: bool | None,
+    livecrawl_timeout_ms: int | None,
+    max_characters: int | None,
+) -> dict[str, Any]:
+    body: dict[str, Any] = {"url": url}
+    if tier is not None:
+        body["tier"] = tier
+    if livecrawl is not None:
+        body["livecrawl"] = livecrawl
+    if livecrawl_timeout_ms is not None:
+        body["livecrawl_timeout_ms"] = livecrawl_timeout_ms
+    if max_characters is not None:
+        body["max_characters"] = max_characters
+    return body
+
+
+def browse_timeout_ms(
+    tier: BrowseTier | None,
+    livecrawl: bool | None,
+    livecrawl_timeout_ms: int | None,
+) -> int:
+    """The server deadline that applies to this browse call, plus the margin."""
+    effective_tier: BrowseTier = tier if tier is not None else DEFAULT_BROWSE_TIER
+    if livecrawl:
+        server_deadline = (
+            livecrawl_timeout_ms
+            if livecrawl_timeout_ms is not None
+            else (
+                DEFAULT_LIVECRAWL_TIMEOUT_AUTO_MS
+                if effective_tier == "auto"
+                else DEFAULT_LIVECRAWL_TIMEOUT_FLEX_MS
+            )
+        )
+        return server_deadline + LIVECRAWL_TIMEOUT_MARGIN_MS
+    return (
+        BROWSE_AUTO_TIMEOUT_MS if effective_tier == "auto" else BROWSE_FLEX_TIMEOUT_MS
+    )
+
+
+def require_credits(credits: Credits | None) -> Credits:
+    if credits is None:
+        raise ArexNetworkError(
+            "invalid_response",
+            "The AREX API response is missing the X-Arex-Credits-Used header.",
+        )
+    return credits
+
+
+class Arex:
+    """Synchronous client for the AREX public API."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        *,
+        base_url: str | None = None,
+        timeout_ms: int | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_unavailable: bool = False,
+        http_client: httpx.Client | None = None,
+    ) -> None:
+        self._transport = Transport(
+            api_key=resolve_api_key(api_key),
+            base_url=base_url or DEFAULT_BASE_URL,
+            timeout_ms=timeout_ms,
+            max_retries=max_retries,
+            retry_unavailable=retry_unavailable,
+            client=http_client,
+        )
+        self.research = Research(self._transport)
+
+    def search(
+        self,
+        query: str,
+        *,
+        limit: int | None = None,
+        idempotency_key: str | None = None,
+        timeout_ms: int | None = None,
+    ) -> SearchResponse:
+        """Run a web search."""
+        validate_search(query, limit)
+        key = resolve_idempotency_key(idempotency_key)
+        body: dict[str, Any] = {"query": query}
+        if limit is not None:
+            body["limit"] = limit
+        result = self._transport.request(
+            "POST",
+            "/v1/search",
+            json=body,
+            idempotency_key=key,
+            timeout_ms=self._transport.resolve_timeout(SEARCH_TIMEOUT_MS, timeout_ms),
+        )
+        return SearchResponse.from_wire(result.data, require_credits(result.credits))
+
+    def browse(
+        self,
+        url: str,
+        *,
+        tier: BrowseTier | None = None,
+        livecrawl: bool | None = None,
+        livecrawl_timeout_ms: int | None = None,
+        max_characters: int | None = None,
+        idempotency_key: str | None = None,
+        timeout_ms: int | None = None,
+    ) -> BrowseDocument:
+        """Fetch and extract one page."""
+        validate_browse(
+            url,
+            tier=tier,
+            livecrawl=livecrawl,
+            livecrawl_timeout_ms=livecrawl_timeout_ms,
+            max_characters=max_characters,
+        )
+        key = resolve_idempotency_key(idempotency_key)
+        default_ms = browse_timeout_ms(tier, livecrawl, livecrawl_timeout_ms)
+        result = self._transport.request(
+            "POST",
+            "/v1/browse",
+            idempotency_key=key,
+            json=browse_body(
+                url, tier, livecrawl, livecrawl_timeout_ms, max_characters
+            ),
+            timeout_ms=self._transport.resolve_timeout(default_ms, timeout_ms),
+        )
+        return BrowseDocument.from_wire(result.data, require_credits(result.credits))
+
+    def usage(self, *, timeout_ms: int | None = None) -> CreditUsage:
+        """Read the current credit windows."""
+        result = self._transport.request(
+            "GET",
+            "/v1/usage",
+            timeout_ms=self._transport.resolve_timeout(USAGE_TIMEOUT_MS, timeout_ms),
+        )
+        return CreditUsage.from_wire(result.data)
+
+    def close(self) -> None:
+        """Close the HTTP client, unless it was injected."""
+        self._transport.close()
+
+    def __enter__(self) -> Arex:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.close()