nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/cli/commands/object_store/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/cli/commands/object_store/object_store.py

@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import importlib
+import logging
+import mimetypes
+import time
+from pathlib import Path
+
+import click
+
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.object_store import ObjectStoreBaseConfig
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+logger = logging.getLogger(__name__)
+
+STORE_CONFIGS = {
+    "s3": {
+        "module": "nat.plugins.s3.object_store", "config_class": "S3ObjectStoreClientConfig"
+    },
+    "mysql": {
+        "module": "nat.plugins.mysql.object_store", "config_class": "MySQLObjectStoreClientConfig"
+    },
+    "redis": {
+        "module": "nat.plugins.redis.object_store", "config_class": "RedisObjectStoreClientConfig"
+    }
+}
+
+
+def get_object_store_config(**kwargs) -> ObjectStoreBaseConfig:
+    """Process common object store arguments and return the config class"""
+    store_type = kwargs.pop("store_type")
+    config = STORE_CONFIGS[store_type]
+    module = importlib.import_module(config["module"])
+    config_class = getattr(module, config["config_class"])
+    return config_class(**kwargs)
+
+
+async def upload_file(object_store: ObjectStore, file_path: Path, key: str):
+    """
+    Upload a single file to object store.
+
+    Args:
+        object_store: The object store instance to use.
+        file_path: The path to the file to upload.
+        key: The key to upload the file to.
+    """
+    try:
+        data = await asyncio.to_thread(file_path.read_bytes)
+
+        item = ObjectStoreItem(data=data,
+                               content_type=mimetypes.guess_type(str(file_path))[0],
+                               metadata={
+                                   "original_filename": file_path.name,
+                                   "file_size": str(len(data)),
+                                   "file_extension": file_path.suffix,
+                                   "upload_timestamp": str(int(time.time()))
+                               })
+
+        # Upload using upsert to allow overwriting
+        await object_store.upsert_object(key, item)
+        click.echo(f"✅ Uploaded: {file_path.name} -> {key}")
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to upload {file_path.name}:\n{e}") from e
+
+
+def object_store_command_decorator(async_func):
+    """
+    Decorator that handles the common object store command pattern.
+
+    The decorated function should take (store: ObjectStore, kwargs) as parameters
+    and return an exit code (0 for success).
+    """
+
+    @click.pass_context
+    def wrapper(ctx: click.Context, **kwargs):
+        config = ctx.obj["store_config"]
+
+        async def work():
+            async with WorkflowBuilder() as builder:
+                await builder.add_object_store(name="store", config=config)
+                store = await builder.get_object_store_client("store")
+                return await async_func(store, **kwargs)
+
+        try:
+            exit_code = asyncio.run(work())
+        except Exception as e:
+            raise click.ClickException(f"Command failed: {e}") from e
+        if exit_code != 0:
+            raise click.ClickException(f"Command failed with exit code {exit_code}")
+        return exit_code
+
+    return wrapper
+
+
+@click.command(name="upload", help="Upload a directory to an object store.")
+@click.argument("local_dir",
+                type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
+                required=True)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def upload_command(store: ObjectStore, local_dir: Path, **_kwargs):
+    """
+    Upload a directory to an object store.
+
+    Args:
+        local_dir: The local directory to upload.
+        store: The object store to use.
+        _kwargs: Additional keyword arguments.
+    """
+    try:
+        click.echo(f"📁 Processing directory: {local_dir}")
+        file_count = 0
+
+        # Process each file recursively
+        for file_path in local_dir.rglob("*"):
+            if file_path.is_file():
+                key = file_path.relative_to(local_dir).as_posix()
+                await upload_file(store, file_path, key)
+                file_count += 1
+
+        click.echo(f"✅ Directory uploaded successfully! {file_count} files uploaded.")
+        return 0
+
+    except Exception as e:
+        raise click.ClickException(f"❌ Failed to upload directory {local_dir}:\n  {e}") from e
+
+
+@click.command(name="delete", help="Delete files from an object store.")
+@click.argument("keys", type=str, required=True, nargs=-1)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def delete_command(store: ObjectStore, keys: list[str], **_kwargs):
+    """
+    Delete files from an object store.
+
+    Args:
+        store: The object store to use.
+        keys: The keys to delete.
+        _kwargs: Additional keyword arguments.
+    """
+    deleted_count = 0
+    failed_count = 0
+    for key in keys:
+        try:
+            await store.delete_object(key)
+            click.echo(f"✅ Deleted: {key}")
+            deleted_count += 1
+        except Exception as e:
+            click.echo(f"❌ Failed to delete {key}: {e}")
+            failed_count += 1
+
+    click.echo(f"✅ Deletion completed! {deleted_count} keys deleted. {failed_count} keys failed to delete.")
+    return 0 if failed_count == 0 else 1
+
+
+@click.group(name="object-store", invoke_without_command=False, help="Manage object store operations.")
+def object_store_command(**_kwargs):
+    """Manage object store operations including uploading files and directories."""
+    pass
+
+
+def register_object_store_commands():
+
+    @click.group(name="s3", invoke_without_command=False, help="S3 object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--endpoint-url", type=str, help="S3 endpoint URL")
+    @click.option("--access-key", type=str, help="S3 access key")
+    @click.option("--secret-key", type=str, help="S3 secret key")
+    @click.option("--region", type=str, help="S3 region")
+    @click.pass_context
+    def s3(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="s3", **kwargs)
+
+    @click.group(name="mysql", invoke_without_command=False, help="MySQL object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="MySQL host")
+    @click.option("--port", type=int, help="MySQL port")
+    @click.option("--db", type=str, help="MySQL database name")
+    @click.option("--username", type=str, help="MySQL username")
+    @click.option("--password", type=str, help="MySQL password")
+    @click.pass_context
+    def mysql(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="mysql", **kwargs)
+
+    @click.group(name="redis", invoke_without_command=False, help="Redis object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="Redis host")
+    @click.option("--port", type=int, help="Redis port")
+    @click.option("--db", type=int, help="Redis db")
+    @click.pass_context
+    def redis(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="redis", **kwargs)
+
+    commands = {"s3": s3, "mysql": mysql, "redis": redis}
+
+    for store_type, config in STORE_CONFIGS.items():
+        try:
+            importlib.import_module(config["module"])
+            command = commands[store_type]
+            object_store_command.add_command(command, name=store_type)
+            command.add_command(upload_command, name="upload")
+            command.add_command(delete_command, name="delete")
+        except ImportError:
+            pass
+
+
+register_object_store_commands()

nat/cli/commands/optimize.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from pathlib import Path
+
+import click
+
+from nat.data_models.optimizer import OptimizerRunConfig
+from nat.profiler.parameter_optimization.optimizer_runtime import optimize_config
+
+logger = logging.getLogger(__name__)
+
+
+@click.group(name=__name__, invoke_without_command=True, help="Optimize a workflow with the specified dataset.")
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=True,
+    help="A JSON/YAML file that sets the parameters for the workflow and evaluation.",
+)
+@click.option(
+    "--dataset",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=False,
+    help="A json file with questions and ground truth answers. This will override the dataset path in the config file.",
+)
+@click.option(
+    "--result_json_path",
+    type=str,
+    default="$",
+    help=("A JSON path to extract the result from the workflow. Use this when the workflow returns "
+          "multiple objects or a dictionary. For example, '$.output' will extract the 'output' field "
+          "from the result."),
+)
+@click.option(
+    "--endpoint",
+    type=str,
+    default=None,
+    help="Use endpoint for running the workflow. Example: http://localhost:8000/generate",
+)
+@click.option(
+    "--endpoint_timeout",
+    type=int,
+    default=300,
+    help="HTTP response timeout in seconds. Only relevant if endpoint is specified.",
+)
+@click.pass_context
+def optimizer_command(ctx, **kwargs) -> None:
+    """ Optimize workflow with the specified dataset"""
+    pass
+
+
+async def run_optimizer(config: OptimizerRunConfig):
+    await optimize_config(config)
+
+
+@optimizer_command.result_callback(replace=True)
+def run_optimizer_callback(
+    processors,  # pylint: disable=unused-argument
+    *,
+    config_file: Path,
+    dataset: Path,
+    result_json_path: str,
+    endpoint: str,
+    endpoint_timeout: int,
+):
+    """Run the optimizer with the provided config file and dataset."""
+    config = OptimizerRunConfig(
+        config_file=config_file,
+        dataset=dataset,
+        result_json_path=result_json_path,
+        endpoint=endpoint,
+        endpoint_timeout=endpoint_timeout,
+    )
+
+    asyncio.run(run_optimizer(config))

nat/cli/commands/registry/publish.py

@@ -40,7 +40,7 @@ async def publish_artifact(registry_handler_config: RegistryHandlerBaseConfig, p
         try:
             artifact = build_artifact(package_root=package_root)
         except Exception as e:
-            logger.exception("Error building artifact: %s", e, exc_info=True)
+            logger.exception("Error building artifact: %s", e)
             return
         await stack.enter_async_context(registry_handler.publish(artifact=artifact))
 
@@ -82,7 +82,7 @@ def publish(channel: str, config_file: str, package_root: str) -> None:
             logger.error("Publish channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(publish_artifact(registry_handler_config=publish_channel_config, package_root=package_root))

nat/cli/commands/registry/pull.py

@@ -66,7 +66,7 @@ async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, pack
             validated_packages = PullRequestPackages(packages=package_list)
 
         except Exception as e:
-            logger.exception("Error processing package names: %s", e, exc_info=True)
+            logger.exception("Error processing package names: %s", e)
             return
 
         await stack.enter_async_context(registry_handler.pull(packages=validated_packages))
@@ -112,7 +112,7 @@ def pull(channel: str, config_file: str, packages: str) -> None:
             logger.error("Pull channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(pull_artifact(pull_channel_config, packages))

nat/cli/commands/registry/remove.py

@@ -41,7 +41,7 @@ async def remove_artifact(registry_handler_config: RegistryHandlerBaseConfig, pa
         try:
             package_name_list = PackageNameVersionList(**{"packages": packages})
         except Exception as e:
-            logger.exception("Invalid package format: '%s'", e, exc_info=True)
+            logger.exception("Invalid package format: '%s'", e)
 
         await stack.enter_async_context(registry_handler.remove(packages=package_name_list))
 
@@ -102,7 +102,7 @@ def remove(channel: str, config_file: str, packages: str) -> None:
             logger.error("Remove channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(remove_artifact(registry_handler_config=remove_channel_config, packages=packages_versions))

nat/cli/commands/registry/search.py

@@ -29,14 +29,13 @@ from nat.utils.data_models.schema_validator import validate_yaml
 logger = logging.getLogger(__name__)
 
 
-async def search_artifacts(  # pylint: disable=R0917
-        registry_handler_config: RegistryHandlerBaseConfig,
-        query: str,
-        search_fields: list[SearchFields],
-        visualize: bool,
-        component_types: list[ComponentEnum],
-        save_path: str | None = None,
-        n_results: int = 10) -> None:
+async def search_artifacts(registry_handler_config: RegistryHandlerBaseConfig,
+                           query: str,
+                           search_fields: list[SearchFields],
+                           visualize: bool,
+                           component_types: list[ComponentEnum],
+                           save_path: str | None = None,
+                           n_results: int = 10) -> None:
 
     from nat.cli.type_registry import GlobalTypeRegistry
     from nat.registry_handlers.schemas.search import SearchQuery
@@ -116,14 +115,13 @@ async def search_artifacts(  # pylint: disable=R0917
     required=False,
     help=("The component types to include in search."),
 )
-def search(  # pylint: disable=R0917
-        config_file: str,
-        channel: str,
-        fields: list[str],
-        query: str,
-        component_types: list[ComponentEnum],
-        n_results: int,
-        output_path: str) -> None:
+def search(config_file: str,
+           channel: str,
+           fields: list[str],
+           query: str,
+           component_types: list[ComponentEnum],
+           n_results: int,
+           output_path: str) -> None:
     """
     Search for NAT artifacts with the specified configuration.
     """
@@ -142,7 +140,7 @@ def search(  # pylint: disable=R0917
             logger.error("Search channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(

nat/cli/commands/start.py

@@ -35,7 +35,6 @@ logger = logging.getLogger(__name__)
 
 class StartCommandGroup(click.Group):
 
-    # pylint: disable=too-many-positional-arguments
     def __init__(
         self,
         name: str | None = None,
@@ -103,12 +102,24 @@ class StartCommandGroup(click.Group):
                 raise ValueError(f"Invalid field '{name}'.Unions are only supported for optional parameters.")
 
             # Handle the types
-            if (issubclass(decomposed_type.root, Path)):
+            # Literal[...] -> map to click.Choice([...])
+            if (decomposed_type.origin is typing.Literal):
+                # typing.get_args returns the literal values; ensure they are strings for Click
+                literal_values = [str(v) for v in decomposed_type.args]
+                param_type = click.Choice(literal_values)
+
+            elif (issubclass(decomposed_type.root, Path)):
                 param_type = click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path)
 
-            elif (issubclass(decomposed_type.root, (list, tuple, set))):
+            elif (issubclass(decomposed_type.root, list | tuple | set)):
                 if (len(decomposed_type.args) == 1):
-                    param_type = decomposed_type.args[0]
+                    inner = DecomposedType(decomposed_type.args[0])
+                    # Support containers of Literal values -> multiple Choice
+                    if (inner.origin is typing.Literal):
+                        literal_values = [str(v) for v in inner.args]
+                        param_type = click.Choice(literal_values)
+                    else:
+                        param_type = inner.root
                 else:
                     param_type = None
 
@@ -225,7 +236,7 @@ class StartCommandGroup(click.Group):
             return asyncio.run(run_plugin())
 
         except Exception as e:
-            logger.error("Failed to initialize workflow", exc_info=True)
+            logger.error("Failed to initialize workflow")
             raise click.ClickException(str(e)) from e
 
     def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None:

nat/cli/commands/uninstall.py

@@ -44,7 +44,7 @@ async def uninstall_packages(packages: list[dict[str, str]]) -> None:
     try:
         package_name_list = PackageNameVersionList(**{"packages": packages})
     except Exception as e:
-        logger.exception("Error validating package format: %s", e, exc_info=True)
+        logger.exception("Error validating package format: %s", e)
         return
 
     async with AsyncExitStack() as stack:

nat/cli/commands/workflow/templates/config.yml.j2

@@ -1,5 +1,4 @@
 general:
-  use_uvloop: true
   logging:
     console:
       _type: console

nat/cli/commands/workflow/templates/pyproject.toml.j2

@@ -3,6 +3,9 @@ build-backend = "setuptools.build_meta"
 {% if editable %}requires = ["setuptools >= 64", "setuptools-scm>=8"]
 
 [tool.setuptools_scm]
+# NAT uses the --first-parent flag to avoid tags from previous releases which have been merged into the develop branch
+# from causing an unexpected version change. This can be safely removed if developing outside of the NAT repository.
+git_describe_command = "git describe --long --first-parent"
 root = "{{ rel_path_to_repo_root}}"{% else %}requires = ["setuptools >= 64"]{% endif %}
 
 [project]
@@ -11,7 +14,7 @@ name = "{{ package_name }}"
 dependencies = [
   "{{ nat_dependency }}",
 ]
-requires-python = ">=3.11,<3.13"
+requires-python = ">=3.11,<3.14"
 description = "Custom NeMo Agent Toolkit Workflow"
 classifiers = ["Programming Language :: Python"]
 

nat/cli/commands/workflow/templates/register.py.j2

@@ -1,4 +1,3 @@
-# pylint: disable=unused-import
 # flake8: noqa
 
 # Import any tools which need to be automatically registered here

nat/cli/commands/workflow/workflow_commands.py

@@ -37,7 +37,7 @@ def _get_nat_dependency(versioned: bool = True) -> str:
     Returns:
         str: The dependency string to use in pyproject.toml
     """
-    # Assume the default dependency is langchain
+    # Assume the default dependency is LangChain/LangGraph
     dependency = "nvidia-nat[langchain]"
 
     if not versioned:
@@ -97,7 +97,7 @@ def find_package_root(package_name: str) -> Path | None:
         try:
             info = json.loads(direct_url)
         except json.JSONDecodeError:
-            logger.error("Malformed direct_url.json for package: %s", package_name)
+            logger.exception("Malformed direct_url.json for package: %s", package_name)
             return None
 
         if not info.get("dir_info", {}).get("editable"):
@@ -161,7 +161,6 @@ def get_workflow_path_from_name(workflow_name: str):
     default="NAT function template. Please update the description.",
     help="""A description of the component being created. Will be used to populate the docstring and will describe the
          component when inspecting installed components using 'nat info component'""")
-# pylint: disable=missing-param-doc
 def create_command(workflow_name: str, install: bool, workflow_dir: str, description: str):
     """
     Create a new NAT workflow using templates.
@@ -172,6 +171,9 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
         workflow_dir (str): The directory to create the workflow package.
         description (str): Description to pre-popluate the workflow docstring.
     """
+    # Fail fast with Click's standard exit code (2) for bad params.
+    if not workflow_name or not workflow_name.strip():
+        raise click.BadParameter("Workflow name cannot be empty.")  # noqa: TRY003
     try:
         # Get the repository root
         try:
@@ -218,15 +220,13 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
         else:
             install_cmd = ['pip', 'install', '-e', str(new_workflow_dir)]
 
-        config_source = configs_dir / 'config.yml'
-
         # List of templates and their destinations
         files_to_render = {
             'pyproject.toml.j2': new_workflow_dir / 'pyproject.toml',
             'register.py.j2': base_dir / 'register.py',
             'workflow.py.j2': base_dir / f'{workflow_name}_function.py',
             '__init__.py.j2': base_dir / '__init__.py',
-            'config.yml.j2': config_source,
+            'config.yml.j2': configs_dir / 'config.yml',
         }
 
         # Render templates
@@ -247,10 +247,6 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
             with open(output_path, 'w', encoding="utf-8") as f:
                 f.write(content)
 
-        # Create symlink for config.yml
-        config_link = new_workflow_dir / 'configs' / 'config.yml'
-        os.symlink(config_source, config_link)
-
         # Create symlinks for config and data directories
         config_dir_source = configs_dir
         config_dir_link = new_workflow_dir / 'configs'
@@ -272,7 +268,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
 
         click.echo(f"Workflow '{workflow_name}' created successfully in '{new_workflow_dir}'.")
     except Exception as e:
-        logger.exception("An error occurred while creating the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while creating the workflow: %s", e)
         click.echo(f"An error occurred while creating the workflow: {e}")
 
 
@@ -308,7 +304,7 @@ def reinstall_command(workflow_name):
 
         click.echo(f"Workflow '{workflow_name}' reinstalled successfully.")
     except Exception as e:
-        logger.exception("An error occurred while reinstalling the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while reinstalling the workflow: %s", e)
         click.echo(f"An error occurred while reinstalling the workflow: {e}")
 
 
@@ -355,7 +351,7 @@ def delete_command(workflow_name: str):
 
         click.echo(f"Workflow '{workflow_name}' deleted successfully.")
     except Exception as e:
-        logger.exception("An error occurred while deleting the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while deleting the workflow: %s", e)
         click.echo(f"An error occurred while deleting the workflow: {e}")
 
 

nat/cli/entrypoint.py

@@ -30,9 +30,14 @@ import time
 import click
 import nest_asyncio
 
+from nat.utils.log_levels import LOG_LEVELS
+
 from .commands.configure.configure import configure_command
 from .commands.evaluate import eval_command
 from .commands.info.info import info_command
+from .commands.mcp.mcp import mcp_command
+from .commands.object_store.object_store import object_store_command
+from .commands.optimize import optimizer_command
 from .commands.registry.registry import registry_command
 from .commands.sizing.sizing import sizing
 from .commands.start import start_command
@@ -43,15 +48,6 @@ from .commands.workflow.workflow import workflow_command
 # Apply at the beginning of the file to avoid issues with asyncio
 nest_asyncio.apply()
 
-# Define log level choices
-LOG_LEVELS = {
-    'DEBUG': logging.DEBUG,
-    'INFO': logging.INFO,
-    'WARNING': logging.WARNING,
-    'ERROR': logging.ERROR,
-    'CRITICAL': logging.CRITICAL
-}
-
 
 def setup_logging(log_level: str):
     """Configure logging with the specified level"""
@@ -107,11 +103,13 @@ cli.add_command(uninstall_command, name="uninstall")
 cli.add_command(validate_command, name="validate")
 cli.add_command(workflow_command, name="workflow")
 cli.add_command(sizing, name="sizing")
+cli.add_command(optimizer_command, name="optimize")
+cli.add_command(object_store_command, name="object-store")
+cli.add_command(mcp_command, name="mcp")
 
 # Aliases
 cli.add_command(start_command.get_command(None, "console"), name="run")  # type: ignore
 cli.add_command(start_command.get_command(None, "fastapi"), name="serve")  # type: ignore
-cli.add_command(start_command.get_command(None, "mcp"), name="mcp")
 
 
 @cli.result_callback()