rapidata 2.36.0__py3-none-any.whl → 2.36.2__py3-none-any.whl

rapidata/rapidata_client/datapoints/{datapoint.py → _datapoint.py}

@@ -1,21 +1,49 @@
 from typing import Sequence, cast
-from rapidata.rapidata_client.datapoints.assets import MediaAsset, TextAsset, MultiAsset, BaseAsset
+from rapidata.rapidata_client.datapoints.assets import (
+    MediaAsset,
+    TextAsset,
+    MultiAsset,
+    BaseAsset,
+)
+from rapidata.rapidata_client.datapoints.assets.constants import (
+    ALLOWED_VIDEO_EXTENSIONS,
+    ALLOWED_IMAGE_EXTENSIONS,
+    ALLOWED_AUDIO_EXTENSIONS,
+)
 from rapidata.rapidata_client.datapoints.metadata import Metadata
-from rapidata.api_client.models.dataset_dataset_id_datapoints_post_request_metadata_inner import DatasetDatasetIdDatapointsPostRequestMetadataInner
-from rapidata.api_client.models.create_datapoint_from_text_sources_model import CreateDatapointFromTextSourcesModel
+from rapidata.api_client.models.dataset_dataset_id_datapoints_post_request_metadata_inner import (
+    DatasetDatasetIdDatapointsPostRequestMetadataInner,
+)
+from rapidata.api_client.models.create_datapoint_from_text_sources_model import (
+    CreateDatapointFromTextSourcesModel,
+)
 from pydantic import StrictStr, StrictBytes
+from rapidata.api_client.models.asset_type import AssetType
+from rapidata.api_client.models.prompt_type import PromptType
+from rapidata.rapidata_client.datapoints.metadata._media_asset_metadata import (
+    MediaAssetMetadata,
+)
+from rapidata.rapidata_client.datapoints.metadata._prompt_metadata import PromptMetadata
+from rapidata.rapidata_client.logging import logger
+
 
 class Datapoint:
-    def __init__(self, asset: MediaAsset | TextAsset | MultiAsset, metadata: Sequence[Metadata] | None = None):
+    def __init__(
+        self,
+        asset: MediaAsset | TextAsset | MultiAsset,
+        metadata: Sequence[Metadata] | None = None,
+    ):
         if not isinstance(asset, (MediaAsset, TextAsset, MultiAsset)):
-            raise TypeError("Asset must be of type MediaAsset, TextAsset, or MultiAsset.")
-        
+            raise TypeError(
+                "Asset must be of type MediaAsset, TextAsset, or MultiAsset."
+            )
+
         if metadata and not isinstance(metadata, Sequence):
             raise TypeError("Metadata must be a list of Metadata objects.")
-        
+
         if metadata and not all(isinstance(m, Metadata) for m in metadata):
             raise TypeError("All metadata objects must be of type Metadata.")
-        
+
         self.asset = asset
         self.metadata = metadata
 
@@ -35,6 +63,45 @@ class Datapoint:
         effective_type = self._get_effective_asset_type()
         return issubclass(effective_type, TextAsset)
 
+    def get_asset_type(self) -> AssetType:
+        """Get the asset type of the datapoint."""
+        if self.is_text_asset():
+            return AssetType.TEXT
+        elif self.is_media_asset():
+            if isinstance(self.asset, MultiAsset):
+                asset = self.asset.assets[0]
+            else:
+                asset = self.asset
+            assert isinstance(asset, MediaAsset)
+            if any(asset.path.endswith(ext) for ext in ALLOWED_IMAGE_EXTENSIONS):
+                return AssetType.IMAGE
+            elif any(asset.path.endswith(ext) for ext in ALLOWED_VIDEO_EXTENSIONS):
+                return AssetType.VIDEO
+            elif any(asset.path.endswith(ext) for ext in ALLOWED_AUDIO_EXTENSIONS):
+                return AssetType.AUDIO
+            else:
+                logger.debug(
+                    f"Cannot get asset type for asset type: {type(self.asset)}"
+                )
+                return AssetType.NONE
+        else:
+            logger.debug(f"Cannot get asset type for asset type: {type(self.asset)}")
+            return AssetType.NONE
+
+    def get_prompt_type(self) -> list[PromptType]:
+        """Get the prompt type of the datapoint."""
+        prompt_types = []
+        for metadata in self.metadata or []:
+            if isinstance(metadata, MediaAssetMetadata):
+                prompt_types.append(PromptType.ASSET)
+            elif isinstance(metadata, PromptMetadata):
+                prompt_types.append(PromptType.TEXT)
+
+        if len(prompt_types) == 0:
+            return [PromptType.NONE]
+
+        return prompt_types
+
     def get_texts(self) -> list[str]:
         """Extract text content from the asset(s)."""
         if isinstance(self.asset, TextAsset):
