mage-ai 0.9.58__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/command_center/blocks/utils.py

@@ -0,0 +1,127 @@
+import urllib.parse
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.command_center.constants import (
+    ApplicationType,
+    ButtonActionType,
+    ItemType,
+    ObjectType,
+)
+from mage_ai.command_center.utils import shorten_directory
+
+
+def add_application_actions(item_dict: Dict) -> Dict:
+    uuid = item_dict.get('uuid')
+    block_metadata = (item_dict.get('metadata') or {}).get('block') or {}
+    file_path = block_metadata.get('file_path')
+
+    actions = []
+    buttons = []
+    pipelines = block_metadata.get('pipelines')
+    if pipelines:
+        pipeline_uuid = pipelines[0].get('uuid')
+        actions.append(dict(
+            page=dict(
+                path=f'/pipelines/{pipeline_uuid}/edit',
+            ),
+            uuid='open_block_in_pipeline',
+        ))
+        buttons.append(dict(
+            label='Open block in pipeline',
+            keyboard_shortcuts=[[13]],
+            action_types=[
+                ButtonActionType.EXECUTE,
+            ],
+        ))
+
+    return dict(
+        applications=[
+            dict(
+                application_type=ApplicationType.DETAIL,
+                actions=[
+                    dict(
+                        request=dict(
+                            operation=OperationType.DETAIL,
+                            resource='blocks',
+                            resource_id=urllib.parse.quote_plus(uuid or ''),
+                            response_resource_key='block',
+                            query=dict(
+                                file_path=urllib.parse.quote_plus(file_path or ''),
+                            ),
+                        ),
+                        uuid='block_detail',
+                    ),
+                ],
+                buttons=buttons,
+                uuid='model_detail',
+            ),
+        ],
+        actions=actions,
+    )
+
+
+async def build_and_score(
+    factory,
+    data_input: Tuple[str, Dict],
+    items: List[Dict],
+    add_application: bool = False,
+):
+    uuid, data = data_input
+    block = data.get('block') or {}
+    pipelines = [d.get('pipeline') for d in (data.get('pipelines') or [])]
+
+    file_path = block.get('file_path') or f'{uuid}'
+    if not Path(file_path).suffix:
+        file_path = f'{file_path}.py'
+
+    path_dict = shorten_directory(file_path)
+    directory = path_dict.get('directory')
+
+    title = block.get('name') or uuid
+    parts = []
+    if block.get('name'):
+        parts = Path(block.get('name')).parts
+    elif uuid:
+        parts = Path(uuid).parts
+
+    if parts:
+        title = parts[len(parts) - 1]
+
+    description = directory
+    if len(pipelines) == 1:
+        description = f'{description} in 1 pipeline'
+    else:
+        description = f'{description} in {len(pipelines)} pipelines'
+
+    item_dict = dict(
+        item_type=ItemType.DETAIL,
+        object_type=ObjectType.BLOCK,
+        title=title,
+        description=description,
+        uuid=uuid,
+        metadata=dict(
+            block=dict(
+                file_path=file_path,
+                language=block.get('language'),
+                pipelines=pipelines,
+                type=block.get('type'),
+            ),
+        ),
+        display_settings_by_attribute=dict(
+            description=dict(
+                text_styles=dict(
+                    monospace=True,
+                    small=True,
+                ),
+            ),
+        ),
+    )
+
+    if add_application:
+        item_dict.update(add_application_actions(item_dict))
+
+    scored = factory.filter_score(item_dict)
+    if scored:
+        items.append(scored)

mage_ai/command_center/constants.py

