deepailab 0.2.0b1__py3-none-any.whl

deepailab/__init__.py

@@ -0,0 +1,142 @@
+"""
+DeepAI Lab Python SDK
+
+Official Python SDK for the DeepAI Lab API platform.
+Supports OpenAI-compatible endpoints, model marketplace, and enterprise features.
+"""
+
+from .client import DeepAILab, AsyncDeepAILab
+from .exceptions import (
+    DeepAILabError,
+    APIError,
+    AuthenticationError,
+    RateLimitError,
+    TimeoutError,
+    ValidationError,
+)
+
+from .auth import (
+    OIDCClient,
+    OIDCError,
+    OIDCTokenSet,
+    DeviceAuthorization,
+    TokenStore,
+    MemoryTokenStore,
+    FileTokenStore,
+    generate_code_verifier,
+    code_challenge,
+    DEFAULT_ISSUER,
+    DEFAULT_SCOPES,
+)
+from .portal import (
+    PortalClient,
+    PortalError,
+    DEFAULT_BRIDGE_BASE_URL,
+)
+
+# Type exports
+from .types import (
+    # Client configuration
+    ClientOptions,
+    RequestOptions,
+    OnBehalfOfOptions,
+    
+    # OpenAI-compatible types
+    ChatCompletionMessage,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionChoice,
+    ChatCompletionUsage,
+    
+    EmbeddingRequest,
+    EmbeddingResponse,
+    Embedding,
+    
+    ModelListResponse,
+    Model,
+    
+    # Model marketplace types
+    InferenceRequest,
+    InferenceResponse,
+    BatchRequest,
+    BatchResponse,
+    ModelStatus,
+    
+    # Observability types
+    RequestMetrics,
+    UsageInfo,
+    CostInfo,
+    
+    # Streaming types
+    StreamingResponse,
+    StreamChunk,
+)
+
+# Version information
+__version__ = "0.2.0b1"
+__author__ = "DeepAI Lab"
+__email__ = "sdk@deepailab.ai"
+__license__ = "MIT"
+
+# Public API
+__all__ = [
+    # Main client
+    "DeepAILab",
+    "AsyncDeepAILab",
+
+    # Auth (Keycloak OIDC)
+    "OIDCClient",
+    "OIDCError",
+    "OIDCTokenSet",
+    "DeviceAuthorization",
+    "TokenStore",
+    "MemoryTokenStore",
+    "FileTokenStore",
+    "generate_code_verifier",
+    "code_challenge",
+    "DEFAULT_ISSUER",
+    "DEFAULT_SCOPES",
+    # Portal aggregation
+    "PortalClient",
+    "PortalError",
+    "DEFAULT_BRIDGE_BASE_URL",
+    
+    # Exceptions
+    "DeepAILabError",
+    "APIError", 
+    "AuthenticationError",
+    "RateLimitError",
+    "TimeoutError",
+    "ValidationError",
+    
+    # Types
+    "ClientOptions",
+    "RequestOptions",
+    "OnBehalfOfOptions",
+    "ChatCompletionMessage",
+    "ChatCompletionRequest", 
+    "ChatCompletionResponse",
+    "ChatCompletionChoice",
+    "ChatCompletionUsage",
+    "EmbeddingRequest",
+    "EmbeddingResponse",
+    "Embedding",
+    "ModelListResponse",
+    "Model",
+    "InferenceRequest",
+    "InferenceResponse",
+    "BatchRequest",
+    "BatchResponse",
+    "ModelStatus",
+    "RequestMetrics",
+    "UsageInfo",
+    "CostInfo",
+    "StreamingResponse",
+    "StreamChunk",
+    
+    # Version info
+    "__version__",
+    "__author__",
+    "__email__",
+    "__license__",
+]

deepailab/auth.py