@@ -59,13 +126,17 @@ class Datapoint:
                     media_assets.append(asset)
             return media_assets
         else:
-            raise ValueError(f"Cannot extract media assets from asset type: {type(self.asset)}")
+            raise ValueError(
+                f"Cannot extract media assets from asset type: {type(self.asset)}"
+            )
 
-    def get_local_file_paths(self) -> list[StrictStr | tuple[StrictStr, StrictBytes] | StrictBytes]:
+    def get_local_file_paths(
+        self,
+    ) -> list[StrictStr | tuple[StrictStr, StrictBytes] | StrictBytes]:
         """Get local file paths for media assets that are stored locally."""
         if not self.is_media_asset():
             return []
-        
+
         media_assets = self.get_media_assets()
         return [asset.to_file() for asset in media_assets if asset.is_local()]
 
@@ -73,34 +144,51 @@ class Datapoint:
         """Get URLs for media assets that are remote."""
         if not self.is_media_asset():
             return []
-        
+
         media_assets = self.get_media_assets()
         return [asset.path for asset in media_assets if not asset.is_local()]
 
-    def get_prepared_metadata(self) -> list[DatasetDatasetIdDatapointsPostRequestMetadataInner]:
+    def get_prepared_metadata(
+        self,
+    ) -> list[DatasetDatasetIdDatapointsPostRequestMetadataInner]:
         """Prepare metadata for API upload."""
         metadata: list[DatasetDatasetIdDatapointsPostRequestMetadataInner] = []
         if self.metadata:
             for meta in self.metadata:
                 meta_model = meta.to_model() if meta else None
                 if meta_model:
-                    metadata.append(DatasetDatasetIdDatapointsPostRequestMetadataInner(meta_model))
+                    metadata.append(
+                        DatasetDatasetIdDatapointsPostRequestMetadataInner(meta_model)
+                    )
         return metadata
 
-    def create_text_upload_model(self, index: int) -> CreateDatapointFromTextSourcesModel:
+    def create_text_upload_model(
+        self, index: int
+    ) -> CreateDatapointFromTextSourcesModel:
         """Create the model for uploading text datapoints."""
         if not self.is_text_asset():
             raise ValueError("Cannot create text upload model for non-text asset")
-        
+
         texts = self.get_texts()
         metadata = self.get_prepared_metadata()
-        
+
         return CreateDatapointFromTextSourcesModel(
             textSources=texts,
             sortIndex=index,
             metadata=metadata,
         )
 
+    def get_datapoint_string(self) -> str:
+        """Get the datapoint string for the datapoint."""
+        if isinstance(self.asset, MediaAsset):
+            return self.asset.path
+        elif isinstance(self.asset, TextAsset):
+            return self.asset.text
+        else:
+            raise ValueError(
+                f"Cannot get datapoint string for asset type: {type(self.asset)}"
+            )
+
     def __str__(self):
         return f"Datapoint(asset={self.asset})"
 

rapidata/rapidata_client/datapoints/assets/_media_asset.py

@@ -18,6 +18,11 @@ import logging
 from functools import cached_property
 from rapidata.rapidata_client.datapoints.assets._sessions import SessionManager
 from rapidata.rapidata_client.logging import logger
+from rapidata.rapidata_client.datapoints.assets.constants import (
+    ALLOWED_IMAGE_EXTENSIONS,
+    ALLOWED_MEDIA_EXTENSIONS,
+)
+
 
 class MediaAsset(BaseAsset):
     """MediaAsset Class with Lazy Loading
@@ -32,34 +37,35 @@ class MediaAsset(BaseAsset):
     Raises:
         FileNotFoundError: If the provided file path does not exist.
     """
