superb-ai-onprem 0.9.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

spb_onprem/models/params/pin_model.py

@@ -0,0 +1,17 @@
+def pin_model_params(
+    dataset_id: str,
+    model_id: str,
+):
+    """Get parameters for pinning a model.
+    
+    Args:
+        dataset_id: The dataset ID
+        model_id: The model ID
+        
+    Returns:
+        dict: Parameters for pinning a model
+    """
+    return {
+        "datasetId": dataset_id,
+        "id": model_id,
+    }

spb_onprem/models/params/unpin_model.py

@@ -0,0 +1,17 @@
+def unpin_model_params(
+    dataset_id: str,
+    model_id: str,
+):
+    """Get parameters for unpinning a model.
+    
+    Args:
+        dataset_id: The dataset ID
+        model_id: The model ID
+        
+    Returns:
+        dict: Parameters for unpinning a model
+    """
+    return {
+        "datasetId": dataset_id,
+        "id": model_id,
+    }

spb_onprem/models/params/update_model.py

@@ -0,0 +1,61 @@
+from typing import List, Union, Any
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.models.entities.model_train_class import ModelTrainClass
+
+
+def update_model_params(
+    dataset_id: str,
+    model_id: str,
+    name: Union[str, UndefinedType] = Undefined,
+    description: Union[str, UndefinedType] = Undefined,
+    training_classes: Union[List[ModelTrainClass], UndefinedType] = Undefined,
+    model_content_id: Union[str, UndefinedType] = Undefined,
+    is_trained: Union[bool, UndefinedType] = Undefined,
+    trained_at: Union[str, UndefinedType] = Undefined,
+    meta: Union[Any, UndefinedType] = Undefined,
+):
+    """Get parameters for updating a model.
+    
+    Args:
+        dataset_id: The dataset ID
+        model_id: The model ID
+        name: Optional new name
+        description: Optional new description
+        training_classes: Optional new training classes
+        model_content_id: Optional new model content ID
+        is_trained: Optional new trained status
+        trained_at: Optional new trained timestamp
+        meta: Optional new metadata
+        
+    Returns:
+        dict: Parameters for updating a model
+    """
+    params = {
+        "datasetId": dataset_id,
+        "id": model_id,
+    }
+    
+    if not isinstance(name, UndefinedType):
+        params["name"] = name
+    
+    if not isinstance(description, UndefinedType):
+        params["description"] = description
+    
+    if not isinstance(training_classes, UndefinedType):
+        params["trainingClasses"] = [
+            tc.model_dump(by_alias=True, exclude_unset=True) for tc in training_classes
+        ] if training_classes is not None else None
+    
+    if not isinstance(model_content_id, UndefinedType):
+        params["modelContentId"] = model_content_id
+    
+    if not isinstance(is_trained, UndefinedType):
+        params["isTrained"] = is_trained
+    
+    if not isinstance(trained_at, UndefinedType):
+        params["trainedAt"] = trained_at
+    
+    if not isinstance(meta, UndefinedType):
+        params["meta"] = meta
+    
+    return params

spb_onprem/models/queries.py

