modelstudio-sdk 0.0.0.dev0__py3-none-any.whl

modelstudio/__init__.py

@@ -0,0 +1,25 @@
+"""Model Studio Python SDK."""
+
+from modelstudio._version import __version__
+from modelstudio.client import ModelStudioClient
+from modelstudio.exceptions import (
+    AuthenticationError,
+    BadRequestError,
+    ConflictError,
+    ModelStudioError,
+    NotFoundError,
+    ServerError,
+    TimeoutError,
+)
+
+__all__ = [
+    "__version__",
+    "ModelStudioClient",
+    "ModelStudioError",
+    "BadRequestError",
+    "AuthenticationError",
+    "NotFoundError",
+    "ConflictError",
+    "ServerError",
+    "TimeoutError",
+]

modelstudio/_http.py

@@ -0,0 +1,116 @@
+"""HTTP transport layer wrapping httpx with auth and error mapping."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from modelstudio.exceptions import (
+    STATUS_CODE_MAP,
+    AuthenticationError,
+    ModelStudioError,
+    ServerError,
+)
+
+
+class HttpTransport:
+    """Thin httpx wrapper that injects auth headers and maps errors to typed exceptions."""
+
+    def __init__(
+        self,
+        base_url: str,
+        jwt_token: str | None = None,
+        timeout: float = 30.0,
+        org_name: str | None = None,
+    ) -> None:
+        headers: dict[str, str] = {
+            "X-Source": "modelstudio-sdk",
+            "Accept": "application/json",
+        }
+        if jwt_token:
+            headers["Authorization"] = f"Bearer {jwt_token}"
+        if org_name:
+            headers["X-Org-Name"] = org_name
+
+        self._client = httpx.Client(
+            base_url=base_url.rstrip("/"),
+            headers=headers,
+            timeout=timeout,
+        )
+
+    def close(self) -> None:
+        self._client.close()
+
+    # --- HTTP methods ---
+
+    def get(self, path: str, params: dict[str, Any] | None = None) -> Any:
+        resp = self._client.get(path, params=params)
+        return self._handle(resp)
+
+    def post(self, path: str, json: Any = None) -> Any:
+        resp = self._client.post(path, json=json)
+        return self._handle(resp)
+
+    def put(self, path: str, json: Any = None) -> Any:
+        resp = self._client.put(path, json=json)
+        return self._handle(resp)
+
+    def patch(self, path: str, json: Any = None) -> Any:
+        resp = self._client.patch(path, json=json)
+        return self._handle(resp)
+
+    def delete(self, path: str, json: Any = None, params: dict[str, Any] | None = None) -> Any:
+        if json is not None:
+            req = self._client.build_request("DELETE", path, json=json, params=params)
+            resp = self._client.send(req)
+        else:
+            resp = self._client.delete(path, params=params)
+        return self._handle(resp)
+
+    def delete_with_body(self, path: str, json: Any = None) -> Any:
+        """DELETE with a JSON request body (e.g. bulk annotation delete)."""
+        req = self._client.build_request("DELETE", path, json=json)
+        resp = self._client.send(req)
+        return self._handle(resp)
+
+    def post_accepted(self, path: str, json: Any = None) -> Any:
+        """POST that expects 202 Accepted (async operations)."""
+        resp = self._client.post(path, json=json)
+        return self._handle(resp)
+
+    # --- Response handling ---
+
+    def _handle(self, resp: httpx.Response) -> Any:
+        if resp.status_code == 204:
+            return None
+
+        if resp.is_success:
+            if not resp.content:
+                return None
+            return resp.json()
+
+        self._raise_for_status(resp)
+
+    def _raise_for_status(self, resp: httpx.Response) -> None:
+        error_code = "Unknown"
+        message = f"HTTP {resp.status_code}"
+
+        try:
+            body = resp.json()
+            error_code = body.get("error", error_code)
+            message = body.get("message", message)
+        except Exception:
+            message = resp.text or message
+
+        exc_class = STATUS_CODE_MAP.get(resp.status_code)
+        if exc_class:
+            raise exc_class(message=message, error_code=error_code)
+
+        if resp.status_code == 401:
+            raise AuthenticationError(message=message)
+        if resp.status_code >= 500:
+            raise ServerError(message=message)
+        raise ModelStudioError(
+            message=message, status_code=resp.status_code, error_code=error_code
+        )

modelstudio/_pandas.py

@@ -0,0 +1,33 @@
+"""Optional pandas DataFrame helpers (guarded import)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def _get_pandas() -> Any:
+    try:
+        import pandas
+
+        return pandas
+    except ImportError:
+        raise ImportError(
+            "pandas is required for DataFrame operations. "
+            "Install it with: pip install 'modelstudio-sdk[pandas]'"
+        ) from None
+
+
+def to_dataframe(items: list[T]) -> pd.DataFrame:
+    """Convert a list of Pydantic models to a pandas DataFrame."""
+    pd = _get_pandas()
+    if not items:
+        return pd.DataFrame()
+    rows = [item.model_dump() for item in items]
+    return pd.DataFrame(rows)

