sume-client 0.1.0__py3-none-any.whl

sume_client/__init__.py

@@ -0,0 +1,74 @@
+from sume_client.exceptions import (
+    CredentialsMissedError,
+    SumeClientError,
+)
+from sume_client.http.client import (
+    AsyncClient,
+    AsyncRequestController,
+    DONE_STATUSES,
+    SyncClient,
+    SyncRequestController,
+)
+from sume_client.types_ import (
+    Cancelled,
+    Completed,
+    Failed,
+    InProgress,
+    NSFW,
+    Queued,
+    Status,
+)
+
+sync_client = SyncClient()
+submit = sync_client.submit
+subscribe = sync_client.subscribe
+status = sync_client.status
+result = sync_client.result
+cancel = sync_client.cancel
+upload = sync_client.upload
+upload_file = sync_client.upload_file
+upload_image = sync_client.upload_image
+
+async_client = AsyncClient()
+submit_async = async_client.submit
+subscribe_async = async_client.subscribe
+status_async = async_client.status
+result_async = async_client.result
+cancel_async = async_client.cancel
+upload_async = async_client.upload
+upload_file_async = async_client.upload_file
+upload_image_async = async_client.upload_image
+
+
+__all__ = [
+    'CredentialsMissedError',
+    'SumeClientError',
+    'SyncRequestController',
+    'AsyncRequestController',
+    'SyncClient',
+    'AsyncClient',
+    'DONE_STATUSES',
+    'Queued',
+    'InProgress',
+    'Completed',
+    'Status',
+    'Failed',
+    'Cancelled',
+    'NSFW',
+    'submit',
+    'subscribe',
+    'status',
+    'result',
+    'cancel',
+    'upload',
+    'upload_file',
+    'upload_image',
+    'submit_async',
+    'subscribe_async',
+    'status_async',
+    'result_async',
+    'cancel_async',
+    'upload_async',
+    'upload_file_async',
+    'upload_image_async',
+]

sume_client/auth.py

@@ -0,0 +1,15 @@
+import os
+
+from sume_client.exceptions import CredentialsMissedError
+
+
+def get_credential_key() -> str:
+    key = os.getenv('SUME_API_KEY', '').strip()
+
+    if key:
+        return key
+
+    raise CredentialsMissedError(
+        'Sume API credentials missing. Please set SUME_API_KEY to a key '
+        'from https://open.sume.com/api-keys.'
+    )

sume_client/exceptions.py

@@ -0,0 +1,22 @@
+from typing import Any, Optional
+
+
+class SumeClientError(Exception):
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: Optional[int] = None,
+        code: Optional[str] = None,
+        request_id: Optional[str] = None,
+        details: Any = None,
+    ) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+        self.code = code
+        self.request_id = request_id
+        self.details = details
+
+
+class CredentialsMissedError(Exception):
+    pass