@@ -1,33 +1,238 @@
-from .params import (
-    get_models_params,
+from spb_onprem.models.params import (
+    models_params,
+    model_params,
+    create_model_params,
+    update_model_params,
+    pin_model_params,
+    unpin_model_params,
     delete_model_params,
 )
 
 
-class Queries:
-    GET_MODELS = {
-        "name": "getModels",
-        "query": '''
-            query GetModels($dataset_id: String!, $filter: ModelFilter, $cursor: String, $length: Int) {
-                models(datasetId: $dataset_id, filter: $filter, cursor: $cursor, length: $length) {
-                    models {
-                        id
-                        name
-                    }
+class Schemas:
+    MODEL_TRAIN_CLASS = '''
+        class
+        annotationType
+        ap
+        trainingAnnotationsCount
+        validationAnnotationsCount
+    '''
+    
+    MODEL = '''
+        id
+        datasetId
+        baselineModel
+        name
+        description
+        trainingClasses {
+            class
+            annotationType
+            ap
+            trainingAnnotationsCount
+            validationAnnotationsCount
+        }
+        trainingDataCount
+        validationDataCount
+        isPinned
+        isTrained
+        trainedAt
+        modelContent {
+            id
+            downloadURL
+        }
+        createdBy
+        createdAt
+        updatedBy
+        updatedAt
+        meta
+        trainingSlices {
+            id
+            datasetId
+            name
+            description
+            isPinned
+            createdAt
+            createdBy
+            updatedAt
+            updatedBy
+        }
+        validationSlices {
+            id
+            datasetId
+            name
+            description
+            isPinned
+            createdAt
+            createdBy
+            updatedAt
+            updatedBy
+        }
+    '''
+
+
+class Queries():
+    MODELS = {
+        "name": "models",
+        "query": f'''
+            query Models(
+                $datasetId: ID!,
+                $filter: ModelFilter,
+                $cursor: String,
+                $length: Int
+            ) {{
+                models(
+                    datasetId: $datasetId,
+                    filter: $filter,
+                    cursor: $cursor,
+                    length: $length
+                ) {{
+                    models {{
+                        {Schemas.MODEL}
+                    }}
                     next
                     totalCount
-                }
-            }
+                }}
+            }}
+        ''',
+        "variables": models_params,
+    }
+    
+    MODEL = {
+        "name": "model",
+        "query": f'''
+            query Model(
+                $datasetId: ID!,
+                $modelId: ID!
+            ) {{
+                model(
+                    datasetId: $datasetId,
+                    id: $modelId
+                ) {{
+                    {Schemas.MODEL}
+                }}
+            }}
+        ''',
+        "variables": model_params,
+    }
+    
+    CREATE_MODEL = {
+        "name": "createModel",
+        "query": f'''
+            mutation CreateModel(
+                $datasetId: ID!,
+                $name: String!,
+                $description: String,
+                $baselineModel: String!,
+                $trainingClasses: [ModelTrainClassInput!],
+                $trainingSliceIds: [ID!]!,
+                $validationSliceIds: [ID!]!,
+                $modelContentId: String,
+                $isTrained: Boolean,
+                $trainedAt: DateTime,
+                $isPinned: Boolean,
+                $meta: JSONObject
+            ) {{
+                createModel(
+                    datasetId: $datasetId,
+                    name: $name,
+                    description: $description,
+                    baselineModel: $baselineModel,
+                    trainingClasses: $trainingClasses,
+                    trainingSliceIds: $trainingSliceIds,
+                    validationSliceIds: $validationSliceIds,
+                    modelContentId: $modelContentId,
+                    isTrained: $isTrained,
+                    trainedAt: $trainedAt,
+                    isPinned: $isPinned,
+                    meta: $meta
+                ) {{
+                    {Schemas.MODEL}
+                }}
+            }}
         ''',
-        "variables": get_models_params
+        "variables": create_model_params,
+    }
+    
+    UPDATE_MODEL = {
+        "name": "updateModel",
+        "query": f'''
+            mutation UpdateModel(
+                $datasetId: ID!,
+                $id: ID!,
+                $name: String,
+                $description: String,
+                $trainingClasses: [ModelTrainClassInput!],
+                $modelContentId: String,
+                $isTrained: Boolean,
+                $trainedAt: DateTime,
+                $meta: JSONObject
+            ) {{
+                updateModel(
+                    datasetId: $datasetId,
+                    id: $id,
+                    name: $name,
+                    description: $description,
+                    trainingClasses: $trainingClasses,
+                    modelContentId: $modelContentId,
+                    isTrained: $isTrained,
+                    trainedAt: $trainedAt,
+                    meta: $meta
+                ) {{
+                    {Schemas.MODEL}
+                }}
+            }}
+        ''',
+        "variables": update_model_params,
+    }
+    
+    PIN_MODEL = {
+        "name": "pinModel",
+        "query": f'''
+            mutation PinModel(
+                $datasetId: ID!,
+                $id: ID!
+            ) {{
+                pinModel(
+                    datasetId: $datasetId,
+                    id: $id
+                ) {{
+                    {Schemas.MODEL}
+                }}
+            }}
+        ''',
+        "variables": pin_model_params,
+    }
+    
+    UNPIN_MODEL = {
+        "name": "unpinModel",
+        "query": f'''
+            mutation UnpinModel(
+                $datasetId: ID!,
+                $id: ID!
+            ) {{
+                unpinModel(
+                    datasetId: $datasetId,
+                    id: $id
+                ) {{
+                    {Schemas.MODEL}
+                }}
+            }}
+        ''',
+        "variables": unpin_model_params,
     }
     
     DELETE_MODEL = {
         "name": "deleteModel",
         "query": '''
