sub2api 0.1.0__py3-none-any.whl

sub2api/_resources.py

@@ -0,0 +1,574 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Iterator, Mapping, Sequence
+from datetime import date, datetime
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
+
+from ._exceptions import ProtocolError
+from ._models import (
+    Announcement,
+    APIKey,
+    Balance,
+    Group,
+    KeyGroupMultiplier,
+    Page,
+    PlatformQuota,
+    Redemption,
+    Resource,
+    Subscription,
+    UsageRecord,
+    User,
+)
+
+if TYPE_CHECKING:
+    from ._client import Sub2API
+
+
+R = TypeVar("R", bound=Resource)
+_UNSET = object()
+
+
+def _mapping(data: Any, operation: str) -> Mapping[str, Any]:
+    if not isinstance(data, Mapping):
+        raise ProtocolError(f"{operation} returned {type(data).__name__}, expected an object")
+    return data
+
+
+def _list(data: Any, operation: str) -> list[Any]:
+    if not isinstance(data, list):
+        raise ProtocolError(f"{operation} returned {type(data).__name__}, expected a list")
+    return data
+
+
+def _page(data: Any, item_type: type[R], operation: str) -> Page[R]:
+    payload = _mapping(data, operation)
+    try:
+        return Page.from_data(payload, item_type)
+    except (TypeError, ValueError) as error:
+        raise ProtocolError(f"{operation} returned invalid pagination data") from error
+
+
+def _positive(value: int, name: str) -> int:
+    if isinstance(value, bool) or value < 1:
+        raise ValueError(f"{name} must be at least 1")
+    return value
+
+
+def _nonnegative(value: int | float, name: str) -> int | float:
+    if isinstance(value, bool) or not math.isfinite(value) or value < 0:
+        raise ValueError(f"{name} must be a finite non-negative number")
+    return value
+
+
+def _date(value: date | str | None) -> str | None:
+    if value is None:
+        return None
+    return value.isoformat() if isinstance(value, date) else value
+
+
+def _datetime(value: datetime | str) -> str:
+    return value.isoformat() if isinstance(value, datetime) else value
+
+
+def _string_list(value: Sequence[str], name: str) -> list[str]:
+    if isinstance(value, (str, bytes)) or any(not isinstance(item, str) for item in value):
+        raise ValueError(f"{name} must be a sequence of strings")
+    return list(value)
+
+
+class AccountResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def profile(self) -> User:
+        return User(_mapping(self._client.request("GET", "auth/me"), "profile"))
+
+    def balance(self) -> Balance:
+        profile = self.profile()
+        try:
+            return Balance(
+                balance=Decimal(str(profile["balance"])),
+                frozen_balance=Decimal(str(profile.get("frozen_balance", 0))),
+                total_recharged=Decimal(str(profile.get("total_recharged", 0))),
+            )
+        except (KeyError, ValueError, TypeError) as error:
+            raise ProtocolError("profile response did not contain valid balance values") from error
+
+    def platform_quotas(self) -> tuple[PlatformQuota, ...]:
+        data = _mapping(
+            self._client.request("GET", "user/platform-quotas"),
+            "platform quotas",
+        )
+        quotas = _list(data.get("platform_quotas", []), "platform quotas")
+        return tuple(PlatformQuota(item) for item in quotas if isinstance(item, Mapping))
+
+
+class KeysResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def __call__(self, **kwargs: Any) -> Page[APIKey]:
+        return self.list(**kwargs)
+
+    def list(
+        self,
+        *,
+        page: int = 1,
+        page_size: int = 20,
+        search: str | None = None,
+        status: str | None = None,
+        group_id: int | None = None,
+        sort_by: str | None = None,
+        sort_order: Literal["asc", "desc"] | None = None,
+    ) -> Page[APIKey]:
+        params: dict[str, Any] = {
+            "page": _positive(page, "page"),
+            "page_size": _positive(page_size, "page_size"),
+            "search": search,
+            "status": status,
+            "group_id": group_id,
+            "sort_by": sort_by,
+            "sort_order": sort_order,
+        }
+        return _page(self._client.request("GET", "keys", params=params), APIKey, "keys")
+
+    def iter(self, **kwargs: Any) -> Iterator[APIKey]:
+        page_number = int(kwargs.pop("page", 1))
+        while True:
+            result = self.list(page=page_number, **kwargs)
+            yield from result
+            if not result.has_next:
+                return
+            page_number += 1
+
+    def all(self, *, page_size: int = 200, **filters: Any) -> tuple[APIKey, ...]:
+        return tuple(self.iter(page_size=page_size, **filters))
+
+    def with_group_multipliers(
+        self,
+        *,
+        page_size: int = 200,
+        **filters: Any,
+    ) -> tuple[KeyGroupMultiplier, ...]:
+        groups: dict[int, Group] = {}
+        for available_group in self._client.groups.list():
+            try:
+                groups[int(available_group["id"])] = available_group
+            except (KeyError, TypeError, ValueError) as error:
+                raise ProtocolError("group contained an invalid ID") from error
+        custom_rates = self._client.groups.rates()
+        resolved: list[KeyGroupMultiplier] = []
+        for api_key in self.all(page_size=page_size, **filters):
+            raw_group_id = api_key.get("group_id")
+            try:
+                group_id = int(raw_group_id) if raw_group_id is not None else None
+            except (TypeError, ValueError) as error:
+                raise ProtocolError("API key contained an invalid group ID") from error
+            group = groups.get(group_id) if group_id is not None else None
+            embedded_group = api_key.get("group")
+            if group is None and isinstance(embedded_group, Mapping):
+                group = Group(embedded_group)
+            base_multiplier: float | None = None
+            if group is not None and group.get("rate_multiplier") is not None:
+                try:
+                    base_multiplier = float(group["rate_multiplier"])
+                except (TypeError, ValueError) as error:
+                    raise ProtocolError("group contained an invalid rate multiplier") from error
+            custom_multiplier = custom_rates.get(group_id) if group_id is not None else None
+            effective_multiplier = (
+                custom_multiplier if custom_multiplier is not None else base_multiplier
+            )
+            resolved.append(
+                KeyGroupMultiplier(
+                    api_key=api_key,
+                    group=group,
+                    group_id=group_id,
+                    base_multiplier=base_multiplier,
+                    custom_multiplier=custom_multiplier,
+                    effective_multiplier=effective_multiplier,
+                )
+            )
+        return tuple(resolved)
+
+    def multiplier_map(
+        self,
+        *,
+        key_by: Literal["key", "id", "name"] = "key",
+        page_size: int = 200,
+        **filters: Any,
+    ) -> dict[str | int, float | None]:
+        result: dict[str | int, float | None] = {}
+        for item in self.with_group_multipliers(page_size=page_size, **filters):
+            value = item.api_key.get(key_by)
+            if not isinstance(value, (str, int)) or isinstance(value, bool):
+                raise ProtocolError(f"API key did not contain a valid {key_by!r} field")
+            if value in result:
+                raise ValueError(f"API key {key_by!r} values are not unique: {value!r}")
+            result[value] = item.effective_multiplier
+        return result
+
+    def get(self, key_id: int) -> APIKey:
+        key_id = _positive(key_id, "key_id")
+        return APIKey(_mapping(self._client.request("GET", f"keys/{key_id}"), "API key"))
+
+    def create(
+        self,
+        name: str,
+        *,
+        group_id: int | None = None,
+        custom_key: str | None = None,
+        ip_whitelist: Sequence[str] | None = None,
+        ip_blacklist: Sequence[str] | None = None,
+        quota: int | float | None = None,
+        expires_in_days: int | None = None,
+        rate_limit_5h: int | float | None = None,
+        rate_limit_1d: int | float | None = None,
+        rate_limit_7d: int | float | None = None,
+        idempotency_key: str | None = None,
+    ) -> APIKey:
+        if not name.strip():
+            raise ValueError("name must not be empty")
+        if group_id is not None:
+            _positive(group_id, "group_id")
+        payload: dict[str, Any] = {"name": name}
+        optional = {
+            "group_id": group_id,
+            "custom_key": custom_key,
+            "ip_whitelist": (
+                _string_list(ip_whitelist, "ip_whitelist") if ip_whitelist is not None else None
+            ),
+            "ip_blacklist": (
+                _string_list(ip_blacklist, "ip_blacklist") if ip_blacklist is not None else None
+            ),
+            "quota": _nonnegative(quota, "quota") if quota is not None else None,
+            "expires_in_days": (
+                _positive(expires_in_days, "expires_in_days")
+                if expires_in_days is not None
+                else None
+            ),
+            "rate_limit_5h": (
+                _nonnegative(rate_limit_5h, "rate_limit_5h") if rate_limit_5h is not None else None
+            ),
+            "rate_limit_1d": (
+                _nonnegative(rate_limit_1d, "rate_limit_1d") if rate_limit_1d is not None else None
+            ),
+            "rate_limit_7d": (
+                _nonnegative(rate_limit_7d, "rate_limit_7d") if rate_limit_7d is not None else None
+            ),
+        }
+        payload.update({key: value for key, value in optional.items() if value is not None})
+        headers = {"Idempotency-Key": idempotency_key} if idempotency_key else None
+        data = self._client.request("POST", "keys", json=payload, headers=headers)
+        return APIKey(_mapping(data, "created API key"))
+
+    def update(
+        self,
+        key_id: int,
+        *,
+        name: Any = _UNSET,
+        group_id: Any = _UNSET,
+        status: Any = _UNSET,
+        ip_whitelist: Any = _UNSET,
+        ip_blacklist: Any = _UNSET,
+        quota: Any = _UNSET,
+        expires_at: Any = _UNSET,
+        clear_expiration: bool = False,
+        reset_quota: Any = _UNSET,
+        rate_limit_5h: Any = _UNSET,
+        rate_limit_1d: Any = _UNSET,
+        rate_limit_7d: Any = _UNSET,
+        reset_rate_limit_usage: Any = _UNSET,
+    ) -> APIKey:
+        key_id = _positive(key_id, "key_id")
+        values = {
+            "name": name,
+            "group_id": group_id,
+            "status": status,
+            "ip_whitelist": ip_whitelist,
+            "ip_blacklist": ip_blacklist,
+            "quota": quota,
+            "expires_at": expires_at,
+            "reset_quota": reset_quota,
+            "rate_limit_5h": rate_limit_5h,
+            "rate_limit_1d": rate_limit_1d,
+            "rate_limit_7d": rate_limit_7d,
+            "reset_rate_limit_usage": reset_rate_limit_usage,
+        }
+        payload = {key: value for key, value in values.items() if value is not _UNSET}
+        if clear_expiration:
+            if expires_at is not _UNSET:
+                raise ValueError("expires_at and clear_expiration cannot be used together")
+            payload["expires_at"] = ""
+        elif expires_at is not _UNSET and expires_at is not None:
+            payload["expires_at"] = _datetime(expires_at)
+        if status is not _UNSET and status not in {"active", "inactive"}:
+            raise ValueError("status must be 'active' or 'inactive'")
+        if name is not _UNSET and (not isinstance(name, str) or not name.strip()):
+            raise ValueError("name must not be empty")
+        if group_id is not _UNSET and group_id is not None:
+            _positive(group_id, "group_id")
+        for field in ("quota", "rate_limit_5h", "rate_limit_1d", "rate_limit_7d"):
+            value = payload.get(field, _UNSET)
+            if value is not _UNSET:
+                payload[field] = _nonnegative(value, field)
+        for field in ("ip_whitelist", "ip_blacklist"):
+            value = payload.get(field, _UNSET)
+            if value is not _UNSET and value is not None:
+                payload[field] = _string_list(value, field)
+        if not payload:
+            raise ValueError("at least one update must be provided")
+        data = self._client.request("PUT", f"keys/{key_id}", json=payload)
+        return APIKey(_mapping(data, "updated API key"))
+
+    def delete(self, key_id: int) -> None:
+        key_id = _positive(key_id, "key_id")
+        self._client.request("DELETE", f"keys/{key_id}")
+
+    def set_status(self, key_id: int, active: bool) -> APIKey:
+        return self.update(key_id, status="active" if active else "inactive")
+
+
+class GroupsResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def __call__(self) -> tuple[Group, ...]:
+        return self.list()
+
+    def list(self) -> tuple[Group, ...]:
+        data = _list(self._client.request("GET", "groups/available"), "groups")
+        return tuple(Group(item) for item in data if isinstance(item, Mapping))
+
+    def rates(self) -> dict[int, float]:
+        data = self._client.request("GET", "groups/rates")
+        if data is None:
+            return {}
+        mapping = _mapping(data, "group rates")
+        try:
+            return {int(key): float(value) for key, value in mapping.items()}
+        except (TypeError, ValueError) as error:
+            raise ProtocolError("group rates contained invalid values") from error
+
+
+class UsageResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def __call__(self, **kwargs: Any) -> Page[UsageRecord]:
+        return self.list(**kwargs)
+
+    def list(
+        self,
+        *,
+        page: int = 1,
+        page_size: int = 20,
+        api_key_id: int | None = None,
+        group_id: int | None = None,
+        model: str | None = None,
+        request_type: str | None = None,
+        stream: bool | None = None,
+        start_date: date | str | None = None,
+        end_date: date | str | None = None,
+        sort_by: str | None = None,
+        sort_order: Literal["asc", "desc"] | None = None,
+        native_compaction_v2: bool | None = None,
+        billing_type: int | None = None,
+        billing_mode: str | None = None,
+    ) -> Page[UsageRecord]:
+        params = {
+            "page": _positive(page, "page"),
+            "page_size": _positive(page_size, "page_size"),
+            "api_key_id": api_key_id,
+            "group_id": group_id,
+            "model": model,
+            "request_type": request_type,
+            "stream": stream,
+            "start_date": _date(start_date),
+            "end_date": _date(end_date),
+            "sort_by": sort_by,
+            "sort_order": sort_order,
+            "native_compaction_v2": native_compaction_v2,
+            "billing_type": billing_type,
+            "billing_mode": billing_mode,
+        }
+        return _page(self._client.request("GET", "usage", params=params), UsageRecord, "usage")
+
+    def iter(self, **kwargs: Any) -> Iterator[UsageRecord]:
+        page_number = int(kwargs.pop("page", 1))
+        while True:
+            result = self.list(page=page_number, **kwargs)
+            yield from result
+            if not result.has_next:
+                return
+            page_number += 1
+
+    def get(self, usage_id: int) -> UsageRecord:
+        usage_id = _positive(usage_id, "usage_id")
+        data = self._client.request("GET", f"usage/{usage_id}")
+        return UsageRecord(_mapping(data, "usage record"))
+
+    def stats(
+        self,
+        *,
+        period: Literal["today", "week", "month"] | None = None,
+        start_date: date | str | None = None,
+        end_date: date | str | None = None,
+        api_key_id: int | None = None,
+        group_id: int | None = None,
+        model: str | None = None,
+    ) -> Resource:
+        params = {
+            "period": period,
+            "start_date": _date(start_date),
+            "end_date": _date(end_date),
+            "api_key_id": api_key_id,
+            "group_id": group_id,
+            "model": model,
+        }
+        return Resource(
+            _mapping(self._client.request("GET", "usage/stats", params=params), "usage stats")
+        )
+
+    def dashboard(self) -> Resource:
+        data = self._client.request("GET", "usage/dashboard/stats")
+        return Resource(_mapping(data, "dashboard stats"))
+
+    def trend(
+        self,
+        *,
+        start_date: date | str | None = None,
+        end_date: date | str | None = None,
+        granularity: Literal["day", "hour"] | None = None,
+        api_key_id: int | None = None,
+        model: str | None = None,
+        group_id: int | None = None,
+    ) -> Resource:
+        params = {
+            "start_date": _date(start_date),
+            "end_date": _date(end_date),
+            "granularity": granularity,
+            "api_key_id": api_key_id,
+            "model": model,
+            "group_id": group_id,
+        }
+        data = self._client.request("GET", "usage/dashboard/trend", params=params)
+        return Resource(_mapping(data, "usage trend"))
+
+    def models(
+        self,
+        *,
+        start_date: date | str | None = None,
+        end_date: date | str | None = None,
+        api_key_id: int | None = None,
+        model: str | None = None,
+        group_id: int | None = None,
+        requested_models: bool = False,
+    ) -> Resource:
+        params = {
+            "start_date": _date(start_date),
+            "end_date": _date(end_date),
+            "api_key_id": api_key_id,
+            "model": model,
+            "group_id": group_id,
+            "model_source": "requested" if requested_models else None,
+        }
+        data = self._client.request("GET", "usage/dashboard/models", params=params)
+        return Resource(_mapping(data, "model usage"))
+
+    def snapshot(
+        self,
+        *,
+        start_date: date | str | None = None,
+        end_date: date | str | None = None,
+        granularity: Literal["day", "hour"] | None = None,
+        include_trend: bool = True,
+        include_model_stats: bool = True,
+        include_group_stats: bool = True,
+        api_key_id: int | None = None,
+        model: str | None = None,
+        group_id: int | None = None,
+    ) -> Resource:
+        params = {
+            "start_date": _date(start_date),
+            "end_date": _date(end_date),
+            "granularity": granularity,
+            "include_trend": include_trend,
+            "include_model_stats": include_model_stats,
+            "include_group_stats": include_group_stats,
+            "api_key_id": api_key_id,
+            "model": model,
+            "group_id": group_id,
+        }
+        data = self._client.request("GET", "usage/dashboard/snapshot-v2", params=params)
+        return Resource(_mapping(data, "usage snapshot"))
+
+    def key_totals(self, api_key_ids: Sequence[int]) -> dict[int, Resource]:
+        ids = [_positive(key_id, "api_key_id") for key_id in api_key_ids]
+        if not ids:
+            return {}
+        data = _mapping(
+            self._client.request(
+                "POST",
+                "usage/dashboard/api-keys-usage",
+                json={"api_key_ids": ids},
+            ),
+            "API key usage totals",
+        )
+        stats = _mapping(data.get("stats", {}), "API key usage totals")
+        try:
+            return {
+                int(key): Resource(value)
+                for key, value in stats.items()
+                if isinstance(value, Mapping)
+            }
+        except ValueError as error:
+            raise ProtocolError("API key usage totals contained an invalid key ID") from error
+
+
+class SubscriptionsResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def __call__(self, *, active: bool = False) -> tuple[Subscription, ...]:
+        return self.list(active=active)
+
+    def list(self, *, active: bool = False) -> tuple[Subscription, ...]:
+        endpoint = "subscriptions/active" if active else "subscriptions"
+        data = _list(self._client.request("GET", endpoint), "subscriptions")
+        return tuple(Subscription(item) for item in data if isinstance(item, Mapping))
+
+
+class AnnouncementsResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def __call__(self) -> tuple[Announcement, ...]:
+        return self.list()
+
+    def list(self) -> tuple[Announcement, ...]:
+        data = _list(self._client.request("GET", "announcements"), "announcements")
+        return tuple(Announcement(item) for item in data if isinstance(item, Mapping))
+
+    def mark_read(self, announcement_id: int) -> None:
+        announcement_id = _positive(announcement_id, "announcement_id")
+        self._client.request("POST", f"announcements/{announcement_id}/read")
+
+
+class RedeemResource:
+    def __init__(self, client: Sub2API) -> None:
+        self._client = client
+
+    def __call__(self, code: str) -> Redemption:
+        return self.apply(code)
+
+    def apply(self, code: str) -> Redemption:
+        if not code.strip():
+            raise ValueError("code must not be empty")
+        data = self._client.request("POST", "redeem", json={"code": code})
+        return Redemption(_mapping(data, "redemption"))
+
+    def history(self) -> tuple[Redemption, ...]:
+        data = _list(self._client.request("GET", "redeem/history"), "redemption history")
+        return tuple(Redemption(item) for item in data if isinstance(item, Mapping))