mage-ai 0.9.7__py3-none-any.whl → 0.9.9__py3-none-any.whl

mage_ai/api/resources/FileResource.py

@@ -3,6 +3,8 @@ from typing import Dict
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.cache.block_action_object import BlockActionObjectCache
+from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.errors import FileExistsError
 from mage_ai.data_preparation.models.file import File
 from mage_ai.orchestration.db import safe_db_query
@@ -21,7 +23,7 @@ class FileResource(GenericResource):
 
     @classmethod
     @safe_db_query
-    def create(self, payload: Dict, user, **kwargs) -> 'FileResource':
+    async def create(self, payload: Dict, user, **kwargs) -> 'FileResource':
         dir_path = payload['dir_path']
         repo_path = get_repo_path()
         content = None
@@ -42,6 +44,11 @@ class FileResource(GenericResource):
                 overwrite=payload.get('overwrite', False),
             )
 
+            block_type = Block.block_type_from_path(dir_path)
+            if block_type:
+                cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+                cache_block_action_object.update_block(block_file_absolute_path=file.file_path)
+
             return self(file, user, **kwargs)
         except FileExistsError as err:
             error = ApiError.RESOURCE_INVALID.copy()
@@ -70,6 +77,19 @@ class FileResource(GenericResource):
         return self.model.delete()
 
     @safe_db_query
-    def update(self, payload, **kwargs):
+    async def update(self, payload, **kwargs):
+        block_type = Block.block_type_from_path(self.model.dir_path)
+        cache_block_action_object = None
+        if block_type:
+            cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+            cache_block_action_object.update_block(
+                block_file_absolute_path=self.model.file_path,
+                remove=True,
+            )
+
         self.model.rename(payload['dir_path'], payload['name'])
+
+        if block_type and cache_block_action_object:
+            cache_block_action_object.update_block(block_file_absolute_path=self.model.file_path)
+
         return self

mage_ai/api/resources/OauthAccessTokenResource.py

@@ -0,0 +1,24 @@
+from datetime import datetime
+from mage_ai.api.resources.DatabaseResource import DatabaseResource
+from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken
+
+
+class OauthAccessTokenResource(DatabaseResource):
+    model_class = Oauth2AccessToken
+
+    @classmethod
+    @safe_db_query
+    def collection(self, query_arg, meta, user, **kwargs):
+        results = Oauth2AccessToken.query.filter(Oauth2AccessToken.user_id == user.id)
+
+        show_all = query_arg.get('show_all', [None])
+        if show_all:
+            show_all = show_all[0]
+
+        if not show_all:
+            results = results.filter(Oauth2AccessToken.expires > datetime.utcnow())
+
+        results = results.order_by(Oauth2AccessToken.expires.desc())
+
+        return results

mage_ai/api/resources/PipelineResource.py

@@ -3,12 +3,19 @@ import asyncio
 from sqlalchemy import or_
 from sqlalchemy.orm import aliased
 
+from mage_ai.ai.constants import LLMUseCase
 from mage_ai.api.operations.constants import DELETE
 from mage_ai.api.resources.BaseResource import BaseResource
+from mage_ai.api.resources.BlockResource import BlockResource
+from mage_ai.api.resources.LlmResource import LlmResource
 from mage_ai.data_preparation.models.block.dbt.utils import (
     add_blocks_upstream_from_refs,
 )
-from mage_ai.data_preparation.models.constants import PipelineStatus
+from mage_ai.data_preparation.models.constants import (
+    BlockLanguage,
+    BlockType,
+    PipelineStatus,
+)
 from mage_ai.data_preparation.models.custom_templates.custom_pipeline_template import (
     CustomPipelineTemplate,
 )
