Switch to Use Node Based Manager

maab/agents/mlzero_default/mlzero_default.sh

@@ -57,9 +57,10 @@ fi
 mlzero \
     -i "$TRAINING_PATH" \
     -o "$OUTPUT_DIR" \
-    -n 8 \
+    -n 3 \
     -v 1 \
-    --initial-instruction "complete the task in 30 minutes"
+    --continuous_improvement \
+    --initial-instruction "complete the task in 10 minutes"
 
 # Check if the process was successful
 if [ $? -ne 0 ]; then

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
   "torchaudio",
   "torchvision",
   "pandas>=2.2",
+  "reportlab>=4.4.3",
   "streamlit>=1.37",
   "streamlit-aggrid>=1.0.2",
   "streamlit-extras>=0.4",

src/autogluon/assistant/agents/tool_selector_agent.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Tuple
+from typing import List, Union
 
 from ..prompts import ToolSelectorPrompt
 from .base_agent import BaseAgent
@@ -10,14 +10,15 @@
 
 class ToolSelectorAgent(BaseAgent):
     """
-    Select the most appropriate tool based on data description and task requirements.
+    Select and rank the most appropriate tools based on data description and task requirements.
 
     Agent Input:
     - data_prompt: Text string containing data prompt
     - description: Description of the task/data from previous analysis
 
     Agent Output:
-    - str: Selected tool name
+    - List[str]: Prioritized list of tool names
+    - str: Selected tool name (for backward compatibility)
     """
 
     def __init__(self, config, manager, llm_config, prompt_template):
@@ -39,8 +40,8 @@ def __init__(self, config, manager, llm_config, prompt_template):
                 multi_turn=self.tool_selector_llm_config.multi_turn,
             )
 
-    def __call__(self) -> Tuple[str, str]:
-        self.manager.log_agent_start("ToolSelectorAgent: choosing the most appropriate ML library for the task.")
+    def __call__(self) -> Union[List[str], str]:
+        self.manager.log_agent_start("ToolSelectorAgent: choosing and ranking ML libraries for the task.")
 
         # Build prompt for tool selection
         prompt = self.tool_selector_prompt.build()
@@ -54,8 +55,9 @@ def __call__(self) -> Tuple[str, str]:
 
         response = self.tool_selector_llm.assistant_chat(prompt)
 
-        selected_tool = self.tool_selector_prompt.parse(response)
+        tools = self.tool_selector_prompt.parse(response)
 
-        self.manager.log_agent_end(f"ToolSelectorAgent: selected {selected_tool}.")
+        tools_str = ", ".join(tools)
+        self.manager.log_agent_end(f"ToolSelectorAgent: selected tools in priority order: {tools_str}")
 
-        return selected_tool
+        return tools

src/autogluon/assistant/coding_agent.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import time
 import uuid
 from datetime import datetime
 from pathlib import Path