-    _logger = logging.getLogger(__name__ + '.MediaAsset')
+
+    _logger = logging.getLogger(__name__ + ".MediaAsset")
 
     ALLOWED_TYPES = [
-        'image/', 
-        'audio/mp3',      # MP3
-        'video/mp4',       # MP4
+        "image/",
+        "audio/mp3",  # MP3
+        "video/mp4",  # MP4
     ]
 
     MIME_TYPES = {
-        'jpg': 'image/jpeg',
-        'jpeg': 'image/jpeg',
-        'png': 'image/png',
-        'gif': 'image/gif',
-        'webp': 'image/webp',
-        'mp3': 'audio/mp3',
-        'mp4': 'video/mp4'
+        "jpg": "image/jpeg",
+        "jpeg": "image/jpeg",
+        "png": "image/png",
+        "gif": "image/gif",
+        "webp": "image/webp",
+        "mp3": "audio/mp3",
+        "mp4": "video/mp4",
     }
 
     FILE_SIGNATURES = {
-        b'\xFF\xD8\xFF': 'image/jpeg',
-        b'\x89PNG\r\n\x1a\n': 'image/png',
-        b'GIF87a': 'image/gif',
-        b'GIF89a': 'image/gif',
-        b'RIFF': 'image/webp',
-        b'ID3': 'audio/mp3',
-        b'\xFF\xFB': 'audio/mp3',
-        b'\xFF\xF3': 'audio/mp3',
-        b'ftyp': 'video/mp4',
+        b"\xFF\xD8\xFF": "image/jpeg",
+        b"\x89PNG\r\n\x1a\n": "image/png",
+        b"GIF87a": "image/gif",
+        b"GIF89a": "image/gif",
+        b"RIFF": "image/webp",
+        b"ID3": "audio/mp3",
+        b"\xFF\xFB": "audio/mp3",
+        b"\xFF\xF3": "audio/mp3",
+        b"ftyp": "video/mp4",
     }
 
     def __init__(self, path: str):
@@ -74,22 +80,24 @@ class MediaAsset(BaseAsset):
             ValueError: If path is not a string.
         """
         if not isinstance(path, str):
-            raise ValueError(f"Media must be a string, either a local file path or a URL, got {type(path)}")
-        
+            raise ValueError(
+                f"Media must be a string, either a local file path or a URL, got {type(path)}"
+            )
+
         self._url = None
         self._content = None
-        self.session: requests.Session  = SessionManager.get_session()
-        
-        if re.match(r'^https?://', path):
+        self.session: requests.Session = SessionManager.get_session()
+
+        if re.match(r"^https?://", path):
             self._url = path
-            self.name = path.split('/')[-1]
+            self.name = path.split("/")[-1]
             self.name = self.__check_name_ending(self.name)
             self.path = path
             return
-        
+
         if not os.path.exists(path):
             raise FileNotFoundError(f"File not found: {path}")
-        
+
         self.path = path
         self.name = path
 
@@ -101,9 +109,9 @@ class MediaAsset(BaseAsset):
         """
         if self._url is None:
             self.path = cast(str, self.path)
-            with open(self.path, 'rb') as f:
+            with open(self.path, "rb") as f:
                 return f.read()
-            
+
         return self.__get_media_bytes(self._url)
 
     def get_duration(self) -> int:
@@ -118,29 +126,31 @@ class MediaAsset(BaseAsset):
             ValueError: If the duration cannot be determined
         """
         path_to_check = self.name.lower()
-        
+
         # Return 0 for static images
-        if any(path_to_check.endswith(ext) for ext in ('.jpg', '.jpeg', '.png', '.webp', '.gif')):
+        if any(path_to_check.endswith(ext) for ext in ALLOWED_IMAGE_EXTENSIONS):
             return 0
 
         try:
             # Create temporary file from content
-            with tempfile.NamedTemporaryFile(suffix=os.path.splitext(self.name)[1], delete=False) as tmp:
+            with tempfile.NamedTemporaryFile(
+                suffix=os.path.splitext(self.name)[1], delete=False
+            ) as tmp:
                 tmp.write(self.content)
                 tmp.flush()
                 tmp_path = tmp.name
-                
+
             try:
                 tag = TinyTag.get(tmp_path)
             finally:
                 # Clean up the temporary file
                 os.unlink(tmp_path)
-            
+
             if tag.duration is None:
                 raise ValueError("Could not read duration from file")
-                
+
             return int(tag.duration * 1000)  # Convert to milliseconds
-            
+
         except Exception as e:
             raise ValueError(f"Could not determine media duration: {str(e)}")
 
@@ -149,53 +159,55 @@ class MediaAsset(BaseAsset):
         Get the dimensions (width, height) of an image file.
         Returns None for non-image files or if dimensions can't be determined.
         """
