test_prediction.py

from prediction import ModelPrediction
import utilities as us


def main():
    model = ModelPrediction()
    response = model.make_prediction(
        question='What is the surnamname of Simone',
        db_schema='CREATE TABLE Player(Name, Age)\nINSERT INTO Player(Name, Age) VALUES (Simone, Papicchio)', 
        model_name="gpt-3.5",
        task='QA'
    )
    print(response)  # dict[response, response_parsed, cost]


if __name__ == "__main__":
    main()