@@ -0,0 +1,66 @@
+from enum import Enum
+
+SETTINGS_FILENAME = '.command_center.yaml'
+
+
+class ItemTagEnum(str, Enum):
+    PINNED = 'pinned'
+    RECENT = 'recent'
+
+
+class ItemType(str, Enum):
+    ACTION = 'action'  # Cast spell
+    CREATE = 'create'  # Conjure
+    DETAIL = 'detail'  # Enchant
+    LIST = 'list'  # Enchant
+    NAVIGATE = 'navigate'  # Teleport
+    OPEN = 'open'  # Summon - open shows a view on the same page vs navigating to another page
+    SUPPORT = 'support'  # Reinforcements
+
+
+class ObjectType(str, Enum):
+    APPLICATION = 'application'
+    BLOCK = 'block'
+    CHAT = 'chat'
+    CODE = 'code'
+    DOCUMENT = 'document'
+    FILE = 'file'
+    FOLDER = 'folder'
+    GIT = 'git'
+    PIPELINE = 'pipeline'
+    PIPELINE_RUN = 'pipeline_run'
+    TRIGGER = 'trigger'
+
+
+class FileExtension(str, Enum):
+    CSV = 'csv'
+    JSON = 'json'
+    MD = 'md'
+    PY = 'py'
+    R = 'r'
+    SH = 'sh'
+    SQL = 'sql'
+    TXT = 'txt'
+    YAML = 'yaml'
+    YML = 'yml'
+
+
+class InteractionType(str, Enum):
+    CLICK = 'click'
+    OPEN_FILE = 'open_file'
+    SCROLL = 'scroll'
+
+
+class ApplicationType(str, Enum):
+    DETAIL = 'detail'
+    DETAIL_LIST = 'detail_list'
+    FORM = 'form'
+
+
+class ButtonActionType(str, Enum):
+    ADD_APPLICATION = 'add_application'
+    CLOSE_APPLICATION = 'close_application'  # Go back out of the current application.
+    CUSTOM_ACTIONS = 'custom_actions'
+    EXECUTE = 'execute'  # Executes the actions associated to the application’s item.
+    REPLACE_APPLICATION = 'replace_application'
+    RESET_FORM = 'reset_form'  # If there is a form filled out, clear out all the values.

mage_ai/command_center/factory.py

