superb-ai-onprem 0.5.9__py3-none-any.whl → 0.5.11__py3-none-any.whl

spb_onprem/datasets/params/delete_dataset.py

@@ -0,0 +1,12 @@
+def delete_dataset_params(dataset_id: str):
+    """Generate variables for delete dataset GraphQL mutation.
+    
+    Args:
+        dataset_id (str): The ID of the dataset to delete.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL mutation.
+    """
+    return {
+        "dataset_id": dataset_id
+    }

spb_onprem/datasets/queries.py

@@ -3,6 +3,7 @@ from spb_onprem.datasets.params import (
     datasets_params,
     create_dataset_params,
     update_dataset_params,
+    delete_dataset_params,
 )
 
 class Schemas:
@@ -77,3 +78,13 @@ class Queries():
         ''',
         "variables": update_dataset_params,
     }
+    
+    DELETE_DATASET = {
+        "name": "deleteDataset",
+        "query": '''
+            mutation DeleteDataset($dataset_id: String!) {
+                deleteDataset(datasetId: $dataset_id)
+            }
+        ''',
+        "variables": delete_dataset_params,
+    }

spb_onprem/datasets/service.py

@@ -130,3 +130,21 @@ class DatasetService(BaseService):
             ),
         )
         return Dataset.model_validate(response)
+    
+    def delete_dataset(self, dataset_id: str) -> bool:
+        """Delete the dataset.
+        
+        Args:
+            dataset_id (str): The ID of the dataset to delete.
+        
+        Returns:
+            bool: True if deletion was successful.
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+
+        response = self.request_gql(
+            Queries.DELETE_DATASET,
+            Queries.DELETE_DATASET["variables"](dataset_id=dataset_id)
+        )
+        return response.get("deleteDataset", False)

spb_onprem/entities.py

@@ -26,6 +26,8 @@ from .activities.entities import (
 )
 from .exports.entities import Export
 from .contents.entities import Content
+from .predictions.entities import PredictionSet
+from .models.entities import Model
 
 __all__ = [
     # Core Entities
@@ -43,6 +45,8 @@ __all__ = [
     "Export",
     "Content",
     "Frame",
+    "PredictionSet",
+    "Model",
 
     # Enums
     "DataType",

spb_onprem/inferences/__init__.py

@@ -0,0 +1,5 @@
+from .service import InferService
+
+__all__ = [
+    "InferService",
+]

spb_onprem/inferences/service.py

@@ -0,0 +1,56 @@
+import os
+from typing import Dict, Any
+
+from spb_onprem.base_service import BaseService
+from spb_onprem.exceptions import BadParameterError
+
+
+class InferService(BaseService):
+    """Service class for handling inference operations."""
+    
+    def __init__(self):
+        super().__init__()
+        self.model_endpoint_url = os.environ.get("MODEL_ENDPOINT_URL")
+        if not self.model_endpoint_url:
+            raise ValueError("MODEL_ENDPOINT_URL environment variable is required")
+    
+    def infer_data(
+        self,
+        model_id: str,
+        base64_image: str
+    ) -> Dict[str, Any]:
+        """Perform inference on image data using the specified model.
+        
+        Args:
+            model_id (str): The model ID to use for inference.
+            base64_image (str): Base64 encoded image data.
+        
+        Returns:
+            Dict[str, Any]: The inference result.
+        
+        Raises:
+            BadParameterError: If required parameters are missing.
+        """
+        if model_id is None:
+            raise BadParameterError("model_id is required.")
+        if base64_image is None:
+            raise BadParameterError("base64_image is required.")
+            
+        payload = {
+            "model_id": model_id,
+            "base64_image": base64_image,
+        }
+        
+        try:
+            response = self.request(
+                method="POST",
+                url=self.model_endpoint_url,
+                json_data=payload,
+                headers={
+                    "Content-Type": "application/json"
+                }
+            )
+            return response
+        except Exception as e:
+            print(f"Failed to fetch inference result: {e}")
+            raise

spb_onprem/models/__init__.py

@@ -0,0 +1,7 @@
+from .service import ModelService
+from .entities import Model
+
+__all__ = [
+    "ModelService",
+    "Model",
+]

spb_onprem/models/entities.py

@@ -0,0 +1,9 @@
+from typing import Optional
+from spb_onprem.base_model import BaseModel
+
+
+class Model(BaseModel):
+    """Model entity representing a model in the dataset."""
+    
+    id: str
+    name: Optional[str] = None

spb_onprem/models/params/__init__.py

@@ -0,0 +1,7 @@
+from .get_models import get_models_params
+from .delete_model import delete_model_params
+
+__all__ = [
+    "get_models_params",
+    "delete_model_params",
+]

spb_onprem/models/params/delete_model.py

@@ -0,0 +1,14 @@
+def delete_model_params(dataset_id: str, model_id: str):
+    """Generate variables for delete model GraphQL mutation.
+    
+    Args:
+        dataset_id (str): The ID of the dataset.
+        model_id (str): The ID of the model to delete.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL mutation.
+    """
+    return {
+        "dataset_id": dataset_id,
+        "id": model_id
+    }

spb_onprem/models/params/get_models.py

@@ -0,0 +1,29 @@
+def get_models_params(
+    dataset_id: str,
+    filter: dict = None,
+    cursor: str = None,
+    length: int = 50
+):
+    """Generate variables for get models GraphQL query.
+    
+    Args:
+        dataset_id (str): The ID of the dataset.
+        filter (dict, optional): Filter for models.
+        cursor (str, optional): Cursor for pagination.
+        length (int): Number of items to retrieve per page.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL query.
+    """
+    params = {
+        "dataset_id": dataset_id,
+        "length": length
+    }
+    
+    if filter is not None:
+        params["filter"] = filter
+        
+    if cursor is not None:
+        params["cursor"] = cursor
+        
+    return params

spb_onprem/models/queries.py

@@ -0,0 +1,33 @@
+from .params import (
+    get_models_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
+                    }
+                    next
+                    totalCount
+                }
+            }
+        ''',
+        "variables": get_models_params
+    }
+    
+    DELETE_MODEL = {
+        "name": "deleteModel",
+        "query": '''
+            mutation DeleteModel($dataset_id: String!, $id: String!) {
+                deleteModel(datasetId: $dataset_id, id: $id)
+            }
+        ''',
+        "variables": delete_model_params
+    }