modelstudio/_polling.py

@@ -0,0 +1,59 @@
+"""Poller for async operations (clone, import, finalize)."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import Any
+
+from modelstudio.exceptions import TimeoutError
+
+
+class OperationPoller:
+    """Polls an async operation endpoint until a terminal status is reached."""
+
+    def __init__(
+        self,
+        transport: Any,
+        poll_url: str,
+        terminal_statuses: set[str] | None = None,
+        interval: float = 2.0,
+        max_wait: float = 300.0,
+    ) -> None:
+        self._transport = transport
+        self._poll_url = poll_url
+        self._terminal = terminal_statuses or {"COMPLETED", "FAILED", "ERROR", "DONE"}
+        self._interval = interval
+        self._max_wait = max_wait
+
+    def poll_once(self) -> dict[str, Any]:
+        """Single poll, returns the status response dict."""
+        return self._transport.get(self._poll_url)  # type: ignore[no-any-return]
+
+    def wait(self, callback: Callable[[dict[str, Any]], None] | None = None) -> dict[str, Any]:
+        """Block until terminal status. Optional progress callback called on each poll."""
+        elapsed = 0.0
+        while elapsed < self._max_wait:
+            result = self.poll_once()
+            if callback:
+                callback(result)
+
+            status = self._extract_status(result)
+            if status and status.upper() in self._terminal:
+                return result
+
+            time.sleep(self._interval)
+            elapsed += self._interval
+
+        raise TimeoutError(
+            f"Operation at {self._poll_url} did not complete within {self._max_wait}s"
+        )
+
+    @staticmethod
+    def _extract_status(data: dict[str, Any]) -> str | None:
+        """Extract status from common response shapes."""
+        for key in ("status", "clone_status", "finalize_status", "import_status"):
+            val = data.get(key)
+            if val is not None:
+                return str(val)
+        return None

modelstudio/_version.py

@@ -0,0 +1 @@
+__version__ = "0.0.0.dev0"

modelstudio/client.py

@@ -0,0 +1,108 @@
+"""ModelStudioClient — main entry point for the SDK."""
+
+from __future__ import annotations
+
+import os
+from uuid import UUID
+
+from modelstudio._http import HttpTransport
+from modelstudio.resources.dataset import Dataset, DatasetsCollection
+
+
+class ModelStudioClient:
+    """Client for the Model Studio REST API.
+
+    Usage::
+
+        # Auto-configured from environment variables
+        client = ModelStudioClient.from_env()
+
+        # Or explicit
+        client = ModelStudioClient(
+            base_url="http://localhost:8081",
+            jwt_token="eyJhbG...",
+        )
+
+        # List datasets
+        datasets = client.datasets.list()
+
+        # Work with a specific dataset
+        ds = client.dataset("uuid-here")
+        overview = ds.overview()
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        jwt_token: str | None = None,
+        timeout: float = 30.0,
+        org_name: str | None = None,
+    ) -> None:
+        self._transport = HttpTransport(
+            base_url=base_url,
+            jwt_token=jwt_token,
+            timeout=timeout,
+            org_name=org_name,
+        )
+        self._base_url = base_url.rstrip("/")
+
+    @classmethod
+    def from_env(cls) -> ModelStudioClient:
+        """Create client from environment variables.
+
+        Reads:
+            MODEL_STUDIO_API_URL: API base URL (required)
+            MODEL_STUDIO_JWT: JWT token (optional)
+            MODEL_STUDIO_ORG: Organization name (optional)
+        """
+        base_url = os.environ.get("MODEL_STUDIO_API_URL")
+        if not base_url:
+            raise ValueError(
+                "MODEL_STUDIO_API_URL environment variable is required. "
+                "Set it to the Model Studio API base URL "
+                "(e.g. http://model-studio-api-service.model-studio.svc.cluster.local)"
+            )
+        return cls(
+            base_url=base_url,
+            jwt_token=os.environ.get("MODEL_STUDIO_JWT"),
+            org_name=os.environ.get("MODEL_STUDIO_ORG"),
+        )
+
+    @property
+    def datasets(self) -> DatasetsCollection:
+        """Access dataset collection methods (list, create, merge)."""
+        return DatasetsCollection(self._transport)
+
+    def dataset(self, dataset_id: str | UUID) -> Dataset:
+        """Get a Dataset resource bound to a specific dataset ID."""
+        return Dataset(self._transport, str(dataset_id))
+
+    def full_image_url(
+        self, dataset_id: str | UUID, image_id: str | UUID, ext: str = "png"
+    ) -> str:
+        """Build a full image URL for direct browser/download use."""
+        return f"{self._base_url}/api/v1/media/images/{dataset_id}/{image_id}.{ext}"
+
+    def thumbnail_url(
+        self,
+        dataset_id: str | UUID,
+        image_id: str | UUID,
+        size: int,
+        hash: str,
+        ext: str = "webp",
+    ) -> str:
+        """Build a thumbnail URL."""
+        return (
+            f"{self._base_url}/api/v1/media/thumbs"
+            f"/{dataset_id}/{image_id}_{size}.{ext}?h={hash}"
+        )
+
+    def close(self) -> None:
+        """Close the underlying HTTP connection."""
+        self._transport.close()
+
+    def __enter__(self) -> ModelStudioClient:
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()

