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

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

@@ -57,6 +57,7 @@ class RepoConfig:
         self.queue_config = None
         self.help_improve_mage = None
         self.openai_api_key = None
+        self._pipelines = None
         self.retry_config = None
         self.ldap_config = None
         self.s3_bucket = None
@@ -117,6 +118,7 @@ class RepoConfig:
             self.project_uuid = repo_config.get('project_uuid')
             self.help_improve_mage = repo_config.get('help_improve_mage')
             self.openai_api_key = repo_config.get('openai_api_key')
+            self.pipelines = repo_config.get('pipelines')
             self.retry_config = repo_config.get('retry_config')
 
             self.ldap_config = repo_config.get('ldap_config')
@@ -149,6 +151,22 @@ class RepoConfig:
         else:
             return get_metadata_path()
 
+    @property
+    def pipelines(self):
+        if isinstance(self._pipelines, dict):
+            self.pipelines = self._pipelines
+
+        return self._pipelines
+
+    @pipelines.setter
+    def pipelines(self, pipelines: Dict = None) -> None:
+        from mage_ai.data_preparation.models.project.models import ProjectPipelines
+
+        if isinstance(pipelines, dict):
+            self._pipelines = ProjectPipelines.load(**(pipelines or {}))
+        else:
+            self._pipelines = pipelines
+
     def to_dict(self, remote: bool = False) -> Dict:
         return dict(
             ai_config=self.ai_config,
@@ -160,6 +178,7 @@ class RepoConfig:
             help_improve_mage=self.help_improve_mage,
             notification_config=self.notification_config,
             openai_api_key=self.openai_api_key,
+            pipelines=self.pipelines.to_dict() if self.pipelines else self.pipelines,
             project_type=self.project_type,
             project_uuid=self.project_uuid,
             queue_config=self.queue_config,
@@ -179,7 +198,10 @@ class RepoConfig:
 
         for key, value in kwargs.items():
             data[key] = value
-            if hasattr(self, key):
+
+            if 'pipelines' == key:
+                self.pipelines = value
+            elif hasattr(self, key):
                 setattr(self, key, value)
 
         with open(self.metadata_path, 'w') as f:
@@ -233,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/data_loaders/streaming/activemq.yaml

@@ -0,0 +1,6 @@
+connector_type: 'activemq'
+connection_host: 'localhost'
+connection_port: 61613
+queue_name: 'default'
+username: 'admin'
+password: 'admin'

mage_ai/data_preparation/templates/data_loaders/streaming/nats.yaml

@@ -0,0 +1,20 @@
+connector_type: nats
+server_url: "nats://localhost:4222"
+subject: subject_name
+stream_name: stream_name
+
+# Optional: Batch size and timeout for reading messages
+batch_size: 100
+timeout: 5
+
+# Optional: TLS configuration
+use_tls: false
+# cafile: "path/to/ca.pem"
+# certfile: "path/to/client-cert.pem"
+# keyfile: "path/to/client-key.pem"
+
+# Optional: NKEY seed string
+# nkeys_seed_str: "SUAPAEYJQOQHJ4"
+
+# Optional: Consumer name for pull-based subscription
+consumer_name: "my_consumer"

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'

mage_ai/io/chroma.py

@@ -0,0 +1,153 @@
+import logging
+from typing import Dict
+
+import chromadb
+from pandas import DataFrame
+
+from mage_ai.io.base import BaseIO
+from mage_ai.io.config import BaseConfigLoader, ConfigKey
+
+logger = logging.getLogger(__name__)
+
+
+class Chroma(BaseIO):
+    def __init__(
+            self,
+            collection: str,
+            path: str = None,
+            verbose: bool = True,
+            **kwargs,) -> None:
+        """
+        Initializes settings to connect to chroma db.
+        """
+        super().__init__(verbose=verbose)
+        self.collection = collection
+        if path is None:
+            self.client = chromadb.Client()
+        else:
+            self.client = chromadb.PersistentClient(path=path)
+
+    @classmethod
+    def with_config(cls, config: BaseConfigLoader) -> 'Chroma':
+        return cls(
+            collection=config[ConfigKey.CHROMA_COLLECTION],
+            path=config[ConfigKey.CHROMA_PATH],
+        )
+
+    def __convert_to_multiple_rows(self, data: Dict, column_name: str, item_length: int):
+        """
+        Convert the data from chroma query into multiple rows.
+
+        Args:
+            data (Dict): Data returned from chroma query.
+            Example format:
+            {
+                'ids': [['id1', '0', 'id2']],
+                'distances': [[0.6036068181428754, 0.6036068181428754, 1.267588382105355]],
+                'metadatas': [[
+                                {'chapter': '3', 'verse': '16'},
+                                None,
+                                {'chapter': '3', 'verse': '5'}
+                              ]],
+                'embeddings': None,
+                'documents': [['abc', 'abc', 'def']],
+                'uris': None,
+                'data': None
+            }
+            column_name (str): Name of the field in data to convert.
+            item_length (int): Number of expected items in the list.
+
+        Returns:
+            DataFrame: Data frame object loaded from the chroma query fuction.
+            Chroma query function requries dictionary contains number of results matched.
+        """
+        print(f"Testing data: {data}")
+        converted_array = []
+        if data[column_name] is None:
+            for _ in range(item_length):
+                converted_array.append('')
+        else:
+            for list_items in data[column_name]:
+                for item in list_items:
+                    if item is None:
+                        converted_array.append('')
+                    else:
+                        converted_array.append(item)
+        return converted_array
+
+    def load(
+        self,
+        n_results: int,
+        collection: str = None,
+        query_embeddings: str = None,
+        query_texts: str = None,
+        **kwargs,
+    ) -> DataFrame:
+        """
+        Loads the data from Chroma with embeddings or query_texts.
+
+        Args:
+            n_results (int): Number of results to return.
+            query_embeddings (str): Embeddings to query.
+            query_texts (str): Texts to query.
+
+        Returns:
+            DataFrame: Data frame object loaded from the chroma query fuction.
+            Chroma query function requries dictionary contains number of results matched.
+        """
+        collection_client = self.client.get_or_create_collection(
+                name=self.collection if collection is None else collection)
+        # Chroma supports query by embeddings or query texts
+        if query_embeddings is not None:
+            data = collection_client.query(
+                query_embeddings=query_embeddings,
+                n_results=n_results
+            )
+        elif query_texts is not None:
+            data = collection_client.query(
+                query_texts=query_texts,
+                n_results=n_results
+            )
+        else:
+            raise ValueError("query_embeddings or query_texts must be specified")
+        # Flatten the query results into multiple rows
+        flatten_data = {}
+        item_length = len(data['ids'][0])
+        flatten_data['ids'] = self.__convert_to_multiple_rows(
+            data, 'ids', item_length)
+        flatten_data['distances'] = self.__convert_to_multiple_rows(
+            data, 'distances', item_length)
+        flatten_data['metadatas'] = self.__convert_to_multiple_rows(
+            data, 'metadatas', item_length)
+        flatten_data['embeddings'] = self.__convert_to_multiple_rows(
+            data, 'embeddings', item_length)
+        flatten_data['documents'] = self.__convert_to_multiple_rows(
+            data, 'documents', item_length)
+        flatten_data['data'] = self.__convert_to_multiple_rows(
+            data, 'data', item_length)
+        return DataFrame.from_dict(flatten_data)
+
+    def export(
+        self,
+        df: DataFrame,
+        document_column: str,
+        collection: str = None,
+        **kwargs,
+    ) -> None:
+        """
+        Exports the input dataframe to the chroma collection.
+        The column required is document_column.
+        It contains the contextual prompt for query and search.
+
+        Args:
+            df (DataFrame): Data frame to export.
+            document_column (str): name of the document.
+        """
+        docs = df[f'{document_column}'].tolist()
+        ids = [str(x) for x in df.index.tolist()]
+        collection_client = self.client.get_or_create_collection(
+                name=self.collection if collection is None else collection)
+        return collection_client.add(
+            documents=docs,
+            ids=ids,
+        )

mage_ai/io/config.py

@@ -27,6 +27,9 @@ class ConfigKey(str, Enum):
     AZURE_STORAGE_ACCOUNT_NAME = 'AZURE_STORAGE_ACCOUNT_NAME'
     AZURE_TENANT_ID = 'AZURE_TENANT_ID'
 
+    CHROMA_COLLECTION = 'CHROMA_COLLECTION'
+    CHROMA_PATH = 'CHROMA_PATH'
+
     CLICKHOUSE_DATABASE = 'CLICKHOUSE_DATABASE'
     CLICKHOUSE_HOST = 'CLICKHOUSE_HOST'
     CLICKHOUSE_INTERFACE = 'CLICKHOUSE_INTERFACE'
@@ -323,6 +326,7 @@ class VerboseConfigKey(str, Enum):
 
     AWS = 'AWS'
     BIGQUERY = 'BigQuery'
+    CHROMA = 'Chroma'
     CLICKHOUSE = 'ClickHouse'
     DRUID = 'Druid'
     DUCKDB = 'Duck DB'
@@ -370,6 +374,10 @@ class ConfigFileLoader(BaseConfigLoader):
             VerboseConfigKey.REDSHIFT,
             'user',
         ),
+        ConfigKey.CHROMA_COLLECTION: (
+            VerboseConfigKey.CHROMA, 'collection'),
+        ConfigKey.CHROMA_PATH: (
+            VerboseConfigKey.CHROMA, 'path'),
         ConfigKey.CLICKHOUSE_DATABASE: (
             VerboseConfigKey.CLICKHOUSE, 'database'),
         ConfigKey.CLICKHOUSE_HOST: (