mage-ai 0.9.46__py3-none-any.whl → 0.9.47__py3-none-any.whl

mage_ai/data_preparation/models/global_hooks/predicates.py

@@ -0,0 +1,316 @@
+import json
+from abc import abstractmethod
+from collections.abc import Iterable
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Union
+
+from mage_ai.api.resources.BaseResource import BaseResource
+from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.global_hooks.constants import (
+    INTERNAL_DEFAULT_PREDICATE_VALUE,
+    PredicateAndOrOperator,
+    PredicateObjectType,
+    PredicateOperator,
+    PredicateValueDataType,
+)
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class PredicateValueType(BaseDataClass):
+    value_data_type: PredicateValueDataType = None
+    value_type: Dict = None
+
+    def __post_init__(self):
+        self.serialize_attribute_enum('value_data_type', PredicateValueDataType)
+
+        if self.value_type and isinstance(self.value_type, dict):
+            self.value_type = self.__class__.load(**self.value_type)
+
+    def to_dict(self, **kwargs) -> Dict:
+        data = super().to_dict(**kwargs)
+
+        if self.value_type and issubclass(self.value_type.__class__, BaseDataClass):
+            data['value_type'] = self.value_type.to_dict(**kwargs)
+
+        return data
+
+
+@dataclass
+class BasePredicate(BaseDataClass):
+    and_or_operator: PredicateAndOrOperator = None
+    left_object_keys: List[str] = None
+    left_object_type: PredicateObjectType = None
+    left_value: Union[bool, Dict, float, int, List, str] = None
+    left_value_type: PredicateValueType = None
+    operator: PredicateOperator = None
+    predicates: List = None
+    right_object_keys: List[str] = None
+    right_object_type: PredicateObjectType = None
+    right_value: Union[bool, Dict, float, int, List, str] = None
+    right_value_type: PredicateValueType = None
+
+    @abstractmethod
+    def validate(
+        self,
+        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        error: Dict = None,  # code, errors, message, type
+        hook: Dict = None,
+        meta: Dict = None,
+        metadata: Dict = None,
+        payload: Dict = None,
+        query: Dict = None,
+        resource: Dict = None,
+        resource_id: Union[int, str] = None,
+        resource_parent_id: Union[int, str] = None,
+        resources: Dict = None,
+        user: Dict = None,
+    ) -> bool:
+        return False
+
+
+@dataclass
+class HookPredicate(BasePredicate):
+    predicates: List[BasePredicate] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_classes('predicates', self.__class__)
+        self.serialize_attribute_class('left_value_type', PredicateValueType)
+        self.serialize_attribute_class('right_value_type', PredicateValueType)
+        self.serialize_attribute_enum('and_or_operator', PredicateAndOrOperator)
+        self.serialize_attribute_enum('left_object_type', PredicateObjectType)
+        self.serialize_attribute_enum('operator', PredicateOperator)
+        self.serialize_attribute_enum('right_object_type', PredicateObjectType)
+
+    def validate(
+        self,
+        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        **kwargs,
+    ) -> bool:
+        if self.predicates:
+            operator_func = all
+
+            if PredicateAndOrOperator.OR == self.and_or_operator:
+                operator_func = any
+
+            return operator_func([predicate.validate(
+                operation_resource,
+                **kwargs,
+            ) for predicate in self.predicates])
+
+        return self.__validate_resource(operation_resource, **kwargs)
+
+    def __validate_resource(
+        self,
+        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        **kwargs,
+    ) -> bool:
+        if isinstance(operation_resource, Iterable) and not isinstance(operation_resource, dict):
+            return all([self.__validate(
+                operation_resource=res,
+                **kwargs,
+            ) for res in operation_resource])
+
+        return self.__validate(operation_resource=operation_resource, **kwargs)
+
+    def __validate(self, **kwargs) -> bool:
+
+        left_object = None
+        right_object = None
+
+        left_value_to_compare = INTERNAL_DEFAULT_PREDICATE_VALUE
+        right_value_to_compare = INTERNAL_DEFAULT_PREDICATE_VALUE
+
+        # Get the value from the object
+        if self.left_object_type:
+            key = self.left_object_type.value
+            if key in kwargs:
+                left_object = kwargs[key]
+
+        if self.right_object_type:
+            key = self.right_object_type.value
+            if key in kwargs:
+                right_object = kwargs[key]
+
+        # Get value from object using keys if self.right_value is not None
+        if self.left_value is not None:
+            left_value_to_compare = self.left_value
+        elif left_object:
+            if self.left_object_keys:
+                left_value_to_compare = get_value(left_object, self.left_object_keys)
+            else:
+                left_value_to_compare = left_object
+
+        # Convert left value to the specified data type
+        if self.left_value_type and INTERNAL_DEFAULT_PREDICATE_VALUE != left_value_to_compare:
+            left_value_to_compare = convert_value(left_value_to_compare, self.left_value_type)
+
+        # Get value from object using keys if self.right_value is not None
+        if self.right_value is not None:
+            right_value_to_compare = self.right_value
+        elif right_object:
+            if self.right_object_keys:
+                right_value_to_compare = get_value(right_object, self.right_object_keys)
+            else:
+                right_value_to_compare = right_object
+
+        # Convert value to the specified data type
+        if self.right_value_type and INTERNAL_DEFAULT_PREDICATE_VALUE != right_value_to_compare:
+            right_value_to_compare = convert_value(right_value_to_compare, self.right_value_type)
+
+        if INTERNAL_DEFAULT_PREDICATE_VALUE == left_value_to_compare:
+            left_value_to_compare = None
+
+        if INTERNAL_DEFAULT_PREDICATE_VALUE == right_value_to_compare:
+            right_value_to_compare = None
+
+        if self.operator:
+            return compare_using_operator(
+                self.operator,
+                left_value_to_compare,
+                right_value_to_compare,
+            )
+
+        return False
+
+
+def get_value(
+    object_arg: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline, int, str],
+    keys: List[str] = None,
+) -> Union[bool, Dict, float, int, List, str]:
+    value_temp = None
+
+    if keys:
+        if not isinstance(object_arg, Iterable) and (
+            issubclass(object_arg.__class__, BaseResource) or
+            isinstance(object_arg, Block) or
+            isinstance(object_arg, Pipeline)
+        ):
+            for idx, key in enumerate(keys):
+                if idx == 0:
+                    value_temp = getattr(object_arg, key)
+                else:
+                    value_temp = get_value(value_temp, [key])
+        elif isinstance(object_arg, Iterable) and \
+                not isinstance(object_arg, dict) and \
+                not isinstance(object_arg, str):
+
+            for idx, key in enumerate(keys):
+                if idx == 0:
+                    if isinstance(key, str) and hasattr(object_arg, key):
+                        value_temp = getattr(object_arg, key)
+                    elif isinstance(key, int) or isinstance(key, float):
+                        index_inner = int(key)
+
+                        if object_arg and index_inner < len(object_arg):
+                            value_temp = object_arg[index_inner]
+                else:
+                    value_temp = get_value(value_temp, [key])
+        elif isinstance(object_arg, dict):
+            for idx, key in enumerate(keys):
+                if idx == 0:
+                    value_temp = object_arg.get(key)
+                else:
+                    value_temp = get_value(value_temp, [key])
+
+    if callable(value_temp):
+        value_temp = value_temp()
+
+    return value_temp
+
+
+def convert_value(
+    value: Union[bool, Dict, float, int, List, str],
+    value_type: PredicateValueType,
+) -> Union[bool, Dict, float, int, List, str]:
+    value_data_type = value_type.value_data_type
+    value_type_sub = value_type.value_type
+
+    if PredicateValueDataType.BOOLEAN == value_data_type:
+        value = bool(value)
+    elif PredicateValueDataType.DICTIONARY == value_data_type:
+        if value is not None:
+            if isinstance(value, str):
+                value = json.loads(value)
+            if value_type_sub:
+                for k, v in value.items():
+                    value[k] = convert_value(v, value_type_sub)
+    elif PredicateValueDataType.FLOAT == value_data_type:
+        if value is not None:
+            value = float(value)
+    elif PredicateValueDataType.INTEGER == value_data_type:
+        if value is not None:
+            value = int(value)
+    elif PredicateValueDataType.LIST == value_data_type:
+        if value is not None:
+            if isinstance(value, str):
+                value = json.loads(value)
+            if value_type_sub:
+                value = [convert_value(v, value_type_sub) for v in value]
+    elif PredicateValueDataType.STRING == value_data_type:
+        if value is None:
+            value = ''
+        if isinstance(value, Enum):
+            value = value.value
+        value = str(value)
+
+    return value
+
+
+def compare_using_operator(
+    operator: PredicateOperator,
+    left_value: Union[bool, Dict, float, int, List, str],
+    right_value: Union[bool, Dict, float, int, List, str],
+) -> bool:
+    try:
+        if operator in [
+            PredicateOperator.GREATER_THAN,
+            PredicateOperator.GREATER_THAN_OR_EQUALS,
+            PredicateOperator.LESS_THAN,
+            PredicateOperator.LESS_THAN_OR_EQUALS,
+        ]:
+            if (
+                left_value is not None and
+                not isinstance(left_value, str) and
+                (isinstance(left_value, dict) or isinstance(left_value, Iterable))
+            ) or (
+                right_value is not None and
+                not isinstance(right_value, str) and
+                (isinstance(right_value, dict) or isinstance(right_value, Iterable))
+            ):
+                if left_value is None:
+                    left_value = 0
+                else:
+                    left_value = len(left_value)
+
+                if right_value is None:
+                    right_value = 0
+                else:
+                    right_value = len(right_value)
+
+        if PredicateOperator.EQUALS == operator:
+            return left_value == right_value
+        if PredicateOperator.GREATER_THAN == operator:
+            return left_value > right_value
+        if PredicateOperator.GREATER_THAN_OR_EQUALS == operator:
+            return left_value >= right_value
+        if PredicateOperator.INCLUDES == operator:
+            return left_value in right_value
+        if PredicateOperator.LESS_THAN == operator:
+            return left_value < right_value
+        if PredicateOperator.LESS_THAN_OR_EQUALS == operator:
+            return left_value <= right_value
+        if PredicateOperator.NOT_EQUALS == operator:
+            return left_value != right_value
+        if PredicateOperator.NOT_INCLUDES == operator:
+            return left_value not in right_value
+        if PredicateOperator.NOT_PRESENT == operator:
+            return False if left_value else True
+        if PredicateOperator.PRESENT == operator:
+            return True if left_value else False
+    except TypeError as err:
+        print(f'[ERROR] Predicate error comparing {left_value} {operator} {right_value}: {err}.')
+
+    return False