-            mutation DeleteModel($dataset_id: String!, $id: String!) {
-                deleteModel(datasetId: $dataset_id, id: $id)
+            mutation DeleteModel(
+                $datasetId: ID!,
+                $id: ID!
+            ) {
+                deleteModel(
+                    datasetId: $datasetId,
+                    id: $id
+                )
             }
         ''',
-        "variables": delete_model_params
-    }
+        "variables": delete_model_params,
+    }

spb_onprem/models/service.py

@@ -1,68 +1,289 @@
-from typing import Optional, Dict, Any, List, Tuple, Union
-
+from typing import Optional, Union, List, Any
 from spb_onprem.base_service import BaseService
-from spb_onprem.base_types import Undefined, UndefinedType
 from spb_onprem.exceptions import BadParameterError
+from spb_onprem.base_types import Undefined, UndefinedType
 from .queries import Queries
-from .entities import Model
+from .entities import Model, ModelPageInfo, ModelTrainClass
+from .params.models import ModelsFilter
 
 
 class ModelService(BaseService):
-    """Service class for handling model operations."""
+    """
+    Service class for handling model-related operations.
+    """
     
     def get_models(
         self,
         dataset_id: str,
-        filter: Union[UndefinedType, Dict[str, Any]] = Undefined,
-        cursor: Union[UndefinedType, str] = Undefined,
-        length: int = 50
-    ) -> Tuple[List[Model], Optional[str], int]:
-        """Get paginated list of models for a dataset.
+        models_filter: Optional[ModelsFilter] = None,
+        cursor: Optional[str] = None,
+        length: Optional[int] = 10
+    ):
+        """
+        Get a list of models based on the provided filter and pagination parameters.
         
         Args:
-            dataset_id (str): The dataset ID.
-            filter (Union[UndefinedType, Dict[str, Any]]): Filter for models.
-            cursor (Union[UndefinedType, str]): Cursor for pagination.
-            length (int): Number of items to retrieve per page.
+            dataset_id (str): The dataset ID
+            models_filter (Optional[ModelsFilter]): Filter criteria for models
+            cursor (Optional[str]): Cursor for pagination
+            length (Optional[int]): Number of items per page (default: 10)
         
         Returns:
-            Tuple[List[Model], Optional[str], int]: A tuple containing models, next cursor, and total count.
+            tuple: A tuple containing:
+                - List[Model]: A list of Model objects
+                - str: Next cursor for pagination
+                - int: Total count of models
         """
         if dataset_id is None:
             raise BadParameterError("dataset_id is required.")
-
+        
+        if length > 50:
+            raise BadParameterError("The maximum length is 50.")
+        
         response = self.request_gql(
-            Queries.GET_MODELS,
-            Queries.GET_MODELS["variables"](
+            Queries.MODELS,
+            Queries.MODELS["variables"](
                 dataset_id=dataset_id,
-                filter=filter,
+                models_filter=models_filter,
                 cursor=cursor,
                 length=length
             )
         )
-        models_list = response.get("models", [])
+        
+        page_info = ModelPageInfo.model_validate(response)
         return (
-            [Model.model_validate(model) for model in models_list],
-            response.get("next"),
-            response.get("totalCount", 0)
+            page_info.models or [],
+            page_info.next,
+            page_info.total_count or 0
+        )
+
+    def get_model(
+        self,
+        dataset_id: str,
+        model_id: str,
+    ):
+        """
+        Retrieve a model by its ID.
+
+        Args:
+            dataset_id (str): The dataset ID
+            model_id (str): The ID of the model to retrieve
+
+        Returns:
+            Model: The retrieved model object
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+        
+        if model_id is None:
+            raise BadParameterError("model_id is required.")
+        
+        response = self.request_gql(
+            Queries.MODEL,
+            Queries.MODEL["variables"](
+                dataset_id=dataset_id,
+                model_id=model_id
+            ),
+        )
+        return Model.model_validate(response)
+    
+    def create_model(
+        self,
+        dataset_id: str,
+        name: str,
+        baseline_model: str,
+        training_slice_ids: List[str],
+        validation_slice_ids: List[str],
+        description: Union[str, UndefinedType] = Undefined,
+        training_classes: Union[List[ModelTrainClass], UndefinedType] = Undefined,
+        model_content_id: Union[str, UndefinedType] = Undefined,
+        is_trained: Union[bool, UndefinedType] = Undefined,
+        trained_at: Union[str, UndefinedType] = Undefined,
+        is_pinned: Union[bool, UndefinedType] = Undefined,
+        meta: Union[Any, UndefinedType] = Undefined,
+    ):
+        """
+        Create a new model.
+
+        Args:
+            dataset_id (str): The dataset ID
+            name (str): The model name
+            baseline_model (str): The baseline model used
+            training_slice_ids (List[str]): The IDs of the training slices
+            validation_slice_ids (List[str]): The IDs of the validation slices
+            description (Optional[str]): The description of the model
+            training_classes (Optional[List[ModelTrainClass]]): The training classes
+            model_content_id (Optional[str]): The model content ID
+            is_trained (Optional[bool]): Whether the model is trained
+            trained_at (Optional[str]): When the model was trained
+            is_pinned (Optional[bool]): Whether the model is pinned
+            meta (Optional[Any]): The metadata of the model
+
+        Returns:
+            Model: The created model object
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+        
+        if name is None:
+            raise BadParameterError("name is required.")
+        
+        if baseline_model is None:
+            raise BadParameterError("baseline_model is required.")
+        
+        if training_slice_ids is None:
+            raise BadParameterError("training_slice_ids is required.")
+        
+        if validation_slice_ids is None:
+            raise BadParameterError("validation_slice_ids is required.")
+        
+        response = self.request_gql(
+            Queries.CREATE_MODEL,
+            Queries.CREATE_MODEL["variables"](
+                dataset_id=dataset_id,
+                name=name,
+                baseline_model=baseline_model,
+                training_slice_ids=training_slice_ids,
+                validation_slice_ids=validation_slice_ids,
+                description=description,
+                training_classes=training_classes,
+                model_content_id=model_content_id,
+                is_trained=is_trained,
+                trained_at=trained_at,
+                is_pinned=is_pinned,
+                meta=meta,
+            ),
+        )
+        return Model.model_validate(response)
+    
+    def update_model(
+        self,
+        dataset_id: str,
+        model_id: str,
+        name: Union[str, UndefinedType] = Undefined,
+        description: Union[str, UndefinedType] = Undefined,
+        training_classes: Union[List[ModelTrainClass], UndefinedType] = Undefined,
+        model_content_id: Union[str, UndefinedType] = Undefined,
+        is_trained: Union[bool, UndefinedType] = Undefined,
+        trained_at: Union[str, UndefinedType] = Undefined,
+        meta: Union[Any, UndefinedType] = Undefined,
+    ):
+        """
+        Update a model.
+
+        Args:
+            dataset_id (str): The dataset ID
+            model_id (str): The ID of the model to update
+            name (Optional[str]): The new name
+            description (Optional[str]): The new description
+            training_classes (Optional[List[ModelTrainClass]]): The new training classes
+            model_content_id (Optional[str]): The new model content ID
+            is_trained (Optional[bool]): The new trained status
+            trained_at (Optional[str]): The new trained timestamp
+            meta (Optional[Any]): The new metadata
+
+        Returns:
+            Model: The updated model object
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+        
+        if model_id is None:
+            raise BadParameterError("model_id is required.")
+        
+        response = self.request_gql(
+            Queries.UPDATE_MODEL,
+            Queries.UPDATE_MODEL["variables"](
+                dataset_id=dataset_id,
+                model_id=model_id,
+                name=name,
+                description=description,
+                training_classes=training_classes,
+                model_content_id=model_content_id,
+                is_trained=is_trained,
+                trained_at=trained_at,
+                meta=meta,
+            ),
         )
+        return Model.model_validate(response)
+    
+    def pin_model(
+        self,
+        dataset_id: str,
+        model_id: str,
+    ):
+        """
+        Pin a model.
+
+        Args:
+            dataset_id (str): The dataset ID
+            model_id (str): The ID of the model to pin
+
+        Returns:
+            Model: The pinned model object
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+        
+        if model_id is None:
+            raise BadParameterError("model_id is required.")
+        
+        response = self.request_gql(
+            Queries.PIN_MODEL,
+            Queries.PIN_MODEL["variables"](
+                dataset_id=dataset_id,
+                model_id=model_id,
+            ),
+        )
+        return Model.model_validate(response)
+    
+    def unpin_model(
+        self,
+        dataset_id: str,
+        model_id: str,
+    ):
+        """
+        Unpin a model.
+
+        Args:
+            dataset_id (str): The dataset ID
+            model_id (str): The ID of the model to unpin
+
+        Returns:
+            Model: The unpinned model object
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+        
+        if model_id is None:
+            raise BadParameterError("model_id is required.")
+        
+        response = self.request_gql(
+            Queries.UNPIN_MODEL,
+            Queries.UNPIN_MODEL["variables"](
+                dataset_id=dataset_id,
+                model_id=model_id,
+            ),
+        )
+        return Model.model_validate(response)
     
     def delete_model(
         self,
         dataset_id: str,
-        model_id: str
+        model_id: str,
     ) -> bool:
-        """Delete a model from the dataset.
+        """Delete a model.
         
         Args:
-            dataset_id (str): The dataset ID.
-            model_id (str): The model ID to delete.
+            dataset_id (str): The dataset ID
+            model_id (str): The ID of the model to delete
         
         Returns:
-            bool: True if deletion was successful.
+            bool: True if deletion was successful
         """
         if dataset_id is None:
             raise BadParameterError("dataset_id is required.")
+        
         if model_id is None:
             raise BadParameterError("model_id is required.")
 
@@ -70,7 +291,7 @@ class ModelService(BaseService):
             Queries.DELETE_MODEL,
             Queries.DELETE_MODEL["variables"](
                 dataset_id=dataset_id,
-                model_id=model_id
+                model_id=model_id,
             )
         )
-        return response
+        return response