@@ -23,7 +30,8 @@ from mage_ai.orchestration.pipeline_scheduler import (
 from mage_ai.server.active_kernel import switch_active_kernel
 from mage_ai.server.kernels import PIPELINE_TO_KERNEL_NAME
 from mage_ai.settings.repo import get_repo_path
-from mage_ai.shared.hash import group_by, ignore_keys
+from mage_ai.shared.array import find_index
+from mage_ai.shared.hash import group_by, ignore_keys, merge_dict
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
@@ -227,6 +235,66 @@ class PipelineResource(BaseResource):
         if update_content:
             update_content = update_content[0]
 
+        llm_payload = payload.get('llm')
+        if llm_payload:
+            llm_use_case = llm_payload.get('use_case')
+            llm_request = llm_payload.get('request')
+
+            if 'pipeline_uuid' not in llm_payload:
+                llm_payload['pipeline_uuid'] = self.model.uuid
+
+            llm_resource = await LlmResource.create(llm_payload, self.current_user, **kwargs)
+            llm_response = llm_resource.model.get('response')
+
+            pipeline_doc = None
+            block_docs = []
+            blocks = self.model.blocks_by_uuid.values()
+
+            async def _add_markdown_block(block_doc: str, block_uuid: str, priority: int):
+                return await BlockResource.create(dict(
+                    content=block_doc.strip() if block_doc else block_doc,
+                    language=BlockLanguage.MARKDOWN,
+                    name=f'Documentation for {block_uuid}',
+                    priority=priority,
+                    type=BlockType.MARKDOWN,
+                ), self.current_user, **merge_dict(kwargs, dict(
+                    parent_model=self.model,
+                )))
+
+            if LLMUseCase.GENERATE_COMMENT_FOR_BLOCK == llm_use_case:
+                pass
+            elif LLMUseCase.GENERATE_DOC_FOR_BLOCK == llm_use_case:
+                block_doc = llm_response.get('block_doc')
+                if block_doc:
+                    block_uuid = llm_request.get('block_uuid')
+                    priority = find_index(lambda x: x.uuid == block_uuid, blocks)
+                    await _add_markdown_block(block_doc, block_uuid, priority)
+            elif LLMUseCase.GENERATE_DOC_FOR_PIPELINE == llm_use_case:
+                block_docs = llm_response.get('block_docs', [])
+                pipeline_doc = llm_response.get('pipeline_doc')
+
+                if pipeline_doc:
+                    lines = []
+                    if payload.get('description'):
+                        lines.append(payload.get('description'))
+                    lines.append(pipeline_doc)
+                    payload['description'] = '\n'.join(lines).strip()
+
+            if block_docs and len(block_docs) >= 1:
+                blocks_with_docs = list(zip(block_docs, blocks))
+                blocks_with_docs.reverse()
+                blocks_with_docs_length = len(blocks_with_docs)
+
+                for idx, tup in enumerate(blocks_with_docs):
+                    priority = blocks_with_docs_length - (idx + 1)
+                    block_doc, block = tup
+
+                    if block_doc:
+                        await _add_markdown_block(block_doc, block.uuid, priority)
+
+            if pipeline_doc:
+                await _add_markdown_block(pipeline_doc, self.model.uuid, 0)
+
         await self.model.update(
             ignore_keys(payload, ['add_upstream_for_block_uuid']),
             update_content=update_content,

mage_ai/api/resources/ProjectResource.py

@@ -11,6 +11,7 @@ async def build_project(repo_config=None, **kwargs):
     project = Project(repo_config=repo_config)
 
     model = merge_dict(project.repo_config.to_dict(), dict(
+        features=project.features,
         latest_version=await project.latest_version(),
         name=project.name,
         project_uuid=project.project_uuid,
@@ -47,7 +48,15 @@ class ProjectResource(GenericResource):
         repo_config = get_repo_config()
 
         data = {}
-        should_log_project = False
+        should_log_project = self.model.get('help_improve_mage') or False
+
+        if 'features' in payload:
+            for k, v in payload.get('features', {}).items():
+                feature = (repo_config.features or {}).get(k)
+                if (feature and not v) or (not feature and v):
+                    if 'features' not in data:
+                        data['features'] = {}
+                    data['features'][k] = v
 
         if 'help_improve_mage' in payload:
             if payload['help_improve_mage']:
@@ -58,6 +67,11 @@ class ProjectResource(GenericResource):
 
             data['help_improve_mage'] = payload['help_improve_mage']
 
+        if 'openai_api_key' in payload:
+            openai_api_key = payload.get('openai_api_key')
+            if repo_config.openai_api_key != openai_api_key:
+                data['openai_api_key'] = payload.get('openai_api_key')
+
         if len(data.keys()) >= 1:
             repo_config.save(**data)
 

mage_ai/api/resources/SearchResultResource.py

@@ -0,0 +1,68 @@
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.cache.block_action_object.constants import (
+    OBJECT_TYPE_BLOCK_FILE,
+    OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE,
+    OBJECT_TYPE_MAGE_TEMPLATE,
+)
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType, PipelineType
+from mage_ai.services.search.constants import SEARCH_TYPE_BLOCK_ACTION_OBJECTS
+from typing import Dict
+
+
+def filter_results(result: Dict) -> bool:
+    block_action_object = result.get('block_action_object')
+    object_type = result.get('object_type')
+
+    block_type = None
+    language = None
+
+    if OBJECT_TYPE_BLOCK_FILE == object_type:
+        block_type = block_action_object.get('type')
+        language = block_action_object.get('language')
+    elif OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE == object_type:
+        block_type = block_action_object.get('block_type')
+        language = block_action_object.get('language')
+    elif OBJECT_TYPE_MAGE_TEMPLATE == object_type:
+        block_type = block_action_object.get('block_type')
+        language = block_action_object.get('language')
+
+    if BlockLanguage.YAML == language and BlockType.DBT != block_type:
+        return False
+
+    if block_type in [
+        BlockType.CALLBACK,
+        BlockType.CONDITIONAL,
+        BlockType.EXTENSION,
+    ]:
+        return False
+
+    return True
+
+
+class SearchResultResource(GenericResource):
+    @classmethod
+    async def create(self, payload: Dict, user, **kwargs):
+        pipeline_type = payload.get('pipeline_type', None)
+        query = payload.get('query', None)
+        ratio = payload.get('ratio', None)
+        search_type = payload.get('type', None)
+
+        results = []
+
+        if query:
+            if SEARCH_TYPE_BLOCK_ACTION_OBJECTS == search_type:
+                from mage_ai.services.search.block_action_objects import search
+
+                results = await search(query, ratio=ratio)
+
+                # TODO (tommy dangerous): remove this when we unify pipeline types
+                if PipelineType.PYTHON == pipeline_type:
+                    results = list(filter(filter_results, results))
+
+                results = results[:12]
+
+        return self(dict(
+            results=results,
+            type=search_type,
+            uuid=query,
+        ), user, **kwargs)

mage_ai/api/resources/VariableResource.py

@@ -42,9 +42,11 @@ class VariableResource(GenericResource):
     def collection(self, query, meta, user, **kwargs):
         pipeline_uuid = kwargs['parent_model'].uuid
 
+        global_only = query.get('global_only', [False])
+        if global_only:
+            global_only = global_only[0]
+
         # Get global variables from project's path
-        variable_manager = VariableManager(variables_dir=get_variables_dir())
-        variables_dict = variable_manager.get_variables_by_pipeline(pipeline_uuid)
         global_variables = [
             dict(
                 uuid=uuid,
@@ -60,18 +62,27 @@ class VariableResource(GenericResource):
                 variables=global_variables,
             )
         ]
-        variables = [
-            dict(
-                block=dict(uuid=uuid),
-                pipeline=dict(uuid=pipeline_uuid),
-                variables=[
-                    get_variable_value(variable_manager, pipeline_uuid, uuid, var) for var in arr
-                    # Not return printed outputs
-                    if var == 'df' or var.startswith('output')
-                  ],
-            )
-            for uuid, arr in variables_dict.items() if uuid != 'global'
-        ] + global_variables_arr
+        variables = global_variables_arr
+        if not global_only:
+            variable_manager = VariableManager(variables_dir=get_variables_dir())
+            variables_dict = variable_manager.get_variables_by_pipeline(pipeline_uuid)
+            variables = [
+                dict(
+                    block=dict(uuid=uuid),
+                    pipeline=dict(uuid=pipeline_uuid),
+                    variables=[
+                        get_variable_value(
+                            variable_manager,
+                            pipeline_uuid,
+                            uuid,
+                            var,
+                        ) for var in arr
+                        # Not return printed outputs
+                        if var == 'df' or var.startswith('output')
+                    ],
+                )
+                for uuid, arr in variables_dict.items() if uuid != 'global'
+            ] + global_variables_arr
 
         return self.build_result_set(
             variables,

mage_ai/cache/base.py

@@ -38,6 +38,9 @@ class BaseCache():
         value_fetched = self.storage.read_json_file(self.__build_path(key), None)
         return value_fetched
 
+    def load_all_data(self):
+        return self.get(self.cache_key)
+
     def set(self, key: str, value: Any) -> None:
         self.storage.write_json_file(self.__build_path(key), value)
 

mage_ai/cache/block_action_object/__init__.py

@@ -0,0 +1,282 @@
+import os
+from mage_ai.cache.base import BaseCache
+from mage_ai.cache.block_action_object.constants import (
+    OBJECT_TYPE_BLOCK_FILE,
+    OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE,
+    OBJECT_TYPE_MAGE_TEMPLATE,
+)
+from mage_ai.cache.constants import CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING
+from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.constants import (
+    BlockLanguage,
+    BlockType,
+    FILE_EXTENSION_TO_BLOCK_LANGUAGE,
+)
+from mage_ai.data_preparation.models.custom_templates.constants import (
+    DIRECTORY_FOR_BLOCK_TEMPLATES,
+)
+from mage_ai.data_preparation.models.custom_templates.custom_block_template import (
+    CustomBlockTemplate,
+)
+from mage_ai.data_preparation.models.custom_templates.utils import (
+    flatten_files,
+    get_templates,
+    group_and_hydrate_files,
+)
+from mage_ai.data_preparation.templates.constants import (
+    TEMPLATES,
+    TEMPLATES_ONLY_FOR_V2,
+)
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.strings import remove_extension_from_filename
+from pathlib import Path
+from typing import Dict
+
+
+def parse_block_file_absolute_path(block_file_absolute_path: str) -> Dict:
+    block_type = Block.block_type_from_path(block_file_absolute_path)
+    file_directory_name = Block.file_directory_name(block_type)
+
+    # This is the block_file_absolute_path without the repo_path
+    file_path = str(block_file_absolute_path).replace(get_repo_path(), '')
+    if file_path.startswith(os.sep):
+        file_path = file_path[1:]
+
+    file_path_parts = file_path.split(os.path.sep)
+
+    if len(file_path_parts) >= 1 and file_directory_name == file_path_parts[0]:
+        file_path_parts = file_path_parts[1:]
+
+    # This is the file_path without the block type directory
+    filename = os.path.join(*file_path_parts)
+
+    filename_parts = filename.split('.')
+    file_extension = None
+    if len(filename_parts) >= 2:
+        file_extension = filename_parts[-1]
+
+    block_language = None
+    if file_extension:
+        block_language = FILE_EXTENSION_TO_BLOCK_LANGUAGE.get(file_extension)
+
+    return dict(
+        block_type=block_type,
+        file_path=file_path,
+        filename=filename,
+        filename_parts=filename_parts,
+        language=block_language,
+    )
+
+
+class BlockActionObjectCache(BaseCache):
+    cache_key = CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING
+
+    @classmethod
+    async def initialize_cache(self, replace: bool = False) -> 'BlockActionObjectCache':
+        cache = self()
+        if replace or not cache.exists():
+            await cache.initialize_cache_for_all_objects()
+
+        return cache
+
+    def build_key_for_block_file(
+        self,
+        file_path: str = None,
+        block_type: BlockType = None,
+        language: BlockLanguage = None,
+        filename: str = None,
+    ) -> str:
+        return ':'.join(list(filter(lambda x: x, [
+            file_path,
+            block_type,
+            language,
+            remove_extension_from_filename(filename).replace('_', ' ') if filename else filename,
+        ])))
+
+    def build_key_for_custom_block_template(
+        self,
+        custom_block_template: CustomBlockTemplate,
+    ) -> str:
+        arr = []
+
+        # This is a specific order
+        for key in [
+            'uuid',
+            'block_type',
+            'language',
+            'name',
+            'description',
+            'tags',
+        ]:
+            val = getattr(custom_block_template, key)
+            if val:
+                if not isinstance(val, list):
+                    val = [val]
+                arr += val
+
+        return ':'.join(arr)
+
+    def build_key_for_mage_template(self, block_action_object: Dict) -> str:
+        arr = []
+
+        # This is a specific order
+        for key in [
+            'path',
+            'block_type',
+            'language',
+            'name',
+            'description',
+            'groups',
+        ]:
+            val = block_action_object.get(key)
+            if val:
+                if not isinstance(val, list):
+                    val = [val]
+                arr += val
+
+        return ':'.join(arr)
+
+    def update_block(
+        self,
+        block: Block = None,
+        block_file_absolute_path: str = None,
+        remove: bool = False,
+    ) -> None:
+        if block:
+            block_type = block.type
+            file_path = os.path.join(block.file.dir_path, block.file.filename)
+            filename = os.path.join(*file_path.split(os.path.sep)[1:])
+            language = block.language
+            uuid = block.uuid
+        else:
+            d = parse_block_file_absolute_path(block_file_absolute_path)
+            block_type = d.get('block_type')
+            file_path = d.get('file_path')
+            filename = d.get('filename')
+            language = d.get('language')
+
+            filename_parts = d.get('filename_parts')
+            uuid = '.'.join(filename_parts[:-1])
+
+        key = self.build_key_for_block_file(
+            file_path,
+            block_type,
+            language,
+            filename,
+        )
+
+        mapping = self.load_all_data()
+        if mapping is None:
+            mapping = {}
+
+        if OBJECT_TYPE_BLOCK_FILE not in mapping:
+            mapping[OBJECT_TYPE_BLOCK_FILE] = {}
+
+        if remove:
+            if key in mapping[OBJECT_TYPE_BLOCK_FILE]:
+                mapping[OBJECT_TYPE_BLOCK_FILE].pop(key, None)
+        else:
+            content = None
+            if block:
+                content = block.content
+            else:
+                with open(block_file_absolute_path, 'r') as f:
+                    content = f.read()
+
+            mapping[OBJECT_TYPE_BLOCK_FILE][key] = dict(
+                content=content,
+                file_path=file_path,
+                language=language,
+                type=block_type,
+                uuid=uuid,
+            )
+
+        self.set(self.cache_key, mapping)
+
+    def update_custom_block_template(
+        self,
+        custom_block_template: CustomBlockTemplate,
+        remove: bool = False,
+    ) -> None:
+        key = self.build_key_for_custom_block_template(custom_block_template)
+
+        mapping = self.load_all_data()
+        if mapping is None:
+            mapping = {}
+
+        if OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE not in mapping:
+            mapping[OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE] = {}
+
+        if remove:
+            if key in mapping[OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE]:
+                mapping[OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE].pop(key, None)
+        else:
+            mapping[OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE][key] = custom_block_template.to_dict(
+                include_content=True,
+            )
+
+        self.set(self.cache_key, mapping)
+
+    async def initialize_cache_for_all_objects(self) -> None:
+        mapping = {
+            OBJECT_TYPE_BLOCK_FILE: {},
+            OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE: {},
+            OBJECT_TYPE_MAGE_TEMPLATE: {},
+        }
+
+        for block_action_object in (TEMPLATES + TEMPLATES_ONLY_FOR_V2):
+            key = self.build_key_for_mage_template(block_action_object)
+            mapping[OBJECT_TYPE_MAGE_TEMPLATE][key] = block_action_object
+
+        file_dicts = get_templates(DIRECTORY_FOR_BLOCK_TEMPLATES)
+        file_dicts_flat = flatten_files(file_dicts)
+        custom_block_templates = group_and_hydrate_files(file_dicts_flat, CustomBlockTemplate)
+
+        for custom_block_template in custom_block_templates:
+            key = self.build_key_for_custom_block_template(custom_block_template)
+            mapping[OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE][key] = custom_block_template.to_dict(
+                include_content=True,
+            )
+
+        for block_type in BlockType:
+            file_directory_name = Block.file_directory_name(block_type)
+            directory_full_path = os.path.join(get_repo_path(), file_directory_name)
+
+            for path in Path(directory_full_path).rglob('*'):
+                if not path.is_file():
+                    continue
+
+                block_file_absolute_path = path.absolute()
+                d = parse_block_file_absolute_path(block_file_absolute_path)
+
+                file_path = d.get('file_path')
+                filename = d.get('filename')
+                filename_parts = d.get('filename_parts')
+                language = d.get('language')
+
+                if '__init__.py' == filename:
+                    continue
+
+                if not language:
+                    continue
+
+                key = self.build_key_for_block_file(
+                    file_path,
+                    block_type,
+                    language,
+                    filename,
+                )
+
+                content = None
+                with open(block_file_absolute_path, 'r') as f:
+                    content = f.read()
+
+                mapping[OBJECT_TYPE_BLOCK_FILE][key] = dict(
+                    content=content,
+                    file_path=file_path,
+                    language=language,
+                    type=block_type,
+                    uuid='.'.join(filename_parts[:-1]),
+                )
+
+        self.set(self.cache_key, mapping)

mage_ai/cache/block_action_object/constants.py

@@ -0,0 +1,4 @@
+OBJECT_TYPE_BLOCK_FILE = 'block_file'
+OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE = 'custom_block_template'
+OBJECT_TYPE_GENERATE_BLOCK = 'generate_block'
+OBJECT_TYPE_MAGE_TEMPLATE = 'mage_template'

mage_ai/cache/constants.py

@@ -1,4 +1,5 @@
 CACHE_KEY_BLOCKS_TO_PIPELINE_MAPPING = 'blocks_to_pipeline_mapping'
+CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING = 'block_action_objects_mapping'
 CACHE_KEY_TAGS_TO_OBJECT_MAPPING = 'tags_to_object_mapping'
 
 MAGE_CACHE_DIRECTORY_DEFAULT = '.cache'

mage_ai/data_integrations/destinations/constants.py

@@ -1,8 +1,10 @@
 DESTINATIONS = [
     dict(name='Amazon S3'),
     dict(name='BigQuery'),
+    dict(name='Clickhouse'),
     dict(name='Delta Lake S3'),
     dict(name='Google Cloud Storage'),
+    dict(name='MongoDB'),
     dict(
         module_name='MSSQL',
         name='Microsoft SQL Server',

mage_ai/data_preparation/executors/k8s_pipeline_executor.py

@@ -27,7 +27,7 @@ class K8sPipelineExecutor(PipelineExecutor):
         **kwargs,
     ) -> None:
         cmd = f'/app/run_app.sh '\
-              f'python mage_ai/cli/main.py run {self.pipeline.repo_config.repo_path} '\
+              f'mage run {self.pipeline.repo_config.repo_path} '\
               f'{self.pipeline.uuid}'
         options = [
             '--executor-type local_python',