@@ -0,0 +1,161 @@
+import asyncio
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Dict, List, Union
+
+from thefuzz import fuzz
+
+from mage_ai.command_center.constants import ObjectType
+from mage_ai.command_center.models import Application, Item
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.shared.array import flatten
+from mage_ai.shared.hash import ignore_keys, merge_dict
+
+DEFAULT_RATIO = 65
+
+
+class BaseFactory:
+    def __init__(
+        self,
+        application: Dict = None,
+        component: str = None,
+        item: Dict = None,
+        page: str = None,
+        page_history: List[Dict] = None,
+        picks: str = None,
+        search: str = None,
+        search_history: List[Dict] = None,
+        search_ratio: int = DEFAULT_RATIO,
+        user: User = None,
+    ):
+        self.application = Application.load(**application) if application else application
+        self.component = component
+        self.item = Item.load(**item) if item else item
+        self.page = page
+        self.page_history = page_history
+        self.picks = picks
+        self.search = search
+        self.search_history = search_history
+        self.search_ratio = search_ratio
+        self.user = user
+
+    @classmethod
+    async def create_items(
+        self,
+        factory_or_items: List[Union['BaseFactory', List[Dict]]],
+        **kwargs,
+    ) -> List[Item]:
+        async def process_factory_items(
+            class_or_items: Union['BaseFactory', List[Dict]],
+            cls=self,
+            kwargs=kwargs,
+        ) -> List[Item]:
+            is_iterable = isinstance(class_or_items, Iterable)
+            factory_class = class_or_items if not is_iterable and issubclass(
+                class_or_items,
+                cls,
+            ) else cls
+            items_dicts = class_or_items if is_iterable else None
+
+            return await factory_class(**kwargs).process(items_dicts)
+
+        application = kwargs.get('application')
+        item = kwargs.get('item')
+        if application and item:
+            object_type = item.get('object_type')
+            if ObjectType.PIPELINE == object_type:
+                from mage_ai.command_center.pipelines.factory import PipelineFactory
+                factory_or_items = [PipelineFactory]
+            elif ObjectType.TRIGGER == object_type:
+                from mage_ai.command_center.triggers.factory import TriggerFactory
+                factory_or_items = [TriggerFactory]
+
+        return flatten(await asyncio.gather(
+            *[process_factory_items(class_or_items) for class_or_items in factory_or_items]
+        ))
+
+    @property
+    def project(self) -> Project:
+        return Project(root_project=True)
+
+    # Subclasses MUST override this
+    async def fetch_items(self, **kwargs) -> List[Dict]:
+        raise Exception
+
+    # Subclasses can override this
+    def get_item_uuid(self, item_dict: Dict, **kwargs) -> str:
+        return item_dict.get('uuid') or ''
+
+    # Subclasses can override this
+    def get_searchable_text(self, item_dict: Dict, **kwargs) -> str:
+        arr = []
+        for key in ['title', 'description', 'uuid', 'item_type', 'object_type']:
+            if item_dict.get(key):
+                value = item_dict.get(key)
+                if value:
+                    value = value.lower()
+                    extension = Path(value).suffix
+                    if extension:
+                        arr.append(Path(value).stem)
+                        arr.append(extension)
+                    arr.append(value)
+
+        text = ' '.join(arr)
+        return ' '.join([
+            text,
+            text.replace('_', ' '),
+        ])
+
+    # Subclasses can override this
+    def score_item(self, _item_dict: Dict, score: int = None) -> int:
+        return score
+
+    def filter_score(self, item_dict: Dict) -> Union[None, Dict]:
+        score = 0
+        if self.search:
+            score = fuzz.partial_token_sort_ratio(
+                self.search,
+                self.get_searchable_text(item_dict),
+            )
+            if score < self.search_ratio:
+                return None
+
+        condition = item_dict.get('condition')
+        if (not condition or condition(dict(project=self.project, user=self.user))):
+            return merge_dict(item_dict, dict(score=self.score_item(item_dict, score=score)))
+
+        return None
+
+    async def process(self, items_dicts: List[Dict] = None) -> List[Item]:
+        items_dicts_scored = []
+        if items_dicts:
+            for item_dict in items_dicts:
+                item_scored = self.filter_score(item_dict)
+                if item_scored:
+                    items_dicts_scored.append(item_scored)
+        else:
+            items_dicts_scored = await self.fetch_items()
+
+        return await self.__post_process_items(self.__build_items(items_dicts_scored))
+
+    async def rank_items(self, items_or_dicts: List[Union[Dict, Item]]) -> List[Union[Dict, Item]]:
+        return [item for item in sorted(
+            items_or_dicts,
+            key=lambda i: (i.score if isinstance(
+                i,
+                Item,
+            ) else i.get('score') if i else -100) or - 100,
+            reverse=True,
+        )]
+
+    def __build_items(self, items_dicts_scored: List[Dict]) -> List[Item]:
+        return [Item.load(**merge_dict(ignore_keys(item_dict, [
+            'condition',
+        ]), dict(
+            title=item_dict.get('title') or item_dict.get('uuid'),
+            uuid=item_dict.get('uuid') or item_dict.get('title'),
+        ))) for item_dict in items_dicts_scored if item_dict]
+
+    async def __post_process_items(self, items: List[Item]) -> List[Item]:
+        return await self.rank_items(items)

mage_ai/command_center/files/constants.py

