MCP and WebUI

.gitignore

@@ -21,6 +21,18 @@ cdk.out/
 output_*.csv
 runs/
 
+# MCP specific
+.flask.pid
+.mcp.pid
+mcp_output/
+mcp/__pycache__/
+mcp/examples/credentials_template.txt
+mcp/examples/example_config.yaml
+mcp/*.pyc
+
+# IDE specific
+.vscode/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -190,10 +202,19 @@ aga-output-*.csv
 # ignore the MLAgent output folder
 output/
 
-# MLZero Specific
-maab/datasets/
-runs/
-maab/runs/
-
 # OS specific
 *.DS_Store
+
+src/autogluon/mcp/server/.flask.pid
+src/autogluon/mcp/server/.mcp.pid
+src/autogluon/mcp/server/test_output/
+src/autogluon/mcp/server/test_data/
+src/autogluon/mcp/server/nohup.out
+src/autogluon/mcp/**/uploads/
+src/autogluon/mcp/**/mlzero-*/
+
+# Credentials
+*.credentials
+*.creds
+credentials.txt
+aws_credentials.txt

README.md

@@ -34,6 +34,7 @@ uv pip install git+https://github.com/autogluon/autogluon-assistant.git
 
 For detailed usage instructions, Anthropic/Azure/OpenAI setup, and advanced configuration options, see our [Getting Started Tutorial](docs/tutorials/getting_started.md).
 
+## API Setup
 ## API Setup
 MLZero uses AWS Bedrock by default. Configure your AWS credentials:
 