-        if not any(self.name.lower().endswith(ext) for ext in ('.jpg', '.jpeg', '.png', '.gif', '.webp')):
+        if not any(self.name.lower().endswith(ext) for ext in ALLOWED_IMAGE_EXTENSIONS):
             return None
-            
+
         try:
             img = Image.open(BytesIO(self.content))
             return img.size
         except Exception:
             return None
 
-    def set_custom_name(self, name: str) -> 'MediaAsset':
+    def set_custom_name(self, name: str) -> "MediaAsset":
         """Set a custom name for the media asset (only works with URLs)."""
         if self._url is not None:
             self.name = self.__check_name_ending(name)
         else:
             raise ValueError("Custom name can only be set for URLs.")
         return self
-    
+
     def __check_name_ending(self, name: str) -> str:
         """Check if the media path is valid."""
-        if not name.endswith(('.jpg', '.jpeg', '.png', '.gif', '.mp3', '.mp4', '.webp')):
-            logger.warning("Warning: Supported file types: jpg, jpeg, png, gif, mp3, mp4. Image might not be displayed correctly.")
-            name = name + '.jpg'
+        if not any(name.endswith(ext) for ext in ALLOWED_MEDIA_EXTENSIONS):
+            logger.warning(
+                f"Warning: Supported file types: {ALLOWED_MEDIA_EXTENSIONS}. Image might not be displayed correctly."
+            )
+            name = name + ".jpg"
         return name
 
     def __get_media_type_from_extension(self, url: str) -> Optional[str]:
         """
         Determine media type from URL file extension.
-        
+
         Args:
             url: The URL to check
-            
+
         Returns:
             Optional[str]: MIME type if valid extension found, None otherwise
         """
         try:
-            ext = url.lower().split('?')[0].split('.')[-1]
+            ext = url.lower().split("?")[0].split(".")[-1]
             return self.MIME_TYPES.get(ext)
         except IndexError:
             return None
-        
+
     def __validate_image_content(self, content: bytes) -> bool:
         """
         Validate image content using PIL.
-        
+
         Args:
             content: Image bytes to validate
-            
+
         Returns:
             bool: True if valid image, False otherwise
         """
@@ -206,14 +218,14 @@ class MediaAsset(BaseAsset):
         except Exception as e:
             self._logger.debug(f"Image validation failed: {str(e)}")
             return False
-        
+
     def __get_media_type_from_signature(self, content: bytes) -> Optional[str]:
         """
         Determine media type from file signature.
-        
+
         Args:
             content: File content bytes
-            
+
         Returns:
             Optional[str]: MIME type if valid signature found, None otherwise
         """
@@ -226,13 +238,13 @@ class MediaAsset(BaseAsset):
     def __get_media_bytes(self, url: str) -> bytes:
         """
         Downloads and validates media files from URL with retry logic and session reuse.
-        
+
         Args:
             url: URL of the media file
-            
+
         Returns:
             bytes: Validated media content
-            
+
         Raises:
             ValueError: If media type is unsupported or content validation fails
             requests.exceptions.RequestException: If download fails after all retries
@@ -243,17 +255,17 @@ class MediaAsset(BaseAsset):
 
         try:
             response = self.session.get(
-                url, 
-                stream=False,
-                timeout=(5, 30)  # (connect timeout, read timeout)
+                url, stream=False, timeout=(5, 30)  # (connect timeout, read timeout)
             )
             response.raise_for_status()
         except requests.exceptions.RequestException as e:
-            self._logger.error(f"Failed to download media from {url} after retries: {str(e)}")
+            self._logger.error(
+                f"Failed to download media from {url} after retries: {str(e)}"
+            )
             raise
 
         content = response.content
-        content_type = response.headers.get('content-type', '').lower()
+        content_type = response.headers.get("content-type", "").lower()
 
         # Case 1: Content-type is already allowed
         if any(content_type.startswith(t) for t in self.ALLOWED_TYPES):
@@ -279,18 +291,18 @@ class MediaAsset(BaseAsset):
 
         # If we get here, validation failed
         error_msg = (
-            f'Could not validate media type from content.\n'
-            f'Content-Type: {content_type}\n'
+            f"Could not validate media type from content.\n"
+            f"Content-Type: {content_type}\n"
             f'URL extension: {url.split("?")[0].split(".")[-1]}\n'
-            f'Allowed types: {self.ALLOWED_TYPES}'
+            f"Allowed types: {self.ALLOWED_TYPES}"
         )
         self._logger.error(error_msg)
         raise ValueError(error_msg)
-    
+
     def is_local(self) -> bool:
         """Check if the media asset is a local file."""
         return self._url is None
-    
+
     def to_file(self) -> StrictStr | tuple[StrictStr, StrictBytes] | StrictBytes:
         """Convert the media asset to a file representation."""
         if self._url is None:
@@ -298,9 +310,9 @@ class MediaAsset(BaseAsset):
             return self.path
         else:
             return (self.name, self.content)
-        
+
     def __str__(self) -> str:
         return f"MediaAsset(path={self.path})"
-    
+
     def __repr__(self) -> str:
         return f"MediaAsset(path={self.path})"

rapidata/rapidata_client/datapoints/assets/_sessions.py

@@ -2,7 +2,7 @@ import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
-from rapidata.rapidata_client.config.config import rapidata_config
+from rapidata.rapidata_client.config.rapidata_config import rapidata_config
 
 
 class SessionManager:
@@ -18,8 +18,8 @@ class SessionManager:
             requests.Session: A singleton requests session with retry logic.
         """
         if cls._session is None:
