superb-ai-onprem 0.1.0__py3-none-any.whl

spb_onprem/__init__.py

@@ -0,0 +1,74 @@
+try:
+    from ._version import version as __version__
+except ImportError:
+    __version__ = "0.0.0.dev0"
+
+# Services
+from .datasets.service import DatasetService
+from .data.service import DataService
+from .slices.service import SliceService
+
+# Core Entities
+from .data.entities import (
+    Data,
+    Scene,
+    Annotation,
+    AnnotationVersion,
+    Prediction,
+    DataMeta,
+)
+from .datasets.entities import Dataset
+from .slices.entities import Slice
+
+# Enums
+from .data.enums import (
+    DataType,
+    SceneType,
+    DataMetaTypes,
+    DataMetaValue,
+)
+
+# Filters
+from .data.params.data_list import (
+    AnnotationFilter,
+    DataListFilter,
+    DataFilterOptions,
+)
+from .datasets.params.datasets import (
+    DatasetsFilter,
+    DatasetsFilterOptions,
+)
+from .slices.params.slices import (
+    SlicesFilterOptions,
+)
+
+__all__ = (
+    # Services
+    "DatasetService",
+    "DataService",
+    "SliceService",
+    
+    # Core Entities
+    "Data",
+    "Scene",
+    "Annotation",
+    "AnnotationVersion",
+    "Prediction",
+    "DataMeta",
+    "Dataset",
+    "Slice",
+    
+    # Enums
+    "DataType",
+    "SceneType",
+    "DataMetaTypes",
+    "DataMetaValue",
+    
+    # Filters
+    "AnnotationFilter",
+    "DataListFilter",
+    "DataFilterOptions",
+    "DatasetsFilter",
+    "DatasetsFilterOptions",
+    "SlicesFilterOptions",
+)

spb_onprem/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)

spb_onprem/base_model.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel, Field
+
+
+class CustomBaseModel(BaseModel):
+    class Config:
+        populate_by_name = True

spb_onprem/base_service.py

@@ -0,0 +1,164 @@
+from typing import Optional, Dict, Any, ClassVar
+import random
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util import Retry
+
+from spb_onprem.users.entities import AuthUser
+from spb_onprem.exceptions import (
+    NotFoundError,
+    UnknownError,
+    BadResponseError,
+    BadRequestError,
+    BaseSDKError,
+    BadRequestParameterError,
+    RequestError,
+    ResponseError,
+)
+
+class RetryWithJitter(Retry):
+    def get_backoff_time(self):
+        base_backoff = super().get_backoff_time()
+        jitter = base_backoff * random.uniform(0.5, 1.5)
+        return jitter
+
+
+class BaseService():
+    """The BaseService class is an abstract base class that defines the interface for services that handle data operations.
+    """
+    _retry_session: ClassVar[Optional[requests.Session]] = None
+    _auth_user: Optional[AuthUser] = None
+    
+    def __init__(self):
+        self._auth_user = AuthUser.get_instance()
+        self.endpoint = f"{self._auth_user.host}/graphql/"
+    
+    @classmethod
+    def requests_retry_session(
+        cls,
+        retries=5,
+        backoff_factor=2,
+        status_forcelist=(500, 502, 504),
+        session=None,
+        allowed_methods=[
+            'GET',
+            'POST',
+            'PUT',
+            'DELETE',
+            'OPTIONS',
+            'HEAD',
+            'PATCH',
+            'TRACE',
+            'CONNECT'
+        ]
+    ) -> requests.Session:
+        if BaseService._retry_session is None:
+            session = session or requests.Session()
+            retry = RetryWithJitter(
+                total=retries,
+                read=retries,
+                connect=retries,
+                backoff_factor=backoff_factor,
+                status_forcelist=status_forcelist,
+                allowed_methods=frozenset(allowed_methods),
+            )
+            adapter = HTTPAdapter(max_retries=retry)
+            session.mount('http://', adapter)
+            session.mount('https://', adapter)
+            BaseService._retry_session = session
+        return BaseService._retry_session
+
+    def request_gql(self, query: Any, variables: Dict[str, Any]):
+        """Request Graphql query to the server."""
+        payload = {
+            "query": query["query"],
+            "variables": variables
+        }
+        
+        # Create a new session for each request
+        session = self.requests_retry_session()
+        
+        try:
+            response = session.post(
+                self.endpoint,
+                json=payload,
+                headers=self._auth_user.auth_headers
+            )
+            response.raise_for_status()
+            
+            result = response.json()
+            if not isinstance(result, dict):
+                raise BadRequestError(f"Invalid response format: {type(result).__name__}, expected dict")
+
+            # Check for GraphQL errors
+            if 'errors' in result and result['errors']:
+                for error in result['errors']:
+                    if error['code'] == 'NOT_FOUND':
+                        raise NotFoundError(error['message'])
+                error_messages = [error.get('message', 'Unknown error') for error in result['errors']]
+                raise UnknownError(f"GraphQL errors: {', '.join(error_messages)}")
+            
+            # Validate response structure
+            if 'data' not in result:
+                raise BadResponseError("Missing 'data' field in response")
+            
+            query_name = query.get("name")
+            if not query_name:
+                raise BadResponseError("Missing query name in query object")
+            
+            # Handle different response structures
+            data = result['data']
+            
+            # For other queries, expect the query name to be directly in data
+            if query_name not in data:
+                raise BadResponseError(f"Missing '{query_name}' in response data")
+            
+            return data[query_name]
+            
+        except requests.exceptions.RequestException as e:
+            raise BadResponseError(f"HTTP request failed: {str(e)}") from e
+        except BaseSDKError as e:
+            raise e
+        except Exception as e:
+            raise ResponseError(f"Unexpected error: {str(e)}") from e
+        finally:
+            # Close the session after use
+            session.close()
+
+    def request(
+        self,
+        method: str,
+        url: str,
+        headers: Optional[dict] = None,
+        params: Optional[dict] = None,
+        data: Optional[dict] = None,
+        json_data: Optional[dict] = None,
+        timeout: int = 30
+    ):
+        session = self.requests_retry_session()
+        try:
+            response = session.request(
+                method=method.upper(),
+                url=url,
+                headers={
+                    **headers,
+                    **self._auth_user.auth_headers
+                },
+                params=params,
+                data=data,
+                json=json_data,
+                timeout=timeout
+            )
+            response.raise_for_status()
+            return response
+
+        except requests.exceptions.RequestException as e:
+            print(f"An error occurred during the HTTP request: {str(e)}")
+            raise BadRequestError(f"HTTP request failed: {str(e)}") from e
+        except ValueError:
+            raise BadRequestParameterError("Failed to parse the HTTP response as JSON.") from e
+        except Exception as e:
+            raise RequestError(f"An error occurred while processing the HTTP response: {str(e)}") from e
+        finally:
+            session.close()

