Skip to content

datarobot_genai.dragent.cli.commands

commands

DRAgentCommandGroup

Bases: StartCommandGroup

NAT StartCommandGroup filtered to dragent frontends with friendly aliases.

Source code in datarobot_genai/dragent/cli/commands.py
class DRAgentCommandGroup(StartCommandGroup):
    """NAT StartCommandGroup filtered to dragent frontends with friendly aliases."""

    def _build_params(self, front_end: RegisteredFrontEndInfo) -> list[click.Parameter]:
        params = super()._build_params(front_end)
        for param in params:
            if not isinstance(param, click.Option):
                continue
            if "--config_file" in param.opts:
                param.required = False
                param.envvar = "DRAGENT_CONFIG_FILE"
            elif "--port" in param.opts:
                param.envvar = "AGENT_PORT"
            elif "--user_prompt" in param.opts:
                param.opts.insert(0, "--input")
            elif "--input_file" in param.opts:
                param.opts.insert(0, "--file")
        return params

    def invoke_subcommand(  # type: ignore[override]
        self,
        ctx: click.Context,
        cmd_name: str,
        config_file: Path | None,
        override: tuple[tuple[str, str], ...],
        **kwargs: object,
    ) -> int | None:
        if config_file is None:
            raise click.ClickException(
                "No config file provided. "
                "Pass --config_file <path> or set the DRAGENT_CONFIG_FILE env var."
            )
        publish_dragent_config_file_env(config_file)
        _bridge_pulumi_otel_env()
        return super().invoke_subcommand(ctx, cmd_name, config_file, override, **kwargs)

    def _load_commands(self) -> dict[str, click.Command]:
        # Depends on StartCommandGroup caching into self._commands (nvidia-nat 1.4.1).
        # If NAT renames/removes this attribute, _load_commands will need updating.
        commands: dict[str, click.Command] | None = self._commands  # type: ignore[has-type,assignment]
        if commands is not None:
            return commands

        # Reset cache so the parent actually populates it, then filter + alias.
        self._commands = None  # type: ignore[assignment]
        loaded = super()._load_commands()

        filtered: dict[str, click.Command] = {}
        for original_name, meta in _FRONTEND_COMMANDS.items():
            if original_name not in loaded:
                raise RuntimeError(
                    f"Frontend '{original_name}' not registered. "
                    f"Ensure nat.front_ends entry points are installed."
                )
            cmd = loaded[original_name]
            cmd.name = meta["alias"]
            cmd.help = meta["help"]
            filtered[meta["alias"]] = cmd

        filtered["query"] = query_command
        self._commands = filtered  # type: ignore[assignment]
        return filtered