diff --git a/llama_stack/providers/inline/eval/meta_reference/eval.py b/llama_stack/providers/inline/eval/meta_reference/eval.py index b555c9f2a..408043db8 100644 --- a/llama_stack/providers/inline/eval/meta_reference/eval.py +++ b/llama_stack/providers/inline/eval/meta_reference/eval.py @@ -18,8 +18,8 @@ from llama_stack.providers.datatypes import EvalTasksProtocolPrivate from llama_stack.providers.utils.common.data_schema_validator import ( ColumnName, - DataSchemaValidatorMixin, get_valid_schemas, + validate_dataset_schema, ) from llama_stack.providers.utils.kvstore import kvstore_impl @@ -31,7 +31,10 @@ from .config import MetaReferenceEvalConfig EVAL_TASKS_PREFIX = "eval_tasks:" -class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate, DataSchemaValidatorMixin): +class MetaReferenceEvalImpl( + Eval, + EvalTasksProtocolPrivate, +): def __init__( self, config: MetaReferenceEvalConfig, @@ -85,7 +88,7 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate, DataSchemaValidatorM candidate = task_config.eval_candidate scoring_functions = task_def.scoring_functions dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id) - self.validate_dataset_schema( + validate_dataset_schema( dataset_def.dataset_schema, get_valid_schemas(Api.eval.value) ) all_rows = await self.datasetio_api.get_rows_paginated( diff --git a/llama_stack/providers/inline/scoring/basic/scoring.py b/llama_stack/providers/inline/scoring/basic/scoring.py index f612abda4..621e217bb 100644 --- a/llama_stack/providers/inline/scoring/basic/scoring.py +++ b/llama_stack/providers/inline/scoring/basic/scoring.py @@ -18,8 +18,8 @@ from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams from llama_stack.distribution.datatypes import Api from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate from llama_stack.providers.utils.common.data_schema_validator import ( - DataSchemaValidatorMixin, get_valid_schemas, + validate_dataset_schema, ) from .config import BasicScoringConfig from .scoring_fn.equality_scoring_fn import EqualityScoringFn @@ -30,7 +30,8 @@ FIXED_FNS = [EqualityScoringFn, SubsetOfScoringFn, RegexParserScoringFn] class BasicScoringImpl( - Scoring, ScoringFunctionsProtocolPrivate, DataSchemaValidatorMixin + Scoring, + ScoringFunctionsProtocolPrivate, ): def __init__( self, @@ -75,7 +76,7 @@ class BasicScoringImpl( save_results_dataset: bool = False, ) -> ScoreBatchResponse: dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id) - self.validate_dataset_schema( + validate_dataset_schema( dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value) ) diff --git a/llama_stack/providers/inline/scoring/braintrust/braintrust.py b/llama_stack/providers/inline/scoring/braintrust/braintrust.py index 4282ef6ec..6cfc94df5 100644 --- a/llama_stack/providers/inline/scoring/braintrust/braintrust.py +++ b/llama_stack/providers/inline/scoring/braintrust/braintrust.py @@ -35,8 +35,9 @@ from llama_stack.distribution.datatypes import Api from llama_stack.distribution.request_headers import NeedsRequestProviderData from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate from llama_stack.providers.utils.common.data_schema_validator import ( - DataSchemaValidatorMixin, get_valid_schemas, + validate_dataset_schema, + validate_row_schema, ) from llama_stack.providers.utils.scoring.aggregation_utils import aggregate_metrics @@ -111,7 +112,6 @@ class BraintrustScoringImpl( Scoring, ScoringFunctionsProtocolPrivate, NeedsRequestProviderData, - DataSchemaValidatorMixin, ): def __init__( self, @@ -171,7 +171,7 @@ class BraintrustScoringImpl( await self.set_api_key() dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id) - self.validate_dataset_schema( + validate_dataset_schema( dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value) ) @@ -194,7 +194,7 @@ class BraintrustScoringImpl( async def score_row( self, input_row: Dict[str, Any], scoring_fn_identifier: Optional[str] = None ) -> ScoringResultRow: - self.validate_row_schema(input_row, get_valid_schemas(Api.scoring.value)) + validate_row_schema(input_row, get_valid_schemas(Api.scoring.value)) await self.set_api_key() assert scoring_fn_identifier is not None, "scoring_fn_identifier cannot be None" expected_answer = input_row["expected_answer"] diff --git a/llama_stack/providers/inline/scoring/llm_as_judge/scoring.py b/llama_stack/providers/inline/scoring/llm_as_judge/scoring.py index 305c13665..a11d0734c 100644 --- a/llama_stack/providers/inline/scoring/llm_as_judge/scoring.py +++ b/llama_stack/providers/inline/scoring/llm_as_judge/scoring.py @@ -19,8 +19,8 @@ from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams from llama_stack.distribution.datatypes import Api from llama_stack.providers.datatypes import ScoringFunctionsProtocolPrivate from llama_stack.providers.utils.common.data_schema_validator import ( - DataSchemaValidatorMixin, get_valid_schemas, + validate_dataset_schema, ) from .config import LlmAsJudgeScoringConfig @@ -31,7 +31,8 @@ LLM_JUDGE_FNS = [LlmAsJudgeScoringFn] class LlmAsJudgeScoringImpl( - Scoring, ScoringFunctionsProtocolPrivate, DataSchemaValidatorMixin + Scoring, + ScoringFunctionsProtocolPrivate, ): def __init__( self, @@ -79,7 +80,7 @@ class LlmAsJudgeScoringImpl( save_results_dataset: bool = False, ) -> ScoreBatchResponse: dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id) - self.validate_dataset_schema( + validate_dataset_schema( dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value) ) diff --git a/llama_stack/providers/utils/common/data_schema_validator.py b/llama_stack/providers/utils/common/data_schema_validator.py index d9e6cb6b5..af58a4592 100644 --- a/llama_stack/providers/utils/common/data_schema_validator.py +++ b/llama_stack/providers/utils/common/data_schema_validator.py @@ -62,26 +62,24 @@ def get_valid_schemas(api_str: str): raise ValueError(f"Invalid API string: {api_str}") -class DataSchemaValidatorMixin: - def validate_dataset_schema( - self, - dataset_schema: Dict[str, Any], - expected_schemas: List[Dict[str, Any]], - ): - if dataset_schema not in expected_schemas: - raise ValueError( - f"Dataset {dataset_schema} does not have a correct input schema in {expected_schemas}" - ) - - def validate_row_schema( - self, - input_row: Dict[str, Any], - expected_schemas: List[Dict[str, Any]], - ): - for schema in expected_schemas: - if all(key in input_row for key in schema): - return - +def validate_dataset_schema( + dataset_schema: Dict[str, Any], + expected_schemas: List[Dict[str, Any]], +): + if dataset_schema not in expected_schemas: raise ValueError( - f"Input row {input_row} does not match any of the expected schemas in {expected_schemas}" + f"Dataset {dataset_schema} does not have a correct input schema in {expected_schemas}" ) + + +def validate_row_schema( + input_row: Dict[str, Any], + expected_schemas: List[Dict[str, Any]], +): + for schema in expected_schemas: + if all(key in input_row for key in schema): + return + + raise ValueError( + f"Input row {input_row} does not match any of the expected schemas in {expected_schemas}" + )