ENG-1110: Onboard Utilities

[thiago-aixplain]

from aixplain.factories import ModelFactory
from aixplain.modules.model.utility_model import UtilityModelInput
from aixplain.enums import DataType

# set the inputs of the python function
inputs = [
    UtilityModelInput(name="inputA", description="input A is the only input", type=DataType.TEXT),
]
# set the python code, which may have a 'main' method with the set input parameters
# The code may be set in a string, a URL or a local path
code = "def main(inputA):\n\treturn inputA"

output_description = "An example is 'test'"

# Onboard the utility model
utility_model = ModelFactory.create_utility_model(
    name="test_script",
    description="This is a test script",
    inputs=inputs,
    code=code,
    output_description=output_description,
)

# Run the Utility as any other model
response = utility_model.run(data={"inputA": "test"})
# ModelResponse(status=SUCCESS, data='test', details={'status': 'Success', 'data': 'test', 'stdout': '', 'stderr': ''}, completed=True, used_credits=1.3470000000000001e-05, run_time=1.402)
assert response.status == "SUCCESS"
assert response.data == "test"

# Update the code as many times as you want
utility_model.code = "def main(inputA):\n\treturn 5"
utility_model.update()
response = utility_model.run(data={"inputA": "test"})
assert response.status == "SUCCESS"
assert str(response.data) == "5"

# Delete the model whenever you want
utility_model.delete()