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

mage_ai/ai/constants.py

@@ -5,3 +5,4 @@ class LLMUseCase(str, Enum):
     GENERATE_COMMENT_FOR_BLOCK = 'generate_comment_for_block'
     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/generator.py

@@ -24,5 +24,11 @@ class Generator:
             return await LLMPipelineWizard().async_generate_pipeline_documentation(
                 pipeline_uuid,
             )
+        elif use_case == LLMUseCase.GENERATE_BLOCK_WITH_DESCRIPTION:
+            from mage_ai.ai.llm_pipeline_wizard import LLMPipelineWizard
+
+            return await LLMPipelineWizard().async_generate_block_with_description(
+                request.get('block_description'),
+            )
 
         raise Exception(f'Use case {use_case} is not supported yet.')

mage_ai/ai/{gen_doc_for_pipeline.py → generator_cmds.py}

@@ -42,5 +42,10 @@ def generate_pipeline_documentation(project_path: str = PROJECT_NAME_DEFAULT,
     ))['pipeline_doc'])
 
 
+@app.command()
+def generate_block_with_description(block_description: str):
+    print(asyncio.run(LLMPipelineWizard().async_generate_block_with_description(block_description)))
+
+
 if __name__ == '__main__':
     app()

mage_ai/ai/llm_pipeline_wizard.py

@@ -1,17 +1,27 @@
 import asyncio
+import json
+import os
 
+import openai
 from langchain.chains import LLMChain
 from langchain.llms import OpenAI
 from langchain.prompts import PromptTemplate
 
