forked from phoenix-oss/llama-stack-mirror
34ab7a3b6c
Lint check in main branch is failing. This fixes the lint check after we moved to ruff in https://github.com/meta-llama/llama-stack/pull/921. We need to move to a `ruff.toml` file as well as fixing and ignoring some additional checks. Signed-off-by: Yuan Tang <terrytangyuan@gmail.com>
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
# Copyright (c) Meta Platforms, IAny, nc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
from llama_stack.apis.common.type_system import (
|
|
ChatCompletionInputType,
|
|
DialogType,
|
|
StringType,
|
|
)
|
|
from llama_stack.apis.datasets import Datasets
|
|
from llama_stack.providers.utils.common.data_schema_validator import (
|
|
ColumnName,
|
|
validate_dataset_schema,
|
|
)
|
|
|
|
EXPECTED_DATASET_SCHEMA = {
|
|
"instruct": [
|
|
{
|
|
ColumnName.chat_completion_input.value: ChatCompletionInputType(),
|
|
ColumnName.expected_answer.value: StringType(),
|
|
}
|
|
],
|
|
"dialog": [
|
|
{
|
|
ColumnName.dialog.value: DialogType(),
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
async def validate_input_dataset_schema(
|
|
datasets_api: Datasets,
|
|
dataset_id: str,
|
|
dataset_type: str,
|
|
) -> None:
|
|
dataset_def = await datasets_api.get_dataset(dataset_id=dataset_id)
|
|
if not dataset_def.dataset_schema or len(dataset_def.dataset_schema) == 0:
|
|
raise ValueError(f"Dataset {dataset_id} does not have a schema defined.")
|
|
|
|
if dataset_type not in EXPECTED_DATASET_SCHEMA:
|
|
raise ValueError(f"Dataset type {dataset_type} is not supported.")
|
|
|
|
validate_dataset_schema(dataset_def.dataset_schema, EXPECTED_DATASET_SCHEMA[dataset_type])
|