modelstudio/exceptions.py

@@ -0,0 +1,80 @@
+"""Typed exception hierarchy mapped to API error envelope."""
+
+from __future__ import annotations
+
+
+class ModelStudioError(Exception):
+    """Base exception for all Model Studio SDK errors."""
+
+    def __init__(
+        self,
+        message: str,
+        status_code: int | None = None,
+        error_code: str | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.status_code = status_code
+        self.error_code = error_code
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"status_code={self.status_code}, "
+            f"error_code={self.error_code!r}, "
+            f"message={self.message!r})"
+        )
+
+
+class BadRequestError(ModelStudioError):
+    """400 Bad Request — invalid input or validation failure."""
+
+    def __init__(self, message: str, error_code: str = "BadRequest") -> None:
+        super().__init__(message, status_code=400, error_code=error_code)
+
+
+class AuthenticationError(ModelStudioError):
+    """401 Unauthorized — missing or invalid JWT token."""
+
+    def __init__(self, message: str = "Unauthorized", error_code: str = "Unauthorized") -> None:
+        super().__init__(message, status_code=401, error_code=error_code)
+
+
+class NotFoundError(ModelStudioError):
+    """404 Not Found — resource does not exist."""
+
+    def __init__(self, message: str, error_code: str = "NotFound") -> None:
+        super().__init__(message, status_code=404, error_code=error_code)
+
+
+class ConflictError(ModelStudioError):
+    """409 Conflict — operation conflicts with current state."""
+
+    def __init__(self, message: str, error_code: str = "Conflict") -> None:
+        super().__init__(message, status_code=409, error_code=error_code)
+
+
+class ServerError(ModelStudioError):
+    """500 Internal Server Error — unexpected server failure."""
+
+    def __init__(
+        self, message: str = "Internal server error", error_code: str = "InternalServerError"
+    ) -> None:
+        super().__init__(message, status_code=500, error_code=error_code)
+
+
+class TimeoutError(ModelStudioError):
+    """Polling timeout — async operation did not complete in time."""
+
+    def __init__(self, message: str = "Operation timed out") -> None:
+        super().__init__(message, status_code=None, error_code="Timeout")
+
+
+# Map HTTP status codes to exception classes
+STATUS_CODE_MAP: dict[int, type[ModelStudioError]] = {
+    400: BadRequestError,
+    401: AuthenticationError,
+    404: NotFoundError,
+    409: ConflictError,
+    500: ServerError,
+}

