ads_dev 1.0.0__py3-none-any.whl

ads_dev/_sync/client.py

@@ -0,0 +1,335 @@
+from __future__ import annotations
+
+import math
+import uuid
+from collections.abc import Iterator, Mapping
+from typing import Any, Generic, TypeVar
+from urllib.parse import quote
+
+import time
+import httpx
+from pydantic import BaseModel
+
+from .._core import (
+    LONGEST_WAIT_S,
+    RETRYABLE,
+    RUN_TIMEOUT_S,
+    GenerationFailed,
+    delay_of,
+    headers_of,
+    is_finished,
+    query_of,
+    refusal_of,
+    settings,
+)
+from ..types import (
+    Ad,
+    AdList,
+    AdSummary,
+    Advertiser,
+    AdvertiserList,
+    Asset,
+    Brand,
+    Character,
+    CharacterList,
+    Generation,
+    GenerationList,
+    LibraryAsset,
+    LibraryAssetList,
+    Model,
+    ModelList,
+    Offer,
+    OfferList,
+    OfferSummary,
+    Voice,
+    VoiceList,
+    WebhookSecret,
+)
+
+Item = TypeVar("Item", bound=BaseModel)
+
+Answer = TypeVar("Answer", bound=BaseModel)
+
+
+class Page(Generic[Item]):
+    def __init__(self, client: Sandstorm, path: str, params: Mapping[str, Any], shape: type[BaseModel]) -> None:
+        self._client = client
+        self._path = path
+        self._params = params
+        self._shape = shape
+
+    def first_page(self) -> Any:
+        return self._client._send("GET", self._path, self._shape, query=self._params)
+
+    def __iter__(self) -> Iterator[Item]:
+        return self._items()
+
+    def _items(self) -> Iterator[Item]:
+        cursor: str | None = None
+        while True:
+            page = self._client._send("GET", self._path, self._shape, query={**self._params, "cursor": cursor})
+            for item in page.data:
+                yield item
+            if not page.has_more or page.next_cursor is None:
+                return
+            cursor = page.next_cursor
+
+
+class Generations:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def create(
+        self,
+        model: str,
+        input: Mapping[str, Any],
+        *,
+        webhook_url: str | None = None,
+        wait: int | None = None,
+        idempotency_key: str | None = None,
+    ) -> Generation:
+        body = {"model": model, "input": dict(input), **({"webhook_url": webhook_url} if webhook_url else {})}
+        return self._client._send(
+            "POST", "/v1/generations", Generation, body=body, wait=wait, idempotency_key=idempotency_key or str(uuid.uuid4())
+        )
+
+    def get(self, id: str, *, wait: int | None = None) -> Generation:
+        return self._client._send("GET", f"/v1/generations/{quote(id, safe='')}", Generation, wait=wait)
+
+    def wait(self, id: str, *, timeout: float = RUN_TIMEOUT_S) -> Generation:
+        deadline = time.monotonic() + timeout
+        while True:
+            left = math.ceil(deadline - time.monotonic())
+            generation = self.get(id, wait=max(1, min(LONGEST_WAIT_S, left)))
+            if is_finished(generation):
+                return generation
+            if left <= 0:
+                raise TimeoutError(f"{id} did not finish within {timeout} s; it is still {generation.status}")
+
+    def run(
+        self,
+        model: str,
+        input: Mapping[str, Any],
+        *,
+        webhook_url: str | None = None,
+        idempotency_key: str | None = None,
+        timeout: float = RUN_TIMEOUT_S,
+    ) -> Generation:
+        created = self.create(model, input, webhook_url=webhook_url, wait=LONGEST_WAIT_S, idempotency_key=idempotency_key)
+        generation = created if is_finished(created) else self.wait(created.id, timeout=timeout)
+        if generation.status != "succeeded":
+            raise GenerationFailed(generation)
+        return generation
+
+    def cancel(self, id: str) -> Generation:
+        return self._client._send("POST", f"/v1/generations/{quote(id, safe='')}/cancel", Generation)
+
+    def list(self, *, status: str | None = None, model: str | None = None, limit: int | None = None) -> Page[Generation]:
+        return Page(self._client, "/v1/generations", {"status": status, "model": model, "limit": limit}, GenerationList)
+
+
+class Models:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(self, *, type: str | None = None) -> Page[Model]:
+        return Page(self._client, "/v1/models", {"type": type}, ModelList)
+
+    def get(self, id: str) -> Model:
+        return self._client._send("GET", f"/v1/models/{quote(id, safe='')}", Model)
+
+
+class Assets:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(
+        self,
+        *,
+        type: str | None = None,
+        offer: str | None = None,
+        liked: bool | None = None,
+        archived: bool | None = None,
+        limit: int | None = None,
+    ) -> Page[LibraryAsset]:
+        return Page(self._client, "/v1/assets", {"type": type, "offer": offer, "liked": liked, "archived": archived, "limit": limit}, LibraryAssetList)
+
+    def get(self, id: str) -> Asset:
+        return self._client._send("GET", f"/v1/assets/{quote(id, safe='')}", Asset)
+
+    def import_url(self, url: str, *, name: str | None = None, idempotency_key: str | None = None) -> Asset:
+        body = {"url": url, **({"name": name} if name else {})}
+        return self._client._send("POST", "/v1/assets", Asset, body=body, idempotency_key=idempotency_key or str(uuid.uuid4()))
+
+
+class Brands:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def get(self) -> Brand:
+        return self._client._send("GET", "/v1/brand", Brand)
+
+
+class Offers:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(self, *, query: str | None = None, limit: int | None = None) -> Page[OfferSummary]:
+        return Page(self._client, "/v1/offers", {"query": query, "limit": limit}, OfferList)
+
+    def get(self, id: str) -> Offer:
+        return self._client._send("GET", f"/v1/offers/{quote(id, safe='')}", Offer)
+
+
+class Characters:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(self, *, query: str | None = None, kind: str | None = None) -> Page[Character]:
+        return Page(self._client, "/v1/characters", {"query": query, "kind": kind}, CharacterList)
+
+
+class Voices:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(
+        self,
+        *,
+        query: str | None = None,
+        gender: str | None = None,
+        age: str | None = None,
+        accent: str | None = None,
+        use_case: str | None = None,
+        limit: int | None = None,
+    ) -> Page[Voice]:
+        params = {"query": query, "gender": gender, "age": age, "accent": accent, "use_case": use_case, "limit": limit}
+        return Page(self._client, "/v1/voices", params, VoiceList)
+
+
+class Advertisers:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(self, *, query: str, limit: int | None = None) -> Page[Advertiser]:
+        return Page(self._client, "/v1/advertisers", {"query": query, "limit": limit}, AdvertiserList)
+
+
+class Ads:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def list(
+        self,
+        *,
+        query: str | None = None,
+        advertiser: str | None = None,
+        platform: str | None = None,
+        format: str | None = None,
+        status: str | None = None,
+        min_active_days: int | None = None,
+        started_after: str | None = None,
+        started_before: str | None = None,
+        limit: int | None = None,
+    ) -> Page[AdSummary]:
+        params = {
+            "query": query,
+            "advertiser": advertiser,
+            "platform": platform,
+            "format": format,
+            "status": status,
+            "min_active_days": min_active_days,
+            "started_after": started_after,
+            "started_before": started_before,
+            "limit": limit,
+        }
+        return Page(self._client, "/v1/ads", params, AdList)
+
+    def get(self, id: str) -> Ad:
+        return self._client._send("GET", f"/v1/ads/{quote(id, safe='')}", Ad)
+
+
+class Webhooks:
+    def __init__(self, client: Sandstorm) -> None:
+        self._client = client
+
+    def secret(self) -> str:
+        return (self._client._send("GET", "/v1/webhooks/secret", WebhookSecret)).secret
+
+
+class Sandstorm:
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        brand: str | None = None,
+        max_retries: int = 2,
+        timeout: float = 60,
+        http: httpx.Client | None = None,
+    ) -> None:
+        key, url = settings(api_key, base_url)
+        self._base_url = url
+        self._headers = headers_of(key, brand)
+        self._max_retries = max_retries
+        self._timeout = timeout
+        self._http = http or httpx.Client()
+        self.generations = Generations(self)
+        self.models = Models(self)
+        self.assets = Assets(self)
+        self.brand = Brands(self)
+        self.offers = Offers(self)
+        self.characters = Characters(self)
+        self.voices = Voices(self)
+        self.advertisers = Advertisers(self)
+        self.ads = Ads(self)
+        self.webhooks = Webhooks(self)
+
+    def __enter__(self) -> Sandstorm:
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._http.close()
+
+    def _send(
+        self,
+        method: str,
+        path: str,
+        shape: type[Answer],
+        *,
+        query: Mapping[str, Any] | None = None,
+        body: Mapping[str, Any] | None = None,
+        wait: int | None = None,
+        idempotency_key: str | None = None,
+    ) -> Answer:
+        headers = {
+            **self._headers,
+            **({"prefer": f"wait={wait}"} if wait is not None else {}),
+            **({"idempotency-key": idempotency_key} if idempotency_key is not None else {}),
+        }
+        attempt = 0
+        while True:
+            try:
+                response = self._http.request(
+                    method,
+                    self._base_url + path,
+                    params=query_of(query or {}),
+                    json=body,
+                    headers=headers,
+                    timeout=self._timeout + (wait or 0),
+                )
+            except httpx.TransportError:
+                if attempt >= self._max_retries:
+                    raise
+                time.sleep(delay_of(attempt, None))
+                attempt += 1
+                continue
+            if response.is_success:
+                return shape.model_validate(response.json())
+            if response.status_code not in RETRYABLE or attempt >= self._max_retries:
+                raise refusal_of(response)
+            time.sleep(delay_of(attempt, response))
+            attempt += 1