@@ -0,0 +1,369 @@
+"""
+Keycloak OIDC authentication for the DeepAILab platform (Python).
+
+Implements the same contract as the JS SDK:
+  - OIDC discovery (cached)
+  - PKCE (S256) helpers
+  - Device Authorization Grant (RFC 8628) for CLI
+  - Authorization Code + PKCE for desktop/web
+  - refresh_token rotation
+  - logout (end-session)
+
+Authoritative contract: docs/architecture/stack-a2/OAUTH2_CLIENT_INTEGRATION.md
+Compatible with Python 3.8+.
+"""
+
+import base64
+import hashlib
+import os
+import secrets
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Tuple
+
+import httpx
+
+DEFAULT_ISSUER = "https://auth.deepailab.ai/realms/deepailab"
+DEFAULT_SCOPES = "openid profile email offline_access"
+
+
+class OIDCError(Exception):
+    """Raised on any OIDC/token error."""
+
+    def __init__(self, message: str, code: str = "oidc_error", status: Optional[int] = None) -> None:
+        super().__init__(message)
+        self.code = code
+        self.status = status
+
+
+@dataclass
+class OIDCTokenSet:
+    access_token: str
+    expires_at: int  # epoch seconds
+    refresh_token: Optional[str] = None
+    id_token: Optional[str] = None
+    scope: Optional[str] = None
+    token_type: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, object]:
+        return {k: v for k, v in self.__dict__.items() if v is not None}
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, object]) -> "OIDCTokenSet":
+        return cls(
+            access_token=str(d["access_token"]),
+            expires_at=int(d["expires_at"]),  # type: ignore[arg-type]
+            refresh_token=_opt_str(d.get("refresh_token")),
+            id_token=_opt_str(d.get("id_token")),
+            scope=_opt_str(d.get("scope")),
+            token_type=_opt_str(d.get("token_type")),
+        )
+
+
+@dataclass
+class DeviceAuthorization:
+    device_code: str
+    user_code: str
+    verification_uri: str
+    expires_in: int
+    interval: int
+    verification_uri_complete: Optional[str] = None
+
+
+# --------------------------------------------------------------------------
+# Token storage
+# --------------------------------------------------------------------------
+
+class TokenStore(ABC):
+    @abstractmethod
+    def load(self) -> Optional[OIDCTokenSet]: ...
+
+    @abstractmethod
+    def save(self, tokens: OIDCTokenSet) -> None: ...
+
+    @abstractmethod
+    def clear(self) -> None: ...
+
+
+class MemoryTokenStore(TokenStore):
+    def __init__(self) -> None:
+        self._tokens: Optional[OIDCTokenSet] = None
+
+    def load(self) -> Optional[OIDCTokenSet]:
+        return self._tokens
+
+    def save(self, tokens: OIDCTokenSet) -> None:
+        self._tokens = tokens
+
+    def clear(self) -> None:
+        self._tokens = None
+
+
+class FileTokenStore(TokenStore):
+    """Persists tokens as JSON with 0600 permissions (CLI/desktop)."""
+
+    def __init__(self, file_path: str) -> None:
+        if not file_path:
+            raise ValueError("FileTokenStore requires a file path")
+        self.file_path = file_path
+
+    def load(self) -> Optional[OIDCTokenSet]:
+        import json
+        try:
+            with open(self.file_path, "r", encoding="utf-8") as fh:
+                data = json.load(fh)
+        except FileNotFoundError:
+            return None
+        if not isinstance(data, dict) or "access_token" not in data:
+            return None
+        return OIDCTokenSet.from_dict(data)
+
+    def save(self, tokens: OIDCTokenSet) -> None:
+        import json
+        directory = os.path.dirname(self.file_path)
+        if directory:
+            os.makedirs(directory, exist_ok=True)
+        tmp = self.file_path + ".tmp"
+        fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+        try:
+            with os.fdopen(fd, "w", encoding="utf-8") as fh:
+                json.dump(tokens.to_dict(), fh)
+        finally:
+            os.chmod(tmp, 0o600)
+        os.replace(tmp, self.file_path)
+        os.chmod(self.file_path, 0o600)
+
+    def clear(self) -> None:
+        try:
+            os.unlink(self.file_path)
+        except FileNotFoundError:
+            pass
+
+
+# --------------------------------------------------------------------------
+# PKCE helpers
+# --------------------------------------------------------------------------
+
+def _b64url(raw: bytes) -> str:
+    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")
+
+
+def generate_code_verifier() -> str:
+    return _b64url(secrets.token_bytes(64))
+
+
+def code_challenge(verifier: str) -> str:
+    return _b64url(hashlib.sha256(verifier.encode("ascii")).digest())
+
+
+def random_state() -> str:
+    return _b64url(secrets.token_bytes(16))
+
+
+def _opt_str(v: object) -> Optional[str]:
+    return None if v is None else str(v)
+
+
+# --------------------------------------------------------------------------
+# OIDC client
+# --------------------------------------------------------------------------
+
+class OIDCClient:
+    def __init__(
+        self,
+        client_id: str,
+        issuer: str = DEFAULT_ISSUER,
+        scope: str = DEFAULT_SCOPES,
+        store: Optional[TokenStore] = None,
+        refresh_skew: int = 60,
+        http: Optional[httpx.Client] = None,
+        now: Optional[Callable[[], float]] = None,
+    ) -> None:
+        if not client_id:
+            raise OIDCError("client_id is required", "config_error")
+        self.client_id = client_id
+        self.issuer = issuer.rstrip("/")
+        self.scope = scope
+        self.store: TokenStore = store or MemoryTokenStore()
+        self.refresh_skew = refresh_skew
+        self._http = http or httpx.Client(timeout=30.0)
+        self._now = now or time.time
+        self._discovery: Optional[Dict[str, object]] = None
+
+    # --- discovery --------------------------------------------------------
+
+    def discover(self) -> Dict[str, object]:
+        if self._discovery is not None:
+            return self._discovery
+        url = self.issuer + "/.well-known/openid-configuration"
+        resp = self._http.get(url, headers={"Accept": "application/json"})
+        if resp.status_code != 200:
+            raise OIDCError("Discovery failed: %s %s" % (resp.status_code, url), "discovery_failed", resp.status_code)
+        self._discovery = resp.json()
+        return self._discovery
+
+    def _endpoint(self, key: str) -> str:
+        disc = self.discover()
+        val = disc.get(key)
+        if not isinstance(val, str):
+            raise OIDCError("Realm discovery has no %s" % key, "endpoint_missing")
+        return val
+
+    def _post_token(self, params: Dict[str, str]) -> OIDCTokenSet:
+        resp = self._http.post(
+            self._endpoint("token_endpoint"),
+            data=params,
+            headers={"Accept": "application/json"},
+        )
+        try:
+            body = resp.json()
+        except Exception:
+            raise OIDCError("Non-JSON token response", "invalid_response", resp.status_code)
+        if resp.status_code >= 400 or "error" in body:
+            raise OIDCError(
+                str(body.get("error_description") or body.get("error") or resp.status_code),
+                str(body.get("error") or "token_error"),
+                resp.status_code,
+            )
+        return self._to_token_set(body)
+
+    def _to_token_set(self, body: Dict[str, object]) -> OIDCTokenSet:
+        expires_in = int(body.get("expires_in", 300))  # type: ignore[arg-type]
+        tokens = OIDCTokenSet(
+            access_token=str(body["access_token"]),
+            expires_at=int(self._now()) + expires_in,
+            refresh_token=_opt_str(body.get("refresh_token")),
+            id_token=_opt_str(body.get("id_token")),
+            scope=_opt_str(body.get("scope")),
+            token_type=_opt_str(body.get("token_type")),
+        )
+        return tokens
+
+    # --- device flow ------------------------------------------------------
+
+    def start_device_authorization(self) -> DeviceAuthorization:
+        endpoint = self._endpoint("device_authorization_endpoint")
+        resp = self._http.post(
+            endpoint,
+            data={"client_id": self.client_id, "scope": self.scope},
+            headers={"Accept": "application/json"},
+        )
+        body = resp.json()
+        if resp.status_code >= 400:
+            raise OIDCError(str(body.get("error_description") or body.get("error") or resp.status_code), "device_start_failed", resp.status_code)
+        return DeviceAuthorization(
+            device_code=str(body["device_code"]),
+            user_code=str(body["user_code"]),
+            verification_uri=str(body["verification_uri"]),
+            expires_in=int(body.get("expires_in", 600)),
+            interval=int(body.get("interval", 5)),
+            verification_uri_complete=_opt_str(body.get("verification_uri_complete")),
+        )
+
+    def poll_device_token(self, device: DeviceAuthorization, sleep: Optional[Callable[[float], None]] = None) -> OIDCTokenSet:
+        do_sleep = sleep or time.sleep
+        interval = device.interval
+        deadline = self._now() + device.expires_in
+        while True:
+            if self._now() > deadline:
+                raise OIDCError("Device code expired", "expired_token")
+            do_sleep(interval)
+            try:
+                tokens = self._post_token({
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                    "device_code": device.device_code,
+                    "client_id": self.client_id,
+                })
+                self.store.save(tokens)
+                return tokens
+            except OIDCError as err:
+                if err.code == "authorization_pending":
+                    continue
+                if err.code == "slow_down":
+                    interval += 5
+                    continue
+                raise
+
+    def login_with_device_flow(self, on_prompt: Callable[[DeviceAuthorization], None], sleep: Optional[Callable[[float], None]] = None) -> OIDCTokenSet:
+        device = self.start_device_authorization()
+        on_prompt(device)
+        return self.poll_device_token(device, sleep=sleep)
+
+    # --- authorization code + PKCE ---------------------------------------
+
+    def create_authorization_request(self, redirect_uri: str) -> Tuple[str, str, str]:
+        """Returns (authorization_url, code_verifier, state)."""
+        from urllib.parse import urlencode
+        verifier = generate_code_verifier()
+        challenge = code_challenge(verifier)
+        state = random_state()
+        query = urlencode({
+            "response_type": "code",
+            "client_id": self.client_id,
+            "redirect_uri": redirect_uri,
+            "scope": self.scope,
+            "state": state,
+            "code_challenge": challenge,
+            "code_challenge_method": "S256",
+        })
+        return self._endpoint("authorization_endpoint") + "?" + query, verifier, state
+
+    def exchange_authorization_code(self, code: str, code_verifier: str, redirect_uri: str) -> OIDCTokenSet:
+        tokens = self._post_token({
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": redirect_uri,
+            "client_id": self.client_id,
+            "code_verifier": code_verifier,
+        })
+        self.store.save(tokens)
+        return tokens
+
+    # --- session ----------------------------------------------------------
+
+    def get_access_token(self) -> str:
+        current = self.store.load()
+        if current is None:
+            raise OIDCError("Not authenticated", "no_session")
+        if current.expires_at - self.refresh_skew > int(self._now()):
+            return current.access_token
+        if not current.refresh_token:
+            raise OIDCError("Access token expired and no refresh_token available", "session_expired")
+        return self.refresh().access_token
+
+    def refresh(self) -> OIDCTokenSet:
+        current = self.store.load()
+        if current is None or not current.refresh_token:
+            raise OIDCError("No refresh_token to refresh", "no_refresh_token")
+        tokens = self._post_token({
+            "grant_type": "refresh_token",
+            "refresh_token": current.refresh_token,
+            "client_id": self.client_id,
+        })
+        if not tokens.refresh_token and current.refresh_token:
+            tokens.refresh_token = current.refresh_token
+        self.store.save(tokens)
+        return tokens
+
+    def is_authenticated(self) -> bool:
+        current = self.store.load()
+        if current is None:
+            return False
+        if current.expires_at > int(self._now()):
+            return True
+        return bool(current.refresh_token)
+
+    def logout(self) -> None:
+        current = self.store.load()
+        try:
+            if current is not None and current.refresh_token:
+                disc = self.discover()
+                endpoint = disc.get("end_session_endpoint")
+                if isinstance(endpoint, str):
+                    try:
+                        self._http.post(endpoint, data={"client_id": self.client_id, "refresh_token": current.refresh_token})
+                    except Exception:
+                        pass
+        finally:
+            self.store.clear()