visionai-sdk-python 2.2.0.dev1__py3-none-any.whl → 2.2.0.dev3__py3-none-any.whl

visionai_sdk_python/auth/resource.py

@@ -1,6 +1,5 @@
 """Authentication resource for sync operations."""
 
-import time
 from typing import TYPE_CHECKING
 
 import jwt
@@ -28,12 +27,17 @@ class AuthResource(AuthMixin):
         """
         self._sdk_client = client
 
-    def get_access_token(self, client_id: str, client_secret: str) -> TokenResponse:
+    def get_access_token(
+        self, client_id: str, client_secret: str, *, force_refresh: bool = False
+    ) -> TokenResponse:
         """Get access token using client credentials flow.
 
-        The token is stored internally and will be used automatically for subsequent
-        API calls. If the token expires, it will be automatically refreshed.
-        Calling this again with the same credentials returns the cached token.
+        The client keeps the resulting credential and uses it for subsequent API
+        calls, minting a new token when the current one expires. Calling this
+        again with the same client credentials returns the cached token;
+        ``force_refresh=True`` mints a new one. A later ``login()`` or
+        ``set_token()`` replaces the credential: the last explicit
+        authentication wins.
 
         New code should build a ``ClientCredentials`` / ``AsyncClientCredentials``
         and pass it as ``credential=`` to the client instead; this method stays for
@@ -42,6 +46,7 @@ class AuthResource(AuthMixin):
         Args:
             client_id: OAuth client ID
             client_secret: OAuth client secret
+            force_refresh: Mint a new token even if a fresh one is cached.
 
         Returns:
             TokenResponse with access_token, expires_in, and token_type
@@ -68,16 +73,15 @@ class AuthResource(AuthMixin):
                 client_secret,
                 http=self._sdk_client._client,
             )
