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

spb_onprem/datasets/params/update_dataset.py

@@ -0,0 +1,39 @@
+from typing import Union
+
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+
+
+def update_dataset_params(
+    dataset_id: str,
+    name: Union[
+        str,
+        UndefinedType,
+    ] = Undefined,
+    description: Union[
+        str,
+        UndefinedType,
+    ] = Undefined,
+):
+    """Update parameters for dataset modification.
+    
+    Args:
+        id: ID of the dataset to update
+        name: Optional new name for the dataset
+        description: Optional new description for the dataset
+        
+    Returns:
+        dict: Parameters for dataset update
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+
+    params = {
+        "id": dataset_id,
+    }
+
+    if name is not Undefined:
+        params["name"] = name
+    if description is not Undefined:
+        params["description"] = description
+    return params

spb_onprem/datasets/queries.py

@@ -0,0 +1,79 @@
+from spb_onprem.datasets.params import (
+    dataset_params,
+    datasets_params,
+    create_dataset_params,
+    update_dataset_params,
+)
+
+class Schemas:
+    DATASET = '''
+        id
+        name
+        description
+        createdAt
+        createdBy
+        updatedAt
+        updatedBy
+    '''
+
+
+class Queries():
+    DATASET = {
+        "name": "dataset",
+        "query": f'''
+            query Dataset(
+                $name: String,
+                $datasetId: ID
+            ) {{
+                dataset(
+                    name: $name,
+                    datasetId: $datasetId
+                ) {{
+                    {Schemas.DATASET}
+                }}
+            }}
+        ''',
+        "variables": dataset_params,
+    }
+    
+    DATASETS = {
+        "name": "datasets",
+        "query": f'''
+            query Datasets(
+                $name: String,
+                $datasetId: ID
+            ) {{
+                datasets(
+                    name: $name,
+                    id: $datasetId
+                ) {{
+                    {Schemas.DATASET}
+                }}
+            }}
+        ''',
+        "variables": datasets_params,
+    }
+    
+    CREATE_DATASET = {
+        "name": "createDataset",
+        "query": f'''
+        mutation CreateDataset($name: String!, $description: String) {{
+            createDataset(name: $name, description: $description) {{
+                {Schemas.DATASET}
+            }}
+        }}
+        ''',
+        "variables": create_dataset_params,
+    }
+
+    UPDATE_DATASET = {
+        "name": "updateDataset",
+        "query": f'''
+            mutation UpdateDataset($updateDatasetId: ID!, $name: String, $description: String) {{
+                updateDataset(id: $updateDatasetId, name: $name, description: $description) {{
+                    {Schemas.DATASET}
+                }}
+            }}
+        ''',
+        "variables": update_dataset_params,
+    }

spb_onprem/datasets/service.py

@@ -0,0 +1,132 @@
+from typing import Optional, Union
+from spb_onprem.base_service import BaseService
+from spb_onprem.exceptions import BadParameterError
+from spb_onprem.base_types import Undefined, UndefinedType
+from .queries import Queries
+from .entities import Dataset
+from .params.datasets import DatasetsFilter
+
+
+class DatasetService(BaseService):
+    """
+    Service class for handling dataset-related operations.
+    """
+    
+    def get_datasets(
+        self,
+        datasets_filter: Optional[DatasetsFilter] = None,
+        cursor: Optional[str] = None,
+        length: Optional[int] = 10
+    ):
+        """
+        Get a list of datasets based on the provided filter and pagination parameters.
+        
+        Args:
+            datasets_filter (Union[DatasetsFilter, UndefinedType]): Filter criteria for datasets
+            cursor (Optional[str]): Cursor for pagination
+            length (Optional[int]): Number of items per page (default: 10)
+        
+        Returns:
+            List[Dataset]: A list of Dataset objects
+        """
+        if length > 50:
+            raise BadParameterError("The maximum length is 50.")
+        
+        response = self.request_gql(
+            Queries.DATASETS,
+            Queries.DATASETS["variables"](
+                datasets_filter=datasets_filter,
+                cursor=cursor,
+                length=length
+            )
+        )
+        datasets = [Dataset.model_validate(dataset) for dataset in response["datasets"]]
+        return (
+            datasets,
+            response.get("next", None),
+            response.get("totalCount", False)
+        )
+
+    def get_dataset(
+        self,
+        dataset_id: Optional[str] = None,
+        name: Optional[str] = None,
+    ):
+        """
+        Retrieve a dataset by its ID or name.
+
+        Args:
+            dataset_id (Optional[str]): The ID of the dataset to retrieve.
+            name (Optional[str]): The name of the dataset to retrieve.
+
+        Returns:
+            Dataset: The retrieved dataset object.
+        """
+        response = self.request_gql(
+            Queries.DATASET,
+            Queries.DATASET["variables"](
+                dataset_id=dataset_id,
+                name=name
+            ),
+        )
+        return Dataset.model_validate(response)
+    
+    def create_dataset(
+        self,
+        name: str,
+        description: Union[
+            str,
+            UndefinedType,
+        ] = Undefined,
+    ):
+        """
+        Create a new dataset.
+
+        Args:
+            name (str): The name of the dataset to create.
+            description (Optional[str]): The description of the dataset to create.
+
+        Returns:
+            Dataset: The created dataset object.
+        """
+        response = self.request_gql(
+            Queries.CREATE_DATASET,
+            Queries.CREATE_DATASET["variables"](
+                name=name,
+                description=description,
+            ),
+        )
+        return Dataset.model_validate(response)
+
+    def update_dataset(
+        self,
+        dataset_id: str,
+        name: Union[
+            str,
+            UndefinedType,
+        ] = Undefined,
+        description: Union[
+            str,
+            UndefinedType,
+        ] = Undefined,
+    ):
+        """
+        Update a dataset.
+
+        Args:
+            dataset_id (str): The ID of the dataset to update.
+            name (Optional[str]): The name of the dataset to update.
+            description (Optional[str]): The description of the dataset to update.
+
+        Returns:
+            Dataset: The updated dataset object.
+        """
+        response = self.request_gql(
+            Queries.UPDATE_DATASET,
+            Queries.UPDATE_DATASET["variables"](
+                dataset_id=dataset_id,
+                name=name,
+                description=description,
+            ),
+        )
+        return Dataset.model_validate(response)

spb_onprem/exceptions.py

@@ -0,0 +1,40 @@
+
+
+class BaseSDKError(Exception):
+    pass
+
+
+class SDKConfigError(BaseSDKError):
+    pass
+
+
+class BadParameterError(BaseSDKError):
+    pass
+
+
+class NotFoundError(BaseSDKError):
+    pass
+
+
+class RequestError(BaseSDKError):
+    pass
+
+
+class BadRequestError(RequestError):
+    pass
+
+
+class BadRequestParameterError(BadRequestError):
+    pass
+
+
+class ResponseError(BaseSDKError):
+    pass
+
+
+class BadResponseError(ResponseError):
+    pass
+
+
+class UnknownError(BaseSDKError):
+    pass

spb_onprem/slices/__init__.py

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

spb_onprem/slices/entities/__init__.py

@@ -0,0 +1,5 @@
+from .slice import Slice
+
+__all__ = (
+    "Slice",
+)

spb_onprem/slices/entities/slice.py

@@ -0,0 +1,17 @@
+from typing import Optional
+from spb_onprem.base_model import CustomBaseModel, Field
+
+
+class Slice(CustomBaseModel):
+    """
+    THE SLICE.
+    """
+    id: Optional[str] = None
+    dataset_id: Optional[str] = Field(None, alias="datasetId")
+    name: Optional[str] = None
+    description: Optional[str] = None
+    is_pinned: Optional[bool] = Field(None, alias="isPinned")
+    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/slices/params/__init__.py

@@ -0,0 +1,23 @@
+from .slices import (
+    SlicesFilter,
+    SlicesFilterOptions,
+    slices_params,
+)
+from .slice import (
+    slice_params,
+)
+from .create_slice import (
+    create_slice_params,
+)
+from .update_slice import (
+    update_slice_params,
+)
+
+__all__ = (
+    "SlicesFilter",
+    "SlicesFilterOptions",
+    "slices_params",
+    "slice_params",
+    "create_slice_params",
+    "update_slice_params",
+)

spb_onprem/slices/params/create_slice.py

@@ -0,0 +1,36 @@
+from typing import Union
+
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+
+def create_slice_params(
+    dataset_id: str,
+    slice_name: str,
+    slice_description: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+):
+    """Create parameters for slice creation.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to create the slice for.
+        slice_name (str): The name of the slice to create.
+        slice_description (Optional[str]): The description of the slice to create.  
+        
+    Returns:
+        dict: Parameters for slice creation
+    """
+
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+    if slice_name is None:
+        raise BadParameterError("Slice name is required")
+
+    params = {
+        "dataset_id": dataset_id,
+        "name": slice_name,
+    }
+    if slice_description is not Undefined:
+        params["description"] = slice_description
+    return params

spb_onprem/slices/params/slice.py

@@ -0,0 +1,42 @@
+from typing import Union
+
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+
+def slice_params(
+    dataset_id: str,
+    slice_id: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    name: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+):
+    """Get parameters for slice lookup.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to lookup the slice for.
+        slice_id (Optional[str]): The ID of the slice to lookup.
+        name (Optional[str]): The name of the slice to lookup.
+        
+    Returns:
+        dict: Parameters for slice lookup
+    """
+
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+
+    params = {
+        "dataset_id": dataset_id,
+    }
+    
+    if slice_id is not Undefined:
+        params["id"] = slice_id
+    elif name is not Undefined:
+        params["name"] = name
+    else:
+        raise BadParameterError("Either slice_id or name must be provided.")
+
+    return params

spb_onprem/slices/params/slices.py

@@ -0,0 +1,62 @@
+from typing import (
+    Union,
+    Optional,
+    List
+)
+from spb_onprem.base_model import CustomBaseModel, Field
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+
+class SlicesFilterOptions(CustomBaseModel):
+    """Options for filtering slices.
+    
+    Attributes:
+        name_contains: Filter slices by name containing this string
+        id_in: Filter slices by list of IDs
+    """
+    name_contains: Optional[str] = Field(None, alias="nameContains")
+    id_in: Optional[List[str]] = Field(None, alias="idIn")
+
+
+class SlicesFilter(CustomBaseModel):
+    """Filter criteria for slice queries.
+    
+    Attributes:
+        must_filter: Conditions that must be met
+        not_filter: Conditions that must not be met
+    """
+    must_filter: Optional[SlicesFilterOptions] = Field(None, alias="must")
+    not_filter: Optional[SlicesFilterOptions] = Field(None, alias="not")
+
+
+def slices_params(
+    dataset_id: str,
+    slices_filter: Union[
+        UndefinedType,
+        SlicesFilter
+    ],
+    cursor: Optional[str] = None,
+    length: Optional[int] = 10
+):
+    """Get the params for the slices query.
+
+    Args:
+        dataset_id (str): the dataset id to get the slices for.
+        slices_filter (Union[ UndefinedType, SlicesFilter ]): the filter for the slices.
+        cursor (Optional[str], optional): The next cursor to get the next page of slices. Defaults to None.
+        length (Optional[int], optional): The number of slices to get. Defaults to 10.
+    """
+
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+    params = {
+        "dataset_id": dataset_id,
+        "cursor": cursor,
+        "length": length,
+    }
+    if slices_filter is not Undefined:
+        params["filter"] = slices_filter.model_dump(
+            by_alias=True, exclude_unset=True
+        )
+
+    return params

spb_onprem/slices/params/update_slice.py

@@ -0,0 +1,45 @@
+from typing import Union
+
+from spb_onprem.base_types import UndefinedType, Undefined
+from spb_onprem.exceptions import BadParameterError
+
+
+def update_slice_params(
+    dataset_id: str,
+    slice_id: str,
+    slice_name: Union[
+        UndefinedType,
+        str
+    ] = Undefined,
+    slice_description: Union[
+        UndefinedType,
+        str
+    ] = Undefined
+):
+    """Update slice parameters.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to update the slice for.
+        slice_id (str): The ID of the slice to update.
+        slice_name (Optional[str]): The name of the slice to update.
+        slice_description (Optional[str]): The description of the slice to update.
+        
+    Returns:
+        dict: Parameters for slice update
+    """
+    if dataset_id is None:
+        raise BadParameterError("Dataset ID is required")
+    if slice_id is None:
+        raise BadParameterError("Slice ID is required")
+
+    variables = {
+        "dataset_id": dataset_id,
+        "id": slice_id,
+    }
+
+    if slice_name is not Undefined:
+        variables["name"] = slice_name
+    if slice_description is not Undefined:
+        variables["description"] = slice_description
+
+    return variables

spb_onprem/slices/queries.py

@@ -0,0 +1,121 @@
+from spb_onprem.slices.params import (
+    slices_params,
+    slice_params,
+    create_slice_params,
+    update_slice_params,
+)
+
+
+class Schemas:
+    """Schemas for slices queries
+    """
+    SLICE_PAGE= '''
+        slices {
+            id
+            name
+        }
+        next
+        totalCount
+    '''
+    
+    SLICE = '''
+        id
+        datasetId
+        name
+        description
+        isPinned
+        createdAt
+        createdBy
+        updatedAt
+        updatedBy
+    '''
+
+
+class Queries:
+    '''
+    Queries for slices
+    '''
+
+    GET_SLICES = {
+        "name": "slices",
+        "query": f'''
+            query slices(
+                $dataset_id: String!,
+                $filter: SliceFilter,
+                $cursor: String,
+                $length: Int
+            ) {{
+                slices(
+                    datasetId: $dataset_id,
+                    filter: $filter,
+                    cursor: $cursor,
+                    length: $length
+                ) {{
+                    {Schemas.SLICE_PAGE}
+                }}
+            }}
+        ''',
+        "variables": slices_params,
+    }
+    
+    GET_SLICE = {
+        "name": "slice",
+        "query": f'''
+            query slice(
+                $dataset_id: String!,
+                $id: ID
+                $name: String
+            ) {{
+                slice(
+                    datasetId: $dataset_id,
+                    id: $id
+                    name: $name
+                ) {{
+                    {Schemas.SLICE}
+                }}
+            }}
+        ''',
+        "variables": slice_params,
+    }
+    
+    CREATE_SLICE = {
+        "name": "createSlice",
+        "query": f'''
+            mutation createSlice(
+                $dataset_id: String!,
+                $name: String!,
+                $description: String
+            ) {{
+                createSlice(
+                    datasetId: $dataset_id,
+                    name: $name,
+                    description: $description
+                ) {{
+                    {Schemas.SLICE}
+                }}
+            }}
+        ''',
+        "variables": create_slice_params,
+    }
+    
+    UPDATE_SLICE = {
+        "name": "updateSlice",
+        "query": f'''
+            mutation updateSlice(
+                $dataset_id: String!,
+                $id: ID!,
+                $name: String,
+                $description: String
+            ) {{
+                updateSlice(
+                    datasetId: $dataset_id,
+                    id: $id,
+                    name: $name,
+                    description: $description
+                ) {{
+                    {Schemas.SLICE}
+                }}
+            }}
+        ''',
+        "variables": update_slice_params,
+    }