sume_client/http/client.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+import uuid
+import weakref
+from functools import cached_property
+from typing import AsyncIterator, Callable, Dict, Iterator, Optional, Tuple, TYPE_CHECKING
+from urllib.parse import quote
+
+import httpx
+
+from sume_client.auth import get_credential_key
+from sume_client.http.error import to_storage_error
+from sume_client.http.retry import create_default_retry_strategy
+from sume_client.http.transport import AsyncHttpTransport, HttpTransport
+from sume_client.mixins import UploadMixin
+from sume_client.types_ import (
+    AnyJSON,
+    Cancelled,
+    Completed,
+    Failed,
+    NSFW,
+    Status,
+)
+from sume_client.utils import get_arguments, get_catalog_url, get_router_url, get_status
+
+if TYPE_CHECKING:
+    from PIL import Image
+
+
+DEFAULT_TIMEOUT = 90.0
+DEFAULT_POLLING_DELAY = 2.0
+
+BASE_URL = 'https://api.open.sume.com'
+USER_AGENT = 'sume-client-py/0.1.0'
+DONE_STATUSES = (Completed, NSFW, Cancelled, Failed)
+UPLOAD_URL = '/v1/files/generate-upload-url'
+
+
+class _NoLoop:
+    pass
+
+
+_NO_LOOP = _NoLoop()
+
+
+def _running_loop() -> object:
+    try:
+        return asyncio.get_running_loop()
+    except RuntimeError:
+        return _NO_LOOP
+
+
+class cached_per_loop:
+    def __init__(self, func: Callable) -> None:
+        self.func = func
+        self.key = f'_{func.__name__}_per_loop'
+
+    def __get__(self, instance, owner=None):
+        if instance is None:
+            return self
+
+        values = instance.__dict__.setdefault(self.key, weakref.WeakKeyDictionary())
+        loop = _running_loop()
+        value = values.get(loop)
+
+        if value is None:
+            value = values.setdefault(loop, self.func(instance))
+
+        return value
+
+
+class RequestController:
+    def __init__(
+        self,
+        request_id: str,
+        response_url: Optional[str] = None,
+        status_url: Optional[str] = None,
+        cancel_url: Optional[str] = None,
+    ) -> None:
+        self.request_id = request_id
+
+        urls = self.build_urls(BASE_URL, request_id)
+
+        self.response_url = response_url or urls['response_url']
+        self.status_url = status_url or urls['status_url']
+        self.cancel_url = cancel_url or urls['cancel_url']
+
+    @staticmethod
+    def build_urls(base_url: str, request_id: str) -> Dict[str, str]:
+        if request_id in ('', '.', '..'):
+            raise ValueError(f'{request_id!r} is not a request id')
+
+        request_base = f"{base_url.rstrip('/')}/v1/jobs/{quote(request_id, safe='')}"
+        return {
+            'status_url': f'{request_base}/status',
+            'response_url': f'{request_base}/result',
+            'cancel_url': f'{request_base}/cancel',
+        }
+
+
+class SyncRequestController(RequestController):
+    def __init__(
+        self,
+        request_id: str,
+        transport: HttpTransport,
+        response_url: Optional[str] = None,
+        status_url: Optional[str] = None,
+        cancel_url: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            request_id=request_id,
+            response_url=response_url,
+            status_url=status_url,
+            cancel_url=cancel_url,
+        )
+        self._transport = transport
+
+    def cancel(self) -> None:
+        self._transport.request('POST', self.cancel_url)
+
+    def status(self) -> Status:
+        return get_status(self._get_status_data())
+
+    def _get_status_data(self) -> AnyJSON:
+        response = self._transport.request('GET', self.status_url)
+        return response.json()['data']
+
+    def _poll(self, delay: float) -> Iterator[Tuple[Status, AnyJSON]]:
+        while True:
+            data = self._get_status_data()
+            status = get_status(data)
+            yield status, data
+
+            if isinstance(status, DONE_STATUSES):
+                break
+
+            time.sleep(delay)
+
+    def poll_request_status(self, *, delay: float = DEFAULT_POLLING_DELAY) -> Iterator[Status]:
+        for status, _ in self._poll(delay):
+            yield status
+
+    def get(self) -> AnyJSON:
+        for status, data in self._poll(DEFAULT_POLLING_DELAY):
+            continue
+
+        if not isinstance(status, Completed):
+            return data
+
+        response = self._transport.request('GET', self.response_url)
+        return response.json()['data']
+
+    def __repr__(self):
+        return f'SyncRequestController(request_id={self.request_id})'
+
+
+class AsyncRequestController(RequestController):
+    def __init__(
+        self,
+        request_id: str,
+        transport: AsyncHttpTransport,
+        response_url: Optional[str] = None,
+        status_url: Optional[str] = None,
+        cancel_url: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            request_id=request_id,
+            response_url=response_url,
+            status_url=status_url,
+            cancel_url=cancel_url,
+        )
+        self._transport = transport
+
+    async def cancel(self) -> None:
+        await self._transport.request('POST', self.cancel_url)
+
+    async def status(self) -> Status:
+        return get_status(await self._get_status_data())
+
+    async def _get_status_data(self) -> AnyJSON:
+        response = await self._transport.request('GET', self.status_url)
+        return response.json()['data']
+
+    async def _poll(self, delay: float) -> AsyncIterator[Tuple[Status, AnyJSON]]:
+        while True:
+            data = await self._get_status_data()
+            status = get_status(data)
+            yield status, data
+
+            if isinstance(status, DONE_STATUSES):
+                break
+
+            await asyncio.sleep(delay)
+
+    async def poll_request_status(
+        self,
+        *,
+        delay: float = DEFAULT_POLLING_DELAY,
+    ) -> AsyncIterator[Status]:
+        async for status, _ in self._poll(delay):
+            yield status
+
+    async def get(self) -> AnyJSON:
+        async for status, data in self._poll(DEFAULT_POLLING_DELAY):
+            continue
+
+        if not isinstance(status, Completed):
+            return data
+
+        response = await self._transport.request('GET', self.response_url)
+        return response.json()['data']
+
+    def __repr__(self):
+        return f'AsyncRequestController(request_id={self.request_id})'
+
+
+class SyncClient(UploadMixin):
+    def __init__(
+        self,
+        base_url: str = BASE_URL,
+        timeout: float = DEFAULT_TIMEOUT,
+        api_key: str | None = None,
+    ):
+        self.base_url = base_url
+        self.timeout = timeout
+        self.api_key = api_key
+        self._model_kinds: Dict[str, str] = {}
+
+    @cached_property
+    def _client(self) -> httpx.Client:
+        api_key = self.api_key or get_credential_key()
+
+        return httpx.Client(
+            headers={
+                'Authorization': f'Bearer {api_key}',
+                'Content-Type': 'application/json',
+                'User-Agent': USER_AGENT,
+            },
+            timeout=self.timeout,
+            base_url=self.base_url,
+        )
+
+    @cached_property
+    def _upload_client(self) -> httpx.Client:
+        return httpx.Client(timeout=self.timeout)
+
+    @cached_property
+    def _transport(self) -> HttpTransport:
+        retry_strategy = create_default_retry_strategy()
+        return HttpTransport(self._client, retry_strategy)
+
+    def get_request_controller(self, request_id: str) -> SyncRequestController:
+        urls = RequestController.build_urls(self.base_url, request_id)
+        return SyncRequestController(request_id, self._transport, **urls)
+
+    def _get_model_kind(self, application: str) -> str:
+        if application not in self._model_kinds:
+            response = self._transport.request('GET', get_catalog_url(application))
+            self._model_kinds[application] = response.json()['data']['kind']
+
+        return self._model_kinds[application]
+
+    def submit(
+        self,
+        application: str,
+        arguments: AnyJSON,
+        *,
+        webhook_url: str | None = None,
+    ) -> SyncRequestController:
+        url = get_router_url(application, self._get_model_kind(application))
+
+        response = self._transport.request(
+            'POST',
+            url,
+            json=get_arguments(application, arguments, webhook_url),
+            headers={'Idempotency-Key': str(uuid.uuid4())},
+            timeout=self.timeout,
+        )
+
+        request_id = response.json()['data']['request_id']
+
+        return self.get_request_controller(request_id)
+
+    def subscribe(
+        self,
+        application: str,
+        arguments: AnyJSON,
+        *,
+        on_enqueue: Optional[Callable[[str], None]] = None,
+        on_queue_update: Optional[Callable[[Status], None]] = None,
+    ) -> AnyJSON:
+        request_controller = self.submit(application=application, arguments=arguments)
+
+        if on_enqueue is not None:
+            on_enqueue(request_controller.request_id)
+
+        if on_queue_update is not None:
+            for status in request_controller.poll_request_status():
+                on_queue_update(status)
+
+        return request_controller.get()
+
+    def status(self, request_id: str) -> Status:
+        request_controller = self.get_request_controller(request_id)
+        return request_controller.status()
+
+    def result(self, request_id: str) -> AnyJSON:
+        request_controller = self.get_request_controller(request_id)
+        return request_controller.get()
+
+    def cancel(self, request_id: str) -> None:
+        request_controller = self.get_request_controller(request_id)
+        request_controller.cancel()
+
+    def _get_upload_url(self, content_type: str, size_bytes: int) -> Tuple[str, str, Dict[str, str]]:
+        response = self._transport.request(
+            'POST',
+            UPLOAD_URL,
+            json={'content_type': content_type, 'size_bytes': size_bytes},
+        )
+        data = response.json()['data']
+
+        upload_headers = data.get('upload_headers') or {'Content-Type': content_type}
+        return data['public_url'], data['upload_url'], upload_headers
+
+    def upload(
+        self,
+        data: bytes | str,
+        content_type: str,
+    ) -> str:
+        data = self.ensure_bytes(data)
+
+        public_url, upload_url, upload_headers = self._get_upload_url(content_type, len(data))
+
+        response = self._upload_client.put(
+            upload_url,
+            content=data,
+            headers=upload_headers,
+        )
+        if response.is_error:
+            raise to_storage_error(response) from None
+
+        return public_url
+
+    def upload_file(self, path: os.PathLike) -> str:
+        mime_type = self.guess_mime_type(path)
+
+        with open(path, 'rb') as file:
+            return self.upload(file.read(), mime_type)
+
+    def upload_image(self, image: Image.Image, format: str = 'jpeg') -> str:
+        image_bytes = self.image_to_bytes(image, format)
+        return self.upload(image_bytes, f'image/{format}')
+
+
+class AsyncClient(UploadMixin):
+    def __init__(
+        self,
+        base_url: str = BASE_URL,
+        timeout: float = DEFAULT_TIMEOUT,
+        api_key: str | None = None,
+    ):
+        self.base_url = base_url
+        self.timeout = timeout
+        self.api_key = api_key
+        self._model_kinds: Dict[str, str] = {}
+
+    @cached_per_loop
+    def _client(self) -> httpx.AsyncClient:
+        api_key = self.api_key or get_credential_key()
+
+        return httpx.AsyncClient(
+            headers={
+                'Authorization': f'Bearer {api_key}',
+                'Content-Type': 'application/json',
+                'User-Agent': USER_AGENT,
+            },
+            timeout=self.timeout,
+            base_url=self.base_url,
+        )
+
+    @cached_per_loop
+    def _upload_client(self) -> httpx.AsyncClient:
+        return httpx.AsyncClient(timeout=self.timeout)
+
+    @cached_property
+    def _transport(self) -> AsyncHttpTransport:
+        retry_strategy = create_default_retry_strategy()
+        return AsyncHttpTransport(lambda: self._client, retry_strategy)
+
+    def get_request_controller(self, request_id: str) -> AsyncRequestController:
+        urls = RequestController.build_urls(self.base_url, request_id)
+        return AsyncRequestController(request_id, self._transport, **urls)
+
+    async def _get_model_kind(self, application: str) -> str:
+        if application not in self._model_kinds:
+            response = await self._transport.request('GET', get_catalog_url(application))
+            self._model_kinds[application] = response.json()['data']['kind']
+
+        return self._model_kinds[application]
+
+    async def submit(
+        self,
+        application: str,
+        arguments: AnyJSON,
+        *,
+        webhook_url: str | None = None,
+    ) -> AsyncRequestController:
+        url = get_router_url(application, await self._get_model_kind(application))
+
+        response = await self._transport.request(
+            'POST',
+            url,
+            json=get_arguments(application, arguments, webhook_url),
+            headers={'Idempotency-Key': str(uuid.uuid4())},
+            timeout=self.timeout,
+        )
+
+        request_id = response.json()['data']['request_id']
+
+        return self.get_request_controller(request_id)
+
+    async def subscribe(
+        self,
+        application: str,
+        arguments: AnyJSON,
+        *,
+        on_enqueue: Optional[Callable[[str], None]] = None,
+        on_queue_update: Optional[Callable[[Status], None]] = None,
+    ) -> AnyJSON:
+        request_controller = await self.submit(application=application, arguments=arguments)
+
+        if on_enqueue is not None:
+            on_enqueue(request_controller.request_id)
+
+        if on_queue_update is not None:
+            async for status in request_controller.poll_request_status():
+                on_queue_update(status)
+
+        return await request_controller.get()
+
+    async def status(self, request_id: str) -> Status:
+        request_controller = self.get_request_controller(request_id)
+        return await request_controller.status()
+
+    async def result(self, request_id: str) -> AnyJSON:
+        request_controller = self.get_request_controller(request_id)
+        return await request_controller.get()
+
+    async def cancel(self, request_id: str) -> None:
+        request_controller = self.get_request_controller(request_id)
+        await request_controller.cancel()
+
+    async def _get_upload_url(self, content_type: str, size_bytes: int) -> Tuple[str, str, Dict[str, str]]:
+        response = await self._transport.request(
+            'POST',
+            UPLOAD_URL,
+            json={'content_type': content_type, 'size_bytes': size_bytes},
+        )
+        data = response.json()['data']
+
+        upload_headers = data.get('upload_headers') or {'Content-Type': content_type}
+        return data['public_url'], data['upload_url'], upload_headers
+
+    async def upload(
+        self,
+        data: bytes | str,
+        content_type: str,
+    ) -> str:
+        data = self.ensure_bytes(data)
+
+        public_url, upload_url, upload_headers = await self._get_upload_url(content_type, len(data))
+
+        response = await self._upload_client.put(
+            upload_url,
+            content=data,
+            headers=upload_headers,
+        )
+        if response.is_error:
+            raise to_storage_error(response) from None
+
+        return public_url
+
+    async def upload_file(self, path: os.PathLike) -> str:
+        mime_type = self.guess_mime_type(path)
+
+        with open(path, 'rb') as file:
+            return await self.upload(file.read(), mime_type)
+
+    async def upload_image(self, image: Image.Image, format: str = 'jpeg') -> str:
+        image_bytes = self.image_to_bytes(image, format)
+        return await self.upload(image_bytes, f'image/{format}')

