Integration with Upgini

.gitignore

@@ -31,6 +31,8 @@ cdk.out/
 .langchain.db
 output_*.csv
 runs/
+*.index
+*.metadata
 
 # MCP specific
 .flask.pid

README.md

@@ -76,6 +76,18 @@ AutoGluon Assistant provides multiple interfaces:
 
 ![Demo](https://github.com/autogluon/autogluon-assistant/blob/main/docs/assets/mcp_demo.gif)
 
+### Integration with Upgini
+
+To enable Upgini integration for enriching your dataset with additional relevant features and selecting the most relevant features from the input dataset, set the `UPGINI_API_KEY` environment variable before invoking MLZero. You can obtain your API key in your [Upgini profile](https://profile.upgini.com).
+
+```bash
+export UPGINI_API_KEY="<your-upgini-api-key>"
+# then run MLZero as usual
+mlzero -i <input_data_folder>
+```
+
+
+
 ## Citation
 If you use Autogluon Assistant (MLZero) in your research, please cite our paper:
 

src/autogluon/assistant/cli/app.py

@@ -95,7 +95,7 @@ def main(
     # 3) Invoke the core run_agent function
     # Override config path if provider is specified and config path is default
     provider_config_path = config_path
-    if llm_provider in ["bedrock", "openai", "anthropic", "sagemaker"] and config_path == DEFAULT_CONFIG_PATH:
+    if llm_provider in ["bedrock", "openai", "anthropic", "sagemaker", "azure"] and config_path == DEFAULT_CONFIG_PATH:
         provider_config_path = Path(DEFAULT_CONFIG_PATH).parent / f"{llm_provider}.yaml"
         if not provider_config_path.exists():
             provider_config_path = DEFAULT_CONFIG_PATH

src/autogluon/assistant/configs/azure.yaml

@@ -0,0 +1,83 @@
+# SageMaker Configuration
+
+per_execution_timeout: 86400
+
+# Data Perception
+max_file_group_size_to_show: 5
+num_example_files_to_show: 1
+
+max_chars_per_file: 768
+num_tutorial_retrievals: 30
+max_num_tutorials: 5
+max_user_input_length: 2048
+max_error_message_length: 2048
+max_tutorial_length: 32768
+create_venv: false
+condense_tutorials: True
+use_tutorial_summary: True
+continuous_improvement: False
+optimize_system_resources: False
+cleanup_unused_env: True
+enable_meta_prompting: False
+
+llm: &default_llm
+  provider: azure
+  model: o3-mini
+  max_tokens: 65535
+  proxy_url: null
+  top_p: 0.9
+  verbose: True
+  multi_turn: False
+  template: null
+  add_coding_format_instruction: false
+  apply_meta_prompting: False
+
+# Ensure all agent types inherit the SageMaker LLM config
+python_coder:
+  <<: *default_llm  # Merge llm_config
+  multi_turn: True
+  apply_meta_prompting: True
+
+bash_coder:
+  <<: *default_llm  # Merge llm_config
+  multi_turn: True
+
+executer:
+  <<: *default_llm  # Merge llm_config
+  max_stdout_length: 8192
+  max_stderr_length: 2048
+
+meta_prompting:
+  <<: *default_llm  # Merge llm_config
+  multi_turn: False
+
+reader:
+  <<: *default_llm  # Merge llm_config
+  details: False
+
+error_analyzer:
+  <<: *default_llm  # Merge llm_config
+
+retriever:
+  <<: *default_llm  # Merge llm_config
+
+reranker:
+  <<: *default_llm  # Merge llm_config
+  temperature: 0.
+  top_p: 1.
+
+description_file_retriever:
+  <<: *default_llm  # Merge llm_config
+  temperature: 0.
+  top_p: 1.
+
+task_descriptor:
+  <<: *default_llm  # Merge llm_config
+  max_description_files_length_to_show: 1024
+  max_description_files_length_for_summarization: 16384
+  apply_meta_prompting: True
+
+tool_selector:
+  <<: *default_llm  # Merge llm_config
+  temperature: 0.
+  top_p: 1.

src/autogluon/assistant/llm/azure_openai_chat.py

@@ -8,6 +8,7 @@
 from .base_chat import BaseAssistantChat
 
 logger = logging.getLogger(__name__)
+NO_TEMPERATURE_MODELS = {"o1", "o1-mini", "o3", "o3-mini"}
 
 
 class AssistantAzureChatOpenAI(AzureChatOpenAI, BaseAssistantChat):
@@ -53,7 +54,7 @@ def create_azure_openai_chat(config, session_name: str) -> AssistantAzureChatOpe
         "max_tokens": config.max_tokens,
     }
 
-    if hasattr(config, "temperature"):
+    if hasattr(config, "temperature") and model not in NO_TEMPERATURE_MODELS:
         kwargs["temperature"] = config.temperature
 
     if hasattr(config, "verbose"):

src/autogluon/assistant/prompts/bash_coder_prompt.py

@@ -97,12 +97,13 @@ def get_env_prompt(self):
         common_env_file = self.manager.common_env_file
         selected_tool_env_file = self.manager.selected_tool_env_file
 
+        install_cmd = f"uv pip install --system -r {selected_tool_env_file} -r {common_env_file}"
         env_prompt = f"""
 Create and configure a conda environment in "{ENV_FOLDER_NAME}" folder under {iteration_folder}:
  - Python version: 3.11
- - Activate the environment
+ - Activate the environment (conda activate "{ENV_FOLDER_NAME}")
  - pip install uv
- - Install required packages from {common_env_file} and {selected_tool_env_file} using uv pip install -r {selected_tool_env_file} -r {common_env_file}"""
+ - Install required packages from {common_env_file} and {selected_tool_env_file} using {install_cmd}"""
 
         if not create_venv:
             env_prompt += f"\n - Do not install or update any package unless there is an error due to the missing package.\n - Do NOT upgrade {selected_tool} which is already installed."

src/autogluon/assistant/prompts/python_coder_prompt.py

@@ -34,25 +34,65 @@ def default_template(self) -> str:
 ONLY save files to the working directory: {per_iteration_output_folder}.
 
 1. Data preprocessing:
-   - Remove training data samples without valid labels (drop NA values from training dataset ONLY, NOT from test dataset) unless explicitly instructed otherwise.
-   - Remove the unneccesary index column (if applicable)
-
-2. Model training:
+   - Remove training data samples without valid labels (drop NA values from the training dataset ONLY, NOT from test dataset) unless explicitly instructed otherwise.
+   - Remove any unneccesary index column (if present).
+   - If the environment variable MLZERO_SAMPLE_SIZE is set and the training dataset contains more than MLZERO_SAMPLE_SIZE rows, draw a random sample:
+        ```python
+        import os
+
+        if "MLZERO_SAMPLE_SIZE" in os.environ:
+            sample_size = int(os.getenv("MLZERO_SAMPLE_SIZE"))
+            if len(train_data) > sample_size:
+                train_data = train_data.sample(n=sample_size, random_state=42)
+            else:
+                print(f"Training dataset has {len(train_data)} rows, which is less than MLZERO_SAMPLE_SIZE {sample_size}. Skipping sampling.")
+        ```
+
+2. Data enrichment:
+    - Perform this step only if UPGINI_API_KEY is set in the environment and the data is tabular. Otherwise, skip this step.
+    - Use Upgini to enrich the data and select the most relevant features from internal and external sources.
+    - Example of usage:
+        ```python
+        from upgini import FeaturesEnricher
+        enricher = FeaturesEnricher()
+        train_data_features = enricher.fit_transform(train_data.drop(columns=label), train_data[label], eval_set=(test_data.drop(columns=label, errors='ignore'), ), calculate_metrics=False)
+        train_data = pd.concat([train_data_features, train_data[[label]]], axis=1)
+        test_data_for_prediction = enricher.transform(test_data.drop(columns=label), errors='ignore')
+        # If the dataset includes a stable identifier column (e.g., id, row_id, key, uid, or any column with 100% unique values), after enrichment, include it in the test data for prediction
+        test_data_for_prediction = pd.concat([test_data[[id_column_name]], test_data_for_prediction], axis=1)
+        ```
+    - If there are multiple target columns, enrich the data separately for each target:
+        ```python
+        train_data_features = dict()
+        test_data_for_prediction = dict()
+        for label in labels:
+            train_data_features[label] = enricher.fit_transform(train_data.drop(columns=labels), train_data[label], eval_set=(test_data.drop(columns=labels, errors='ignore'), ), calculate_metrics=False)
+            test_data_features = enricher.transform(test_data.drop(columns=labels, errors='ignore'), errors='ignore')
+            test_data_for_prediction[label] = pd.concat([test_data[[id_column_name]], test_data_features], axis=1)
+
+        predictions = dict()
+        for label in labels:
+            predictor = ...
+            predictor.fit(train_data=train_data_features[label], ...)
+            predictions[label] = predictor.predict(test_data_for_prediction[label])
+        ```
+
+3. Model training:
    - Use {selected_tool} with appropriate parameters for the task
    - If a model is trained, save it in a folder with random timestamp within {per_iteration_output_folder}
 
-3. Prediction:
+4. Prediction:
    - Make predictions on the test data. Always preserve and use the ORIGINAL INDICES from the test data to maintain exact row correspondence - DO NOT generate new indices or rely on assumed ordering.
    - Save the predicted results to {per_iteration_output_folder}, result file name should be "results", the format and extension should be same as the test data file
    - Output column names must exactly match those in the training or sample submission files without adding "predicted_" prefixes or creating any new columns.
    - At the end, implement validation checks that assert the prediction file maintains exact test data indices, verify correct column names match requirements, and confirm proper output format.
 
-4. Documentation:
+5. Documentation:
    - Add a brief docstring at the beginning of the script explaining its purpose
    - Include additional installation steps with comments at the beginning of the script
    - Include comments explaining any complex operations or design decisions
 
-5. Others:
+6. Others:
    - To avoid DDP errors, wrap the code in: if __name__ == "__main__":
    - Ensure errors are propagated up and not silently caught - do not use try/except blocks unless you explicitly re-raise the exception.
 

src/autogluon/assistant/tools_registry/_common/requirements.txt

@@ -12,3 +12,4 @@ python-dateutil
 datasets>=3.6
 six
 pyarrow
+upgini