diff --git a/tests/cli/test_cli_artifacts.py b/tests/cli/test_cli_artifacts.py new file mode 100644 index 0000000..d7afb7f --- /dev/null +++ b/tests/cli/test_cli_artifacts.py @@ -0,0 +1,94 @@ +"""Test artifacts subcommand.""" + +from __future__ import annotations + +from vlmrun.cli.cli import app + + +def test_get_artifact_by_session_id(runner, mock_client, config_file): + """Test getting an artifact with --session-id.""" + result = runner.invoke( + app, + [ + "artifacts", + "get", + "test-object-456", + "--session-id", + "550e8400-e29b-41d4-a716-446655440000", + "--raw", + ], + ) + assert result.exit_code == 0 + assert "bytes" in result.stdout.lower() or "Received" in result.stdout + + +def test_get_artifact_by_execution_id(runner, mock_client, config_file): + """Test getting an artifact with --execution-id.""" + result = runner.invoke( + app, + [ + "artifacts", + "get", + "test-object-456", + "--execution-id", + "exec-001", + "--raw", + ], + ) + assert result.exit_code == 0 + + +def test_get_artifact_to_output_file(runner, mock_client, config_file, tmp_path): + """Test saving an artifact to a specific output path.""" + out = tmp_path / "artifact.bin" + result = runner.invoke( + app, + [ + "artifacts", + "get", + "test-object-456", + "--session-id", + "550e8400-e29b-41d4-a716-446655440000", + "--raw", + "--output", + str(out), + ], + ) + assert result.exit_code == 0 + assert out.exists() + assert out.read_bytes() == b"mock artifact content" + + +def test_get_artifact_missing_ids(runner, mock_client, config_file): + """Test that omitting both --session-id and --execution-id fails.""" + result = runner.invoke( + app, + ["artifacts", "get", "test-object-456"], + ) + assert result.exit_code == 1 + assert "session-id" in result.stdout.lower() or "Error" in result.stdout + + +def test_get_artifact_both_ids(runner, mock_client, config_file): + """Test that providing both --session-id and --execution-id fails.""" + result = runner.invoke( + app, + [ + "artifacts", + "get", + "test-object-456", + "--session-id", + "sess-001", + "--execution-id", + "exec-001", + ], + ) + assert result.exit_code == 1 + assert "only one" in result.stdout.lower() or "Error" in result.stdout + + +def test_artifacts_no_args(runner, mock_client, config_file): + """Test that running `vlmrun artifacts` with no args shows help.""" + result = runner.invoke(app, ["artifacts"]) + # Typer returns exit code 0 or 2 for no_args_is_help + assert result.exit_code in (0, 2) diff --git a/vlmrun/cli/_cli/artifacts.py b/vlmrun/cli/_cli/artifacts.py new file mode 100644 index 0000000..c5e09b5 --- /dev/null +++ b/vlmrun/cli/_cli/artifacts.py @@ -0,0 +1,168 @@ +"""Artifacts CLI commands — retrieve artifacts by session or execution ID.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import typer +from rich.console import Console +from rich.panel import Panel +from rich.text import Text + +if TYPE_CHECKING: + from vlmrun.client import VLMRun + +app = typer.Typer( + help="Retrieve artifacts generated by chat completions or agent executions.", + add_completion=False, + no_args_is_help=True, +) + +console = Console() + + +@app.command() +def get( + ctx: typer.Context, + object_id: str = typer.Argument( + ..., + help="Object ID for the artifact (format: _<6-hex-chars>, e.g. img_a1b2c3)", + ), + session_id: str | None = typer.Option( + None, + "--session-id", + "-s", + help="Session ID from chat completions (mutually exclusive with --execution-id)", + ), + execution_id: str | None = typer.Option( + None, + "--execution-id", + "-e", + help="Execution ID from agent executions (mutually exclusive with --session-id)", + ), + output: Path | None = typer.Option( + None, + "--output", + "-o", + help="Output file path. If omitted, the artifact is saved to the default cache directory.", + ), + raw: bool = typer.Option( + False, + "--raw", + help="Return raw bytes instead of converting to the appropriate type.", + ), +) -> None: + """Download an artifact by object ID and session/execution ID.""" + client: VLMRun = ctx.obj + + if session_id is None and execution_id is None: + console.print( + "[red bold]Error:[/] Either --session-id or --execution-id is required." + ) + raise typer.Exit(1) + + if session_id is not None and execution_id is not None: + console.print( + "[red bold]Error:[/] Only one of --session-id or --execution-id is allowed, not both." + ) + raise typer.Exit(1) + + try: + result = client.artifacts.get( + object_id=object_id, + session_id=session_id, + execution_id=execution_id, + raw_response=raw, + ) + except Exception as exc: + console.print(f"[red bold]Error:[/] {exc}") + raise typer.Exit(1) from exc + + id_label = ( + f"session_id={session_id}" if session_id else f"execution_id={execution_id}" + ) + + if output is not None: + _write_to_output(result, output) + console.print( + Panel( + Text.from_markup( + f"[bold]Object ID:[/] {object_id}\n" + f"[bold]{id_label}[/]\n" + f"[bold]Saved to:[/] [green]{output}[/]" + ), + title="[bold]Artifact Downloaded[/bold]", + title_align="left", + border_style="blue", + padding=(0, 1), + ) + ) + return + + # No explicit output — display what the SDK returned + if isinstance(result, bytes): + if output is None: + console.print( + f"[yellow]Received {len(result)} bytes. " + "Use --output to save to a file.[/]" + ) + elif isinstance(result, Path): + console.print( + Panel( + Text.from_markup( + f"[bold]Object ID:[/] {object_id}\n" + f"[bold]{id_label}[/]\n" + f"[bold]Saved to:[/] [green]{result}[/]" + ), + title="[bold]Artifact Downloaded[/bold]", + title_align="left", + border_style="blue", + padding=(0, 1), + ) + ) + else: + # PIL Image or other types — save to cache and report + try: + from vlmrun.constants import VLMRUN_ARTIFACTS_CACHE_DIR + + sess_id: str = session_id or execution_id # type: ignore[assignment] + cache_dir = VLMRUN_ARTIFACTS_CACHE_DIR / sess_id + cache_dir.mkdir(parents=True, exist_ok=True) + dest = cache_dir / f"{object_id}.jpg" + result.save(str(dest)) + console.print( + Panel( + Text.from_markup( + f"[bold]Object ID:[/] {object_id}\n" + f"[bold]{id_label}[/]\n" + f"[bold]Saved to:[/] [green]{dest}[/]" + ), + title="[bold]Artifact Downloaded[/bold]", + title_align="left", + border_style="blue", + padding=(0, 1), + ) + ) + except Exception: + console.print(f"[green]Artifact retrieved:[/] {type(result).__name__}") + + +def _write_to_output(result: object, output: Path) -> None: + """Write an artifact result to the given output path. + + Args: + result: The artifact content (bytes, Path, or PIL Image). + output: Destination file path. + """ + output.parent.mkdir(parents=True, exist_ok=True) + + if isinstance(result, bytes): + output.write_bytes(result) + elif isinstance(result, Path): + import shutil + + shutil.copy2(result, output) + else: + # Assume PIL Image + result.save(str(output)) diff --git a/vlmrun/cli/cli.py b/vlmrun/cli/cli.py index 1bc26e0..5594c0b 100644 --- a/vlmrun/cli/cli.py +++ b/vlmrun/cli/cli.py @@ -12,6 +12,7 @@ from rich.text import Text from vlmrun.client import VLMRun +from vlmrun.cli._cli.artifacts import app as artifacts_app from vlmrun.cli._cli.chat import CHAT_HELP, chat from vlmrun.cli._cli.config import app as config_app, resolve_config from vlmrun.cli._cli.execute import EXECUTE_HELP, execute @@ -129,6 +130,7 @@ def main( app.add_typer(hub_app, name="hub") app.add_typer(models_app, name="models") app.add_typer(skills_app, name="skills") +app.add_typer(artifacts_app, name="artifacts") app.add_typer(config_app, name="config") if __name__ == "__main__":