modelstudio/models/__init__.py

@@ -0,0 +1,81 @@
+"""Pydantic v2 response/request models matching the Model Studio API."""
+
+from modelstudio.models.annotations import (
+    AnnotationModel,
+    CreateAnnotationRequest,
+    DeleteAnnotationsResponse,
+    PaginatedAnnotationModel,
+    UpdateAnnotationRequest,
+)
+from modelstudio.models.common import PagedResponse
+from modelstudio.models.images import (
+    DatasetImageModel,
+    ImageModel,
+    MarkSyntheticResponse,
+    PaginatedImageModel,
+    RemoveEmptyImagesResponse,
+    SyntheticStatsModel,
+)
+from modelstudio.models.imports import (
+    ImportJobListModel,
+    ImportJobModel,
+    ImportJobQueuedModel,
+    ImportLogEntryModel,
+    ImportQueuedModel,
+    ImportStatusModel,
+    PreImportValidationModel,
+    ValidationStartedModel,
+)
+from modelstudio.models.media import (
+    MediaJobActionModel,
+    MediaJobErrorModel,
+    MediaJobStatusModel,
+)
+from modelstudio.models.splits import (
+    AlgorithmStats,
+    ClassAwareRedistributeResponse,
+    ClassError,
+    CreateSplitsResponse,
+    DistributionError,
+    LeakageResponse,
+    RedistributeResponse,
+    SmartRedistributeResponse,
+    SplitModel,
+    SplitStats,
+)
+
+__all__ = [
+    "AnnotationModel",
+    "CreateAnnotationRequest",
+    "DeleteAnnotationsResponse",
+    "PaginatedAnnotationModel",
+    "UpdateAnnotationRequest",
+    "PagedResponse",
+    "DatasetImageModel",
+    "ImageModel",
+    "MarkSyntheticResponse",
+    "PaginatedImageModel",
+    "RemoveEmptyImagesResponse",
+    "SyntheticStatsModel",
+    "ImportJobListModel",
+    "ImportJobModel",
+    "ImportJobQueuedModel",
+    "ImportLogEntryModel",
+    "ImportQueuedModel",
+    "ImportStatusModel",
+    "PreImportValidationModel",
+    "ValidationStartedModel",
+    "MediaJobActionModel",
+    "MediaJobErrorModel",
+    "MediaJobStatusModel",
+    "AlgorithmStats",
+    "ClassAwareRedistributeResponse",
+    "ClassError",
+    "CreateSplitsResponse",
+    "DistributionError",
+    "LeakageResponse",
+    "RedistributeResponse",
+    "SmartRedistributeResponse",
+    "SplitModel",
+    "SplitStats",
+]

modelstudio/models/annotations.py

@@ -0,0 +1,80 @@
+"""Annotation-related models."""
+
+from __future__ import annotations
+
+from typing import Any
+from uuid import UUID
+
+from pydantic import BaseModel
+
+
+class AnnotationModel(BaseModel):
+    """COCO annotation."""
+
+    annotation_id: int
+    image_id: UUID
+    category_id: int
+    segmentation: Any | None = None
+    bbox: list[float] = []
+    ignore: bool = False
+    iscrowd: bool = False
+    area: float = 0.0
+    keypoints: Any | None = None
+    num_keypoints: int | None = None
+
+
+class CreateAnnotationRequest(BaseModel):
+    """Request for creating an annotation."""
+
+    annotation_id: int | None = None
+    image_id: UUID
+    category_id: int
+    segmentation: Any | None = None
+    bbox: list[float] = []
+    ignore: bool = False
+    iscrowd: bool = False
+    area: float = 0.0
+    keypoints: Any | None = None
+    num_keypoints: int | None = None
+
+
+class UpdateAnnotationRequest(BaseModel):
+    """Request for updating an annotation."""
+
+    category_id: int | None = None
+    segmentation: Any | None = None
+    bbox: list[float] | None = None
+    area: float | None = None
+    iscrowd: bool | None = None
+    ignore: bool | None = None
+
+
+class DeleteAnnotationsResponse(BaseModel):
+    """Response from DELETE /datasets/{id}/annotations."""
+
+    success: bool = False
+    message: str | None = None
+    annotations_deleted: int = 0
+    images_deleted: int = 0
+    categories_deleted: int = 0
+
+
+class PaginatedAnnotationModel(BaseModel):
+    """An annotation in the paginated dataset-level annotation listing."""
+
+    annotation_id: int
+    image_id: UUID
+    category_id: int
+    segmentation: Any | None = None
+    bbox: list[float] = []
+    ignore: bool = False
+    iscrowd: bool = False
+    area: float = 0.0
+    keypoints: Any | None = None
+    num_keypoints: int | None = None
+    category_name: str | None = None
+    segmentation_type: str | None = None
+    tags: list[str] = []
+    confidence_score: float | None = None
+    annotator_id: str | None = None
+    status: str | None = None