-            max_retries: int = rapidata_config.upload_max_retries
-            max_workers: int = rapidata_config.max_upload_workers
+            max_retries: int = rapidata_config.uploadMaxRetries
+            max_workers: int = rapidata_config.maxUploadWorkers
             cls._session = requests.Session()
             retries = Retry(
                 total=max_retries,

rapidata/rapidata_client/datapoints/assets/constants.py

@@ -0,0 +1,7 @@
+ALLOWED_IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".gif", ".webp"]
+ALLOWED_AUDIO_EXTENSIONS = [".mp3"]
+ALLOWED_VIDEO_EXTENSIONS = [".mp4"]
+
+ALLOWED_MEDIA_EXTENSIONS = (
+    ALLOWED_IMAGE_EXTENSIONS + ALLOWED_AUDIO_EXTENSIONS + ALLOWED_VIDEO_EXTENSIONS
+)

rapidata/rapidata_client/exceptions/failed_upload_exception.py

@@ -1,21 +1,27 @@
 from typing import cast
 from rapidata.api_client.models.file_asset_model import FileAssetModel
-from rapidata.api_client.models.get_failed_datapoints_result import GetFailedDatapointsResult
+from rapidata.api_client.models.get_failed_datapoints_result import (
+    GetFailedDatapointsResult,
+)
 from rapidata.api_client.models.multi_asset_model import MultiAssetModel
-from rapidata.api_client.models.original_filename_metadata_model import OriginalFilenameMetadataModel
+from rapidata.api_client.models.original_filename_metadata_model import (
+    OriginalFilenameMetadataModel,
+)
 from rapidata.api_client.models.source_url_metadata_model import SourceUrlMetadataModel
 from rapidata.rapidata_client.datapoints.assets import MediaAsset, MultiAsset
-from rapidata.rapidata_client.datapoints.datapoint import Datapoint
+from rapidata.rapidata_client.datapoints._datapoint import Datapoint
 from rapidata.rapidata_client.order._rapidata_dataset import RapidataDataset
 from rapidata.rapidata_client.order.rapidata_order import RapidataOrder
 
+
 class FailedUploadException(Exception):
     """Custom error class for Failed Uploads to the Rapidata order."""
+
     def __init__(
-        self, 
+        self,
         dataset: RapidataDataset,
         order: RapidataOrder,
-        failed_uploads: list[Datapoint]
+        failed_uploads: list[Datapoint],
     ):
         self.dataset = dataset
         self.order = order