sume_client/http/error.py

@@ -0,0 +1,91 @@
+import re
+from typing import Any, Dict, List
+
+import httpx
+
+from sume_client.exceptions import SumeClientError
+
+STORAGE_ERROR_FIELD = re.compile(r'<(Code|Message)>([^<]{1,500})</\1>')
+
+
+def read_json(response: httpx.Response) -> Any:
+    try:
+        return response.json()
+    except ValueError:
+        return None
+
+
+def read_error(response: httpx.Response) -> Dict[str, Any]:
+    data = read_json(response)
+    error = data.get('error') if isinstance(data, dict) else None
+    return error if isinstance(error, dict) else {}
+
+
+def field_errors(error: Dict[str, Any]) -> List[str]:
+    details = error.get('details')
+    items = details.get('errors') if isinstance(details, dict) else None
+
+    if not isinstance(items, list):
+        return []
+
+    messages = []
+    for item in items:
+        message = item.get('message') if isinstance(item, dict) else None
+        if not isinstance(message, str) or not message:
+            continue
+        path = item.get('path')
+        messages.append(f'{path}: {message}' if isinstance(path, str) and path else message)
+
+    return messages
+
+
+def extract_error_message(response: httpx.Response) -> str:
+    data = read_json(response)
+    fallback = response.text or f'HTTP {response.status_code}'
+
+    if not isinstance(data, dict):
+        return fallback
+
+    error = data.get('error')
+    if isinstance(error, dict):
+        message = error.get('message') or error.get('code') or fallback
+        fields = field_errors(error)
+        return f"{message} {'; '.join(fields)}" if fields else message
+
+    for key in ('message', 'detail', 'error'):
+        value = data.get(key)
+        if isinstance(value, str) and value:
+            return value
+
+    return fallback
+
+
+def to_client_error(response: httpx.Response) -> SumeClientError:
+    error = read_error(response)
+    return SumeClientError(
+        extract_error_message(response),
+        status_code=response.status_code,
+        code=error.get('code'),
+        request_id=error.get('request_id') or response.headers.get('x-sume-request-id'),
+        details=error.get('details'),
+    )
+
+
+def raise_for_status(response: httpx.Response) -> None:
+    try:
+        response.raise_for_status()
+    except httpx.HTTPStatusError as exception:
+        raise to_client_error(response) from exception
+
+
+def to_storage_error(response: httpx.Response) -> SumeClientError:
+    fields = dict(STORAGE_ERROR_FIELD.findall(response.text))
+    code = fields.get('Code')
+    message = f'Upload to storage failed with HTTP {response.status_code}'
+
+    if code:
+        message += f': {code}'
+    if fields.get('Message'):
+        message += f" ({fields['Message']})"
+
+    return SumeClientError(message, status_code=response.status_code, code=code)