spb_onprem/models/service.py

@@ -0,0 +1,76 @@
+from typing import Optional, Dict, Any, List, Tuple, Union
+
+from spb_onprem.base_service import BaseService
+from spb_onprem.base_types import Undefined, UndefinedType
+from spb_onprem.exceptions import BadParameterError
+from .queries import Queries
+from .entities import Model
+
+
+class ModelService(BaseService):
+    """Service class for handling model 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.
+        
+        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.
+        
+        Returns:
+            Tuple[List[Model], Optional[str], int]: A tuple containing models, next cursor, and total count.
+        """
+        if dataset_id is None:
+            raise BadParameterError("dataset_id is required.")
+
+        response = self.request_gql(
+            Queries.GET_MODELS,
+            Queries.GET_MODELS["variables"](
+                dataset_id=dataset_id,
+                filter=filter,
+                cursor=cursor,
+                length=length
+            )
+        )
+        models_list = response.get("models", [])
+        return (
+            [Model.model_validate(model) for model in models_list],
+            response.get("next"),
+            response.get("totalCount", 0)
+        )
+    
+    def delete_model(
+        self,
+        dataset_id: str,
+        model_id: str
+    ) -> bool:
+        """Delete a model from the dataset.
+        
+        Args:
+            dataset_id (str): The dataset ID.
+            model_id (str): The model ID to delete.
+        
+        Returns:
+            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.")
+
+        response = self.request_gql(
+            Queries.DELETE_MODEL,
+            Queries.DELETE_MODEL["variables"](
+                dataset_id=dataset_id,
+                model_id=model_id
+            )
+        )
+        return response.get("deleteModel", False)

spb_onprem/predictions/__init__.py

@@ -0,0 +1,7 @@
+from .service import PredictionService
+from .entities import PredictionSet
+
+__all__ = [
+    "PredictionService",
+    "PredictionSet",
+]

spb_onprem/predictions/entities.py

@@ -0,0 +1,11 @@
+from typing import List, Optional
+from spb_onprem.base_model import CustomBaseModel, Field
+
+
+class PredictionSet(CustomBaseModel):
+    """PredictionSet entity representing a set of predictions in the dataset."""
+    
+    id: str
+    name: Optional[str] = None
+    annotations_contents: Optional[List[str]] = Field(None, alias="annotationsContents")
+    evaluation_result_content: Optional[dict] = Field(None, alias="evaluationResultContent")

spb_onprem/predictions/params/__init__.py

@@ -0,0 +1,15 @@
+from .get_prediction_sets import get_prediction_sets_params
+from .get_prediction_set import get_prediction_set_params
+from .delete_prediction_set import delete_prediction_set_params
+from .delete_prediction_from_data import delete_prediction_from_data_params
+from .create_prediction_set import create_prediction_set_params
+from .update_prediction_set_data_info import update_prediction_set_data_info_params
+
+__all__ = [
+    "get_prediction_sets_params",
+    "get_prediction_set_params", 
+    "delete_prediction_set_params",
+    "delete_prediction_from_data_params",
+    "create_prediction_set_params",
+    "update_prediction_set_data_info_params",
+]

spb_onprem/predictions/params/create_prediction_set.py

@@ -0,0 +1,44 @@
+from typing import Optional, Dict, Any
+
+
+def create_prediction_set_params(
+    dataset_id: str,
+    model_id: str,
+    name: str,
+    type: str,
+    description: Optional[str] = None,
+    annotations_count: Optional[int] = None,
+    data_count: Optional[int] = None,
+) -> Dict[str, Any]:
+    """Convert parameters for create_prediction_set GraphQL mutation.
+    
+    Args:
+        dataset_id (str): The dataset ID
+        model_id (str): The model ID
+        name (str): Name of the prediction set
+        type (str): Type of the prediction set
+        description (Optional[str]): Description of the prediction set
+        annotations_count (Optional[int]): Number of annotations
+        data_count (Optional[int]): Number of data items
+        
+    Returns:
+        Dict[str, Any]: Parameters formatted for GraphQL mutation
+    """
+    params = {
+        "datasetId": dataset_id,
+        "modelId": model_id,
+        "name": name,
+        "type": type,
+    }
+    
+    if description is not None:
+        params["description"] = description
+        
+    if annotations_count is not None or data_count is not None:
+        params["dataInfo"] = {}
+        if annotations_count is not None:
+            params["dataInfo"]["annotationsCount"] = annotations_count
+        if data_count is not None:
+            params["dataInfo"]["dataCount"] = data_count
+    
+    return params

spb_onprem/predictions/params/delete_prediction_from_data.py

@@ -0,0 +1,20 @@
+def delete_prediction_from_data_params(
+    dataset_id: str, 
+    data_id: str, 
+    prediction_set_id: str
+):
+    """Generate variables for delete prediction from data GraphQL mutation.
+    
+    Args:
+        dataset_id (str): The ID of the dataset.
+        data_id (str): The ID of the data.
+        prediction_set_id (str): The ID of the prediction set.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL mutation.
+    """
+    return {
+        "dataset_id": dataset_id,
+        "data_id": data_id,
+        "prediction_set_id": prediction_set_id
+    }

spb_onprem/predictions/params/delete_prediction_set.py

@@ -0,0 +1,14 @@
+def delete_prediction_set_params(dataset_id: str, prediction_set_id: str):
+    """Generate variables for delete prediction set GraphQL mutation.
+    
+    Args:
+        dataset_id (str): The ID of the dataset.
+        prediction_set_id (str): The ID of the prediction set to delete.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL mutation.
+    """
+    return {
+        "dataset_id": dataset_id,
+        "id": prediction_set_id
+    }

spb_onprem/predictions/params/get_prediction_set.py

@@ -0,0 +1,14 @@
+def get_prediction_set_params(dataset_id: str, prediction_set_id: str):
+    """Generate variables for get prediction set GraphQL query.
+    
+    Args:
+        dataset_id (str): The ID of the dataset.
+        prediction_set_id (str): The ID of the prediction set.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL query.
+    """
+    return {
+        "dataset_id": dataset_id,
+        "id": prediction_set_id
+    }

spb_onprem/predictions/params/get_prediction_sets.py

@@ -0,0 +1,29 @@
+def get_prediction_sets_params(
+    dataset_id: str,
+    filter: dict = None,
+    cursor: str = None,
+    length: int = 50
+):
+    """Generate variables for get prediction sets GraphQL query.
+    
+    Args:
+        dataset_id (str): The ID of the dataset.
+        filter (dict, optional): Filter for prediction sets.
+        cursor (str, optional): Cursor for pagination.
+        length (int): Number of items to retrieve per page.
+        
+    Returns:
+        dict: Variables dictionary for the GraphQL query.
+    """
+    params = {
+        "dataset_id": dataset_id,
+        "length": length
+    }
+    
+    if filter is not None:
+        params["filter"] = filter
+        
+    if cursor is not None:
+        params["cursor"] = cursor
+        
+    return params

spb_onprem/predictions/params/update_prediction_set_data_info.py

@@ -0,0 +1,28 @@
+from typing import Dict, Any
+
+
+def update_prediction_set_data_info_params(
+    dataset_id: str,
+    id: str,
+    annotation_count: int,
+    data_count: int
+) -> Dict[str, Any]:
+    """Convert parameters for update_prediction_set_data_info GraphQL mutation.
+    
+    Args:
+        dataset_id (str): The dataset ID
+        id (str): The prediction set ID to update
+        annotation_count (int): Number of annotations
+        data_count (int): Number of data items
+        
+    Returns:
+        Dict[str, Any]: Parameters formatted for GraphQL mutation
+    """
+    return {
+        "datasetId": dataset_id,
+        "updatePredictionSetId": id,
+        "dataInfo": {
+            "annotationsCount": annotation_count,
+            "dataCount": data_count,
+        }
+    }

spb_onprem/predictions/queries.py

@@ -0,0 +1,110 @@
+from .params import (
+    get_prediction_sets_params,
+    get_prediction_set_params,
+    delete_prediction_set_params,
+    delete_prediction_from_data_params,
+    create_prediction_set_params,
+    update_prediction_set_data_info_params,
+)
+
+
+class Queries:
+    GET_PREDICTION_SETS = {
+        "name": "getPredictionSets",
+        "query": '''
+            query GetPredictionSets($dataset_id: String!, $filter: PredictionSetFilter, $cursor: String, $length: Int) {
+                predictionSets(datasetId: $dataset_id, filter: $filter, cursor: $cursor, length: $length) {
+                    predictionSets {
+                        id
+                        name
+                    }
+                    next
+                    totalCount
+                }
+            }
+        ''',
+        "variables": get_prediction_sets_params
+    }
+    
+    GET_PREDICTION_SET = {
+        "name": "getPredictionSet",
+        "query": '''
+            query GetPredictionSet($dataset_id: String!, $id: String!) {
+                predictionSet(datasetId: $dataset_id, id: $id) {
+                    id
+                    name
+                    annotationsContents
+                    evaluationResultContent {
+                        id
+                    }
+                }
+            }
+        ''',
+        "variables": get_prediction_set_params
+    }
+    
+    DELETE_PREDICTION_SET = {
+        "name": "deletePredictionSet",
+        "query": '''
+            mutation DeletePredictionSet($dataset_id: String!, $id: String!) {
+                deletePredictionSet(datasetId: $dataset_id, id: $id)
+            }
+        ''',
+        "variables": delete_prediction_set_params
+    }
+    
+    DELETE_PREDICTION_FROM_DATA = {
+        "name": "deletePredictionFromData",
+        "query": '''
+            mutation DeletePrediction($dataset_id: String!, $data_id: String!, $set_id: String!) {
+                deletePrediction(datasetId: $dataset_id, dataId: $data_id, setId: $set_id)
+            }
+        ''',
+        "variables": delete_prediction_from_data_params
+    }
+    
+    CREATE_PREDICTION_SET = {
+        "name": "createPredictionSet",
+        "query": '''
+            mutation CreatePredictionSet(
+                $datasetId: ID!,
+                $modelId: ID!,
+                $name: String!,
+                $description: String,
+                $type: PredictionSetTypes!,
+                $dataInfo: PredictionSetDataInfoInput
+            ) {
+                createPredictionSet(
+                    datasetId: $datasetId,
+                    modelId: $modelId,
+                    name: $name,
+                    description: $description,
+                    type: $type,
+                    dataInfo: $dataInfo
+                ) {
+                    id
+                }
+            }
+        ''',
+        "variables": create_prediction_set_params
+    }
+    
+    UPDATE_PREDICTION_SET_DATA_INFO = {
+        "name": "updatePredictionSet",
+        "query": '''
+            mutation UpdatePredictionSet(
+                $datasetId: ID!,
+                $updatePredictionSetId: ID!,
+                $dataInfo: PredictionSetDataInfoInput
+            ) {
+                updatePredictionSet(
+                    datasetId: $datasetId,
+                    id: $updatePredictionSetId,
+                    dataInfo: $dataInfo
+                ) {
+                    id
+                }
+            }
+        ''',
+        "variables": update_prediction_set_data_info_params
+    }