@@ -47,25 +48,45 @@ We also support Anthropic, Azure, and OpenAI. Support for more LLM providers (e.
 
 ## Basic Usage
 
-### CLI UI
+## Basic Usage for CLI UI
 
 ![Demo](https://github.com/autogluon/autogluon-assistant/blob/main/docs/assets/cli_demo.gif)
 
 ```bash
 mlzero -i <input_data_folder> [-u <optional_user_instructions>]
 ```
 
-## WEB UI
+## Basic Usage for WEB UI
 
 ![Demo](https://github.com/autogluon/autogluon-assistant/blob/main/docs/assets/web_demo.gif)
 
-```bash
-mlzero-backend # command to start backend
-mlzero-frontend # command to start frontend on 8509(default)
-```
+### Model Execution Settings
+
+The settings above the divider line control how the model runs, while the settings below the divider line relate to the model being used (including provider, credentials, and model parameters).
+
+### Model Execution Configuration
+
+**Max Iterations**: The number of rounds the model will run. The program automatically stops when this limit is reached. Default is 5, adjustable as needed.
+
+**Manual Prompts Between Iterations**: Choose whether to add iteration-specific prompts between iterations or not.
+
+**Log Verbosity**: Select the level of detail for the logs you want to see. Three options are available: brief, info, and detail. Brief is recommended.
+- **Brief**: Contains key essential information
+- **Info**: Includes brief information plus detailed information such as file save locations
+- **Detail**: Includes info-level information plus all model training related information
+
+### Model Configuration
+
+You can select the LLM provider, model, and credentials to use. If using Bedrock as the provider, you can use EC2 defaults. You can also upload your own config file, which will override the provider and model name settings. Provided credentials will be validated.
+
+### Chat Input Box
+
+1. **Initial Task Submission**: When starting a task for the first time, drag the input folder into this chat input box, enter any description or requirements about the task, then press Enter or click the submit button on the right. Note: Submitting larger files may sometimes fail - you can try multiple times if needed.
+
+2. **Manual Prompts**: If you selected "Manual prompts between iterations" in settings, you can input prompts here.
+
+3. **Task Cancellation**: After submitting a task, if you want to cancel it, submit "cancel" in this input box.
 
-1. **Configure**: Set your model provider and credentials in settings
-2. **Upload & Describe**: Drag your data folder into the chat input box, then type what you want to accomplish and press Enter
 
 ## Citation
 If you use Autogluon Assistant (MLZero) in your research, please cite our paper:

docs/tutorials/getting_started.md

@@ -70,7 +70,7 @@ You can select the LLM provider, model, and credentials to use. If using Bedrock
 ### CLI
 
 ```bash
-mlzero -i INPUT_DATA_FOLDER [-o OUTPUT_DIR] [-c CONFIG_PATH] [-n MAX_ITERATIONS] [--need-user-input] [-u INITIAL_USER_INPUT] [-e EXTRACT_TO] [-v VERBOSITY_LEVEL]
+mlzero -i INPUT_DATA_FOLDER [-o OUTPUT_DIR] [-c CONFIG_PATH] [-n MAX_ITERATIONS] [--ENABLE-PER-ITERATION-INSTRUCTION] [--INITIAL-INSTRUCTION] [-e EXTRACT_TO] [-v VERBOSITY_LEVEL]
 ```
 
 #### Required Arguments
@@ -89,10 +89,10 @@ mlzero -i INPUT_DATA_FOLDER [-o OUTPUT_DIR] [-c CONFIG_PATH] [-n MAX_ITERATIONS]
 - `-n, --max-iterations`:  
   Maximum number of iterations. Default is `5`.
 
-- `--need-user-input`:  
+- `--enable-per-iteration-instruction`:  
   Whether to prompt user input at each iteration. Defaults to `False`.
 
-- `-u, --user-input`:  
+- `--initial-instruction`:  
   Initial user input to use in the first iteration. Optional.
 
 - `-e, --extract-to`:  

maab/agents/mlzero_default/mlzero_default.sh

@@ -59,7 +59,7 @@ mlzero \
     -o "$OUTPUT_DIR" \
     -n 10 \
     -v 1 \
-    -u "complete the task in 10 minutes"
+    --initial-instruction "complete the task in 10 minutes"
 
 # Check if the process was successful
 if [ $? -ne 0 ]; then

pyproject.toml

@@ -42,12 +42,16 @@ dependencies = [
   "streamlit-extras>=0.4",
   "psutil>=5.9.8",
   "peft>=0.15.2",
+  "fastmcp>=2.0.0",
+  "aiohttp>=3.8.0",
+  "requests>=2.28.0",
+  "boto3>=1.28.0",
+  "mcp>=0.1.0",
 ]
 
 [project.scripts]
 aga          = "autogluon.assistant.cli.app:app"
 mlzero       = "autogluon.assistant.cli.app:app"
-mlzero-webui = "autogluon.assistant.webui.Home:main"
 mlzero-backend = "autogluon.assistant.webui.backend.app:main"
 mlzero-frontend = "autogluon.assistant.webui.runner:run_frontend"
 
@@ -56,6 +60,7 @@ dev = [
   "black>=24",
   "GitRepo>=1",
   "pytest>=7",
+  "pytest-asyncio>=0.21.0",
   "isort",
   "ruff==0.4",
 ]
@@ -65,7 +70,7 @@ Homepage = "https://github.com/autogluon/autogluon-assistant"
 
 [tool.setuptools.packages.find]
 where = ["src"]
-include = ["autogluon.assistant*"]
+include = ["autogluon.assistant*", "autogluon.mcp*"]
 namespaces = true
 
 [tool.setuptools.package-data]

src/autogluon/assistant/cli/app.py

@@ -1,13 +1,24 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 
-from pathlib import Path
+import multiprocessing.resource_tracker
 
-import typer
 
-from autogluon.assistant.coding_agent import run_agent
+def _noop(*args, **kwargs):
+    pass
 
-from .. import __file__ as assistant_file
+
+multiprocessing.resource_tracker.register = _noop
+multiprocessing.resource_tracker.unregister = _noop
+multiprocessing.resource_tracker.ensure_running = _noop
+
+from pathlib import Path  # noqa: E402
+
+import typer  # noqa: E402
+
+from autogluon.assistant.coding_agent import run_agent  # noqa: E402
+
+from .. import __file__ as assistant_file  # noqa: E402
 
 PACKAGE_ROOT = Path(assistant_file).parent
 DEFAULT_CONFIG_PATH = PACKAGE_ROOT / "configs" / "default.yaml"
@@ -31,14 +42,39 @@ def main(
         "--config",
         help=f"YAML config file (default: {DEFAULT_CONFIG_PATH})",
     ),
-    max_iterations: int = typer.Option(5, "-n", "--max-iterations", help="Max iteration count"),
-    need_user_input: bool = typer.Option(False, "--need-user-input", help="Whether to prompt user each iteration"),
-    initial_user_input: str | None = typer.Option(None, "-u", "--user-input", help="Initial user input"),
+    max_iterations: int = typer.Option(
+        5,
+        "-n",
+        "--max-iterations",
+        help="Max iteration count. If the task hasn’t succeeded after this many iterations, it will terminate.",
+    ),
+    need_user_input: bool = typer.Option(
+        False,
+        "--enable-per-iteration-instruction",
+        help="If enabled, you can provide a prompt for the next iteration at each iteration",
+    ),
+    initial_user_input: str | None = typer.Option(
+        None, "--initial-instruction", help="You can provide the initial instruction here."
+    ),
     extract_archives_to: str | None = typer.Option(
-        None, "-e", "--extract-to", help="Directory in which to unpack any archives"
+        None,
+        "-e",
+        "--extract-to",
+        help="Copy input data to specified directory and automatically extract all .zip archives. ",
     ),
     # === Logging parameters ===
-    verbosity: int = typer.Option(1, "-v", "--verbosity", help="Verbosity level (0–4)"),
+    verbosity: int = typer.Option(
+        1,
+        "-v",
+        "--verbosity",
+        help=(
+            "-v 0: Only includes error messages\n"
+            "-v 1: Contains key essential information\n"
+            "-v 2: Includes brief information plus detailed information such as file save locations\n"
+            "-v 3: Includes info-level information plus all model training related information\n"
+            "-v 4: Includes full debug information"
+        ),
+    ),
 ):
     """
     mlzero: a CLI for running the AutoMLAgent pipeline.

src/autogluon/assistant/coding_agent.py

@@ -6,7 +6,7 @@
 
 from omegaconf import OmegaConf
 
-from .managers import Manager
+# from .managers import Manager
 from .rich_logging import configure_logging
 from .utils import extract_archives
 
@@ -45,6 +45,7 @@ def run_agent(
     output_dir.mkdir(parents=False, exist_ok=True)
 
     configure_logging(verbosity=verbosity, output_dir=output_dir)
+    from .managers import Manager
 
     if extract_archives_to is not None:
         if extract_archives_to and extract_archives_to != input_data_folder:

src/autogluon/assistant/rich_logging.py

@@ -47,6 +47,7 @@ def _configure_logging(console_level: int, output_dir: Path = None) -> None:
         console = Console(file=sys.stderr)
         console_handler = RichHandler(console=console, markup=True, rich_tracebacks=True)
         console_handler.setLevel(console_level)
+        console_handler.name = CONSOLE_HANDLER
         handlers = [console_handler]
     else:
         stdout_handler = logging.StreamHandler(sys.stdout)

src/autogluon/assistant/tools_registry/indexing.py

@@ -1,8 +1,13 @@
+import contextlib
+import io
 import logging
+import os
 import pickle
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+
 import faiss
 import numpy as np
 from FlagEmbedding import FlagAutoModel
@@ -32,6 +37,10 @@ def __del__(self):
         """Cleanup method to properly close the embedding model."""
         self.cleanup()
 
+    def __log_wrapper(self, input):
+        with contextlib.redirect_stderr(io.StringIO()):
+            return self.model.encode(input)
+
     def cleanup(self):
         """Cleanup the embedding model to avoid multiprocessing issues."""
         if self.model is not None:
@@ -41,6 +50,8 @@ def cleanup(self):
                     self.model.close()
                 elif hasattr(self.model, "stop_multi_process_pool"):
                     self.model.stop_multi_process_pool()
+                else:
+                    del self.model
             except Exception as e:
                 logger.debug(f"Error during model cleanup: {e}")
             finally:
@@ -141,7 +152,7 @@ def _build_tool_index(self, tool_name: str, tutorial_type: str) -> Tuple[faiss.I
 
             for i in range(0, len(summaries), batch_size):
                 batch_summaries = summaries[i : i + batch_size]
-                batch_embeddings = self.model.encode(batch_summaries)
+                batch_embeddings = self.__log_wrapper(batch_summaries)
 
                 # Ensure proper format
                 if not isinstance(batch_embeddings, np.ndarray):
@@ -316,7 +327,7 @@ def search(self, query: str, tool_name: str, condensed: bool = False, top_k: int
             return []
 
         # Generate query embedding
-        query_embedding = self.model.encode([query])
+        query_embedding = self.__log_wrapper([query])
 
         # Ensure proper data type and memory layout
         if not isinstance(query_embedding, np.ndarray):