diff --git a/maxim/dataset/dataset.py b/maxim/dataset/dataset.py index dc07038..2878cbb 100644 --- a/maxim/dataset/dataset.py +++ b/maxim/dataset/dataset.py @@ -2,10 +2,10 @@ from typing import Any, Dict, Optional from ..models.dataset import ( - ContextToEvaluateColumn, + CONTEXT_TO_EVALUATE_COLUMN, DataStructure, - ExpectedOutputColumn, - InputColumn, + EXPECTED_OUTPUT_COLUMN, + INPUT_COLUMN, ) @@ -20,7 +20,7 @@ def sanitize_data_structure(data_structure: Optional[DataStructure]) -> None: encountered_context_to_evaluate = False if data_structure: for value in data_structure.values(): - if value == InputColumn: + if value == INPUT_COLUMN: if encountered_input: raise Exception( "Data structure contains more than one input", @@ -28,7 +28,7 @@ def sanitize_data_structure(data_structure: Optional[DataStructure]) -> None: ) else: encountered_input = True - elif value == ExpectedOutputColumn: + elif value == EXPECTED_OUTPUT_COLUMN: if encountered_expected_output: raise Exception( "Data structure contains more than one expectedOutput", @@ -36,7 +36,7 @@ def sanitize_data_structure(data_structure: Optional[DataStructure]) -> None: ) else: encountered_expected_output = True - elif value == ContextToEvaluateColumn: + elif value == CONTEXT_TO_EVALUATE_COLUMN: if encountered_context_to_evaluate: raise Exception( "Data structure contains more than one contextToEvaluate", diff --git a/maxim/models/dataset.py b/maxim/models/dataset.py index f3bc112..40a0668 100644 --- a/maxim/models/dataset.py +++ b/maxim/models/dataset.py @@ -63,6 +63,16 @@ def to_json(self): return return_dict +# Column name constants used throughout the dataset utilities. +INPUT_COLUMN: str = "INPUT" +EXPECTED_OUTPUT_COLUMN: str = "EXPECTED_OUTPUT" +CONTEXT_TO_EVALUATE_COLUMN: str = "CONTEXT_TO_EVALUATE" +VARIABLE_COLUMN: str = "VARIABLE" +NULLABLE_VARIABLE_COLUMN: str = "NULLABLE_VARIABLE" +OUTPUT_COLUMN: str = "OUTPUT" + +# Type aliases for column names. These are used for type checking but also +# need to be concrete values at runtime for equality comparisons. InputColumn = Literal["INPUT"] ExpectedOutputColumn = Literal["EXPECTED_OUTPUT"] ContextToEvaluateColumn = Literal["CONTEXT_TO_EVALUATE"]