+from mage_ai.data_cleaner.transformer_actions.constants import ActionType, Axis
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.constants import (
     NON_PIPELINE_EXECUTABLE_BLOCK_TYPES,
     BlockLanguage,
     BlockType,
+    PipelineType,
 )
 from mage_ai.data_preparation.models.pipeline import Pipeline
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_repo_config, get_repo_path
+from mage_ai.data_preparation.templates.template import fetch_template_source
+from mage_ai.io.base import DataSource
+from mage_ai.server.logger import Logger
+
+logger = Logger().new_server_logger(__name__)
 
 BLOCK_LANGUAGE_TO_FILE_TYPE_VARIABLE = {
     BlockLanguage.MARKDOWN: 'markdown script',
@@ -39,11 +49,58 @@ Write a detailed summarization of the data pipeline based on the content provide
 ```{block_content}```
 """
 TRANSFORMERS_FOLDER = 'transformers'
+CLASSIFICATION_FUNCTION_NAME = "classify_description"
+TEMPLATE_CLASSIFICATION_FUNCTION = [
+    {
+        "name": CLASSIFICATION_FUNCTION_NAME,
+        "description": "Classify the code description provided into following properties.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                BlockType.__name__: {
+                    "type": "string",
+                    "description": "Type of the code block. It either "
+                                   "loads data from a source, export data to a source "
+                                   "or transform data from one format to another.",
+                    "enum": ["data_exporter", "data_loader", "transformer"]
+                },
+                BlockLanguage.__name__: {
+                    "type": "string",
+                    "description": "Programming language of the code block.",
+                    "enum": [type.name.lower() for type in BlockLanguage]
+                },
+                PipelineType.__name__: {
+                    "type": "string",
+                    "description": "Type of pipeline description to build.",
+                    "enum": [type.name.lower() for type in PipelineType]
+                },
+                ActionType.__name__: {
+                    "type": "string",
+                    "description": f"If {BlockType.__name__} is transformer, "
+                                   f"{ActionType.__name__} specifies what kind "
+                                   "of action the code performs.",
+                    "enum": [type.name.lower() for type in ActionType]
+                },
+                DataSource.__name__: {
+                    "type": "string",
+                    "description": f"If {BlockType.__name__} is data_loader or "
+                                   f"data_exporter, {DataSource.__name__} field specify "
+                                   "where the data loads from or exports to.",
+                    "enum": [type.name.lower() for type in DataSource]
+                },
+            },
+            "required": [BlockType.__name__, BlockLanguage.__name__, PipelineType.__name__],
+        },
+    }
+]
 
 
 class LLMPipelineWizard:
     def __init__(self):
-        self.llm = OpenAI(temperature=0)
+        repo_config = get_repo_config()
+        openai_api_key = repo_config.openai_api_key or os.getenv('OPENAI_API_KEY')
+        openai.api_key = openai_api_key
+        self.llm = OpenAI(openai_api_key=openai_api_key, temperature=0)
 
     async def __async_llm_generate_documentation(
         self,
@@ -68,6 +125,54 @@ class LLMPipelineWizard:
                                 purpose=purpose,
                                 add_on_prompt=add_on_prompt)
 
+    def __load_template_params(self, function_args: json):
+        block_type = BlockType(function_args[BlockType.__name__].lower())
+        block_language = BlockLanguage(
+                            function_args.get(BlockLanguage.__name__, "python").lower())
+        pipeline_type = PipelineType(function_args.get(PipelineType.__name__, "python").lower())
+        config = {}
+        config['action_type'] = function_args.get(ActionType.__name__, None)
+        if config['action_type']:
+            if config['action_type'] in [
+                ActionType.FILTER,
+                ActionType.DROP_DUPLICATE,
+                ActionType.REMOVE,
+                ActionType.SORT
+            ]:
+                config['axis'] = Axis.ROW
+            else:
+                config['axis'] = Axis.COLUMN
+        config['data_source'] = function_args.get(DataSource.__name__, None)
+        return block_type, block_language, pipeline_type, config
+
+    async def async_generate_block_with_description(self, block_description: str) -> dict:
+        messages = [{"role": "user", "content": block_description}]
+        response = await openai.ChatCompletion.acreate(
+            model="gpt-3.5-turbo-0613",
+            messages=messages,
+            functions=TEMPLATE_CLASSIFICATION_FUNCTION,
+            function_call={"name": CLASSIFICATION_FUNCTION_NAME},  # explicitly set function call
+        )
+        response_message = response["choices"][0]["message"]
+        if response_message.get("function_call"):
+            function_args = json.loads(response_message["function_call"]["arguments"])
+            block_type, block_language, pipeline_type, config = self.__load_template_params(
+                function_args)
+            return dict(
+                block_type=block_type,
+                configuration=config,
+                content=fetch_template_source(
+                    block_type=block_type,
+                    config=config,
+                    language=block_language,
+                    pipeline_type=pipeline_type,
+                ),
+                language=block_language,
+            )
+        else:
+            logger.error("Failed to interpret the description as a block template.")
+            return None
+
     async def async_generate_pipeline_documentation(
         self,
         pipeline_uuid: str,
@@ -92,7 +197,7 @@ class LLMPipelineWizard:
         chain = LLMChain(llm=self.llm, prompt=prompt_template)
         pipeline_doc = chain.run(block_content=block_docs_content)
         return dict(
-            block_content=block_docs,
+            block_docs=block_docs,
             pipeline_doc=pipeline_doc,
         )
 

mage_ai/api/policies/BlockPolicy.py

@@ -56,6 +56,7 @@ BlockPolicy.allow_read([
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
 BlockPolicy.allow_write([
+    'block_action_object',
     'color',
     'config',
     'configuration',

mage_ai/api/policies/BlockTemplatePolicy.py

@@ -25,3 +25,12 @@ BlockTemplatePolicy.allow_read(BlockTemplatePresenter.default_attributes + [], s
     constants.DETAIL,
     constants.LIST,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+BlockTemplatePolicy.allow_query([
+    'show_all',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/LlmPolicy.py

@@ -23,9 +23,10 @@ LlmPolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
+
 LlmPolicy.allow_read([
+    'response',
     'use_case',
-    'request',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
@@ -34,8 +35,8 @@ LlmPolicy.allow_read([
 
 
 LlmPolicy.allow_write([
+    'request',
     'use_case',
-    'response',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[

mage_ai/api/policies/OauthAccessTokenPolicy.py

@@ -0,0 +1,31 @@
+from mage_ai.api.oauth_scope import OauthScope
+from mage_ai.api.operations import constants
+from mage_ai.api.policies.BasePolicy import BasePolicy
+from mage_ai.api.presenters.OauthAccessTokenPresenter import OauthAccessTokenPresenter
+
+
+class OauthAccessTokenPolicy(BasePolicy):
+    pass
+
+
+OauthAccessTokenPolicy.allow_actions([
+    constants.LIST,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+OauthAccessTokenPolicy.allow_read(OauthAccessTokenPresenter.default_attributes, scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+OauthAccessTokenPolicy.allow_query([
+    'show_all',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/PipelinePolicy.py

@@ -96,6 +96,7 @@ PipelinePolicy.allow_write([
     'callbacks',
     'conditionals',
     'extensions',
+    'llm',
     'schedules',
 ] + PipelinePresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/policies/ProjectPolicy.py

@@ -35,3 +35,12 @@ ProjectPolicy.allow_write([
 ], on_action=[
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
+
+ProjectPolicy.allow_write([
+    'features',
+    'openai_api_key',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.UPDATE,
+], condition=lambda policy: policy.has_at_least_editor_role())

mage_ai/api/policies/SearchResultPolicy.py

@@ -0,0 +1,32 @@
+from mage_ai.api.oauth_scope import OauthScope
+from mage_ai.api.operations import constants
+from mage_ai.api.policies.BasePolicy import BasePolicy
+from mage_ai.api.presenters.SearchResultPresenter import SearchResultPresenter
+
+
+class SearchResultPolicy(BasePolicy):
+    pass
+
+
+SearchResultPolicy.allow_actions([
+    constants.CREATE,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+SearchResultPolicy.allow_read(SearchResultPresenter.default_attributes, scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+SearchResultPolicy.allow_write([
+    'pipeline_type',
+    'query',
+    'ratio',
+    'type',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/VariablePolicy.py

@@ -55,3 +55,11 @@ VariablePolicy.allow_write([
     constants.CREATE,
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access())
+
+VariablePolicy.allow_query([
+    'global_only',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/presenters/BlockTemplatePresenter.py

@@ -9,4 +9,5 @@ class BlockTemplatePresenter(BasePresenter):
         'language',
         'name',
         'path',
+        'template_variables',
     ]

mage_ai/api/presenters/OauthAccessTokenPresenter.py

@@ -0,0 +1,9 @@
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class OauthAccessTokenPresenter(BasePresenter):
+    default_attributes = [
+        'expires',
+        'token',
+        'user_id',
+    ]

mage_ai/api/presenters/ProjectPresenter.py

@@ -3,9 +3,11 @@ from mage_ai.api.presenters.BasePresenter import BasePresenter
 
 class ProjectPresenter(BasePresenter):
     default_attributes = [
+        'features',
         'help_improve_mage',
         'latest_version',
         'name',
+        'openai_api_key',
         'project_uuid',
         'version',
     ]

mage_ai/api/presenters/SearchResultPresenter.py

@@ -0,0 +1,73 @@
+from mage_ai.api.operations import constants
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+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.services.search.constants import SEARCH_TYPE_BLOCK_ACTION_OBJECTS
+from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.strings import remove_extension_from_filename
+
+
+class SearchResultPresenter(BasePresenter):
+    default_attributes = [
+        'results',
+        'type',
+        'uuid',
+    ]
+
+    def present(self, **kwargs):
+        display_format = kwargs['format']
+
+        if constants.CREATE != display_format:
+            return self.model
+
+        results = self.model.get('results', [])
+        search_type = self.model.get('type')
+        results_transformed = []
+
+        if SEARCH_TYPE_BLOCK_ACTION_OBJECTS == search_type:
+            for result in results:
+                block_action_object = result.get('block_action_object')
+                object_type = result.get('object_type')
+                uuid = result.get('uuid')
+
+                block_type = None
+                description = block_action_object.get('description')
+                language = block_action_object.get('language')
+                title = None
+
+                if OBJECT_TYPE_BLOCK_FILE == object_type:
+                    block_type = block_action_object.get('type')
+                    u = block_action_object.get('uuid')
+                    title = ' '.join(list(filter(lambda x: x, [
+                        remove_extension_from_filename(u).replace('_', ' ') if u else u,
+                    ])))
+                elif OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE == object_type:
+                    block_type = block_action_object.get('block_type')
+                    template_name = block_action_object.get('name')
+                    template_uuid = block_action_object.get('template_uuid')
+                    u = template_name or template_uuid
+                    title = ' '.join(list(filter(lambda x: x, [
+                        remove_extension_from_filename(u).replace('_', ' ') if u else u,
+                    ])))
+                elif OBJECT_TYPE_MAGE_TEMPLATE == object_type:
+                    block_type = block_action_object.get('block_type')
+                    u = block_action_object.get('name')
+                    title = ' '.join(list(filter(lambda x: x, [
+                        remove_extension_from_filename(u).replace('_', ' ') if u else u,
+                    ])))
+
+                results_transformed.append(dict(
+                    block_type=block_type,
+                    description=description,
+                    language=language,
+                    object_type=object_type,
+                    title=title,
+                    uuid=uuid,
+                ))
+
+        return merge_dict(self.model, dict(
+            results=results_transformed,
+        ))

mage_ai/api/resources/BlockResource.py

@@ -3,6 +3,12 @@ import urllib.parse
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.cache.block import BlockCache
+from mage_ai.cache.block_action_object import BlockActionObjectCache
+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.block import Block
 from mage_ai.data_preparation.models.block.dbt import DBTBlock
 from mage_ai.data_preparation.models.block.utils import clean_name
@@ -29,6 +35,31 @@ class BlockResource(GenericResource):
         content = payload.get('content')
         language = payload.get('language')
         name = payload.get('name')
+        block_name = name or payload.get('uuid')
+
+        payload_config = payload.get('config') or {}
+
+        block_action_object = payload.get('block_action_object')
+        if block_action_object:
+            object_type = block_action_object.get('object_type')
+
+            cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+            mapping = cache_block_action_object.load_all_data()
+            objects_mapping = mapping.get(object_type)
+            object_uuid = block_action_object.get('uuid')
+            object_from_cache = objects_mapping.get(object_uuid)
+
+            if OBJECT_TYPE_BLOCK_FILE == object_type:
+                block_name = object_from_cache.get('uuid')
+                block_type = object_from_cache.get('type')
+                language = object_from_cache.get('language')
+            elif OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE == object_type:
+                payload_config['custom_template_uuid'] = object_from_cache.get('template_uuid')
+            elif OBJECT_TYPE_MAGE_TEMPLATE == object_type:
+                block_type = object_from_cache.get('block_type')
+                language = object_from_cache.get('language')
+                payload_config['template_path'] = object_from_cache.get('path')
+                payload_config['template_variables'] = object_from_cache.get('template_variables')
 
         """
         New DBT models include "content" in its block create payload,
@@ -47,9 +78,6 @@ class BlockResource(GenericResource):
                     Please choose a different model name, or add a DBT model by \
                     selecting single model from file.')
 
-        payload_config = payload.get('config')
-
-        block_name = name or payload.get('uuid')
         block_attributes = dict(
             color=payload.get('color'),
             config=payload_config,
@@ -104,6 +132,9 @@ class BlockResource(GenericResource):
         cache = await BlockCache.initialize_cache()
         cache.add_pipeline(block, pipeline)
 
+        cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+        cache_block_action_object.update_block(block)
+
         return self(block, user, **kwargs)
 
     @classmethod
@@ -189,10 +220,16 @@ class BlockResource(GenericResource):
         if pipeline:
             cache.remove_pipeline(self.model, pipeline.uuid)
 
+        cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+        cache_block_action_object.update_block(self.model, remove=True)
+
         return self.model.delete(force=force)
 
     @safe_db_query
-    def update(self, payload, **kwargs):
+    async def update(self, payload, **kwargs):
+        cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+        cache_block_action_object.update_block(self.model, remove=True)
+
         query = kwargs.get('query', {})
         update_state = query.get('update_state', [False])
         if update_state:
@@ -202,6 +239,8 @@ class BlockResource(GenericResource):
             update_state=update_state,
         )
 
+        cache_block_action_object.update_block(self.model)
+
     async def get_pipelines_from_cache(self):
         await BlockCache.initialize_cache()
 

mage_ai/api/resources/BlockTemplateResource.py

@@ -1,6 +1,10 @@
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.templates.constants import TEMPLATES, TEMPLATES_BY_UUID
+from mage_ai.data_preparation.templates.constants import (
+    TEMPLATES,
+    TEMPLATES_BY_UUID,
+    TEMPLATES_ONLY_FOR_V2,
+)
 from mage_ai.orchestration.db import safe_db_query
 
 
@@ -8,8 +12,17 @@ class BlockTemplateResource(GenericResource):
     @classmethod
     @safe_db_query
     def collection(self, query, meta, user, **kwargs):
+        show_all = query.get('show_all', [None])
+        if show_all:
+            show_all = show_all[0]
+
+        arr = TEMPLATES.copy()
+
+        if show_all:
+            arr += TEMPLATES_ONLY_FOR_V2.copy()
+
         return self.build_result_set(
-            TEMPLATES,
+            arr,
             user,
             **kwargs,
         )

mage_ai/api/resources/CustomTemplateResource.py

@@ -1,6 +1,7 @@
 import urllib.parse
 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.custom_templates.constants import (
     DIRECTORY_FOR_BLOCK_TEMPLATES,
     DIRECTORY_FOR_PIPELINE_TEMPLATES,
@@ -51,7 +52,7 @@ class CustomTemplateResource(GenericResource):
         )
 
     @classmethod
-    def create(self, payload, user, **kwargs):
+    async def create(self, payload, user, **kwargs):
         custom_template = None
         object_type = payload.get(OBJECT_TYPE_KEY)
         template_uuid = payload.get('template_uuid')
@@ -76,6 +77,9 @@ class CustomTemplateResource(GenericResource):
                 )
 
                 custom_template.save()