modelstudio/models/categories.py

@@ -0,0 +1,106 @@
+"""Category-related models."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class CategoryModel(BaseModel):
+    """A dataset category."""
+
+    category_id: int
+    name: str
+    supercategory: str | None = None
+    annotation_count: int = 0
+    image_count: int = 0
+    pixel_count: int | None = None
+    pixel_percentage: float | None = None
+    area_m2: float | None = None
+
+
+class CategoryListModel(BaseModel):
+    """Response from GET /datasets/{id}/categories."""
+
+    categories: list[CategoryModel] = []
+    total_categories: int = 0
+    total_annotations: int = 0
+
+
+class MergeCategoriesResponse(BaseModel):
+    """Response from POST /datasets/{id}/categories/merge."""
+
+    merged: bool = False
+    target_id: int | None = None
+    target_name: str | None = None
+    affected_annotations: int = 0
+    deleted_categories: list[str] = []
+
+
+class RenameCategoryResponse(BaseModel):
+    """Response from PUT /datasets/{id}/categories/{cid}/rename."""
+
+    old_name: str | None = None
+    new_name: str | None = None
+    category_id: int | None = None
+    affected_annotations: int = 0
+
+
+class RemoveCategoryResponse(BaseModel):
+    """Response from DELETE /datasets/{id}/categories/{cid}."""
+
+    removed: Any = None
+    category_id: int | None = None
+    annotations_affected: int = 0
+    action_taken: str | None = None
+
+
+class ConsolidateLabelsResponse(BaseModel):
+    """Response from POST /datasets/{id}/categories/consolidate."""
+
+    consolidated: Any = None
+    mapping_applied: dict[str, Any] | None = None
+    total_annotations_updated: int = 0
+
+
+class MergeBySupercategoryResponse(BaseModel):
+    """Response from POST /datasets/{id}/categories/merge-by-supercategory."""
+
+    supercategory: str | None = None
+    target_category_name: str | None = None
+    target_category_id: int | None = None
+    categories_merged: int = 0
+    merged_categories: list[str] = []
+    annotations_affected: int = 0
+
+
+class SplitCategoryResponse(BaseModel):
+    """Response from POST /datasets/{id}/categories/split."""
+
+    operation_id: str | None = None
+    source_category_name: str | None = None
+    target_categories: list[str] = []
+    annotations_to_assign: int = 0
+
+
+class CancelSplitCategoryResponse(BaseModel):
+    """Response from DELETE /datasets/{id}/categories/split/{opId}."""
+
+    operation_id: str | None = None
+    source_category_name: str | None = None
+    target_categories_expected: int = 0
+    target_categories_deleted: int = 0
+    status: str | None = None
+
+
+class MergeSourcesResponse(BaseModel):
+    """Response from GET /datasets/{id}/categories/{cid}/merge-sources."""
+
+    dataset_id: str | None = None
+    category_id: int | None = None
+    category_name: str | None = None
+    has_merge_history: bool = False
+    leaf_categories: list[dict[str, Any]] = []
+    merge_chain: list[dict[str, Any]] = []
+    suggested_split_request: dict[str, Any] | None = None