@@ -30,24 +36,47 @@ def _parse_failed_uploads(failed_uploads: GetFailedDatapointsResult) -> list[Dat
     if not failed_datapoints:
         return []
     if isinstance(failed_datapoints[0].asset.actual_instance, FileAssetModel):
-        failed_assets = [MediaAsset(__get_asset_name(cast(FileAssetModel, datapoint.asset.actual_instance))) for datapoint in failed_datapoints]
+        failed_assets = [
+            MediaAsset(
+                __get_asset_name(cast(FileAssetModel, datapoint.asset.actual_instance))
+            )
+            for datapoint in failed_datapoints
+        ]
     elif isinstance(failed_datapoints[0].asset.actual_instance, MultiAssetModel):
         failed_assets = []
-        backend_assets = [cast(MultiAssetModel, failed_upload.asset.actual_instance).assets for failed_upload in failed_datapoints]
+        backend_assets = [
+            cast(MultiAssetModel, failed_upload.asset.actual_instance).assets
+            for failed_upload in failed_datapoints
+        ]
         for assets in backend_assets:
-            failed_assets.append(MultiAsset([MediaAsset(__get_asset_name(cast(FileAssetModel, asset.actual_instance))) for asset in assets if isinstance(asset.actual_instance, FileAssetModel)]))
+            failed_assets.append(
+                MultiAsset(
+                    [
+                        MediaAsset(
+                            __get_asset_name(
+                                cast(FileAssetModel, asset.actual_instance)
+                            )
+                        )
+                        for asset in assets
+                        if isinstance(asset.actual_instance, FileAssetModel)
+                    ]
+                )
+            )
     else:
-        raise ValueError(f"Unsupported asset type: {type(failed_datapoints[0].asset.actual_instance)}")
-    
+        raise ValueError(
+            f"Unsupported asset type: {type(failed_datapoints[0].asset.actual_instance)}"
+        )
+
     return [Datapoint(asset=asset) for asset in failed_assets]
 
+
 def __get_asset_name(failed_datapoint: FileAssetModel) -> str:
     metadata = failed_datapoint.metadata
     if "sourceUrl" in metadata:
         return cast(SourceUrlMetadataModel, metadata["sourceUrl"].actual_instance).url
     elif "originalFilename" in metadata:
-        return cast(OriginalFilenameMetadataModel, metadata["originalFilename"].actual_instance).original_filename
+        return cast(
+            OriginalFilenameMetadataModel, metadata["originalFilename"].actual_instance
+        ).original_filename
     else:
         return ""
-
-

rapidata/rapidata_client/filter/response_count_filter.py

@@ -3,6 +3,8 @@ from rapidata.rapidata_client.filter._base_filter import RapidataFilter
 from rapidata.api_client.models.response_count_user_filter_model import (
     ResponseCountUserFilterModel,
 )
+from rapidata.api_client.models.comparison_operator import ComparisonOperator
+
 
 class ResponseCountFilter(RapidataFilter):
     """ResponseCountFilter Class
@@ -11,12 +13,12 @@ class ResponseCountFilter(RapidataFilter):
         response_count (int): The number of user responses to filter by.
         dimension (str): The dimension to apply the filter on (e.g. "default", "electrical", etc.).
         operator (str): The comparison operator to use. Must be one of:
-            - "Equal"
-            - "NotEqual"
-            - "LessThan"
-            - "LessThanOrEqual"
-            - "GreaterThan"
-            - "GreaterThanOrEqual"
+            - ComparisonOperator.EQUAL
+            - ComparisonOperator.NOTEQUAL
+            - ComparisonOperator.LESSTHAN
+            - ComparisonOperator.LESSTHANOREQUAL
+            - ComparisonOperator.GREATERTHAN
+            - ComparisonOperator.GREATERTHANOREQUAL
 
     Raises:
         ValueError: If `response_count` is not an integer.
@@ -27,14 +29,14 @@ class ResponseCountFilter(RapidataFilter):
         ```python
         from rapidata import ResponseCountFilter
 
-        filter = ResponseCountFilter(response_count=10, dimension="electrical", operator="GreaterThan")
+        filter = ResponseCountFilter(response_count=10, dimension="electrical", operator=ComparisonOperator.GREATERTHAN)
         ```
         This will filter users who have a response count greater than 10 for the "electrical" dimension.
     """
 
-    def __init__(self, response_count: int, dimension: str, operator: str):
-        if operator not in ["Equal", "NotEqual", "LessThan", "LessThanOrEqual", "GreaterThan", "GreaterThanOrEqual"]:
-            raise ValueError("Operator must be one of Equal, NotEqual, LessThan, LessThanOrEqual, GreaterThan, GreaterThanOrEqual")
+    def __init__(
+        self, response_count: int, dimension: str, operator: ComparisonOperator
+    ):
 
         self.response_count = response_count
         self.dimension = dimension
@@ -42,5 +44,8 @@ class ResponseCountFilter(RapidataFilter):
 
     def _to_model(self):
         return ResponseCountUserFilterModel(
-            _t="ResponseCountFilter", responseCount=self.response_count, dimension=self.dimension, operator=self.operator
+            _t="ResponseCountFilter",
+            responseCount=self.response_count,
+            dimension=self.dimension,
+            operator=self.operator,
         )