spb_onprem/base_types.py

@@ -0,0 +1,11 @@
+
+
+class UndefinedType():
+    """A singleton type used to represent an undefined value."""
+    def __repr__(self):
+        return "Undefined"
+
+    def __bool__(self):
+        return False  # Ensures it evaluates as False in boolean contexts
+
+Undefined = UndefinedType()

spb_onprem/contents/__init__.py

@@ -0,0 +1,6 @@
+from .service import ContentService
+
+
+__all__ = (
+    "ContentService",
+)

spb_onprem/contents/entities/__init__.py

@@ -0,0 +1,8 @@
+from .base_content import BaseContent
+from .content import Content
+
+
+__all__ = (
+    "BaseContent",
+    "Content",
+)

spb_onprem/contents/entities/base_content.py

@@ -0,0 +1,13 @@
+from typing import Optional
+
+from spb_onprem.base_model import CustomBaseModel, Field
+
+class BaseContent(CustomBaseModel):
+    """The content.
+    This is the actual file that is stored in the file storage.
+
+    Args:
+        BaseModel (_type_): _description_
+    """
+    id: str
+    download_url:Optional[str] = Field(None, alias="downloadURL")

spb_onprem/contents/entities/content.py

@@ -0,0 +1,17 @@
+from typing import Optional
+
+from spb_onprem.base_model import CustomBaseModel, Field
+
+class Content(CustomBaseModel):
+    """The content.
+    This is the actual file that is stored in the file storage.
+
+    Args:
+        BaseModel (_type_): _description_
+    """
+    id: str
+    download_url:Optional[str] = Field(None, alias="downloadURL")
+    key: Optional[str] = None
+    location: Optional[dict] = None
+    created_at: Optional[str] = Field(None, alias="createdAt")
+    created_by: Optional[str] = Field(None, alias="createdBy")

spb_onprem/contents/queries.py

@@ -0,0 +1,39 @@
+from typing import (
+    Union,
+)
+
+from spb_onprem.base_types import (
+    Undefined,
+    UndefinedType,
+)
+
+class Queries:
+    """ Content Queries """
+    
+    @staticmethod
+    def create_variables(
+        key: Union[str, UndefinedType] = Undefined,
+    ):
+        if key is not Undefined:
+            return {"key": key}
+        else:
+            return {}
+
+    CREATE = {
+        "name": "createContent",
+        "query": '''
+            mutation CreateContent($key: String) {
+                createContent(key: $key) {
+                    content {
+                        id
+                        key
+                        location
+                        createdAt
+                        createdBy
+                    }
+                    uploadURL
+                }
+            }
+        ''',
+        "variables": create_variables
+    }

