mage-ai 0.9.12__py3-none-any.whl → 0.9.14__py3-none-any.whl

mage_ai/ai/generator.py

@@ -33,7 +33,7 @@ class Generator:
         elif use_case == LLMUseCase.GENERATE_PIPELINE_WITH_DESCRIPTION:
             from mage_ai.ai.llm_pipeline_wizard import LLMPipelineWizard
 
-            return await LLMPipelineWizard().async_generate_pipeline_with_description(
+            return await LLMPipelineWizard().async_generate_pipeline_from_description(
                 request.get('pipeline_description'),
             )
         elif use_case == LLMUseCase.GENERATE_COMMENT_FOR_CODE:

mage_ai/ai/generator_cmds.py

@@ -49,7 +49,7 @@ def generate_block_with_description(block_description: str):
 
 @app.command()
 def generate_pipeline_with_description(pipeline_description: str):
-    print(asyncio.run(LLMPipelineWizard().async_generate_pipeline_with_description(
+    print(asyncio.run(LLMPipelineWizard().async_generate_pipeline_from_description(
                                             pipeline_description)))
 
 

mage_ai/ai/llm_pipeline_wizard.py

@@ -2,6 +2,7 @@ import ast
 import asyncio
 import json
 import os
+import re
 from typing import Dict
 
 import openai
@@ -53,23 +54,24 @@ PROMPT_TO_SPLIT_BLOCKS = """
 A BLOCK does one action either reading data from one data source, transforming the data from
 one format to another or exporting data into a data source.
 Based on the code description delimited by triple backticks, your task is to identify
-how many BLOCKS required and function for each BLOCK.
+how many BLOCKS required, function for each BLOCK and upstream blocks between BLOCKs.
 
 Use the following format:
-BLOCK 1: <block function>
-BLOCK 2: <block function>
-BLOCK 3: <block function>
+BLOCK 1: function: <block function>. upstream: <upstream blocks>
+BLOCK 2: function: <block function>. upstream: <upstream blocks>
+BLOCK 3: function: <block function>. upstream: <upstream blocks>
 ...
 
 Example:
 <code description>: ```
-Read data from MySQL, filter out rows with book_price > 100, and save data to BigQuery.
+Read data from MySQL and Postgres, filter out rows with book_price > 100, and save data to BigQuery.
 ```
 
 Answer:
-BLOCK 1: load data from MySQL
-BLOCK 2: filter out rows with book_price > 100
-BLOCK 3: export data to BigQuery
+BLOCK 1: function: load data from MySQL. upstream:
+BLOCK 2: function: load data from Postgres. upstream:
+BLOCK 3: function: filter out rows with book_price > 100. upstream: 1, 2
+BLOCK 4: function: export data to BigQuery. upstream: 3
 
 <code description>: ```{code_description}```"""
 PROMPT_FOR_FUNCTION_COMMENT = """
@@ -81,6 +83,7 @@ Your task is to write comments for each function inside.
 The comment should follow Google Docstring format.
 Return your response in JSON format with function name as key and the comment as value.
 """
+BLOCK_SPLIT_PATTERN = r"BLOCK\s+(\w+):\s+function:\s+(.*?)\.\s+upstream:\s*(.*?)$"
 TRANSFORMERS_FOLDER = 'transformers'
 CLASSIFICATION_FUNCTION_NAME = "classify_description"
 TEMPLATE_CLASSIFICATION_FUNCTION = [
@@ -197,7 +200,10 @@ class LLMPipelineWizard:
                                     function_args.get(DataSource.__name__))
         return block_type, block_language, pipeline_type, config
 
-    async def async_generate_block_with_description(self, block_description: str) -> dict:
+    async def async_generate_block_with_description(
+            self,
+            block_description: str,
+            upstream_blocks: [str]) -> dict:
         messages = [{"role": "user", "content": block_description}]
         response = await openai.ChatCompletion.acreate(
             model="gpt-3.5-turbo-0613",
@@ -220,6 +226,7 @@ class LLMPipelineWizard:
                     pipeline_type=pipeline_type,
                 ),
                 language=block_language,
+                upstream_blocks=upstream_blocks,
             )
         else:
             logger.error("Failed to interpret the description as a block template.")
@@ -238,11 +245,12 @@ class LLMPipelineWizard:
     async def __async_generate_blocks(self,
                                       block_dict: dict,
                                       block_id: int,
-                                      block_description: str) -> dict:
-        block = await self.async_generate_block_with_description(block_description)
+                                      block_description: str,
+                                      upstream_blocks: [str]) -> dict:
+        block = await self.async_generate_block_with_description(block_description, upstream_blocks)
         block_dict[block_id] = block
 
-    async def async_generate_pipeline_with_description(self, pipeline_description: str) -> dict:
+    async def async_generate_pipeline_from_description(self, pipeline_description: str) -> dict:
         splited_block_descriptions = await self.__async_split_description_by_blocks(
             pipeline_description)
         blocks = {}
@@ -250,11 +258,17 @@ class LLMPipelineWizard:
         for line in splited_block_descriptions.strip().split('\n'):
             if line.startswith("BLOCK") and ":" in line:
                 # Extract the block_id and block_description from the line
-                raw_block_id, block_description = line.split(":", 1)
-                block_id = raw_block_id.strip().split(" ")[-1]
-                block_description = block_description.strip()
-                block_tasks.append(
-                    self.__async_generate_blocks(blocks, block_id, block_description))
+                match = re.search(BLOCK_SPLIT_PATTERN, line)
+                if match:
+                    block_id = match.group(1)
+                    block_description = match.group(2).strip()
+                    upstream_blocks = match.group(3).split(", ")
+                    block_tasks.append(
+                        self.__async_generate_blocks(
+                            blocks,
+                            block_id,
+                            block_description,
+                            upstream_blocks))
         await asyncio.gather(*block_tasks)
         return blocks
 

mage_ai/api/policies/BlockPolicy.py

@@ -70,6 +70,7 @@ BlockPolicy.allow_write([
     'pipelines',
     'priority',
     'replicated_block',
+    'require_unique_name',
     'type',
     'upstream_blocks',
     'uuid',

mage_ai/api/policies/PipelineRunPolicy.py

@@ -47,6 +47,7 @@ PipelineRunPolicy.allow_read(PipelineRunPresenter.default_attributes + [
 PipelineRunPolicy.allow_read(PipelineRunPresenter.default_attributes + [
     'block_runs',
     'block_runs_count',
+    'completed_block_runs_count',
     'pipeline_schedule_name',
     'pipeline_schedule_token',
     'pipeline_schedule_type',

mage_ai/api/policies/PipelineSchedulePolicy.py

@@ -39,6 +39,7 @@ PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes +
 
 PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes + [
     'event_matchers',
+    'tags',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
@@ -48,6 +49,7 @@ PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes +
 PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes + [
     'event_matchers',
     'runtime_average',
+    'tags',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
@@ -58,6 +60,7 @@ PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes +
     'event_matchers',
     'last_pipeline_run_status',
     'pipeline_runs_count',
+    'tags',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
@@ -89,6 +92,7 @@ PipelineSchedulePolicy.allow_write([
     'sla',
     'start_time',
     'status',
+    'tags',
     'variables',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
@@ -102,3 +106,11 @@ PipelineSchedulePolicy.allow_query([
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
+
+PipelineSchedulePolicy.allow_query([
+    'tag[]',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/TagPolicy.py

@@ -14,8 +14,26 @@ TagPolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
+TagPolicy.allow_actions([
+    constants.CREATE,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_editor_role())
+
 TagPolicy.allow_read(TagPresenter.default_attributes + [], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.LIST,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
+
+TagPolicy.allow_read(TagPresenter.default_attributes + [], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_editor_role())
+
+TagPolicy.allow_write(TagPresenter.default_attributes + [], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_editor_role())

mage_ai/api/presenters/BlockRunPresenter.py

@@ -17,4 +17,7 @@ class BlockRunPresenter(BasePresenter):
     ]
 
     def present(self, **kwargs):
-        return self.model.to_dict()
+        if isinstance(self.model, dict):
+            return self.model
+        else:
+            return self.model.to_dict()

mage_ai/api/presenters/PipelineRunPresenter.py

@@ -27,6 +27,7 @@ class PipelineRunPresenter(BasePresenter):
             additional_attributes = [
                 'block_runs',
                 'block_runs_count',
+                'completed_block_runs_count',
                 'pipeline_schedule_name',
                 'pipeline_schedule_token',
                 'pipeline_schedule_type',
@@ -66,6 +67,7 @@ PipelineRunPresenter.register_format(
     PipelineRunPresenter.default_attributes + [
         'block_runs',
         'block_runs_count',
+        'completed_block_runs_count',
         'pipeline_schedule_name',
         'pipeline_schedule_token',
         'pipeline_schedule_type',

mage_ai/api/presenters/PipelineSchedulePresenter.py

@@ -26,17 +26,24 @@ class PipelineSchedulePresenter(BasePresenter):
         data = self.model.to_dict()
 
         if constants.LIST == display_format:
-            return self.model.to_dict(include_attributes=[
+            data = self.model.to_dict(include_attributes=[
                 'event_matchers',
                 'last_pipeline_run_status',
                 'pipeline_runs_count',
             ])
+            data['tags'] = sorted([tag.name for tag in self.get_tag_associations])
+
+            return data
         elif display_format in [constants.DETAIL, constants.UPDATE]:
-            return self.model.to_dict(include_attributes=[
+            data = self.model.to_dict(include_attributes=[
                 'event_matchers',
             ])
+            data['tags'] = sorted([tag.name for tag in self.get_tag_associations])
+
+            return data
         elif 'with_runtime_average' == display_format:
             data['runtime_average'] = self.model.runtime_average()
+            data['tags'] = sorted([tag.name for tag in self.get_tag_associations])
 
         return data
 

mage_ai/api/presenters/TagPresenter.py

@@ -3,5 +3,8 @@ from mage_ai.api.presenters.BasePresenter import BasePresenter
 
 class TagPresenter(BasePresenter):
     default_attributes = [
+        'description',
+        'id',
+        'name',
         'uuid',
     ]

mage_ai/api/resources/BackfillResource.py

@@ -1,10 +1,13 @@
 from datetime import datetime
+
+import pytz
+from sqlalchemy import desc
+
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
-from mage_ai.orchestration.backfills.service import start_backfill, cancel_backfill
+from mage_ai.orchestration.backfills.service import cancel_backfill, start_backfill
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.schedules import Backfill
 from mage_ai.shared.hash import extract, merge_dict
-from sqlalchemy import desc
 
 ALLOWED_PAYLOAD_KEYS = [
     'block_uuid',
@@ -51,7 +54,7 @@ class BackfillResource(DatabaseResource):
             if Backfill.Status.INITIAL == payload['status']:
                 pipeline_runs += start_backfill(self.model)
                 return super().update(dict(
-                    started_at=datetime.now(),
+                    started_at=datetime.now(tz=pytz.UTC),
                     status=payload['status'],
                 ))
             elif Backfill.Status.CANCELLED == payload['status']:

mage_ai/api/resources/BlockResource.py

@@ -36,6 +36,7 @@ class BlockResource(GenericResource):
         language = payload.get('language')
         name = payload.get('name')
         block_name = name or payload.get('uuid')
+        require_unique_name = payload.get('require_unique_name')
 
         payload_config = payload.get('config') or {}
 
@@ -120,6 +121,7 @@ class BlockResource(GenericResource):
                 block_name,
                 block_type,
                 get_repo_path(),
+                require_unique_name=require_unique_name,
                 **block_attributes,
             )
 

mage_ai/api/resources/BlockRunResource.py

@@ -1,7 +1,12 @@
+from sqlalchemy.orm import aliased
+
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
-from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun, PipelineSchedule
 from mage_ai.orchestration.db import safe_db_query
-from sqlalchemy.orm import aliased
+from mage_ai.orchestration.db.models.schedules import (
+    BlockRun,
+    PipelineRun,
+    PipelineSchedule,
+)
 
 
 class BlockRunResource(DatabaseResource):
@@ -24,16 +29,17 @@ class BlockRunResource(DatabaseResource):
                 pipeline_schedule_name,
             ) = tup
 
-            block_run = BlockRun()
-            block_run.block_uuid = block_uuid
-            block_run.completed_at = completed_at
-            block_run.created_at = created_at
-            block_run.id = id
-            block_run.pipeline_run_id = pipeline_run_id
-            block_run.status = status
-            block_run.updated_at = updated_at
-            block_run.pipeline_schedule_id = pipeline_schedule_id
-            block_run.pipeline_schedule_name = pipeline_schedule_name
+            block_run = dict(
+                block_uuid=block_uuid,
+                completed_at=completed_at,
+                created_at=created_at,
+                id=id,
+                pipeline_run_id=pipeline_run_id,
+                status=status,
+                updated_at=updated_at,
+                pipeline_schedule_id=pipeline_schedule_id,
+                pipeline_schedule_name=pipeline_schedule_name,
+            )
 
             block_runs.append(block_run)
 

mage_ai/api/resources/DatabaseResource.py

@@ -130,6 +130,29 @@ class DatabaseResource(BaseResource):
     def delete(self, **kwargs):
         return self.model.delete()
 
+    async def process_update(self, payload, **kwargs):
+        self.on_update_callback = None
+        self.on_update_failure_callback = None
+
+        try:
+            res = self.update(payload, **kwargs)
+            if res and inspect.isawaitable(res):
+                res = await res
+
+            db_connection.session.commit()
+
+            if self.on_update_callback:
+                self.on_update_callback(resource=res)
+
+            return res
+        except Exception as err:
+            db_connection.session.rollback()
+
+            if self.on_update_failure_callback:
+                self.on_update_failure_callback(resource=res)
+
+            raise err
+
     @safe_db_query
     def update(self, payload, **kwargs):
         for k, v in payload.items():

mage_ai/api/resources/GitBranchResource.py

@@ -265,7 +265,14 @@ class GitBranchResource(GenericResource):
         untracked_files: List[str],
         limit: int = None
     ) -> Dict:
-        arr = modified_files + staged_files + untracked_files
+        arr = []
+
+        if modified_files and isinstance(modified_files, list):
+            arr += modified_files
+        if staged_files and isinstance(staged_files, list):
+            arr += staged_files
+        if untracked_files and isinstance(untracked_files, list):
+            arr += untracked_files
 
         if limit:
             arr = arr[:limit]

mage_ai/api/resources/GlobalDataProductResource.py

@@ -35,9 +35,12 @@ class GlobalDataProductResource(GenericResource):
 
     def update(self, payload, **kwargs):
         uuid = payload.get('uuid')
-        if self.model.uuid != uuid and GlobalDataProduct.get(uuid):
+        if self.model and self.model.uuid != uuid and GlobalDataProduct.get(uuid):
             error = ApiError.RESOURCE_INVALID.copy()
             error.update(dict(message=f'A global data product with UUID {uuid} already exists.'))
             raise ApiError(error)
 
-        self.model.update(payload)
+        if self.model:
+            self.model.update(payload)
+        else:
+            self.model = self.create(payload, self.current_user, **kwargs).model

mage_ai/api/resources/PipelineResource.py

@@ -1,4 +1,5 @@
 import asyncio
+from typing import Dict, List
 
 from sqlalchemy import or_
 from sqlalchemy.orm import aliased
@@ -312,17 +313,28 @@ class PipelineResource(BaseResource):
                 schedule.update(status=status)
 
         @safe_db_query
-        def cancel_pipeline_runs(status, pipeline_uuid):
-            pipeline_runs = (
-                PipelineRun.
-                query.
-                filter(PipelineRun.pipeline_uuid == pipeline_uuid).
-                filter(PipelineRun.status.in_([
-                    PipelineRun.PipelineRunStatus.INITIAL,
-                    PipelineRun.PipelineRunStatus.RUNNING,
-                ]))
-            )
-            for pipeline_run in pipeline_runs:
+        def cancel_pipeline_runs(
+            pipeline_uuid: str = None,
+            pipeline_runs: List[Dict] = None,
+        ):
+            if pipeline_runs is not None:
+                pipeline_run_ids = [run.get('id') for run in pipeline_runs]
+                pipeline_runs_to_cancel = (
+                    PipelineRun.
+                    query.
+                    filter(PipelineRun.id.in_(pipeline_run_ids))
+                )
+            else:
+                pipeline_runs_to_cancel = (
+                    PipelineRun.
+                    query.
+                    filter(PipelineRun.pipeline_uuid == pipeline_uuid).
+                    filter(PipelineRun.status.in_([
+                        PipelineRun.PipelineRunStatus.INITIAL,
+                        PipelineRun.PipelineRunStatus.RUNNING,
+                    ]))
+                )
+            for pipeline_run in pipeline_runs_to_cancel:
                 PipelineScheduler(pipeline_run).stop()
 
         def retry_pipeline_runs(pipeline_runs):
@@ -334,15 +346,19 @@ class PipelineResource(BaseResource):
 
         def _update_callback(resource):
             if status:
+                pipeline_runs = payload.get('pipeline_runs')
                 if status in [
                     ScheduleStatus.ACTIVE.value,
                     ScheduleStatus.INACTIVE.value,
                 ]:
                     update_schedule_status(status, pipeline_uuid)
                 elif status == PipelineRun.PipelineRunStatus.CANCELLED.value:
-                    cancel_pipeline_runs(status, pipeline_uuid)
-                elif status == 'retry' and payload.get('pipeline_runs'):
-                    retry_pipeline_runs(payload.get('pipeline_runs'))
+                    if pipeline_runs is None:
+                        cancel_pipeline_runs(pipeline_uuid=pipeline_uuid)
+                    else:
+                        cancel_pipeline_runs(pipeline_runs=pipeline_runs)
+                elif status == 'retry' and pipeline_runs:
+                    retry_pipeline_runs(pipeline_runs)
 
         self.on_update_callback = _update_callback