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

spb_onprem/models/params/models.py

@@ -0,0 +1,60 @@
+from typing import Optional, List, Union
+
+from spb_onprem.base_model import CustomBaseModel, Field
+from spb_onprem.base_types import Undefined, UndefinedType
+
+
+class ModelsFilterOptions(CustomBaseModel):
+    """Options for filtering models.
+    
+    Attributes:
+        name_contains: Filter models by name containing this string
+        id_in: Filter models by list of IDs
+        is_pinned: Filter models by pinned status
+        is_trained: Filter models by trained status
+    """
+    name_contains: Optional[str] = Field(None, alias="nameContains")
+    id_in: Optional[List[str]] = Field(None, alias="idIn")
+    is_pinned: Optional[bool] = Field(None, alias="isPinned")
+    is_trained: Optional[bool] = Field(None, alias="isTrained")
+
+
+class ModelsFilter(CustomBaseModel):
+    """Filter criteria for model queries.
+    
+    Attributes:
+        must_filter: Conditions that must be met
+        not_filter: Conditions that must not be met
+    """
+    must_filter: Optional[ModelsFilterOptions] = Field(None, alias="must")
+    not_filter: Optional[ModelsFilterOptions] = Field(None, alias="not")
+
+
+def models_params(
+    dataset_id: str,
+    models_filter: Union[
+        ModelsFilter,
+        UndefinedType
+    ] = Undefined,
+    cursor: Optional[str] = None,
+    length: Optional[int] = 10
+):
+    """Get parameters for listing models.
+    
+    Args:
+        dataset_id: Required dataset ID
+        models_filter: Optional filter criteria for models
+        cursor: Optional cursor for pagination
+        length: Optional number of items per page (default: 10)
+        
+    Returns:
+        dict: Parameters for listing models
+    """
+    return {
+        "datasetId": dataset_id,
+        "filter": models_filter.model_dump(
+            by_alias=True, exclude_unset=True
+        ) if models_filter and not isinstance(models_filter, UndefinedType) else None,
+        "cursor": cursor,
+        "length": length
+    }

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,
+    }