Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 41 additions & 14 deletions plugins/nemo-agents/src/nemo_agents_plugin/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,18 @@
import httpx
import typer
import yaml
from nemo_agents_plugin.cli_context import (
DEFAULT_BASE_URL as _DEFAULT_BASE_URL,
)
from nemo_agents_plugin.cli_context import (
BaseUrlOption,
)
from nemo_agents_plugin.cli_context import (
resolve_base_url as _resolve_base_url,
)
from nemo_agents_plugin.cli_context import (
resolve_context_headers as _resolve_context_headers,
)
from nemo_agents_plugin.leaderboard.cli import register_leaderboard_commands
from nemo_agents_plugin.usage.cli import register_usage_commands
from nemo_platform.cli.core.formatters import Column, format_output
Expand All @@ -54,7 +66,6 @@

logger = logging.getLogger(__name__)

_DEFAULT_BASE_URL = "http://localhost:8080"
_DEFAULT_WORKSPACE = "default"
_LIST_OUTPUT_FORMAT = Literal["table", "json", "yaml", "csv", "markdown", "raw"]
_AGENT_LIST_COLUMNS = [
Expand Down Expand Up @@ -145,7 +156,7 @@ def invoke(
help="Name of a specific deployment to invoke (platform required).",
),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
timeout: float = typer.Option(
300,
"--timeout",
Expand All @@ -160,6 +171,7 @@ def invoke(
),
) -> None:
"""Invoke an agent — locally (with --agent-config) or via the platform (with --agent or --agent-deployment)."""
base_url = _resolve_base_url(base_url)
if agent_config:
_local_invoke(agent_config, input, input_file, workspace=workspace, base_url=base_url)
elif agent or agent_deployment:
Expand Down Expand Up @@ -628,9 +640,10 @@ def create(
),
description: str = typer.Option("", "--description"),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
) -> None:
"""Register an agent on the platform."""
base_url = _resolve_base_url(base_url)
from nemo_agents_plugin.utils import inject_default_model