@@ -0,0 +1,235 @@
+import urllib.parse
+from typing import Dict
+
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.api.policies.FilePolicy import FilePolicy
+from mage_ai.command_center.constants import (
+    ApplicationType,
+    ButtonActionType,
+    InteractionType,
+    ItemType,
+    ObjectType,
+)
+from mage_ai.presenters.interactions.constants import InteractionInputType
+
+ITEMS = [
+    dict(
+        item_type=ItemType.CREATE,
+        object_type=ObjectType.FILE,
+        title='Create a new file',
+        applications=[
+            dict(
+                application_type=ApplicationType.FORM,
+                buttons=[
+                    dict(
+                        label='Cancel',
+                        tooltip='Discard changes and go back.',
+                        keyboard_shortcuts=[['metaKey', 27]],
+                        action_types=[
+                            ButtonActionType.RESET_FORM,
+                            ButtonActionType.CLOSE_APPLICATION,
+                        ],
+                    ),
+                    dict(
+                        label='Create new file',
+                        tooltip='Save changes and create the new file.',
+                        keyboard_shortcuts=[[13]],
+                        action_types=[
+                            ButtonActionType.EXECUTE,
+                            ButtonActionType.RESET_FORM,
+                        ],
+                    ),
+                ],
+                settings=[
+                    dict(
+                        label='Directory',
+                        description=(
+                            'The location to create the file in. '
+                            'Must be relative to the top level directory.'
+                        ),
+                        placeholder='e.g. utils',
+                        display_settings=dict(
+                            icon_uuid='FolderOutline',
+                        ),
+                        name='request.payload.file.dir_path',
+                        type=InteractionInputType.TEXT_FIELD,
+                        required=True,
+                        monospace=True,
+                        action_uuid='create_file',
+                    ),
+                    dict(
+                        label='File name',
+                        description='A descriptive name of the file to be created.',
+                        display_settings=dict(
+                            icon_uuid='File',
+                        ),
+                        placeholder='e.g. magic_powers.py',
+                        name='request.payload.file.name',
+                        type=InteractionInputType.TEXT_FIELD,
+                        required=True,
+                        monospace=True,
+                        action_uuid='create_file',
+                    ),
+                ],
+                uuid='new_file',
+            ),
+        ],
+        actions=[
+            dict(
+                request=dict(
+                    operation=OperationType.CREATE,
+                    payload=dict(
+                        file=dict(
+                            overwrite=False,
+                        ),
+                        file_json_only=True,
+                    ),
+                    resource='files',
+                    response_resource_key='file'
+                ),
+                uuid='create_file',
+            ),
+            dict(
+                interaction=dict(
+                    type=InteractionType.OPEN_FILE,
+                ),
+                upstream_action_value_key_mapping=dict(
+                    create_file={
+                        'data.file.path': 'interaction.options.file_path',
+                    }
+                ),
+                uuid='open_file',
+            ),
+        ],
+        condition=lambda opts: FilePolicy(
+            None,
+            opts.get('user'),
+        ).has_at_least_editor_role(),
+    ),
+    dict(
+        item_type=ItemType.CREATE,
+        object_type=ObjectType.FOLDER,
+        title='Create a new folder',
+        applications=[
+            dict(
+                application_type=ApplicationType.FORM,
+                buttons=[
+                    dict(
+                        label='Cancel',
+                        tooltip='Discard changes and go back.',
+                        keyboard_shortcuts=[['metaKey', 27]],
+                        action_types=[
+                            ButtonActionType.RESET_FORM,
+                            ButtonActionType.CLOSE_APPLICATION,
+                        ],
+                    ),
+                    dict(
+                        label='Create new folder',
+                        tooltip='Save changes and create the new folder.',
+                        keyboard_shortcuts=[[13]],
+                        action_types=[
+                            ButtonActionType.EXECUTE,
+                            ButtonActionType.RESET_FORM,
+                        ],
+                    ),
+                ],
+                settings=[
+                    dict(
+                        label='Directory',
+                        description='The parent directory of this folder.',
+                        placeholder='e.g. utils',
+                        display_settings=dict(
+                            icon_uuid='FolderOutline',
+                        ),
+                        name='request.payload.folder.path',
+                        type=InteractionInputType.TEXT_FIELD,
+                        required=True,
+                        monospace=True,
+                        action_uuid='create_folder',
+                    ),
+                    dict(
+                        label='Folder name',
+                        description='A descriptive name of the folder to be created.',
+                        display_settings=dict(
+                            icon_uuid='Alphabet',
+                        ),
+                        placeholder='e.g. transformers',
+                        name='request.payload.folder.name',
+                        type=InteractionInputType.TEXT_FIELD,
+                        required=True,
+                        monospace=True,
+                        action_uuid='create_folder',
+                    ),
+                ],
+                uuid='new_folder',
+            ),
+        ],
+        actions=[
+            dict(
+                request=dict(
+                    operation=OperationType.CREATE,
+                    payload=dict(
+                        folder=dict(
+                            overwrite=False,
+                        ),
+                    ),
+                    resource='folders',
+                    response_resource_key='folder'
+                ),
+                uuid='create_folder',
+            ),
+            dict(
+                page=dict(
+                    path='/files',
+                ),
+                uuid='navigate_folder',
+            ),
+        ],
+        condition=lambda opts: FilePolicy(
+            None,
+            opts.get('user'),
+        ).has_at_least_editor_role(),
+    ),
+]
+
+
+def add_application_actions(item_dict: Dict) -> Dict:
+    uuid = item_dict.get('uuid')
+
+    return dict(
+        applications=[
+            dict(
+                application_type=ApplicationType.DETAIL,
+                actions=[
+                    dict(
+                        request=dict(
+                            operation=OperationType.DETAIL,
+                            resource='file_contents',
+                            resource_id=urllib.parse.quote_plus(uuid),
+                            response_resource_key='file_content',
+                        ),
+                        uuid='file_content_detail',
+                    ),
+                ],
+                buttons=[
+                    dict(
+                        label='Open file',
+                        keyboard_shortcuts=[[13]],
+                        action_types=[
+                            ButtonActionType.EXECUTE,
+                        ],
+                    ),
+                ],
+                uuid='model_detail',
+            ),
+        ],
+        actions=[
+            dict(
+                interaction=dict(
+                    options=dict(file_path=uuid),
+                    type=InteractionType.OPEN_FILE,
+                ),
+                uuid='open_file',
+            ),
+        ],
+    )

