mage-ai 0.9.7__tar.gz → 0.9.9__tar.gz

{mage-ai-0.9.7/mage_ai.egg-info → mage-ai-0.9.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.7
+Version: 0.9.9
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage

{mage-ai-0.9.7 → mage-ai-0.9.9}/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-0.9.7 → mage-ai-0.9.9}/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-0.9.7/mage_ai/ai/gen_doc_for_pipeline.py → mage-ai-0.9.9/mage_ai/ai/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-0.9.9/mage_ai/ai/llm_pipeline_wizard.py

@@ -0,0 +1,229 @@
+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_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',
+    BlockLanguage.PYTHON: 'python script',
+    BlockLanguage.R: 'R script',
+    BlockLanguage.SQL: 'sql script',
+    BlockLanguage.YAML: 'yaml configuration',
+}
+BLOCK_TYPE_TO_PURPOSE_VARIABLE = {
+    BlockType.DATA_LOADER: "load data from a specific data source.",
+    BlockType.DATA_EXPORTER: "export data to a specific data source.",
+    BlockType.TRANSFORMER: "transform data from one format into different format.",
+}
+DATA_LOADERS_FOLDER = 'data_loaders'
+DATA_EXPORTERS_FOLDER = 'data_exporters'
+PROMPT_FOR_BLOCK = """
+The {file_type} delimited by triple backticks is used to {purpose}.
+Write a documentation based on the {file_type}. {add_on_prompt}
+Ignore the imported libraries and the @test decorator.
+```{block_content}```"""
+PROMPT_FOR_SUMMARIZE_BLOCK_DOC = """
+A data pipeline reads data from source, transform the data and export data into another source.
+The content delimited by triple backticks contains explains of all components in one data pipeline.
+Write a detailed summarization of the data pipeline based on the content provided.
+```{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):
+        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,
+        block_content: str,
+        file_type: str,
+        purpose: str,
+        template: str,
+        add_on_prompt: str = '',
+    ):
+        prompt_template = PromptTemplate(
+            input_variables=[
+                'block_content',
+                'file_type',
+                'purpose',
+                'add_on_prompt'
+            ],
+            template=template,
+        )
+        chain = LLMChain(llm=self.llm, prompt=prompt_template)
+        return await chain.arun(block_content=block_content,
+                                file_type=file_type,
+                                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,
+        project_path: str = None,
+        print_block_doc: bool = False,
+    ) -> dict:
+        pipeline = Pipeline.get(
+            uuid=pipeline_uuid,
+            repo_path=project_path or get_repo_path(),
+        )
+        async_block_docs = []
+        for block in pipeline.blocks_by_uuid.values():
+            if block.type in NON_PIPELINE_EXECUTABLE_BLOCK_TYPES:
+                continue
+            async_block_docs.append(self.__async_generate_block_documentation(block))
+        block_docs = await asyncio.gather(*async_block_docs)
+        block_docs_content = '\n'.join(block_docs)
+        if print_block_doc:
+            print(block_docs_content)
+        prompt_template = PromptTemplate(input_variables=['block_content'],
+                                         template=PROMPT_FOR_SUMMARIZE_BLOCK_DOC)
+        chain = LLMChain(llm=self.llm, prompt=prompt_template)
+        pipeline_doc = chain.run(block_content=block_docs_content)
+        return dict(
+            block_docs=block_docs,
+            pipeline_doc=pipeline_doc,
+        )
+
+    async def async_generate_block_documentation_with_name(
+        self,
+        pipeline_uuid: str,
+        block_uuid: str,
+        project_path: str = None,
+    ) -> str:
+        pipeline = Pipeline.get(uuid=pipeline_uuid,
+                                repo_path=project_path or get_repo_path())
+        return asyncio.run(
+            self.__async_generate_block_documentation(pipeline.get_block(block_uuid)))
+
+    async def __async_generate_block_documentation(
+        self,
+        block: Block,
+    ) -> str:
+        add_on_prompt = ""
+        if block.type == BlockType.TRANSFORMER:
+            add_on_prompt = "Focus on the customized business logic in execute_transformer_action \
+                             function."
+        return await self.__async_llm_generate_documentation(
+                    block.content,
+                    BLOCK_LANGUAGE_TO_FILE_TYPE_VARIABLE[block.language],
+                    BLOCK_TYPE_TO_PURPOSE_VARIABLE.get(block.type, ""),
+                    PROMPT_FOR_BLOCK,
+                    add_on_prompt
+                )

{mage-ai-0.9.7 → mage-ai-0.9.9}/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-0.9.7 → mage-ai-0.9.9}/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-0.9.7 → mage-ai-0.9.9}/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-0.9.9/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-0.9.7 → mage-ai-0.9.9}/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-0.9.7 → mage-ai-0.9.9}/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-0.9.9/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-0.9.7 → mage-ai-0.9.9}/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-0.9.7 → mage-ai-0.9.9}/mage_ai/api/presenters/BlockTemplatePresenter.py

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

mage-ai-0.9.9/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-0.9.7 → mage-ai-0.9.9}/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-0.9.9/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-0.9.7 → mage-ai-0.9.9}/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-0.9.7 → mage-ai-0.9.9}/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,
         )