@@ -17,7 +18,7 @@ def run_agent(
     input_data_folder,
     output_folder=None,
     config_path=None,
-    max_iterations=5,
+    max_iterations=10,  # Default higher for MCTS search
     continuous_improvement=None,
     enable_meta_prompting=None,
     enable_per_iteration_instruction=False,
@@ -26,6 +27,24 @@ def run_agent(
     manager=None,
     verbosity=1,
 ):
+    """
+    Run the AutoGluon Assistant with MCTS-based search strategy.
+
+    Args:
+        input_data_folder: Path to input data directory
+        output_folder: Path to output directory
+        config_path: Path to configuration file
+        max_iterations: Maximum number of iterations
+        continuous_improvement: Whether to continue after finding a valid solution
+        enable_meta_prompting: Whether to enable meta-prompting
+        enable_per_iteration_instruction: Whether to ask for user input at each iteration
+        initial_user_input: Initial user instruction
+        extract_archives_to: Path to extract archives to
+        verbosity: Verbosity level
+
+    Returns:
+        None
+    """
     # Get the directory of the current file
     current_file_dir = Path(__file__).parent
 
@@ -36,7 +55,7 @@ def run_agent(
         # Generate a random UUID4
         random_uuid = uuid.uuid4()
         # Create the folder name using the pattern
-        folder_name = f"mlzero-{current_datetime}-{random_uuid}"
+        folder_name = f"mlzero-mcts-{current_datetime}-{random_uuid}"
 
         # Create the full path for the new folder
         output_folder = os.path.join(working_dir, folder_name)
@@ -47,7 +66,7 @@ def run_agent(
     output_dir.mkdir(parents=False, exist_ok=True)
 
     configure_logging(verbosity=verbosity, output_dir=output_dir)
-    from .managers import Manager
+    from .managers.node_manager import NodeManager
 
     if extract_archives_to is not None:
         if extract_archives_to and extract_archives_to != input_data_folder:
@@ -100,38 +119,59 @@ def run_agent(
         config.enable_meta_prompting = enable_meta_prompting
 
     if manager is None:
-        manager = Manager(
+        # Create a new NodeManager instance
+        manager = NodeManager(
             input_data_folder=input_data_folder,
             output_folder=output_folder,
             config=config,
             enable_per_iteration_instruction=enable_per_iteration_instruction,
             initial_user_input=initial_user_input,
         )
 
-    manager.set_initial_user_input(
-        enable_per_iteration_instruction=enable_per_iteration_instruction, initial_user_input=initial_user_input
-    )
+    # Initialize the manager (generate initial prompts)
+    manager.initialize()
 
-    while manager.time_step + 1 < max_iterations:
-        logger.brief(f"Starting iteration {manager.time_step + 1}!")
+    # Execute the MCTS search
+    iteration = 0
+    start_time = time.time()
 
-        manager.step()
+    while iteration < max_iterations:
+        # Log the current iteration
+        logger.brief(f"Starting MCTS iteration {iteration + 1}/{max_iterations}")
 
-        # Generate code
-        manager.update_python_code()
-        manager.update_bash_script()
+        # Perform one step of the Monte Carlo Tree Search
+        success = manager.step()
 
-        successful = manager.execute_code()
-
-        if successful:
+        if success:
+            # Create a best run copy when we find a successful solution
             manager.create_best_run_copy()
+
+            # If not in continuous improvement mode, we can stop
             if not config.continuous_improvement:
+                logger.brief("Stopping search - solution found and continuous improvement is disabled")
                 break
+        elif success is None:
+            logger.brief("Stopping search - all nodes are terminal.")
+            break
+        else:
+            pass
 
-        if manager.time_step + 1 >= max_iterations:
+        # Increment iteration counter
+        iteration += 1
+
+        # Check if we've exceeded the maximum iterations
+        if iteration >= max_iterations:
             logger.warning(f"[bold red]Warning: Reached maximum iterations ({max_iterations})[/bold red]")
 
+    manager.visualize_results()
     manager.report_token_usage()
-    manager.get_validation_score_summary()
-    logger.brief(f"output saved in {output_dir}.")
+    # Report token usage and validation score summary
     manager.cleanup()
+
+    # Log summary
+    elapsed_time = time.time() - start_time
+    logger.brief(f"MCTS search completed in {elapsed_time:.2f} seconds")
+    logger.brief(f"Total nodes explored: {manager.time_step + 1}")
+    logger.brief(f"Best validation score: {manager.best_validation_score}")
+    logger.brief(f"Tools used: {', '.join(manager.used_tools)}")
+    logger.brief(f"Output saved in {output_dir}")

src/autogluon/assistant/configs/anthropic.yaml

@@ -12,7 +12,7 @@ max_num_tutorials: 5
 max_user_input_length: 2048
 max_error_message_length: 2048
 max_tutorial_length: 32768
-create_venv: false
+configure_env: false
 condense_tutorials: True
 use_tutorial_summary: True
 continuous_improvement: False

src/autogluon/assistant/configs/bedrock.yaml

@@ -1,6 +1,15 @@
-# Bedrock Configuration
+# Tutorial Prompt Generator Configuration
 
-per_execution_timeout: 86400
+per_execution_timeout: 7200
+
+# MCTS (Monte Carlo Tree Search) parameters
+exploration_constant: 1.414    # Controls exploration vs exploitation trade-off in UCT formula (higher = more exploration)
+max_debug_depth: 3         # Maximum depth of debug nodes in the search tree
+failure_offset: 2             # Number of failures to ignore before applying failure penalty
+failure_penalty_weight: 0.5   # Weight of the penalty for failed executions in UCT calculation
+initial_root_children: 5      # Maximum number of child nodes from root before considering fully expanded
+max_debug_children: 2         # Maximum number of debug child nodes for a single parent node
+max_evolve_children: 4        # Maximum number of evolution child nodes for a single parent node
 
 # Data Perception
 max_file_group_size_to_show: 5
@@ -10,17 +19,20 @@ 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
+configure_env: false
 condense_tutorials: True
 use_tutorial_summary: True
 continuous_improvement: False
 optimize_system_resources: False
 cleanup_unused_env: True
 enable_meta_prompting: False
 
+# Default LLM Configuration
+# For each agent (coder, etc.) you can use a different one
 llm: &default_llm
+  # Note: bedrock is only supported in limited AWS regions
+  #       and requires AWS credentials
   provider: bedrock
   model: "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
   max_tokens: 65535
@@ -44,9 +56,7 @@ bash_coder:
 
 executer:
   <<: *default_llm  # Merge llm_config
-  max_stdout_length: 8192
-  max_stderr_length: 2048
-
+
 meta_prompting:
   <<: *default_llm  # Merge llm_config
   multi_turn: False
@@ -80,4 +90,4 @@ task_descriptor:
 tool_selector:
   <<: *default_llm  # Merge llm_config
   temperature: 0.
-  top_p: 1.
+  top_p: 1.

src/autogluon/assistant/configs/data_visualizer.yaml

@@ -11,7 +11,7 @@ num_tutorial_retrievals: 30
 max_num_tutorials: 5
 max_user_input_length: 2048
 max_tutorial_length: 32768
-create_venv: false
+configure_env: false
 condense_tutorials: True
 use_tutorial_summary: True
 continuous_improvement: False
@@ -82,7 +82,7 @@ python_coder:
 
     {tool_prompt}
 
-    {best_code_prompt}
+    {code_improvement_prompt}
 
     Please provide the complete Python script that accomplishes these tasks, ensuring it's ready to run given the appropriate data inputs.
 
@@ -96,7 +96,7 @@ python_coder:
     {user_input_truncate_end_2048}
 
     ### Previous Errors
-    {all_error_analyses}
+    {all_previous_error_analyses}
 
     ### Tutorials for Reference
     {tutorial_prompt}

src/autogluon/assistant/configs/default.yaml

@@ -1,6 +1,15 @@
 # Tutorial Prompt Generator Configuration
 
-per_execution_timeout: 86400
+per_execution_timeout: 7200
+
+# MCTS (Monte Carlo Tree Search) parameters
+exploration_constant: 1.414    # Controls exploration vs exploitation trade-off in UCT formula (higher = more exploration)
+max_debug_depth: 3         # Maximum depth of debug nodes in the search tree
+failure_offset: 2             # Number of failures to ignore before applying failure penalty
+failure_penalty_weight: 0.5   # Weight of the penalty for failed executions in UCT calculation
+initial_root_children: 5      # Maximum number of child nodes from root before considering fully expanded
+max_debug_children: 2         # Maximum number of debug child nodes for a single parent node
+max_evolve_children: 4        # Maximum number of evolution child nodes for a single parent node
 
 # Data Perception
 max_file_group_size_to_show: 5
@@ -11,7 +20,7 @@ num_tutorial_retrievals: 30
 max_num_tutorials: 5
 max_user_input_length: 2048
 max_tutorial_length: 32768
-create_venv: false
+configure_env: false
 condense_tutorials: True
 use_tutorial_summary: True
 continuous_improvement: False
@@ -26,10 +35,6 @@ llm: &default_llm
   #       and requires AWS credentials
   provider: bedrock
   model: "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
-  #provider: openai
-  #model: gpt-4o-2024-08-06
-  #provider: anthropic
-  # model: claude-3-7-sonnet-20250219
   max_tokens: 65535
   proxy_url: null
   temperature: 0.1

src/autogluon/assistant/configs/openai.yaml

@@ -12,7 +12,7 @@ max_num_tutorials: 5
 max_user_input_length: 2048
 max_error_message_length: 2048
 max_tutorial_length: 32768
-create_venv: false
+configure_env: false
 condense_tutorials: True
 use_tutorial_summary: True
 continuous_improvement: False

src/autogluon/assistant/configs/sagemaker.yaml

@@ -11,7 +11,7 @@ num_tutorial_retrievals: 30
 max_num_tutorials: 5
 max_user_input_length: 2048
 max_tutorial_length: 32768
-create_venv: false
+configure_env: false
 condense_tutorials: True
 use_tutorial_summary: True
 continuous_improvement: False

src/autogluon/assistant/constants.py

@@ -8,6 +8,9 @@
 LOGO_NIGHT_PATH = PACKAGE_ROOT / "webui" / "static" / "sidebar_icon.png"
 LOGO_PATH = PACKAGE_ROOT / "webui" / "static" / "page_icon.png"
 
+### Default Library
+DEFAULT_LIBRARY = "machine learning"
+
 ### WebUI
 
 VALID_CODING_LANGUAGES = ["python", "bash"]

src/autogluon/assistant/managers/__init__.py

@@ -1 +1 @@
-from .manager import Manager
+from .node_manager import Node, NodeManager