spb_onprem/contents/service.py

@@ -0,0 +1,132 @@
+import mimetypes
+import requests
+import json
+
+from io import BytesIO
+from typing import Optional, Union
+
+from spb_onprem.base_service import BaseService
+from spb_onprem.base_types import (
+    Undefined,
+    UndefinedType,
+)
+from .entities import Content, BaseContent
+from .queries import Queries
+
+
+
+class ContentService(BaseService):
+    """The content service for the SDK.
+    Content service is the service that handles the content operations.
+    """
+
+    def upload_content(
+        self,
+        file_path: str,
+        key: Union[
+            str,
+            UndefinedType    
+        ] = Undefined,
+    ):
+        '''
+        Uploads the content to the server.
+        Args:
+            file_path (str):
+                The path of the file to be uploaded.
+                You must provide the full path of the file (with extensions).
+        '''
+        with open(file_path, 'rb') as f:
+            file = f.read()
+        response = self.request_gql(
+            query=Queries.CREATE,
+            variables=Queries.CREATE["variables"](key)
+        )
+        upload_url = response['uploadURL']
+        
+        self.request(
+            method="PUT",
+            url=upload_url,
+            headers={
+                'Content-Type': mimetypes.guess_type(file_path)[0]
+            },
+            data=file,
+        )
+        content = response['content']
+        return BaseContent.model_validate(content)
+    
+    def upload_json_content(
+        self,
+        data: dict,
+        key: Union[
+            str,
+            UndefinedType    
+        ] = Undefined,
+    ):
+        '''
+        Uploads the JSON content to the server.
+
+        Args:
+            data (dict):
+                The JSON data to be uploaded.
+            key (Optional[str]):
+                An optional key to associate with the uploaded content.
+        '''
+        response = self.request_gql(
+            query=Queries.CREATE,
+            variables=Queries.CREATE["variables"](key) if key else None
+        )
+        upload_url = response['uploadURL']
+        self.request(
+            method="PUT",
+            url=upload_url,
+            headers={
+                'Content-Type': 'application/json'
+            },
+            json_data=data,
+        )
+        content = response['content']
+        return BaseContent.model_validate(content)
+
+    def upload_content_with_data(
+        self,
+        file_data: BytesIO,
+        content_type: str,
+        key: Union[
+            str,
+            UndefinedType    
+        ] = Undefined,
+    ):
+        '''
+        Uploads the content to the server.
+
+        Args:
+            file_data (BytesIO):
+                The file data to be uploaded.
+            content_type (str):
+                The MIME type of the file being uploaded (e.g., "image/jpeg").
+            key (Optional[str]):
+                An optional key to associate with the uploaded content.
+        '''
+        # Reset the BytesIO pointer to the beginning
+        file_data.seek(0)
+
+        # Request to get the upload URL
+        response = self.request_gql(
+            query=Queries.CREATE,
+            variables=Queries.CREATE["variables"](key)
+        )
+        upload_url = response['uploadURL']
+
+        # Upload the file data using the PUT request
+        self.request(
+            method="PUT",
+            url=upload_url,
+            headers={
+                'Content-Type': content_type
+            },
+            data=file_data.read(),
+        )
+
+        # Retrieve the uploaded content details
+        content = response['content']
+        return BaseContent.model_validate(content)

spb_onprem/data/__init__.py

@@ -0,0 +1,6 @@
+from .service import DataService
+
+
+__all__ = (
+    "DataService",
+)

spb_onprem/data/entities/__init__.py

@@ -0,0 +1,15 @@
+from .annotation import Annotation, AnnotationVersion
+from .data_meta import DataMeta
+from .data import Data
+from .prediction import Prediction
+from .scene import Scene
+
+
+__all__ = (
+    "Data",
+    "Scene",
+    "Annotation",
+    "AnnotationVersion",
+    "Prediction",
+    "DataMeta",
+)

spb_onprem/data/entities/annotation.py

@@ -0,0 +1,25 @@
+from typing import Optional, List
+
+from spb_onprem.base_model import CustomBaseModel
+from spb_onprem.contents.entities import BaseContent
+
+
+class AnnotationVersion(CustomBaseModel):
+    """
+    The version of the annotation.
+    Annotation version is the version of the data annotation.
+    This has the content of the data annotation.
+    """
+    id: Optional[str] = None
+    content: Optional[BaseContent] = None
+    meta: Optional[dict] = None
+
+
+class Annotation(CustomBaseModel):
+    """
+    The annotation of the data.
+    Annotation has the versions of the data annotation.
+    """
+    versions: Optional[List[AnnotationVersion]] = None
+    meta: Optional[dict] = None
+

