From 8df708b91e9bbc51f96f50e1456b27abd059ad60 Mon Sep 17 00:00:00 2001 From: Josh Salomon Date: Mon, 14 Apr 2025 13:50:51 +0300 Subject: [PATCH] CLI fix: Fix CLI for datasets register Fixed the `llama-stack-cli datasets register` to match the modified API (some mandatory parameters were removed and a new one was added). Added the dataset purpose to the output of the `llama-stack-cli datasets list` command. Signed-off-by: Josh Salomon --- .../lib/cli/datasets/list.py | 2 +- .../lib/cli/datasets/register.py | 26 +++++++------------ 2 files changed, 11 insertions(+), 17 deletions(-) diff --git a/src/llama_stack_client/lib/cli/datasets/list.py b/src/llama_stack_client/lib/cli/datasets/list.py index 8e21ad0d..61d625c9 100644 --- a/src/llama_stack_client/lib/cli/datasets/list.py +++ b/src/llama_stack_client/lib/cli/datasets/list.py @@ -19,7 +19,7 @@ def list_datasets(ctx): """Show available datasets on distribution endpoint""" client = ctx.obj["client"] console = Console() - headers = ["identifier", "provider_id", "metadata", "type"] + headers = ["identifier", "provider_id", "metadata", "type", "purpose"] datasets_list_response = client.datasets.list() if datasets_list_response: diff --git a/src/llama_stack_client/lib/cli/datasets/register.py b/src/llama_stack_client/lib/cli/datasets/register.py index 54c164d1..d990e30c 100644 --- a/src/llama_stack_client/lib/cli/datasets/register.py +++ b/src/llama_stack_client/lib/cli/datasets/register.py @@ -7,7 +7,7 @@ import json import mimetypes import os -from typing import Optional +from typing import Optional, Literal import click import yaml @@ -32,34 +32,30 @@ def data_url_from_file(file_path: str) -> str: @click.command("register") @click.help_option("-h", "--help") @click.option("--dataset-id", required=True, help="Id of the dataset") -@click.option("--provider-id", help="Provider ID for the dataset", default=None) -@click.option("--provider-dataset-id", help="Provider's dataset ID", default=None) +@click.option( + "--purpose", + type=click.Choice(["post-training/messages", "eval/question-answer", "eval/messages-answer"]), + help="Purpose of the dataset", + required=True, +) @click.option("--metadata", type=str, help="Metadata of the dataset") @click.option("--url", type=str, help="URL of the dataset", required=False) @click.option( "--dataset-path", required=False, help="Local file path to the dataset. If specified, upload dataset via URL" ) -@click.option("--schema", type=str, help="JSON schema of the dataset", required=True) @click.pass_context @handle_client_errors("register dataset") def register( ctx, dataset_id: str, - provider_id: Optional[str], - provider_dataset_id: Optional[str], + purpose: Literal["post-training/messages", "eval/question-answer", "eval/messages-answer"], metadata: Optional[str], url: Optional[str], dataset_path: Optional[str], - schema: str, ): """Create a new dataset""" client = ctx.obj["client"] - try: - dataset_schema = json.loads(schema) - except json.JSONDecodeError as err: - raise click.BadParameter("Schema must be valid JSON") from err - if metadata: try: metadata = json.loads(metadata) @@ -74,11 +70,9 @@ def register( response = client.datasets.register( dataset_id=dataset_id, - dataset_schema=dataset_schema, - url={"uri": url}, - provider_id=provider_id, - provider_dataset_id=provider_dataset_id, + source={"uri": url}, metadata=metadata, + purpose=purpose, ) if response: click.echo(yaml.dump(response.dict()))