config_dict = _load_yaml(agent_config)
Expand All @@ -653,7 +666,7 @@ def create(
def list_agents(
ctx: typer.Context,
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
output_format: Optional[_LIST_OUTPUT_FORMAT] = typer.Option(
None,
"--format",
Expand All @@ -671,6 +684,7 @@ def list_agents(
),
) -> None:
"""List agents on the platform."""
base_url = _resolve_base_url(base_url)
resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/agents")
_print_list_response(
ctx,
Expand All @@ -684,20 +698,22 @@ def list_agents(
def get(
name: str = typer.Argument(..., help="Agent name."),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
) -> None:
"""Get an agent by name."""
base_url = _resolve_base_url(base_url)
resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/agents/{name}")
typer.echo(json.dumps(resp, indent=2))

@app.command(rich_help_panel="Agent Resources (requires running cluster)")
def delete(
name: str = typer.Argument(..., help="Agent name."),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."),
) -> None:
"""Delete an agent from the platform."""
base_url = _resolve_base_url(base_url)
if not yes:
typer.confirm(f"Delete agent '{name}'?", abort=True)
_api_request("DELETE", base_url, f"/apis/agents/v2/workspaces/{workspace}/agents/{name}")
Expand Down Expand Up @@ -725,7 +741,7 @@ def deploy(
help="Maximum seconds to wait for a terminal status (only with --wait).",
),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
) -> None:
"""Deploy an agent on the platform.

Expand All @@ -736,6 +752,7 @@ def deploy(
scripted pipelines that prefer to poll separately via ``nemo agents
deployments wait``.
"""
base_url = _resolve_base_url(base_url)
payload: dict = {"agent": agent}
if name:
payload["name"] = name
Expand Down Expand Up @@ -792,7 +809,7 @@ def logs(
help="Print only the absolute log file path and exit (useful for scripting).",
),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
) -> None:
"""Show logs for an agent deployment.

Expand All @@ -816,6 +833,7 @@ def logs(
raise typer.Exit(code=1)

if agent and not name:
base_url = _resolve_base_url(base_url)
candidates = [
d
for d in _unwrap_list(
Expand Down Expand Up @@ -859,10 +877,11 @@ def undeploy(
None, "--agent", "--all", "-a", help="Remove all deployments for this agent."
),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."),
) -> None:
"""Stop and remove a deployment (or all deployments for an agent)."""
base_url = _resolve_base_url(base_url)
if name:
if not yes:
typer.confirm(f"Undeploy '{name}'?", abort=True)
Expand All @@ -888,7 +907,7 @@ def undeploy(
def deployments_list(
ctx: typer.Context,
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
output_format: Optional[_LIST_OUTPUT_FORMAT] = typer.Option(
None,
"--format",
Expand All @@ -906,6 +925,7 @@ def deployments_list(
),
) -> None:
"""List deployments."""
base_url = _resolve_base_url(base_url)
resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/deployments")
_print_list_response(
ctx,
Expand All @@ -919,20 +939,22 @@ def deployments_list(
def deployments_get(
name: str = typer.Argument(..., help="Deployment name."),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
) -> None:
"""Get a deployment by name."""
base_url = _resolve_base_url(base_url)
resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/deployments/{name}")
typer.echo(json.dumps(resp, indent=2))

@deps_app.command(name="delete")
def deployments_delete(
name: str = typer.Argument(..., help="Deployment name."),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."),
) -> None:
"""Delete a deployment by name."""
base_url = _resolve_base_url(base_url)
if not yes:
typer.confirm(f"Delete deployment '{name}'?", abort=True)
_api_request("DELETE", base_url, f"/apis/agents/v2/workspaces/{workspace}/deployments/{name}")
Expand All @@ -950,7 +972,7 @@ def deployments_wait(
timeout: int = typer.Option(300, "--timeout", "-t", help="Maximum seconds to wait."),
interval: float = typer.Option(2.0, "--interval", help="Poll interval in seconds."),
workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"),
base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"),
base_url: BaseUrlOption = None,
) -> None:
"""Wait for a deployment to reach 'running' or 'failed' status.

Expand All @@ -960,6 +982,7 @@ def deployments_wait(
Provide either a deployment name directly or --agent to resolve the
latest active deployment for that agent automatically.
"""
base_url = _resolve_base_url(base_url)
if not name and not agent:
typer.echo("Error: provide a deployment name or --agent.", err=True)
raise typer.Exit(code=1)
Expand Down Expand Up @@ -1200,13 +1223,14 @@ def _platform_invoke(
path = f"/apis/agents/v2/workspaces/{workspace}/deployments/{deployment}/-/v1/chat/completions"

url = base_url.rstrip("/") + path
headers = _resolve_context_headers()
target_label = agent or deployment
for query in queries:
payload = {"messages": [{"role": "user", "content": query}], "stream": False}
try:
with request_progress(f"Waiting for agent '{target_label}'...", disabled=no_progress):
with httpx.Client(timeout=timeout) as client:
resp = client.post(url, json=payload)
resp = client.post(url, json=payload, headers=headers or None)
resp.raise_for_status()
body = resp.json()
typer.echo(json.dumps(body, indent=2))
Expand Down Expand Up @@ -1299,6 +1323,9 @@ def _api_request(method: str, base_url: str, path: str, *, json_body: dict[str,
request_kwargs: dict[str, Any] = {}
if json_body is not None:
request_kwargs["json"] = json_body
headers = _resolve_context_headers()
if headers:
request_kwargs["headers"] = headers
try:
with httpx.Client(timeout=30) as client:
resp = client.request(method, url, **request_kwargs)
Expand Down
111 changes: 111 additions & 0 deletions plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Shared CLI-context resolution for the ``nemo agents`` command group.

The agents plugin makes platform calls from two places — the platform
commands in :mod:`nemo_agents_plugin.cli` (raw ``httpx``) and
``nemo agents usage show`` in :mod:`nemo_agents_plugin.usage.cli` (the
NeMoPlatform SDK client). Both must resolve the platform base URL and the
auth token the same way every other ``nemo`` command does: through the
shared CLI context object stored on ``typer.Context.obj``.

These helpers read the *ambient* Click context so callers deep in a command's
call stack can resolve configuration without threading the context object
through every function signature. They live in their own module (rather than
in ``cli.py``) because ``cli.py`` imports the usage CLI at module load, so a
back-import would create a cycle.
"""

from __future__ import annotations

import logging
from typing import Annotated, Any, Optional

import click
import typer

logger = logging.getLogger(__name__)

DEFAULT_BASE_URL = "http://localhost:8080"

BASE_URL_HELP = (
"Platform base URL. Resolution order: "
"(1) this --base-url flag or NEMO_BASE_URL; "
"(2) shared CLI config (`nemo config set --base-url`) or NMP_BASE_URL; "
f"(3) {DEFAULT_BASE_URL} (default)."
)

# Reusable ``--base-url`` option shared across every ``nemo agents`` command
# so the option and its help text are defined once. ``None`` means "unset" so
# ``resolve_base_url`` can fall back to the shared CLI context / config.
BaseUrlOption = Annotated[
Optional[str],
typer.Option("--base-url", envvar="NEMO_BASE_URL", help=BASE_URL_HELP),
]


def current_cli_state() -> Any:
"""Return the shared CLI context object (``typer.Context.obj``) if present.

Returns ``None`` when the plugin is exercised outside a Click invocation
(e.g. a direct unit test), so callers fall back to their own defaults.
"""
ctx = click.get_current_context(silent=True)
return ctx.obj if ctx is not None else None


def base_url_from_context() -> str | None:
"""Return the base URL configured in the shared CLI context, if any."""
state = current_cli_state()
if state is None or not hasattr(state, "get_base_url"):
return None
try:
return state.get_base_url(default=None)
except Exception:
logger.debug("Failed to resolve base URL from CLI context", exc_info=True)
return None


def resolve_base_url(base_url: str | None) -> str:
"""Resolve the platform base URL and announce the target on stderr.

Precedence:
1. Explicit ``--base-url`` / ``NEMO_BASE_URL`` on the command.
2. The shared CLI context — ``nemo config set --base-url`` and the
``NMP_BASE_URL`` env var — so ``nemo agents`` targets the same
platform as every other ``nemo`` command.
3. The built-in localhost default.

The resolved target is echoed to stderr (never stdout, so piped/JSON
output stays clean) so a mis-pointed command is visible instead of
silently hitting the wrong platform.
"""
resolved = base_url or base_url_from_context() or DEFAULT_BASE_URL
click.echo(f"Targeting {resolved}", err=True)
return resolved


def resolve_context_headers() -> dict[str, str]:
"""Return auth (and other) default headers from the shared CLI context.

Mirrors ``nemo_platform_plugin.commands._resolve_submit_auth_headers``:
reads the SDK client config off the shared context so ``nemo agents``
attaches the same ``Authorization: Bearer`` token as the rest of the CLI
(i.e. the token established by ``nemo auth login``). Returns an empty
mapping when no context or token is available — leaving requests
unauthenticated exactly as before, so local unauthenticated dev keeps
working.
"""
state = current_cli_state()
if state is None or not hasattr(state, "get_sdk_context"):
return {}
try:
client_config = state.get_sdk_context().user.get_client_config()
except Exception:
logger.debug("Failed to resolve auth headers from CLI context", exc_info=True)
return {}
headers = client_config.get("default_headers") if isinstance(client_config, dict) else None
if isinstance(headers, dict):
return {str(key): str(value) for key, value in headers.items()}
return {}
Loading