-            self._sdk_client._credential = credential
-        token = credential.get_token()
-        self._sdk_client._store_token(
+        token = credential.get_token(force_refresh=force_refresh)
+        # Installed only once it works, and over a login() / set_token() token:
+        # the last explicit authentication wins.
+        self._sdk_client._credential = credential
+        return TokenResponse(
             access_token=token.token,
-            expires_in=max(int(token.expires_on - time.monotonic()), 0),
-            credentials={"client_id": client_id, "client_secret": client_secret},
-            credentials_type="client",
+            expires_in=token.expires_in(),
+            token_type=token.token_type,
         )
-        assert credential.last_response is not None
-        return credential.last_response
 
     def login(self, email: str, password: str) -> TokenResponse:
         """Login with email and password to get JWT token.

visionai_sdk_python/client.py

@@ -4,7 +4,7 @@ from ._base import _BaseClient
 from ._request_attribution import merge_request_attribution
 from .auth.credentials import ClientCredentials
 from .auth.resource import AuthResource
-from .exceptions import AuthenticationError
+from .exceptions import AuthenticationError, from_transport_error
 from .vlm.resource import VLMResource
 
 
@@ -19,7 +19,10 @@ class Client(_BaseClient):
         max_connections: int = 100,
         max_keepalive_connections: int = 20,
         credential: ClientCredentials | None = None,
+        http_client: httpx.Client | None = None,
     ) -> None:
+        """See ``_BaseClient``; ``http_client`` is an httpx.Client to send through
+        (tests, shared pools); the client closes only the one it created."""
         super().__init__(
             auth_url=auth_url,
             vlm_url=vlm_url,
@@ -30,7 +33,7 @@ class Client(_BaseClient):
             max_keepalive_connections=max_keepalive_connections,
             credential=credential,
         )
-        self._client = httpx.Client(
+        self._client = http_client or httpx.Client(
             verify=self.verify_ssl,
             timeout=self.timeout,
             limits=httpx.Limits(
@@ -38,6 +41,7 @@ class Client(_BaseClient):
                 max_keepalive_connections=self.max_keepalive_connections,
             ),
         )
+        self._http_owner = http_client is None  # a client we built is ours to close
         # Register the different function resource
         self.auth = AuthResource(self)
         self.vlm = VLMResource(self)
@@ -48,12 +52,42 @@ class Client(_BaseClient):
         try:
             response = self._client.request(method, url, **kwargs)
         except httpx.RequestError as e:
-            raise self._map_transport_error(e) from e
+            raise from_transport_error(e) from e
         return self._handle_response(response)
 
+    def request(self, method: str, url: str, **kwargs) -> httpx.Response:
+        """Send any request with this client's bearer token.
+
+        ``url`` is a full URL (the client has two bases); ``kwargs`` go to httpx.
+        With a ``credential``, a 401 buys one retry with the credential's current
+        token, which it replaces when the rejected one is old enough to have been
+        revoked (``REPLACE_COOLDOWN_S``); a 401 on the retry raises. Without a
+        credential, or with a body that cannot be sent twice (a stream or a
+        file), a 401 raises.
+
+        Raises:
+            ValueError: ``url`` is not a full URL.
+            AuthenticationError: No token, or the fresh one was refused too.
+            RateLimitedError, NetworkError, ServerError, ClientError: As the
+                resource methods.
+        """
+        url = self._full_url(url)
+        token = self._ensure_token()
+        try:
+            return self._request(method, url, **self._with_auth(kwargs, token))
+        except AuthenticationError:
+            if self._credential is None or not self._resendable(kwargs):
+                raise
+            self._credential.invalidate(token)
+            replacement = (self._credential.get_token()).token
+            if replacement == token:
+                raise
+        return self._request(method, url, **self._with_auth(kwargs, replacement))
+
     def close(self) -> None:
-        """Close the HTTP client and release connections."""
-        self._client.close()
+        """Close the HTTP client this client created and release its connections."""
+        if self._http_owner:
+            self._client.close()
 
     def __enter__(self) -> "Client":
         """Context manager entry."""
@@ -64,7 +98,9 @@ class Client(_BaseClient):
         self.close()
 
     def _refresh_token(self) -> None:
-        """Refresh token using stored credentials.
+        """Refresh the token with the stored login credentials.
+
+        Client credentials refresh through the credential object instead.
 
         Raises:
             AuthenticationError: If no credentials are stored or refresh fails.
@@ -72,26 +108,19 @@ class Client(_BaseClient):
         if self._credentials is None or self._credentials_type is None:
             raise AuthenticationError("No credentials available for token refresh")
 
-        if self._credentials_type == "login":
-            self.auth.login(
-                email=self._credentials["email"],
-                password=self._credentials["password"],
-            )
-        elif self._credentials_type == "client":
-            self.auth.get_access_token(
-                client_id=self._credentials["client_id"],
-                client_secret=self._credentials["client_secret"],
-            )
+        self.auth.login(
+            email=self._credentials["email"],
+            password=self._credentials["password"],
+        )
 
-    def _ensure_token(self) -> None:
-        """Ensure a valid token is available, refreshing if necessary.
+    def _ensure_token(self) -> str:
+        """The token to send: the credential's, else the stored one, refreshed if expiring.
 
         Raises:
             AuthenticationError: If no token is available or token expired without credentials.
         """
         if self._credential is not None:
-            self._access_token = (self._credential.get_token()).token
-            return
+            return (self._credential.get_token()).token
         if self._access_token is None:
             raise AuthenticationError(
                 "Not authenticated. Call login() or get_access_token() first."
@@ -103,3 +132,4 @@ class Client(_BaseClient):
                     "Token expired and no credentials available for refresh"
                 )
             self._refresh_token()
+        return self._access_token

visionai_sdk_python/exceptions.py

@@ -1,3 +1,6 @@
+import httpx
+
+
 class VisionaiSDKError(Exception):
     """Base exception for all SDK errors."""
 
@@ -29,7 +32,7 @@ class PermissionDeniedError(ClientError):
 
 
 class RateLimitedError(ClientError):
-    """429 - Too many requests; on client-token, the monthly mint quota is used up."""
+    """429 - Too many requests; on client-token, Platform is rate limiting mints."""
 
     def __init__(self, message: str) -> None:
         super().__init__(429, message)
@@ -45,3 +48,50 @@ class NetworkError(VisionaiSDKError):
 
 class JwksDiscoveryError(VisionaiSDKError):
     """Failed to fetch or parse the OIDC discovery document or JWKS endpoint."""
+
+
+def from_transport_error(error: httpx.RequestError) -> VisionaiSDKError:
+    """The SDK exception for an httpx request that never got a response."""
+    if isinstance(error, httpx.TimeoutException):
+        return NetworkError("Request timed out")
+    if isinstance(error, httpx.NetworkError):
+        return NetworkError(f"Network error: {error}")
+    return VisionaiSDKError(f"Request failed: {error}")
+
+
+def raise_for_status(response: httpx.Response) -> httpx.Response:
+    """Raise SDK-specific exceptions for non-2xx responses.
+
+    Raises:
+        AuthenticationError: If the server returns 401
+        PermissionDeniedError: If the server returns 403
+        ClientError: If the server returns any other 4xx
+        ServerError: If the server returns 5xx
+        APIError: If the server returns any other non-2xx status
+    """
+    try:
+        response.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        body = None
+        if e.response.content:
+            try:
+                body = e.response.json()
+            except Exception:
+                pass
+        if isinstance(body, dict):
+            detail: str = body.get("detail") or body.get("message") or str(e)
+        else:
+            detail = str(e)
+        status = e.response.status_code
+        if status == 401:
+            raise AuthenticationError(detail) from e
+        if status == 403:
+            raise PermissionDeniedError(detail) from e
+        if status == 429:
+            raise RateLimitedError(detail) from e
+        if 400 <= status < 500:
+            raise ClientError(status, detail) from e
+        if 500 <= status < 600:
+            raise ServerError(status, detail) from e
+        raise APIError(status, detail) from e
+    return response

visionai_sdk_python/relay/__init__.py

@@ -0,0 +1,67 @@
+"""The platform-relay wire protocol, shared by platform-relay and platform-edge-agent.
+
+Internal to the platform; the spec is platform-service's
+``docs/adr/0001-relay-protocol.md``. See :mod:`.protocol`.
+"""
+
+from .protocol import (
+    AGENT_TIMEOUT_S,
+    HEARTBEAT_INTERVAL_S,
+    MAX_BODY_BYTES,
+    MAX_FRAME_BYTES,
+    REGISTER_DEADLINE_S,
+    RELAY_TIMEOUT_S,
+    STALE_AFTER_S,
+    SUBPROTOCOL,
+    AgentFrame,
+    CloseCode,
+    Error,
+    ErrorCode,
+    Frame,
+    Heartbeat,
+    InvalidFrame,
+    ModelEntry,
+    Register,
+    Registered,
+    Rejected,
+    RejectReason,
+    RelayFrame,
+    Request,
+    Response,
+    WorkspaceEntry,
+    dump_frame,
+    is_workspace_id,
+    parse_agent_frame,
+    parse_relay_frame,
+)
+
+__all__ = [
+    "AGENT_TIMEOUT_S",
+    "HEARTBEAT_INTERVAL_S",
+    "MAX_BODY_BYTES",
+    "MAX_FRAME_BYTES",
+    "REGISTER_DEADLINE_S",
+    "RELAY_TIMEOUT_S",
+    "STALE_AFTER_S",
+    "SUBPROTOCOL",
+    "AgentFrame",
+    "CloseCode",
+    "Error",
+    "ErrorCode",
+    "Frame",
+    "Heartbeat",
+    "InvalidFrame",
+    "ModelEntry",
+    "Register",
+    "Registered",
+    "Rejected",
+    "RejectReason",
+    "RelayFrame",
+    "Request",
+    "Response",
+    "WorkspaceEntry",
+    "dump_frame",
+    "is_workspace_id",
+    "parse_agent_frame",
+    "parse_relay_frame",
+]

visionai_sdk_python/relay/protocol.py

@@ -0,0 +1,178 @@
+"""Wire format between platform-relay and platform-edge-agent.
+
+The spec is platform-service's ``docs/adr/0001-relay-protocol.md``; this module is
+the one implementation both sides import. Frames are JSON objects told apart by
+``type``. Receivers ignore unknown fields, and the string-typed ``reason`` /
+``code`` fields accept values they do not know, so either side can be upgraded
+first.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from enum import IntEnum, StrEnum
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
+
+SUBPROTOCOL = "platform-relay.v1"
+# The first frame after the upgrade must be `register`, this soon; else the relay closes 1008.
+REGISTER_DEADLINE_S = 10.0
+HEARTBEAT_INTERVAL_S = 30.0
+# No frame for this long: the relay closes 4001 and forgets the connection.
+STALE_AFTER_S = 90.0
+# Both sides refuse to send a larger frame; the relay refuses a larger chat body (413).
+MAX_FRAME_BYTES = 16 * 1024 * 1024
+MAX_BODY_BYTES = 15 * 1024 * 1024
+# Timeouts nest: the agent waits for VS, the relay waits a little longer for the agent.
+AGENT_TIMEOUT_S = 100
+RELAY_TIMEOUT_S = 110.0
+
+# Platform workspace ids: lowercase hyphenated UUIDs.
+UUID_RE = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}")
+
+
+def is_workspace_id(value: str) -> bool:
+    return UUID_RE.fullmatch(value) is not None
+
+
+class CloseCode(IntEnum):
+    NORMAL = 1000
+    GOING_AWAY = 1001
+    POLICY_VIOLATION = 1008
+    TOO_BIG = 1009
+    REPLACED = 4000
+    STALE = 4001
+
+
+class RejectReason(StrEnum):
+    INVALID_ID = "invalid_id"
+    CLUSTER_MISMATCH = "cluster_mismatch"
+    PLATFORM_UNAVAILABLE = "platform_unavailable"
+
+
+class ErrorCode(StrEnum):
+    INVALID_MESSAGE = "invalid_message"
+    UNKNOWN_WORKSPACE = "unknown_workspace"
+    PATH_NOT_ALLOWED = "path_not_allowed"
+    UPSTREAM_UNREACHABLE = "upstream_unreachable"
+    UPSTREAM_TIMEOUT = "upstream_timeout"
+
+
+class Frame(BaseModel):
+    # Receivers ignore unknown fields.
+    model_config = ConfigDict(extra="ignore")
+
+
+class ModelEntry(Frame):
+    name: str
+    version: str
+    # Required by the ADR. A receiver counts None as 0 and logs it: the sender has a bug.
+    replicas: int | None = Field(default=None, ge=0)
+
+
+class WorkspaceEntry(Frame):
+    id: str
+    models: list[ModelEntry] = Field(default_factory=list)
+
+
+class Register(Frame):
+    type: Literal["register"] = "register"
+    agent_id: str = Field(min_length=1)
+    cluster_id: str = Field(min_length=1)
+    workspaces: list[WorkspaceEntry] = Field(default_factory=list)
+
+
+class Rejected(Frame):
+    id: str
+    reason: str  # a RejectReason; a plain str so an older agent accepts new reasons
+
+
+class Registered(Frame):
+    type: Literal["registered"] = "registered"
+    accepted: list[str] = Field(default_factory=list)
+    rejected: list[Rejected] = Field(default_factory=list)
+
+
+class Heartbeat(Frame):
+    type: Literal["heartbeat"] = "heartbeat"
+
+
+class Request(Frame):
+    type: Literal["request"] = "request"
+    request_id: str
+    workspace_id: str
+    method: str
+    path: str
+    headers: dict[str, str] = Field(default_factory=dict)
+    body: dict | list | str | None = None
+    timeout_s: int = Field(gt=0)
+
+
+class Response(Frame):
+    type: Literal["response"] = "response"
+    request_id: str
+    status: int
+    headers: dict[str, str] = Field(default_factory=dict)
+    body: dict | list | str | None = None
+
+
+class Error(Frame):
+    type: Literal["error"] = "error"
+    request_id: str | None = None
+    code: str  # an ErrorCode; a plain str so an older receiver accepts new codes
+    message: str
+
+
+# What each side receives.
+AgentFrame = Register | Heartbeat | Response | Error  # agent -> relay
+RelayFrame = Registered | Request | Error  # relay -> agent
+_agent_frame = TypeAdapter(Annotated[AgentFrame, Field(discriminator="type")])
+_relay_frame = TypeAdapter(Annotated[RelayFrame, Field(discriminator="type")])
+
+
+class InvalidFrame(ValueError):
+    """A frame the receiver cannot read; the text is safe to send back in an ``error``."""
+
+    def __init__(self, message: str, request_id: str | None = None) -> None:
+        super().__init__(message)
+        self.request_id = request_id
+
+
+def parse_agent_frame(text: str) -> AgentFrame:
+    """A frame the relay received from an agent."""
+    return _parse(text, _agent_frame)
+
+
+def parse_relay_frame(text: str) -> RelayFrame:
+    """A frame the agent received from the relay."""
+    return _parse(text, _relay_frame)
+
+
+def _parse(text: str, adapter: TypeAdapter):
+    try:
+        data = json.loads(text)
+    except ValueError:
+        raise InvalidFrame("frame is not JSON") from None
+    if not isinstance(data, dict):
+        raise InvalidFrame("frame is not a JSON object")
+    kind = data.get("type")
+    if not isinstance(kind, str):
+        raise InvalidFrame("frame has no type")
+    request_id = data.get("request_id")
+    request_id = request_id if isinstance(request_id, str) else None
+    try:
+        return adapter.validate_python(data)
+    except ValidationError as exc:
+        if any(err["type"] == "union_tag_invalid" for err in exc.errors()):
+            raise InvalidFrame(f"unknown type {kind!r}", request_id) from None
+        fields = ", ".join(
+            ".".join(str(part) for part in err["loc"][1:]) or "body"
+            for err in exc.errors()
+        )
+        raise InvalidFrame(f"{kind} fails validation: {fields}", request_id) from None
+
+
+def dump_frame(frame: BaseModel) -> str:
+    return frame.model_dump_json(exclude_none=True)

visionai_sdk_python/vlm/async_resource.py

@@ -28,8 +28,8 @@ class AsyncVLMResource(VLMMixin):
     ) -> ResponseNormalModel | ResponseErrorModel:
         """Submit an inference request to the VLM service.
 
-        Uses the internally stored access token obtained from login() or get_access_token().
-        If the token is expiring soon, it will be automatically refreshed.
+        Sent with the client's token (``credential``, ``login()`` or ``get_access_token()``);
+        a 401 gets one fresh token and one retry.
 
         Args:
             payload: Inference parameters as a NIMRequestModel instance or a dict
@@ -45,17 +45,15 @@ class AsyncVLMResource(VLMMixin):
             NetworkError: If the request times out or a network error occurs.
             VisionaiSDKError: If the request fails for any other reason.
         """
-        # Ensure token is valid
-        await self._sdk_client._ensure_token()
-
         # Prepare request (from Mixin)
         body = self._prepare_chat_request(payload)
 
-        # I/O operation (async)
-        response = await self._sdk_client._request(
+        # I/O operation (async); the client adds the token and retries once on 401
+        response = await self._sdk_client.request(
             "POST",
-            self._sdk_client._build_url(self._sdk_client._vlm_base(), VLMEndpoint.CHAT),
-            headers=self._sdk_client._build_auth_header(self._sdk_client._access_token),
+            self._sdk_client._build_url(
+                self._sdk_client._require_vlm_url(), VLMEndpoint.CHAT
+            ),
             json=body,
         )
 
@@ -67,8 +65,8 @@ class AsyncVLMResource(VLMMixin):
     ) -> ResponseNormalModel | ResponseErrorModel:
         """Poll the result of a previously submitted inference request.
 
-        Uses the internally stored access token obtained from login() or get_access_token().
-        If the token is expiring soon, it will be automatically refreshed.
+        Sent with the client's token (``credential``, ``login()`` or ``get_access_token()``);
+        a 401 gets one fresh token and one retry.
 
         Args:
             result_id: Chat result ID returned from a prior chat() call.
@@ -82,16 +80,12 @@ class AsyncVLMResource(VLMMixin):
             NetworkError: If the request times out or a network error occurs.
             VisionaiSDKError: If the request fails for any other reason.
         """
-        # Ensure token is valid
-        await self._sdk_client._ensure_token()
-
-        # I/O operation (async)
-        response = await self._sdk_client._request(
+        # I/O operation (async); the client adds the token and retries once on 401
+        response = await self._sdk_client.request(
             "GET",
             self._sdk_client._build_url(
-                self._sdk_client._vlm_base(), f"{VLMEndpoint.CHAT}/{result_id}"
+                self._sdk_client._require_vlm_url(), f"{VLMEndpoint.CHAT}/{result_id}"
             ),
-            headers=self._sdk_client._build_auth_header(self._sdk_client._access_token),
         )
 
         # Parse response (from Mixin)

visionai_sdk_python/vlm/resource.py

@@ -28,8 +28,8 @@ class VLMResource(VLMMixin):
     ) -> ResponseNormalModel | ResponseErrorModel:
         """Submit an inference request to the VLM service.
 
-        Uses the internally stored access token obtained from login() or get_access_token().
-        If the token is expiring soon, it will be automatically refreshed.
+        Sent with the client's token (``credential``, ``login()`` or ``get_access_token()``);
+        a 401 gets one fresh token and one retry.
 
         Args:
             payload: Inference parameters as a NIMRequestModel instance or a dict
@@ -45,17 +45,15 @@ class VLMResource(VLMMixin):
             NetworkError: If the request times out or a network error occurs.
             VisionaiSDKError: If the request fails for any other reason.
         """
-        # Ensure token is valid
-        self._sdk_client._ensure_token()
-
         # Prepare request (from Mixin)
         body = self._prepare_chat_request(payload)
 
-        # I/O operation (sync)
-        response = self._sdk_client._request(
+        # I/O operation (sync); the client adds the token and retries once on 401
+        response = self._sdk_client.request(
             "POST",
-            self._sdk_client._build_url(self._sdk_client._vlm_base(), VLMEndpoint.CHAT),
-            headers=self._sdk_client._build_auth_header(self._sdk_client._access_token),
+            self._sdk_client._build_url(
+                self._sdk_client._require_vlm_url(), VLMEndpoint.CHAT
+            ),
             json=body,
         )
 
@@ -65,8 +63,8 @@ class VLMResource(VLMMixin):
     def get_chat(self, result_id: str) -> ResponseNormalModel | ResponseErrorModel:
         """Poll the result of a previously submitted inference request.
 
-        Uses the internally stored access token obtained from login() or get_access_token().
-        If the token is expiring soon, it will be automatically refreshed.
+        Sent with the client's token (``credential``, ``login()`` or ``get_access_token()``);
+        a 401 gets one fresh token and one retry.
 
         Args:
             result_id: Chat result ID returned from a prior chat() call.
@@ -80,16 +78,12 @@ class VLMResource(VLMMixin):
             NetworkError: If the request times out or a network error occurs.
             VisionaiSDKError: If the request fails for any other reason.
         """
-        # Ensure token is valid
-        self._sdk_client._ensure_token()
-
-        # I/O operation (sync)
-        response = self._sdk_client._request(
+        # I/O operation (sync); the client adds the token and retries once on 401
+        response = self._sdk_client.request(
             "GET",
             self._sdk_client._build_url(
-                self._sdk_client._vlm_base(), f"{VLMEndpoint.CHAT}/{result_id}"
+                self._sdk_client._require_vlm_url(), f"{VLMEndpoint.CHAT}/{result_id}"
             ),
-            headers=self._sdk_client._build_auth_header(self._sdk_client._access_token),
         )
 
         # Parse response (from Mixin)