mage_ai/command_center/files/factory.py

@@ -0,0 +1,86 @@
+import asyncio
+import os
+from datetime import datetime
+from typing import Dict, List
+
+from mage_ai.cache.file import FileCache
+from mage_ai.command_center.constants import ItemType, ObjectType
+from mage_ai.command_center.factory import BaseFactory
+from mage_ai.command_center.files.constants import ITEMS, add_application_actions
+from mage_ai.command_center.utils import shorten_directory
+from mage_ai.shared.hash import merge_dict
+
+
+class FileFactory(BaseFactory):
+    async def fetch_items(self, **kwargs) -> List[Dict]:
+        items = []
+
+        for item_dict in ITEMS:
+            item_scored = self.filter_score(item_dict)
+            if item_scored:
+                items.append(item_dict)
+
+        files = []
+
+        async def build_and_score(line: str, files=files, factory=self):
+            full_path, file_size, modified_timestamp = line.split(',')
+            filename = os.path.basename(full_path)
+            modified_at = datetime.fromtimestamp(float(modified_timestamp)).isoformat()
+
+            path_dict = shorten_directory(full_path)
+            extension = path_dict.get('extension')
+            directory = path_dict.get('directory')
+
+            item_dict = dict(
+                item_type=ItemType.DETAIL,
+                object_type=ObjectType.FILE,
+                title=filename,
+                description=directory,
+                uuid=full_path,
+                metadata=dict(
+                    file=dict(
+                        extension=extension,
+                        full_path=full_path,
+                        modified_at=modified_at,
+                        modified_timestamp=modified_timestamp,
+                        size=file_size,
+                    ),
+                ),
+                display_settings_by_attribute=dict(
+                    description=dict(
+                        text_styles=dict(
+                            monospace=True,
+                            small=True,
+                        ),
+                    ),
+                ),
+            )
+
+            scored = factory.filter_score(item_dict)
+            if scored:
+                files.append(scored)
+
+        if self.search:
+            now = datetime.utcnow().timestamp()
+            cache = FileCache.initialize_cache_with_settings()
+            lines = await cache.load() or []
+            print(f'[FileFactory] Load files: {len(lines)} - {datetime.utcnow().timestamp() - now}')
+
+            now = datetime.utcnow().timestamp()
+            await asyncio.gather(
+                *[build_and_score(item_dict) for item_dict in lines]
+            )
+            print(
+                f'[FileFactory] Search {self.search}: '
+                f'{len(files)} - {datetime.utcnow().timestamp() - now}',
+            )
+
+            now = datetime.utcnow().timestamp()
+            files = await self.rank_items(files)
+            files = [merge_dict(
+                item_dict,
+                add_application_actions(item_dict),
+            ) for item_dict in files]
+            print(f'[FileFactory] Rank items: {datetime.utcnow().timestamp() - now}')
+
+        return files[:10] + items