mage_ai/data_preparation/models/pipeline.py

@@ -643,6 +643,7 @@ class Pipeline:
             executor_type=self.executor_type,
             name=self.name,
             notification_config=self.notification_config,
+            remote_variables_dir=self.remote_variables_dir,
             retry_config=self.retry_config,
             run_pipeline_in_one_process=self.run_pipeline_in_one_process,
             settings=self.settings.to_dict() if self.settings else self.settings,
@@ -650,6 +651,7 @@ class Pipeline:
             type=self.type.value if type(self.type) is not str else self.type,
             updated_at=updated_at,
             uuid=self.uuid,
+            variables_dir=self.variables_dir,
         )
 
         if self.variables is not None:

mage_ai/data_preparation/models/pipelines/integration_pipeline.py

@@ -272,7 +272,6 @@ class IntegrationPipeline(Pipeline):
                             sample_data_json = data.get('sample_data')
                             sample_data = pd.DataFrame.from_dict(json.loads(sample_data_json))
                             stream_id = data.get('stream_id')
-
                             variables = {
                                 f'output_sample_data_{stream_id}': sample_data,
                             }

mage_ai/data_preparation/preferences.py

@@ -20,6 +20,7 @@ GIT_SYNC_ON_PIPELINE_RUN_VAR = 'GIT_SYNC_ON_PIPELINE_RUN'
 GIT_SYNC_ON_START_VAR = 'GIT_SYNC_ON_START'
 GIT_SYNC_ON_EXECUTOR_START_VAR = 'GIT_SYNC_ON_EXECUTOR_START'
 GIT_SYNC_SUBMODULES = 'GIT_SYNC_SUBMODULES'