spb_onprem/data/entities/data.py

@@ -0,0 +1,28 @@
+from typing import List, Optional
+from spb_onprem.base_model import CustomBaseModel, Field
+from spb_onprem.data.enums import DataType
+from .scene import Scene
+from .annotation import Annotation
+from .prediction import Prediction
+from .data_meta import DataMeta
+from spb_onprem.contents.entities import BaseContent
+
+class Data(CustomBaseModel):
+    """
+    THE DATA.
+    """
+    id: Optional[str] = None
+    dataset_id: Optional[str] = Field(None, alias="datasetId")
+    slice_ids: Optional[List[str]] = Field(None, alias="sliceIds")
+    key: Optional[str] = None
+    type: Optional[DataType] = None
+    scene: Optional[List[Scene]] = None
+    thumbnail: Optional[BaseContent] = None
+    annotation: Optional[Annotation] = None
+    predictions: Optional[List[Prediction]] = None
+    meta: Optional[List[DataMeta]] = None
+    system_meta: Optional[List[DataMeta]] = Field(None, alias="systemMeta")
+    created_at: Optional[str] = Field(None, alias="createdAt")
+    created_by: Optional[str] = Field(None, alias="createdBy")
+    updated_at: Optional[str] = Field(None, alias="updatedAt")
+    updated_by: Optional[str] = Field(None, alias="updatedBy")

spb_onprem/data/entities/data_meta.py

@@ -0,0 +1,31 @@
+from typing import Optional, Union
+from datetime import datetime
+from spb_onprem.data.enums import DataMetaTypes, DataMetaValue
+from spb_onprem.base_model import CustomBaseModel
+
+
+class DataMeta(CustomBaseModel):
+    """
+    The meta of the data.
+    Meta is the metadata of the data.
+    """
+    key: Optional[str] = None
+    type: Optional[DataMetaTypes] = None
+    value: Optional[DataMetaValue] = None
+
+    @classmethod
+    def from_dict(cls, meta: dict[str, DataMetaValue]) -> list["DataMeta"]:
+        return [
+            cls(
+                key=key,
+                value=val,
+                type=(
+                    DataMetaTypes.BOOLEAN if isinstance(val, bool) else
+                    DataMetaTypes.NUMBER if isinstance(val, (int, float)) else
+                    DataMetaTypes.DATETIME if isinstance(val, datetime) else
+                    DataMetaTypes.STRING if isinstance(val, str) else
+                    DataMetaTypes.ANNOTATION  # list or dict
+                ),
+            )
+            for key, val in meta.items()
+        ]

spb_onprem/data/entities/prediction.py

@@ -0,0 +1,13 @@
+from typing import Optional
+from spb_onprem.base_model import CustomBaseModel, Field
+from spb_onprem.contents.entities import BaseContent
+
+
+class Prediction(CustomBaseModel):
+    """
+    The prediction of the data.
+    Prediction has the predictions of the data.
+    """
+    set_id: Optional[str] = Field(None, alias="setId")
+    content: Optional[BaseContent] = None
+    meta: Optional[dict] = None

spb_onprem/data/entities/scene.py

@@ -0,0 +1,14 @@
+from typing import Optional
+from spb_onprem.base_model import CustomBaseModel
+from spb_onprem.data.enums import SceneType
+from spb_onprem.contents.entities import BaseContent
+
+class Scene(CustomBaseModel):
+    """
+    The scene of the data.
+    Scene is the representation of the file of the data.
+    """
+    id: Optional[str] = None
+    type: Optional[SceneType] = None
+    content: Optional[BaseContent] = None
+    meta: Optional[dict] = None

spb_onprem/data/enums/__init__.py

@@ -0,0 +1,10 @@
+from .data_meta_type import DataMetaTypes, DataMetaValue
+from .data_type import DataType
+from .scene_type import SceneType
+
+__all__ = (
+    "DataType",
+    "SceneType",
+    "DataMetaTypes",
+    "DataMetaValue",
+)

spb_onprem/data/enums/data_meta_type.py

@@ -0,0 +1,15 @@
+from enum import Enum
+from typing import Union
+from datetime import datetime
+
+
+class DataMetaTypes(str, Enum):
+    """The data meta types."""
+    STRING = "String"
+    NUMBER = "Number"
+    BOOLEAN = "Boolean"
+    DATETIME = "DateTime"
+    ANNOTATION = "Annotation"
+
+
+DataMetaValue = Union[str, int, float, bool, datetime, dict, list]

spb_onprem/data/enums/data_type.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class DataType(str, Enum):
+    """
+    The type of the data.
+    This is used to determine the type of the data.
+    """
+    SUPERB_IMAGE = "SUPERB_IMAGE"