mage-ai 0.9.73__py3-none-any.whl → 0.9.74__py3-none-any.whl

mage_ai/ai/constants.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
-class LLMUseCase(str, Enum):
+class LLMUseCase(StrEnum):
     GENERATE_DOC_FOR_BLOCK = 'generate_doc_for_block'
     GENERATE_DOC_FOR_PIPELINE = 'generate_doc_for_pipeline'
     GENERATE_BLOCK_WITH_DESCRIPTION = 'generate_block_with_description'

mage_ai/ai/llm_pipeline_wizard.py

@@ -30,6 +30,7 @@ from mage_ai.data_preparation.templates.template import (
 from mage_ai.io.base import DataSource
 from mage_ai.orchestration.ai.config import AIConfig
 from mage_ai.server.logger import Logger
+from mage_ai.settings import ENABLE_HUGGING_FACE, ENABLE_OPEN_AI
 
 logger = Logger().new_server_logger(__name__)
 
@@ -190,12 +191,13 @@ TEMPLATE_CLASSIFICATION_FUNCTION = [
 class LLMPipelineWizard:
     def __init__(self):
         ai_config = AIConfig.load(config=get_repo_config().ai_config)
-        if ai_config.mode == AIMode.OPEN_AI:
+        if ENABLE_OPEN_AI and ai_config.mode == AIMode.OPEN_AI:
             self.client = OpenAIClient(ai_config.open_ai_config)
-        elif ai_config.mode == AIMode.HUGGING_FACE:
+        elif ENABLE_HUGGING_FACE and ai_config.mode == AIMode.HUGGING_FACE:
             self.client = HuggingFaceClient(ai_config.hugging_face_config)
         else:
-            raise Exception('AI Mode is not available.')
+            raise Exception('AI Mode is not available. ENABLE_OPEN_AI: '
+                            f'{ENABLE_OPEN_AI}, ENABLE_HUGGING_FACE: {ENABLE_HUGGING_FACE}')
 
     async def __async_llm_call(
         self,

mage_ai/ai/openai_client.py

@@ -187,7 +187,18 @@ class OpenAIClient(AIClient):
             self,
             block_description: str):
         messages = [{'role': 'user', 'content': block_description}]
+        # Fetch response form API call with retries
+        # retry __chat_completion_request twice.
+        # If it still returns error, raise error in find_block_params.
+        # If not error anymore, proceed with rest of the code change.
+        max_retries = 2
+        attempt = 0
         response = self.__chat_completion_request(messages)
+        while attempt <= max_retries and isinstance(response, Exception):
+            response = self.__chat_completion_request(messages)
+            attempt += 1
+        if isinstance(response, Exception):
+            raise Exception("Error in __chat_completion_request after retries: " + str(response))
         arguments = response.choices[0].message.tool_calls[0].function.arguments
         if arguments:
             function_args = json.loads(arguments)

mage_ai/api/constants.py

@@ -1,20 +1,19 @@
-from enum import Enum
-
 from mage_ai.api.operations.constants import OperationType
 from mage_ai.orchestration.db.models.oauth import Permission
+from mage_ai.shared.enum import StrEnum
 
 
-class AttributeOperationType(str, Enum):
+class AttributeOperationType(StrEnum):
     QUERY = 'query'
     READ = 'read'
     WRITE = 'write'
 
 
-class AttributeType(str, Enum):
+class AttributeType(StrEnum):
     ALL = '__*MAGE*__'
 
 
-class AuthorizeStatusType(str, Enum):
+class AuthorizeStatusType(StrEnum):
     ALL = 'all'
     FAILED = 'failed'
     SUCCEEDED = 'succeeded'

mage_ai/api/oauth_scope.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
 class OauthScope():
@@ -18,7 +18,7 @@ class OauthScope():
     TOKEN_SCOPES = []
 
 
-class OauthScopeType(str, Enum):
+class OauthScopeType(StrEnum):
     CLIENT_ALL = 'all'
     CLIENT_INTERNAL = 'internal'
     CLIENT_PRIVATE = 'private'

mage_ai/api/operations/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 ALL = 'all'
 CREATE = 'create'
@@ -20,7 +20,7 @@ META_KEY_ORDER_BY = '_order_by[]'
 COOKIE_PREFIX = '__COOKIE__'
 
 
-class OperationType(str, Enum):
+class OperationType(StrEnum):
     ALL = ALL
     CREATE = CREATE
     DELETE = DELETE

mage_ai/api/presenters/SyncPresenter.py

@@ -1,3 +1,5 @@
+from typing import Any, Union
+
 from mage_ai.api.presenters.BasePresenter import BasePresenter
 from mage_ai.data_preparation.shared.secrets import get_secret_value
 from mage_ai.settings.repo import get_repo_path
@@ -16,6 +18,7 @@ class SyncPresenter(BasePresenter):
         'ssh_private_key_secret_name',
         'ssh_public_key',
         'ssh_public_key_secret_name',
+        'sync_on_executor_start',
         'sync_on_pipeline_run',
         'sync_on_start',
         'sync_submodules',
@@ -24,7 +27,7 @@ class SyncPresenter(BasePresenter):
         'username',
     ]
 
-    def present(self, **kwargs):
+    async def prepare_present(self, **kwargs) -> Union[Any, None]:
         data = self.model
 
         def filter_invalid_secret_values(data):

mage_ai/authentication/oauth/constants.py

@@ -1,8 +1,8 @@
-from enum import Enum
 from typing import Optional
 
 from mage_ai.settings import get_settings_value
 from mage_ai.settings.keys import GHE_HOSTNAME
+from mage_ai.shared.enum import StrEnum
 
 ACTIVE_DIRECTORY_CLIENT_ID = '51aec820-9d49-40a9-b046-17c1f28f620d'
 
@@ -10,7 +10,7 @@ GITHUB_CLIENT_ID = '8577f13ddc81e2848b07'
 GITHUB_STATE = '1337'
 
 
-class ProviderName(str, Enum):
+class ProviderName(StrEnum):
     ACTIVE_DIRECTORY = 'active_directory'
     AZURE_DEVOPS = 'azure_devops'
     BITBUCKET = 'bitbucket'

mage_ai/authentication/operation_history/constants.py

@@ -1,8 +1,8 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 MAGE_OPERATION_HISTORY_DIRECTORY_DEFAULT = '.operation_history'
 MAGE_OPERATION_HISTORY_DIRECTORY_ENVIRONMENT_VARIABLE_NAME = 'MAGE_OPERATION_HISTORY_DIRECTORY'
 
 
-class ResourceType(str, Enum):
+class ResourceType(StrEnum):
     PIPELINE = 'pipeline'

mage_ai/authentication/permissions/constants.py

@@ -1,10 +1,9 @@
-from enum import Enum
-
 from mage_ai.api.operations.constants import OperationType
 from mage_ai.data_preparation.models.constants import BlockType, PipelineType
+from mage_ai.shared.enum import IntEnum, StrEnum
 
 
-class EntityName(str, Enum):
+class EntityName(StrEnum):
     ALL = 'ALL'
     ALL_EXCEPT_RESERVED = 'ALL_EXCEPT_RESERVED'
     AutocompleteItem = 'AutocompleteItem'
@@ -109,7 +108,7 @@ RESERVED_ENTITY_NAMES = [
 ]
 
 
-class BaseEntityType(str, Enum):
+class BaseEntityType(StrEnum):
     pass
 
 
@@ -137,7 +136,7 @@ class PipelineEntityType(BaseEntityType):
     STREAMING = PipelineType.STREAMING.value
 
 
-class PermissionAccess(int, Enum):
+class PermissionAccess(IntEnum):
     OWNER = 1
     ADMIN = 2
     # Editor: list, detail, create, update, delete
@@ -177,7 +176,7 @@ class PermissionAccess(int, Enum):
     DISABLE_UNLESS_CONDITIONS = 1073741824
 
 
-class PermissionCondition(str, Enum):
+class PermissionCondition(StrEnum):
     HAS_NOTEBOOK_EDIT_ACCESS = 'HAS_NOTEBOOK_EDIT_ACCESS'
     HAS_PIPELINE_EDIT_ACCESS = 'HAS_PIPELINE_EDIT_ACCESS'
     USER_OWNS_ENTITY = 'USER_OWNS_ENTITY'

mage_ai/cache/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 CACHE_KEY_BLOCKS_TO_PIPELINE_MAPPING = 'blocks_to_pipeline_mapping'
 CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING = 'block_action_objects_mapping'
@@ -11,5 +11,5 @@ MAGE_CACHE_DIRECTORY_DEFAULT = '.cache'
 MAGE_CACHE_DIRECTORY_ENVIRONMENT_VARIABLE_NAME = 'MAGE_CACHE_DIRECTORY'
 
 
-class CacheItemType(str, Enum):
+class CacheItemType(StrEnum):
     DBT = 'dbt'

mage_ai/cache/dbt/constants.py

@@ -1,10 +1,9 @@
-from enum import Enum
-
 from mage_ai.data_preparation.models.constants import (
     BLOCK_TYPE_DIRECTORY_NAME,
     PIPELINES_FOLDER,
     BlockType,
 )
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import merge_dict
 
 IGNORE_DIRECTORY_NAMES = merge_dict(
@@ -19,7 +18,7 @@ PROJECT_FILENAME = 'dbt_project.yml'
 PROJECT_FILENAMES = [PROJECT_FILENAME, 'dbt_project.yaml']
 
 
-class FileType(str, Enum):
+class FileType(StrEnum):
     MODEL = 'model'
     PROFILES = 'profiles'
     PROJECT = 'project'

mage_ai/cluster_manager/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 # ECS environment variables
 ECS_CLUSTER_NAME = 'ECS_CLUSTER_NAME'
@@ -29,7 +29,7 @@ GCP_BACKEND_CONFIG_ANNOTATION = 'cloud.google.com/backend-config'
 NODE_PORT_SERVICE_TYPE = 'NodePort'
 
 
-class ClusterType(str, Enum):
+class ClusterType(StrEnum):
     EMR = 'emr'
     ECS = 'ecs'
     CLOUD_RUN = 'cloud_run'

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -43,7 +43,7 @@ from mage_ai.services.k8s.constants import (
 )
 from mage_ai.settings import MAGE_SETTINGS_ENVIRONMENT_VARIABLES
 from mage_ai.shared.array import find
-from mage_ai.shared.hash import safe_dig
+from mage_ai.shared.hash import merge_dict, safe_dig
 
 
 class WorkloadManager:
@@ -391,6 +391,7 @@ class WorkloadManager:
             )
 
         pod_spec = self.pod_config.spec.to_dict() if self.pod_config else dict()
+        pod_labels = self.pod_config.metadata.labels or dict()
         stateful_set_template_spec = dict(
             imagePullSecrets=pod_spec.get('image_pull_secrets'),
             initContainers=init_containers,
@@ -410,7 +411,7 @@ class WorkloadManager:
                 'replicas': 1,
                 'minReadySeconds': 10,
                 'template': {
-                    'metadata': {'labels': {'app': name}},
+                    'metadata': merge_dict(pod_labels, {'labels': {'app': name}}),
                     'spec': stateful_set_template_spec,
                 },
                 'volumeClaimTemplates': [

mage_ai/command_center/constants.py

@@ -1,14 +1,14 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 SETTINGS_FILENAME = '.command_center.yaml'
 
 
-class ItemTagEnum(str, Enum):
+class ItemTagEnum(StrEnum):
     PINNED = 'pinned'
     RECENT = 'recent'
 
 
-class ItemType(str, Enum):
+class ItemType(StrEnum):
     ACTION = 'action'  # Cast spell
     CREATE = 'create'  # Conjure
     DELETE = 'delete'
@@ -22,7 +22,7 @@ class ItemType(str, Enum):
     UPDATE = 'update'
 
 
-class ObjectType(str, Enum):
+class ObjectType(StrEnum):
     APPLICATION = 'application'
     APPLICATION_EXPANSION = 'application_expansion'
     AUTHENTICATION = 'authentication'
@@ -42,11 +42,11 @@ class ObjectType(str, Enum):
     VERSION_CONTROL_FILE = 'version_control_file'
 
 
-class ModeType(str, Enum):
+class ModeType(StrEnum):
     VERSION_CONTROL = 'version_control'
 
 
-class FileExtension(str, Enum):
+class FileExtension(StrEnum):
     CSV = 'csv'
     JSON = 'json'
     MD = 'md'
@@ -59,7 +59,7 @@ class FileExtension(str, Enum):
     YML = 'yml'
 
 
-class ButtonActionType(str, Enum):
+class ButtonActionType(StrEnum):
     ADD_APPLICATION = 'add_application'
     CLOSE_APPLICATION = 'close_application'  # Go back out of the current application.
     CLOSE_COMMAND_CENTER = 'close_command_center'
@@ -70,7 +70,7 @@ class ButtonActionType(str, Enum):
     SELECT_ITEM_FROM_REQUEST = 'select_item_from_request'
 
 
-class InteractionType(str, Enum):
+class InteractionType(StrEnum):
     CLICK = 'click'
     CLOSE_APPLICATION = ButtonActionType.CLOSE_APPLICATION.value
     CLOSE_COMMAND_CENTER = ButtonActionType.CLOSE_COMMAND_CENTER.value
@@ -81,7 +81,7 @@ class InteractionType(str, Enum):
     SELECT_ITEM = 'select_item'
 
 
-class ApplicationType(str, Enum):
+class ApplicationType(StrEnum):
     DETAIL = 'detail'
     DETAIL_LIST = 'detail_list'
     EXPANSION = 'expansion'
@@ -89,22 +89,22 @@ class ApplicationType(str, Enum):
     LIST = 'list'
 
 
-class ValidationType(str, Enum):
+class ValidationType(StrEnum):
     CONFIRMATION = 'confirmation'
     CUSTOM_VALIDATION_PARSERS = 'custom_validation_parsers'
 
 
-class RenderLocationType(str, Enum):
+class RenderLocationType(StrEnum):
     ITEMS_CONTAINER_AFTER = 'items_container_after'
 
 
-class ApplicationExpansionUUID(str, Enum):
+class ApplicationExpansionUUID(StrEnum):
     ArcaneLibrary = 'ArcaneLibrary'
     PortalTerminal = 'PortalTerminal'
     VersionControlFileDiffs = 'VersionControlFileDiffs'
 
 
-class ApplicationExpansionStatus(str, Enum):
+class ApplicationExpansionStatus(StrEnum):
     ACTIVE = 'ACTIVE'
     CLOSED = 'CLOSED'
     INACTIVE = 'INACTIVE'

mage_ai/data_cleaner/column_types/constants.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
-class ColumnType(str, Enum):
+class ColumnType(StrEnum):
     CATEGORY = 'category'
     CATEGORY_HIGH_CARDINALITY = 'category_high_cardinality'
     DATETIME = 'datetime'

mage_ai/data_cleaner/transformer_actions/column.py

@@ -1,3 +1,8 @@
+import logging
+
+import numpy as np
+import pandas as pd
+
 from mage_ai.data_cleaner.column_types.column_type_detector import find_syntax_errors
 from mage_ai.data_cleaner.column_types.constants import NUMBER_TYPES, ColumnType
 from mage_ai.data_cleaner.estimators.outlier_removal import OutlierRemover
@@ -15,11 +20,10 @@ from mage_ai.data_cleaner.transformer_actions.helpers import (
     get_time_window_str,
 )
 from mage_ai.data_cleaner.transformer_actions.udf.base import execute_udf
-from mage_ai.data_cleaner.transformer_actions.utils import clean_column_name, generate_string_cols
-import logging
-import pandas as pd
-import numpy as np
-
+from mage_ai.data_cleaner.transformer_actions.utils import (
+    clean_column_name,
+    generate_string_cols,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -115,7 +119,8 @@ def impute(df, action, **kwargs):
                 mode = df[column].mode().iloc[0]
                 if dtype == ColumnType.LIST:
                     df[column] = df[column].apply(
-                        lambda element: element if element not in [None, np.nan] else mode
+                        lambda element, mode=mode: element if element not in [
+                            None, np.nan] else mode
                     )
                 else:
                     df[columns] = df[columns].fillna(mode)
@@ -180,7 +185,6 @@ def reformat(df, action, **kwargs):
     columns = action['action_arguments']
     options = action['action_options']
     reformat_action = options['reformat']
-    df.loc[:, columns] = df[columns].replace(r'^\s*$', np.nan, regex=True)
 
     if reformat_action == 'caps_standardization':
         capitalization = options['capitalization']
@@ -189,6 +193,8 @@ def reformat(df, action, **kwargs):
                 df.loc[:, column] = df[columns][column].str.upper()
             else:
                 df.loc[:, column] = df[columns][column].str.lower()
+        # Convert empty strings to NaN for this action
+        df.loc[:, columns] = df[columns].replace(r'^\s*$', np.nan, regex=True)
     elif reformat_action == 'currency_to_num':
         for column in generate_string_cols(df, columns):
             clean_col = df[column].replace(CURRENCY_SYMBOLS, '', regex=True)
@@ -215,6 +221,12 @@ def reformat(df, action, **kwargs):
             df.loc[:, column] = pd.to_datetime(
                 clean_col, infer_datetime_format=True, errors='coerce'
             )
+    elif reformat_action == 'trim':
+        for column in columns:
+            df[column] = df[column].str.strip()
+    else:
+        # Apply NaN replacement only for other actions
+        df.loc[:, columns] = df[columns].replace(r'^\s*$', np.nan, regex=True)
 
     return df
 

mage_ai/data_cleaner/transformer_actions/constants.py

@@ -1,9 +1,10 @@
-from mage_ai.data_cleaner.column_types.constants import ColumnType
-from enum import Enum
-import pandas as pd
-import numpy as np
 import re
 
+import numpy as np
+import pandas as pd
+
+from mage_ai.data_cleaner.column_types.constants import ColumnType
+from mage_ai.shared.enum import StrEnum
 
 CONSTANT_IMPUTATION_DEFAULTS = {
     ColumnType.CATEGORY: 'missing',
@@ -34,7 +35,7 @@ INVALID_VALUE_PLACEHOLDERS = {
 }
 
 
-class ActionType(str, Enum):
+class ActionType(StrEnum):
     ADD = 'add'
     AVERAGE = 'average'
     CLEAN_COLUMN_NAME = 'clean_column_name'
@@ -73,18 +74,18 @@ class ActionType(str, Enum):
     STANDARDIZE = 'standardize'
 
 
-class Axis(str, Enum):
+class Axis(StrEnum):
     COLUMN = 'column'
     ROW = 'row'
 
 
-class VariableType(str, Enum):
+class VariableType(StrEnum):
     FEATURE = 'feature'
     FEATURE_SET = 'feature_set'
     FEATURE_SET_VERSION = 'feature_set_version'
 
 
-class Operator(str, Enum):
+class Operator(StrEnum):
     CONTAINS = 'contains'
     NOT_CONTAINS = 'not contains'
     EQUALS = '=='
@@ -95,7 +96,7 @@ class Operator(str, Enum):
     LESS_THAN_OR_EQUALS = '<='
 
 
-class ImputationStrategy(str, Enum):
+class ImputationStrategy(StrEnum):
     AVERAGE = 'average'
     COLUMN = 'column'
     CONSTANT = 'constant'

mage_ai/data_integrations/sources/constants.py

@@ -17,6 +17,7 @@ SQL_SOURCES = [
 SQL_SOURCES_MAPPING = index_by(get_uuid, SQL_SOURCES)
 
 SOURCES = sorted([
+    dict(name='Airtable'),
     dict(name='Amazon S3'),
     dict(name='Amplitude'),
     dict(name='Api'),
@@ -33,6 +34,7 @@ SOURCES = sorted([
     dict(name='GitHub'),
     dict(name='Google Ads'),
     dict(name='Google Analytics'),
+    dict(name='Google Cloud Storage'),
     dict(name='Google Search Console'),
     dict(name='Google Sheets'),
     dict(name='HubSpot'),

mage_ai/data_preparation/git/utils.py

@@ -172,7 +172,10 @@ def get_access_token(git_config, repo_path: str) -> str:
 def build_authenticated_remote_url(remote_url: str, username: str, token: str) -> str:
     # https://[username]:[token]@github.com/[remote_url]
     url = urlsplit(remote_url)
-    url = url._replace(netloc=f'{username}:{token}@{url.netloc}')
+    if remote_url.startswith('https://bitbucket.org/'):
+        url = url._replace(netloc=f'x-token-auth:{token}@{url.netloc}')
+    else:
+        url = url._replace(netloc=f'{username}:{token}@{url.netloc}')
     return urlunsplit(url)
 
 

mage_ai/data_preparation/logging/__init__.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Dict
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.logger import LoggingLevel
 
 
-class LoggerType(str, Enum):
+class LoggerType(StrEnum):
     DEFAULT = 'file'
     S3 = 's3'
     GCS = 'gcs'

mage_ai/data_preparation/logging/logger.py

@@ -61,9 +61,9 @@ class DictLogger():
         )
 
         if error:
-            data['error'] = traceback.format_exc(),
-            data['error_stack'] = traceback.format_stack(),
-            data['error_stacktrace'] = str(error),
+            data['error'] = traceback.format_exc()
+            data['error_stack'] = traceback.format_stack()
+            data['error_stacktrace'] = str(error)
 
         msg = simplejson.dumps(
             merge_dict(self.logging_tags or dict(), merge_dict(kwargs, data)),

mage_ai/data_preparation/models/block/__init__.py

@@ -878,7 +878,7 @@ class Block(
         relative_path: bool = False,
     ) -> str:
         file_extension = BLOCK_LANGUAGE_TO_FILE_EXTENSION[language]
-        block_directory = f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
+        block_directory = block_type + 's' if block_type != BlockType.CUSTOM else block_type
 
         parts = []
         if not relative_path:
@@ -1006,6 +1006,12 @@ class Block(
 
         return extract_full_table_name(self.content)
 
+    def get_typed_content(
+        self,
+        content: str,
+    ) -> str:
+        return '@' + self.type + '\n' + content
+
     @classmethod
     def after_create(cls, block: 'Block', **kwargs) -> None:
         widget = kwargs.get('widget')
@@ -1193,7 +1199,7 @@ class Block(
 
     @classmethod
     def file_directory_name(self, block_type: BlockType) -> str:
-        return f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
+        return block_type + 's' if block_type != BlockType.CUSTOM else block_type
 
     @classmethod
     def block_type_from_path(
@@ -1212,7 +1218,7 @@ class Block(
         for block_type in BlockType:
             if BlockType.CUSTOM == block_type and dir_name == block_type:
                 return BlockType.CUSTOM
-            elif dir_name == f'{block_type}s':
+            elif dir_name == block_type + 's':
                 return block_type
 
     @classmethod

mage_ai/data_preparation/models/block/data_integration/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 BLOCK_CATALOG_FILENAME = 'catalog.json'
 REPLICATION_METHOD_INCREMENTAL = 'INCREMENTAL'
@@ -81,6 +81,6 @@ MAX_QUERY_STRING_SIZE = 10 * MB_1
 VARIABLE_BOOKMARK_VALUES_KEY = '__bookmark_values__'
 
 
-class IngestMode(str, Enum):
+class IngestMode(StrEnum):
     DISK = 'disk'
     MEMORY = 'memory'

mage_ai/data_preparation/models/block/dbt/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 DBT_DIRECTORY_NAME = 'dbt'
 
@@ -13,12 +13,12 @@ SKIP_LIMIT_ADAPTER_NAMES = [
 ]
 
 
-class Flag(str, Enum):
+class Flag(StrEnum):
     PROFILES_DIR = 'profiles-dir'
     PROJECT_DIR = 'project-dir'
 
 
-class LogLevel(str, Enum):
+class LogLevel(StrEnum):
     DEBUG = 'debug'
     INFO = 'info'
     WARN = 'warn'