+GIT_ENABLE_GIT_INTEGRATION_VAR = 'GIT_ENABLE_GIT_INTEGRATION'
 
 
 class Preferences:
@@ -30,8 +31,7 @@ class Preferences:
         user: User = None,
     ):
         self.repo_path = repo_path or get_repo_path()
-        self.preferences_file_path = \
-            os.path.join(self.repo_path, PREFERENCES_FILE)
+        self.preferences_file_path = os.path.join(self.repo_path, PREFERENCES_FILE)
         self.user = user
         project_preferences = dict()
         try:
@@ -47,33 +47,43 @@ class Preferences:
             pass
 
         # Git settings
-        env_sync_config = dict()
         if os.getenv(GIT_REPO_LINK_VAR):
-            env_sync_config = dict(
+            self.sync_config = dict(
                 remote_repo_link=os.getenv(GIT_REPO_LINK_VAR),
                 repo_path=os.getenv(GIT_REPO_PATH_VAR, os.getcwd()),
                 auth_type=os.getenv(GIT_AUTH_TYPE_VAR),
                 username=os.getenv(GIT_USERNAME_VAR),
                 email=os.getenv(GIT_EMAIL_VAR),
                 branch=os.getenv(GIT_BRANCH_VAR),
-                sync_on_pipeline_run=bool(int(os.getenv(GIT_SYNC_ON_PIPELINE_RUN_VAR) or 0)),
+                sync_on_pipeline_run=bool(
+                    int(os.getenv(GIT_SYNC_ON_PIPELINE_RUN_VAR) or 0)
+                ),
                 sync_on_start=bool(int(os.getenv(GIT_SYNC_ON_START_VAR) or 0)),
-                sync_on_executor_start=bool(int(os.getenv(GIT_SYNC_ON_EXECUTOR_START_VAR) or 0)),
+                sync_on_executor_start=bool(
+                    int(os.getenv(GIT_SYNC_ON_EXECUTOR_START_VAR) or 0)
+                ),
                 sync_submodules=bool(int(os.getenv(GIT_SYNC_SUBMODULES) or 0)),
+                enable_git_integration=bool(
+                    int(os.getenv(GIT_ENABLE_GIT_INTEGRATION_VAR) or 0)
+                ),
             )
-        project_sync_config = project_preferences.get('sync_config', dict())
-        if user:
-            user_git_settings = user.git_settings or {}
-            project_sync_config = merge_dict(project_sync_config, user_git_settings)
+        else:
+            project_sync_config = project_preferences.get('sync_config', dict())
+            if user:
+                user_git_settings = user.git_settings or {}
+                project_sync_config = merge_dict(project_sync_config, user_git_settings)
 
-        self.sync_config = merge_dict(env_sync_config, project_sync_config)
+            self.sync_config = project_sync_config
 
     def is_git_integration_enabled(self) -> bool:
-        return self.sync_config.get('remote_repo_link') and \
-            self.sync_config.get('repo_path') and \
-            self.sync_config.get('username') and \
-            self.sync_config.get('email') and \
-            self.sync_config.get('branch') is None
+        return (
+            self.sync_config.get('remote_repo_link')
+            and self.sync_config.get('repo_path')
+            and self.sync_config.get('username')
+            and self.sync_config.get('email')
+            and self.sync_config.get('branch') is None
+            or self.sync_config.get('enable_git_integration')
+        )
 
     def update_preferences(self, updates: Dict):
         preferences = self.to_dict()
@@ -90,8 +100,9 @@ class Preferences:
 def get_preferences(repo_path=None, user: User = None) -> Preferences:
     default_preferences = Preferences(repo_path=repo_path)
     if user:
-        if user.preferences is None \
-                and os.path.exists(default_preferences.preferences_file_path):
+        if user.preferences is None and os.path.exists(
+            default_preferences.preferences_file_path
+        ):
             return default_preferences
         else:
             return Preferences(user=user)

mage_ai/data_preparation/repo_manager.py

@@ -255,16 +255,26 @@ def get_repo_config(repo_path=None) -> RepoConfig:
 
 
 def get_project_type(repo_path=None) -> ProjectType:
+    from mage_ai.settings.repo import MAGE_PROJECT_TYPE_ENV_VAR
     try:
-        return get_repo_config(repo_path=repo_path).project_type
+        project_type_from_env = os.getenv(MAGE_PROJECT_TYPE_ENV_VAR)
+        if project_type_from_env:
+            return ProjectType(project_type_from_env)
+        else:
+            return get_repo_config(repo_path=repo_path).project_type
     except Exception:
         # default to standalone project type
         return ProjectType.STANDALONE
 
 
 def get_cluster_type(repo_path=None) -> Optional[ClusterType]:
+    from mage_ai.settings.repo import MAGE_CLUSTER_TYPE_ENV_VAR
     try:
-        return get_repo_config(repo_path=repo_path).cluster_type
+        cluster_type_from_env = os.getenv(MAGE_CLUSTER_TYPE_ENV_VAR)
+        if cluster_type_from_env:
+            return ClusterType(cluster_type_from_env)
+        else:
+            return get_repo_config(repo_path=repo_path).cluster_type
     except Exception:
         # default to None
         return None

mage_ai/data_preparation/sync/__init__.py

@@ -32,6 +32,8 @@ class GitConfig(BaseConfig):
     ssh_private_key_secret_name: str = GIT_SSH_PRIVATE_KEY_SECRET_NAME
     ssh_public_key_secret_name: str = GIT_SSH_PUBLIC_KEY_SECRET_NAME
     access_token_secret_name: str = GIT_ACCESS_TOKEN_SECRET_NAME
+    # Force Mage to show git integration panel in the UI
+    enable_git_integration: bool = False
     # This is not necessary anymore, but leaving it for backwards compatibility
     type: str = 'git'
 

mage_ai/data_preparation/templates/constants.py

@@ -192,6 +192,13 @@ TEMPLATES_ONLY_FOR_V2 = [
         path='data_loaders/snowflake.py',
     ),
     #   Databases
+    dict(
+        block_type=BlockType.DATA_LOADER,
+        groups=[GROUP_DATABASES],
+        language=BlockLanguage.PYTHON,
+        name='Chroma',
+        path='data_loaders/chroma.py',
+    ),
     dict(
         block_type=BlockType.DATA_LOADER,
         groups=[GROUP_DATABASES],
@@ -552,6 +559,13 @@ TEMPLATES_ONLY_FOR_V2 = [
         path='data_exporters/snowflake.py',
     ),
     #   Databases
+    dict(
+        block_type=BlockType.DATA_EXPORTER,
+        groups=[GROUP_DATABASES],
+        language=BlockLanguage.PYTHON,
+        name='Chroma',
+        path='data_exporters/chroma.py',
+    ),
     dict(
         block_type=BlockType.DATA_EXPORTER,
         groups=[GROUP_DATABASES],

mage_ai/data_preparation/templates/data_exporters/chroma.py

@@ -0,0 +1,24 @@
+from os import path
+
+from pandas import DataFrame
+
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.chroma import Chroma
+
+if 'data_exporter' not in globals():
+    from mage_ai.data_preparation.decorators import data_exporter
+
+
+@data_exporter
+def export_data_to_chroma(df: DataFrame, **kwargs) -> None:
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+    document_column = 'documents'
+    new_collection = 'new_colletion'
+
+    Chroma.with_config(ConfigFileLoader(config_path, config_profile)).export(
+        df,
+        collection=new_collection,
+        document_column=document_column
+    )

mage_ai/data_preparation/templates/data_exporters/streaming/rabbitmq.yaml

@@ -0,0 +1,7 @@
+connector_type: 'rabbitmq'
+connection_host: 'localhost'
+connection_port: 5672
+queue_name: 'default'
+username: 'guest'
+password: 'guest'
+amqp_url_virtual_host: '%2f'

mage_ai/data_preparation/templates/data_loaders/chroma.py

@@ -0,0 +1,27 @@
+{% extends "data_loaders/default.jinja" %}
+{% block imports %}
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.chroma import Chroma
+from os import path
+{{ super() -}}
+{% endblock %}
+
+
+{% block content %}
+@data_loader
+def load_data_from_chroma(*args, **kwargs):
+    """
+    Template for loading data from Chroma.
+    """
+    query_embeddings = 'Query embeddings'
+    query_texts = 'Query texts'
+    n_results = 1
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+
+    return Chroma.with_config(ConfigFileLoader(config_path, config_profile)).load(
+        n_results=n_results,
+        query_embeddings=query_embeddings,
+        query_texts=query_texts)
+{% endblock %}

mage_ai/data_preparation/templates/repo/io_config.yaml

@@ -12,6 +12,9 @@ default:
   AZURE_CLIENT_SECRET: "{{ env_var('AZURE_CLIENT_SECRET') }}"
   AZURE_STORAGE_ACCOUNT_NAME: "{{ env_var('AZURE_STORAGE_ACCOUNT_NAME') }}"
   AZURE_TENANT_ID: "{{ env_var('AZURE_TENANT_ID') }}"
+  # Chroma
+  CHROMA_COLLECTION: collection_name
+  CHROMA_PATH: path
   # Clickhouse
   CLICKHOUSE_DATABASE: default
   CLICKHOUSE_HOST: host.docker.internal

mage_ai/io/base.py

@@ -16,6 +16,7 @@ QUERY_ROW_LIMIT = 10_000_000
 class DataSource(str, Enum):
     API = 'api'
     BIGQUERY = 'bigquery'
+    CHROMA = 'chroma'
     CLICKHOUSE = 'clickhouse'
     DRUID = 'druid'
     DUCKDB = 'duckdb'