+
+                cache = await BlockActionObjectCache.initialize_cache()
+                cache.update_custom_block_template(custom_template)
         elif DIRECTORY_FOR_PIPELINE_TEMPLATES == object_type:
             custom_template = CustomPipelineTemplate.load(template_uuid=template_uuid)
 
@@ -120,13 +124,21 @@ class CustomTemplateResource(GenericResource):
             print(f'[WARNING] CustomTemplateResource.member: {err}')
             raise ApiError(ApiError.RESOURCE_NOT_FOUND)
 
-    def delete(self, **kwargs):
+    async def delete(self, **kwargs):
+        cache = await BlockActionObjectCache.initialize_cache()
+        cache.update_custom_block_template(self.model, remove=True)
         self.model.delete
 
-    def update(self, payload, **kwargs):
+    async def update(self, payload, **kwargs):
+        cache = await BlockActionObjectCache.initialize_cache()
+
+        cache.update_custom_block_template(self.model, remove=True)
+
         for key, value in ignore_keys(payload, [
             'uuid',
             OBJECT_TYPE_KEY,
         ]).items():
             setattr(self.model, key, value)
         self.model.save()
+
+        cache.update_custom_block_template(self.model)

mage_ai/api/resources/FileContentResource.py

@@ -2,6 +2,8 @@ import urllib.parse
 
 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.file import File
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.settings.repo import get_repo_path
@@ -16,7 +18,7 @@ class FileContentResource(GenericResource):
         file = File.from_path(file_path, get_repo_path())
         return self(file, user, **kwargs)
 
-    def update(self, payload, **kwargs):
+    async def update(self, payload, **kwargs):
         error = ApiError.RESOURCE_INVALID.copy()
 
         content = payload.get('content')
@@ -31,3 +33,8 @@ class FileContentResource(GenericResource):
             content = file_version.content()
 
         self.model.update_content(content)
+
+        block_type = Block.block_type_from_path(self.model.dir_path)
+        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)