eng 986: cli supply agent parameters with runs (#36)

* new: implemented run context

* fix: fixed typing bug

* new: showing run context if set (ENG-989)

* improved run context format

* docs: documented run context data

* fix: lint driven fix

Author: evilsocket

CLI.md

@@ -85,6 +85,9 @@ $ dreadnode agent deploy [OPTIONS]
 
 * `-m, --model TEXT`: The inference model to use for this run
 * `-d, --dir DIRECTORY`: The agent directory  [default: .]
+* `-e, --env-var TEXT`: Environment vars to override for this run (key=value)
+* `-p, --param TEXT`: Define custom parameters for this run (key = value in toml syntax or @filename.toml for multiple values)
+* `-c, --command TEXT`: Override the container command for this run.
 * `-s, --strike TEXT`: The strike to use for this run
 * `-w, --watch`: Watch the run status  [default: True]
 * `--help`: Show this message and exit.

README.md

@@ -165,6 +165,18 @@ dreadnode agent push
 # start a new run using the latest agent version.
 dreadnode agent deploy
 
+# start a new run using the latest agent version with custom environment variables
+dreadnode agent deploy --env-var TEST_ENV=test --env-var ANOTHER_ENV=another_value
+
+# start a new run using the latest agent version with custom parameters (using toml syntax)
+dreadnode agent deploy --param "foo = 'bar'" --param "baz = 123.0"
+
+# start a new run using the latest agent version with custom parameters from a toml file
+dreadnode agent deploy --param @parameters.toml
+
+# start a new run using the latest agent version and override the container command
+dreadnode agent deploy --command "echo 'Hello, world!'"
+
 # show the latest run of the currently active agent
 dreadnode agent latest
 

dreadnode_cli/agent/cli.py

@@ -4,6 +4,7 @@
 import time
 import typing as t
 
+import toml
 import typer
 from rich import box, print
 from rich.live import Live
@@ -318,6 +319,31 @@ def push(
     print(":tada: Agent pushed. use [bold]dreadnode agent deploy[/] to start a new run.")
 
 
+def prepare_run_context(
+    env_vars: list[str] | None, parameters: list[str] | None, command: str | None
+) -> Client.StrikeRunContext | None:
+    if not env_vars and not parameters and not command:
+        return None
+
+    context = Client.StrikeRunContext()
+
+    if env_vars:
+        context.environment = {env_var.split("=")[0]: env_var.split("=")[1] for env_var in env_vars}
+
+    if parameters:
+        context.parameters = {}
+        for param in parameters:
+            if param.startswith("@"):
+                context.parameters.update(toml.load(open(param[1:])))
+            else:
+                context.parameters.update(toml.loads(param))
+
+    if command:
+        context.command = command
+
+    return context
+
+
 @cli.command(help="Start a new run using the latest active agent version")
 @pretty_cli
 def deploy(
@@ -328,6 +354,22 @@ def deploy(
         pathlib.Path,
         typer.Option("--dir", "-d", help="The agent directory", file_okay=False, resolve_path=True),
     ] = pathlib.Path("."),
+    env_vars: t.Annotated[
+        list[str] | None,
+        typer.Option("--env-var", "-e", help="Environment vars to override for this run (key=value)"),
+    ] = None,
+    parameters: t.Annotated[
+        list[str] | None,
+        typer.Option(
+            "--param",
+            "-p",
+            help="Define custom parameters for this run (key = value in toml syntax or @filename.toml for multiple values)",
+        ),
+    ] = None,
+    command: t.Annotated[
+        str | None,
+        typer.Option("--command", "-c", help="Override the container command for this run."),
+    ] = None,
     strike: t.Annotated[str | None, typer.Option("--strike", "-s", help="The strike to use for this run")] = None,
     watch: t.Annotated[bool, typer.Option("--watch", "-w", help="Watch the run status")] = True,
 ) -> None:
@@ -346,6 +388,8 @@ def deploy(
     if strike is None:
         raise Exception("No strike specified, use -s/--strike or set the strike in the agent config")
 
+    context = prepare_run_context(env_vars, parameters, command)
+
     user_models = UserModels.read()
     user_model: Client.UserModel | None = None
 
@@ -376,7 +420,9 @@ def deploy(
                 f"Model '{model}' is not user-defined nor is it available in strike '{strike_response.name}'"
             )
 
-    run = client.start_strike_run(agent.latest_version.id, strike=strike, model=model, user_model=user_model)
+    run = client.start_strike_run(
+        agent.latest_version.id, strike=strike, model=model, user_model=user_model, context=context
+    )
     agent_config.add_run(run.id).write(directory)
     formatted = format_run(run, server_url=server_config.url)
 

dreadnode_cli/agent/format.py

@@ -291,6 +291,19 @@ def format_run(
     table.add_row("start", format_time(run.start))
     table.add_row("end", format_time(run.end))
 
+    if run.context and (run.context.environment or run.context.parameters or run.context.command):
+        table.add_row("", "")
+        if run.context.environment:
+            table.add_row(
+                "environment", " ".join(f"[magenta]{k}[/]=[yellow]{v}[/]" for k, v in run.context.environment.items())
+            )
+        if run.context.parameters:
+            table.add_row(
+                "parameters", " ".join(f"[magenta]{k}[/]=[yellow]{v}[/]" for k, v in run.context.parameters.items())
+            )
+        if run.context.command:
+            table.add_row("command", f"[bold][red]{run.context.command}[/red][/bold]")
+
     components: list[RenderableType] = [
         table,
         format_zones_verbose(run.zones, include_logs=include_logs) if verbose else format_zones_summary(run.zones),

dreadnode_cli/api.py

@@ -351,6 +351,11 @@ class StrikeRunZone(_StrikeRunZone):
         outputs: list["Client.StrikeRunOutput"]
         inferences: list[dict[str, t.Any]]
 
+    class StrikeRunContext(BaseModel):
+        environment: dict[str, str] | None = None
+        parameters: dict[str, t.Any] | None = None
+        command: str | None = None
+
     class _StrikeRun(BaseModel):
         id: UUID
         strike_id: UUID
@@ -364,6 +369,7 @@ class _StrikeRun(BaseModel):
         agent_name: str | None = None
         agent_revision: int
         agent_version: "Client.StrikeAgentVersion"
+        context: "Client.StrikeRunContext | None" = None
         status: "Client.StrikeRunStatus"
         start: datetime | None
         end: datetime | None
@@ -440,6 +446,7 @@ def start_strike_run(
         *,
         model: str | None = None,
         user_model: UserModel | None = None,
+        context: StrikeRunContext | None = None,
         strike: UUID | str | None = None,
     ) -> StrikeRunResponse:
         response = self.request(
@@ -450,6 +457,7 @@ def start_strike_run(
                 "model": model,
                 "user_model": user_model.model_dump(mode="json") if user_model else None,
                 "strike": str(strike) if strike else None,
+                "context": context.model_dump(mode="json") if context else None,
             },
         )
         return self.StrikeRunResponse(**response.json())

poetry.lock

@@ -20,6 +20,8 @@ httpx = "^0.27.2"
 ruamel-yaml = "^0.18.6"
 docker = "^7.1.0"
 pydantic-yaml = "^1.4.0"
+toml = "^0.10.2"
+types-toml = "^0.10.8.20240310"
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"