onceonly-sdk 1.2.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

onceonly/__init__.py

@@ -1,3 +1,4 @@
+from .version import __version__
 from .client import OnceOnly, create_client
 from .models import CheckLockResult
 from .exceptions import (
@@ -9,8 +10,6 @@ from .exceptions import (
     ApiError,
 )
 
-__version__ = "1.2.0"
-
 __all__ = [
     "OnceOnly",
     "create_client",
@@ -21,4 +20,5 @@ __all__ = [
     "RateLimitError",
     "ValidationError",
     "ApiError",
+    "__version__",
 ]

onceonly/_http.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, Dict, Optional, Union, Callable, Awaitable
+
+import httpx
+
+from .exceptions import ApiError, UnauthorizedError, OverLimitError, RateLimitError, ValidationError
+
+
+def try_extract_detail(resp: httpx.Response) -> Optional[Union[Dict[str, Any], str, Any]]:
+    try:
+        j = resp.json()
+        if isinstance(j, dict) and "detail" in j:
+            return j.get("detail")
+        return j
+    except Exception:
+        return None
+
+
+def error_text(resp: httpx.Response, default: str) -> str:
+    d = try_extract_detail(resp)
+    if isinstance(d, dict):
+        return d.get("error") or d.get("message") or default
+    if isinstance(d, str) and d.strip():
+        return d
+    return default
+
+
+def _parse_retry_after(resp: httpx.Response) -> Optional[float]:
+    # headers are case-insensitive in httpx
+    ra = resp.headers.get("Retry-After")
+    if not ra:
+        return None
+    ra = ra.strip()
+    try:
+        return float(ra)
+    except Exception:
+        return None
+
+
+def parse_json_or_raise(resp: httpx.Response) -> Dict[str, Any]:
+    # typed errors
+    if resp.status_code in (401, 403):
+        raise UnauthorizedError(error_text(resp, "Invalid API Key (Unauthorized)."))
+
+    if resp.status_code == 402:
+        d = try_extract_detail(resp)
+        raise OverLimitError("Usage limit reached. Please upgrade your plan.", detail=d if isinstance(d, dict) else {})
+
+    if resp.status_code == 429:
+        retry_after = _parse_retry_after(resp)
+        raise RateLimitError(error_text(resp, "Rate limit exceeded. Please slow down."), retry_after_sec=retry_after)
+
+    if resp.status_code == 422:
+        raise ValidationError(error_text(resp, f"Validation Error: {resp.text}"))
+
+    if resp.status_code < 200 or resp.status_code >= 300:
+        d = try_extract_detail(resp)
+        raise ApiError(
+            error_text(resp, f"API Error ({resp.status_code})"),
+            status_code=resp.status_code,
+            detail=d if isinstance(d, dict) else {},
+        )
+
+    try:
+        data = resp.json()
+    except Exception:
+        data = {}
+    return data if isinstance(data, dict) else {"data": data}
+
+
+def request_with_retries_sync(
+    fn: Callable[[], httpx.Response],
+    *,
+    max_retries: int,
+    base_backoff: float,
+    max_backoff: float,
+) -> httpx.Response:
+    attempt = 0
+    while True:
+        resp = fn()
+        if resp.status_code != 429 or attempt >= max_retries:
+            return resp
+
+        ra = _parse_retry_after(resp)
+        sleep_s = ra if ra is not None else min(max_backoff, base_backoff * (2**attempt))
+        time.sleep(max(0.0, float(sleep_s)))
+        attempt += 1
+
+
+async def request_with_retries_async(
+    fn: Callable[[], Awaitable[httpx.Response]],
+    *,
+    max_retries: int,
+    base_backoff: float,
+    max_backoff: float,
+) -> httpx.Response:
+    attempt = 0
+    while True:
+        resp = await fn()
+        if resp.status_code != 429 or attempt >= max_retries:
+            return resp
+
+        ra = _parse_retry_after(resp)
+        sleep_s = ra if ra is not None else min(max_backoff, base_backoff * (2**attempt))
+        await asyncio.sleep(max(0.0, float(sleep_s)))
+        attempt += 1

onceonly/_util.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from dataclasses import is_dataclass, asdict
+from typing import Any, Dict, Mapping, Optional, Union
+
+MetadataLike = Union[Mapping[str, Any], Any]  # Mapping | pydantic model | dataclass | any
+
+
+def to_metadata_dict(metadata: Optional[MetadataLike]) -> Optional[Dict[str, Any]]:
+    """
+    Accepts:
+    - Mapping[str, Any]
+    - Pydantic model (duck-typed: has model_dump())
+    - dataclass
+    - anything else => {"value": str(obj)}
+
+    Returns plain JSON-ready dict (best effort).
+    """
+    if metadata is None:
+        return None
+
+    # Pydantic v2
+    md = getattr(metadata, "model_dump", None)
+    if callable(md):
+        try:
+            out = md()
+            return out if isinstance(out, dict) else {"data": out}
+        except Exception:
+            return {"value": str(metadata)}
+
+    # dataclass
+    if is_dataclass(metadata):
+        try:
+            out = asdict(metadata)
+            return out if isinstance(out, dict) else {"data": out}
+        except Exception:
+            return {"value": str(metadata)}
+
+    # mapping
+    if isinstance(metadata, Mapping):
+        try:
+            return dict(metadata)
+        except Exception:
+            return {"value": str(metadata)}
+
+    return {"value": str(metadata)}

onceonly/ai.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import asyncio
+import time
+import logging
+from typing import Any, Dict, Optional, Awaitable, Callable
+
+import httpx
+
+from ._http import (
+    parse_json_or_raise,
+    request_with_retries_sync,
+    request_with_retries_async,
+)
+from ._util import to_metadata_dict, MetadataLike
+from .ai_models import AiRun, AiStatus, AiResult
+
+logger = logging.getLogger("onceonly")
+
+
+class AiClient:
+    """
+    AI helpers for long-running backend tasks.
+
+    Typical usage for agents:
+        result = client.ai.run_and_wait(key="job:123", metadata={...})
+
+    Endpoints:
+    - POST /ai/run    => start/attach to a run (idempotent by key)
+    - GET  /ai/status => poll status
+    - GET  /ai/result => fetch final result (completed/failed)
+    """
+
+    def __init__(
+        self,
+        sync_client: httpx.Client,
+        async_client_getter: Callable[[], Awaitable[httpx.AsyncClient]],
+        *,
+        max_retries_429: int = 0,
+        retry_backoff: float = 0.5,
+        retry_max_backoff: float = 5.0,
+    ):
+        self._c = sync_client
+        self._get_ac = async_client_getter
+
+        self._max_retries_429 = int(max_retries_429)
+        self._retry_backoff = float(retry_backoff)
+        self._retry_max_backoff = float(retry_max_backoff)
+
+    # ---- sync ----
+
+    def run(self, key: str, ttl: Optional[int] = None, metadata: Optional[MetadataLike] = None) -> AiRun:
+        payload: Dict[str, Any] = {"key": key}
+        if ttl is not None:
+            payload["ttl"] = int(ttl)
+        md = to_metadata_dict(metadata)
+        if md is not None:
+            payload["metadata"] = md
+
+        resp = request_with_retries_sync(
+            lambda: self._c.post("/ai/run", json=payload),
+            max_retries=self._max_retries_429,
+            base_backoff=self._retry_backoff,
+            max_backoff=self._retry_max_backoff,
+        )
+        data = parse_json_or_raise(resp)
+
+        logger.debug("ai.run key=%s status=%s version=%s", key, data.get("status"), data.get("version"))
+        return AiRun.from_dict(data)
+
+    def status(self, key: str) -> AiStatus:
+        resp = request_with_retries_sync(
+            lambda: self._c.get("/ai/status", params={"key": key}),
+            max_retries=self._max_retries_429,
+            base_backoff=self._retry_backoff,
+            max_backoff=self._retry_max_backoff,
+        )
+        data = parse_json_or_raise(resp)
+
+        logger.debug("ai.status key=%s status=%s version=%s", key, data.get("status"), data.get("version"))
+        return AiStatus.from_dict(data)
+
+    def result(self, key: str) -> AiResult:
+        resp = request_with_retries_sync(
+            lambda: self._c.get("/ai/result", params={"key": key}),
+            max_retries=self._max_retries_429,
+            base_backoff=self._retry_backoff,
+            max_backoff=self._retry_max_backoff,
+        )
+        data = parse_json_or_raise(resp)
+
+        logger.debug("ai.result key=%s status=%s", key, data.get("status"))
+        return AiResult.from_dict(data)
+
+    def wait(self, key: str, timeout: float = 60.0, poll_min: float = 0.5, poll_max: float = 5.0) -> AiResult:
+        t0 = time.time()
+        while True:
+            st = self.status(key)
+            if st.status in ("completed", "failed"):
+                return self.result(key)
+
+            if time.time() - t0 >= timeout:
+                return AiResult(ok=False, status="timeout", key=key, error_code="timeout")
+
+            sleep_s = st.retry_after_sec if isinstance(st.retry_after_sec, int) else poll_min
+            sleep_s = max(poll_min, min(poll_max, float(sleep_s)))
+            time.sleep(sleep_s)
+
+    def run_and_wait(
+        self,
+        key: str,
+        *,
+        ttl: Optional[int] = None,
+        metadata: Optional[MetadataLike] = None,
+        timeout: float = 60.0,
+        poll_min: float = 0.5,
+        poll_max: float = 5.0,
+    ) -> AiResult:
+        self.run(key=key, ttl=ttl, metadata=metadata)
+        return self.wait(key=key, timeout=timeout, poll_min=poll_min, poll_max=poll_max)
+
+    # ---- async ----
+
+    async def run_async(self, key: str, ttl: Optional[int] = None, metadata: Optional[MetadataLike] = None) -> AiRun:
+        payload: Dict[str, Any] = {"key": key}
+        if ttl is not None:
+            payload["ttl"] = int(ttl)
+        md = to_metadata_dict(metadata)
+        if md is not None:
+            payload["metadata"] = md
+
+        c = await self._get_ac()
+        resp = await request_with_retries_async(
+            lambda: c.post("/ai/run", json=payload),
+            max_retries=self._max_retries_429,
+            base_backoff=self._retry_backoff,
+            max_backoff=self._retry_max_backoff,
+        )
+        data = parse_json_or_raise(resp)
+
+        logger.debug("ai.run_async key=%s status=%s version=%s", key, data.get("status"), data.get("version"))
+        return AiRun.from_dict(data)
+
+    async def status_async(self, key: str) -> AiStatus:
+        c = await self._get_ac()
+        resp = await request_with_retries_async(
+            lambda: c.get("/ai/status", params={"key": key}),
+            max_retries=self._max_retries_429,
+            base_backoff=self._retry_backoff,
+            max_backoff=self._retry_max_backoff,
+        )
+        data = parse_json_or_raise(resp)
+
+        logger.debug("ai.status_async key=%s status=%s version=%s", key, data.get("status"), data.get("version"))
+        return AiStatus.from_dict(data)
+
+    async def result_async(self, key: str) -> AiResult:
+        c = await self._get_ac()
+        resp = await request_with_retries_async(
+            lambda: c.get("/ai/result", params={"key": key}),
+            max_retries=self._max_retries_429,
+            base_backoff=self._retry_backoff,
+            max_backoff=self._retry_max_backoff,
+        )
+        data = parse_json_or_raise(resp)
+
+        logger.debug("ai.result_async key=%s status=%s", key, data.get("status"))
+        return AiResult.from_dict(data)
+
+    async def wait_async(self, key: str, timeout: float = 60.0, poll_min: float = 0.5, poll_max: float = 5.0) -> AiResult:
+        t0 = time.time()
+        while True:
+            st = await self.status_async(key)
+            if st.status in ("completed", "failed"):
+                return await self.result_async(key)
+
+            if time.time() - t0 >= timeout:
+                return AiResult(ok=False, status="timeout", key=key, error_code="timeout")
+
+            sleep_s = st.retry_after_sec if isinstance(st.retry_after_sec, int) else poll_min
+            sleep_s = max(poll_min, min(poll_max, float(sleep_s)))
+            await asyncio.sleep(sleep_s)
+
+    async def run_and_wait_async(
+        self,
+        key: str,
+        *,
+        ttl: Optional[int] = None,
+        metadata: Optional[MetadataLike] = None,
+        timeout: float = 60.0,
+        poll_min: float = 0.5,
+        poll_max: float = 5.0,
+    ) -> AiResult:
+        await self.run_async(key=key, ttl=ttl, metadata=metadata)
+        return await self.wait_async(key=key, timeout=timeout, poll_min=poll_min, poll_max=poll_max)

onceonly/ai_models.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+
+def _dict_or_none(v: Any) -> Optional[Dict[str, Any]]:
+    return v if isinstance(v, dict) else None
+
+
+@dataclass(frozen=True)
+class AiRun:
+    ok: bool
+    status: str
+    key: str
+    lease_id: Optional[str] = None
+    version: int = 0
+    charged: Optional[int] = None
+    usage: Optional[int] = None
+    limit: Optional[int] = None
+    retry_after_sec: Optional[int] = None
+    done_at: Optional[str] = None
+    error_code: Optional[str] = None
+    result_hash: Optional[str] = None
+    result: Optional[Dict[str, Any]] = None
+
+    @staticmethod
+    def from_dict(d: Dict[str, Any]) -> "AiRun":
+        return AiRun(
+            ok=bool(d.get("ok", False)),
+            status=str(d.get("status") or ""),
+            key=str(d.get("key") or ""),
+            lease_id=d.get("lease_id"),
+            version=int(d.get("version") or 0),
+            charged=d.get("charged"),
+            usage=d.get("usage"),
+            limit=d.get("limit"),
+            retry_after_sec=d.get("retry_after_sec"),
+            done_at=d.get("done_at"),
+            error_code=d.get("error_code"),
+            result_hash=d.get("result_hash"),
+            result=_dict_or_none(d.get("result")),
+        )
+
+
+@dataclass(frozen=True)
+class AiStatus:
+    ok: bool
+    status: str
+    key: str
+    lease_id: Optional[str] = None
+    version: int = 0
+    ttl_left: Optional[int] = None
+    first_seen_at: Optional[str] = None
+    done_at: Optional[str] = None
+    result_hash: Optional[str] = None
+    error_code: Optional[str] = None
+    retry_after_sec: Optional[int] = None
+
+    @staticmethod
+    def from_dict(d: Dict[str, Any]) -> "AiStatus":
+        return AiStatus(
+            ok=bool(d.get("ok", False)),
+            status=str(d.get("status") or ""),
+            key=str(d.get("key") or ""),
+            lease_id=d.get("lease_id"),
+            version=int(d.get("version") or 0),
+            ttl_left=d.get("ttl_left"),
+            first_seen_at=d.get("first_seen_at"),
+            done_at=d.get("done_at"),
+            result_hash=d.get("result_hash"),
+            error_code=d.get("error_code"),
+            retry_after_sec=d.get("retry_after_sec"),
+        )
+
+
+@dataclass(frozen=True)
+class AiResult:
+    ok: bool
+    status: str
+    key: str
+    result: Optional[Dict[str, Any]] = None
+    result_hash: Optional[str] = None
+    error_code: Optional[str] = None
+    done_at: Optional[str] = None
+
+    @staticmethod
+    def from_dict(d: Dict[str, Any]) -> "AiResult":
+        return AiResult(
+            ok=bool(d.get("ok", False)),
+            status=str(d.get("status") or ""),
+            key=str(d.get("key") or ""),
+            result=_dict_or_none(d.get("result")),
+            result_hash=d.get("result_hash"),
+            error_code=d.get("error_code"),
+            done_at=d.get("done_at"),
+        )