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

arex/models.py

@@ -0,0 +1,414 @@
+"""Public result models and SDK-owned defaults."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, get_args
+
+from arex import _wire
+from arex._wire import (
+    WIRE_LIMITS,
+    BrowseTier,
+    ResearchFailureCode,
+    ResearchMode,
+    ResearchModeRequest,
+    ResearchStatus,
+)
+
+__all__ = [
+    "API_KEY_PATTERN",
+    "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",
+    "CreditUsage",
+    "CreditWeek",
+    "CreditWindow",
+    "Credits",
+    "ResearchFailureCode",
+    "ResearchMode",
+    "ResearchModeRequest",
+    "ResearchResult",
+    "ResearchRun",
+    "ResearchRunError",
+    "ResearchSource",
+    "ResearchStatus",
+    "SearchResponse",
+    "SearchResult",
+    "is_arex_api_key",
+]
+
+
+# Authentication
+API_KEY_PREFIX = "sk_arex_"
+API_KEY_PATTERN = re.compile(r"^sk_arex_[A-Za-z0-9_-]{32}$")
+DEFAULT_BASE_URL = "https://api.arex-research.com"
+
+DEFAULT_SEARCH_LIMIT = WIRE_LIMITS["SearchRequest"]["limit"]["default"]
+BROWSE_TIERS = get_args(BrowseTier)
+DEFAULT_BROWSE_TIER: BrowseTier = WIRE_LIMITS["BrowseRequest"]["tier"]["default"]
+DEFAULT_LIVECRAWL_TIMEOUT_AUTO_MS = 20000
+DEFAULT_LIVECRAWL_TIMEOUT_FLEX_MS = 90000
+DEFAULT_BROWSE_MAX_CHARACTERS = WIRE_LIMITS["BrowseRequest"]["max_characters"][
+    "default"
+]
+RESEARCH_MODES = get_args(ResearchModeRequest)
+DEFAULT_RESEARCH_MODE: ResearchModeRequest = WIRE_LIMITS["ResearchCreateRequest"][
+    "mode"
+]["default"]
+
+# SDK-owned timeout, polling and retry defaults
+SEARCH_TIMEOUT_MS = 30000
+USAGE_TIMEOUT_MS = 15000
+BROWSE_AUTO_TIMEOUT_MS = 30000
+BROWSE_FLEX_TIMEOUT_MS = 70000
+LIVECRAWL_TIMEOUT_MARGIN_MS = 10000
+RESEARCH_REQUEST_TIMEOUT_MS = 30000
+DEFAULT_POLL_INTERVAL_MS = 2000
+FAST_POLL_TIMEOUT_MS = 150000
+DEEP_POLL_TIMEOUT_MS = 180000
+DEFAULT_MAX_RETRIES = 2
+DEFAULT_RETRY_DELAY_MS = 1000
+
+
+def is_arex_api_key(value: str) -> bool:
+    """Whether `value` has the shape of an AREX API key."""
+    return API_KEY_PATTERN.match(value) is not None
+
+
+def _invalid(field: str) -> Exception:
+    from arex.errors import ArexNetworkError
+
+    return ArexNetworkError(
+        "invalid_response",
+        f"The AREX API response is missing or mistyped the field {field!r}.",
+    )
+
+
+def _object(data: Any, field: str) -> Mapping[str, Any]:
+    if not isinstance(data, Mapping):
+        raise _invalid(field)
+    return data
+
+
+def _str(data: Mapping[str, Any], key: str) -> str:
+    value = data.get(key)
+    if not isinstance(value, str):
+        raise _invalid(key)
+    return value
+
+
+def _opt_str(data: Mapping[str, Any], key: str) -> str | None:
+    value = data.get(key)
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise _invalid(key)
+    return value
+
+
+def _int(data: Mapping[str, Any], key: str) -> int:
+    value = data.get(key)
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise _invalid(key)
+    return value
+
+
+def _nullable_int(data: Mapping[str, Any], key: str) -> int | None:
+    value = data.get(key)
+    if value is None:
+        return None
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise _invalid(key)
+    return value
+
+
+def _bool(data: Mapping[str, Any], key: str) -> bool:
+    value = data.get(key)
+    if not isinstance(value, bool):
+        raise _invalid(key)
+    return value
+
+
+def _literal(data: Mapping[str, Any], key: str, allowed: tuple[str, ...]) -> str:
+    value = data.get(key)
+    if not isinstance(value, str) or value not in allowed:
+        raise _invalid(key)
+    return value
+
+
+@dataclass(frozen=True, slots=True)
+class Credits:
+    """Credit counters reported alongside a charged response."""
+
+    used: int
+    limit: int | None = None
+    remaining: int | None = None
+    reset_at: str | None = None
+
+    @classmethod
+    def from_headers(cls, headers: Mapping[str, str]) -> Credits | None:
+        """Read the `X-Arex-Credits-*` headers, or `None` when they are absent."""
+        lowered = {key.lower(): value for key, value in headers.items()}
+        used = _header_int(lowered, "x-arex-credits-used")
+        if used is None:
+            return None
+        return cls(
+            used=used,
+            limit=_header_int(lowered, "x-arex-credits-limit"),
+            remaining=_header_int(lowered, "x-arex-credits-remaining"),
+            reset_at=lowered.get("x-arex-credits-reset"),
+        )
+
+
+def _header_int(headers: Mapping[str, str], name: str) -> int | None:
+    raw = headers.get(name)
+    if raw is None:
+        return None
+    try:
+        return int(raw.strip())
+    except ValueError:
+        return None
+
+
+@dataclass(frozen=True, slots=True)
+class SearchResult:
+    """One search hit."""
+
+    title: str
+    url: str
+    snippet: str | None = None
+
+    @classmethod
+    def from_wire(cls, data: _wire.SearchResponseResultsItem) -> SearchResult:
+        obj = _object(data, "results[]")
+        return cls(
+            title=_str(obj, "title"),
+            url=_str(obj, "url"),
+            snippet=_opt_str(obj, "snippet"),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class SearchResponse:
+    """The result of `Arex.search`."""
+
+    results: tuple[SearchResult, ...]
+    credits: Credits
+
+    @classmethod
+    def from_wire(cls, data: _wire.SearchResponse, credits: Credits) -> SearchResponse:
+        obj = _object(data, "body")
+        results = obj.get("results")
+        if not isinstance(results, list):
+            raise _invalid("results")
+        return cls(
+            results=tuple(SearchResult.from_wire(item) for item in results),
+            credits=credits,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class BrowseDocument:
+    """The result of `Arex.browse`."""
+
+    url: str
+    title: str | None
+    published_at: str | None
+    image_url: str | None
+    image_urls: tuple[str, ...] | None
+    content: str
+    truncated: bool
+    credits: Credits
+
+    @classmethod
+    def from_wire(cls, data: _wire.BrowseResponse, credits: Credits) -> BrowseDocument:
+        obj = _object(data, "body")
+        raw_urls = obj.get("image_urls")
+        if raw_urls is None:
+            image_urls = None
+        elif (
+            isinstance(raw_urls, list)
+            and len(raw_urls) <= WIRE_LIMITS["BrowseResponse"]["image_urls"]["maxItems"]
+            and all(isinstance(item, str) for item in raw_urls)
+        ):
+            image_urls = tuple(raw_urls)
+        else:
+            raise _invalid("image_urls")
+        return cls(
+            url=_str(obj, "url"),
+            title=_opt_str(obj, "title"),
+            published_at=_opt_str(obj, "published_at"),
+            image_url=_opt_str(obj, "image_url"),
+            image_urls=image_urls,
+            content=_str(obj, "content"),
+            truncated=_bool(obj, "truncated"),
+            credits=credits,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class CreditWindow:
+    """Credit consumption inside one accounting window."""
+
+    used: int
+    limit: int | None
+    remaining: int | None
+    reset_at: str
+
+
+@dataclass(frozen=True, slots=True)
+class CreditWeek(CreditWindow):
+    """The weekly credit window, which also reports when it started."""
+
+    start: str
+
+    @classmethod
+    def from_wire(cls, data: _wire.UsageResponseWeek) -> CreditWeek:
+        obj = _object(data, "week")
+        return cls(
+            used=_int(obj, "used"),
+            limit=_nullable_int(obj, "limit"),
+            remaining=_nullable_int(obj, "remaining"),
+            reset_at=_str(obj, "reset_at"),
+            start=_str(obj, "start"),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class CreditUsage:
+    """The result of `Arex.usage`: the daily window plus the weekly one."""
+
+    date: str
+    used: int
+    limit: int | None
+    remaining: int | None
+    reset_at: str
+    week: CreditWeek | None = None
+
+    @classmethod
+    def from_wire(cls, data: _wire.UsageResponse) -> CreditUsage:
+        obj = _object(data, "body")
+        week = obj.get("week")
+        return cls(
+            date=_str(obj, "date"),
+            used=_int(obj, "used"),
+            limit=_nullable_int(obj, "limit"),
+            remaining=_nullable_int(obj, "remaining"),
+            reset_at=_str(obj, "reset_at"),
+            week=None if week is None else CreditWeek.from_wire(week),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class ResearchSource:
+    """One source cited by a research answer."""
+
+    title: str
+    url: str
+
+    @classmethod
+    def from_wire(cls, data: _wire.ResearchRunResultSourcesItem) -> ResearchSource:
+        obj = _object(data, "result.sources[]")
+        return cls(title=_str(obj, "title"), url=_str(obj, "url"))
+
+
+@dataclass(frozen=True, slots=True)
+class ResearchResult:
+    """The answer of a completed research run."""
+
+    answer: str
+    sources: tuple[ResearchSource, ...]
+
+    @classmethod
+    def from_wire(cls, data: _wire.ResearchRunResult) -> ResearchResult:
+        obj = _object(data, "result")
+        sources = obj.get("sources")
+        if not isinstance(sources, list):
+            raise _invalid("result.sources")
+        return cls(
+            answer=_str(obj, "answer"),
+            sources=tuple(ResearchSource.from_wire(item) for item in sources),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class ResearchRunError:
+    """Why a research run failed."""
+
+    code: ResearchFailureCode
+    retryable: bool
+
+    @classmethod
+    def from_wire(cls, data: _wire.ResearchRunError) -> ResearchRunError:
+        obj = _object(data, "error")
+        return cls(
+            code=_literal(obj, "code", get_args(ResearchFailureCode)),  # type: ignore[arg-type]
+            retryable=_bool(obj, "retryable"),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class ResearchRun:
+    """A research run at one point in its lifecycle."""
+
+    id: str
+    status: ResearchStatus
+    mode: ResearchMode
+    requested_mode: ResearchModeRequest
+    created_at: str
+    started_at: str | None = None
+    completed_at: str | None = None
+    result: ResearchResult | None = None
+    error: ResearchRunError | None = None
+    credits: Credits | None = None
+
+    @classmethod
+    def from_wire(
+        cls, data: _wire.ResearchRun, credits: Credits | None = None
+    ) -> ResearchRun:
+        obj = _object(data, "body")
+        status = _literal(obj, "status", get_args(ResearchStatus))
+        mode = _literal(obj, "mode", get_args(ResearchMode))
+        # Payloads are meaningful only in their matching terminal state.
+        result = obj.get("result") if status == "completed" else None
+        error = obj.get("error") if status == "failed" else None
+        if status == "completed" and result is None:
+            raise _invalid("result")
+        if status == "failed" and error is None:
+            raise _invalid("error")
+        return cls(
+            id=_str(obj, "id"),
+            status=status,  # type: ignore[arg-type]
+            mode=mode,  # type: ignore[arg-type]
+            requested_mode=_literal(
+                obj, "requested_mode", get_args(ResearchModeRequest)
+            ),  # type: ignore[arg-type]
+            created_at=_str(obj, "created_at"),
+            started_at=_opt_str(obj, "started_at"),
+            completed_at=_opt_str(obj, "completed_at"),
+            result=None if result is None else ResearchResult.from_wire(result),
+            error=None if error is None else ResearchRunError.from_wire(error),
+            credits=credits,
+        )