diff --git a/aixplain/modules/model/utils.py b/aixplain/modules/model/utils.py index 68131c96..f3691928 100644 --- a/aixplain/modules/model/utils.py +++ b/aixplain/modules/model/utils.py @@ -120,13 +120,19 @@ def parse_code(code: Union[Text, Callable]) -> Tuple[Text, List, Text]: if input_type in ["int", "float"]: input_type = "number" - inputs.append(UtilityModelInput(name=input_name, type=DataType.NUMBER, description="")) + inputs.append( + UtilityModelInput(name=input_name, type=DataType.NUMBER, description=f"The {input_name} input is a number") + ) elif input_type == "bool": input_type = "boolean" - inputs.append(UtilityModelInput(name=input_name, type=DataType.BOOLEAN, description="")) + inputs.append( + UtilityModelInput(name=input_name, type=DataType.BOOLEAN, description=f"The {input_name} input is a boolean") + ) elif input_type == "str": input_type = "text" - inputs.append(UtilityModelInput(name=input_name, type=DataType.TEXT, description="")) + inputs.append( + UtilityModelInput(name=input_name, type=DataType.TEXT, description=f"The {input_name} input is a text") + ) else: raise Exception(f"Utility Model Error: Unsupported input type: {input_type}") diff --git a/tests/unit/utility_test.py b/tests/unit/utility_test.py index 265e31b3..2f3ba6a1 100644 --- a/tests/unit/utility_test.py +++ b/tests/unit/utility_test.py @@ -24,7 +24,9 @@ def test_utility_model(): assert utility_model.name == "utility_model_test" assert utility_model.description == "utility_model_test" assert utility_model.code == "utility_model_test" - assert utility_model.inputs == [UtilityModelInput(name="originCode", description="", type=DataType.TEXT)] + assert utility_model.inputs == [ + UtilityModelInput(name="originCode", description="The originCode input is a text", type=DataType.TEXT) + ] assert utility_model.output_examples == "output_description" @@ -136,7 +138,9 @@ def test_parse_code(): with patch("aixplain.factories.file_factory.FileFactory.upload", return_value="code_link"): code = "def main(originCode: str) -> str:\n return originCode" code_link, inputs, description = parse_code(code) - assert inputs == [UtilityModelInput(name="originCode", description="", type=DataType.TEXT)] + assert inputs == [ + UtilityModelInput(name="originCode", description="The originCode input is a text", type=DataType.TEXT) + ] assert description == "" assert code_link == "code_link" @@ -152,8 +156,8 @@ def main(a: int, b: int): code = main code_link, inputs, description = parse_code(code) assert inputs == [ - UtilityModelInput(name="a", description="", type=DataType.NUMBER), - UtilityModelInput(name="b", description="", type=DataType.NUMBER), + UtilityModelInput(name="a", description="The a input is a number", type=DataType.NUMBER), + UtilityModelInput(name="b", description="The b input is a number", type=DataType.NUMBER), ] assert description == "This function adds two numbers" assert code_link == "code_link"