def start_training(bot: str, user: str):
""" Prevents training of the bot if the training session is in progress otherwise start training """
exception = None
model_file = None
training_status = None
ModelProcessor.set_training_status(
bot=bot,
user=user,
status=MODEL_TRAINING_STATUS.INPROGRESS.value,
)
try:
model_file = train_model_for_bot(bot)
training_status = MODEL_TRAINING_STATUS.DONE.value
except Exception as e:
logging.exception(e)
training_status = MODEL_TRAINING_STATUS.FAIL.value
exception = str(e)
raise AppException(exception)
finally:
ModelProcessor.set_training_status(
bot=bot,
user=user,
status=training_status,
model_path=model_file,
exception=exception,
)
AgentProcessor.reload(bot)
return model_file
def test_set_training_status_Fail(self,
test_set_training_status_inprogress):
assert test_set_training_status_inprogress.__len__() == 1
assert test_set_training_status_inprogress.first().bot == "tests"
assert test_set_training_status_inprogress.first().user == "testUser"
assert test_set_training_status_inprogress.first(
).status == "Inprogress"
training_status_inprogress_id = test_set_training_status_inprogress.first(
).id
ModelProcessor.set_training_status(
bot="tests",
user="******",
status="Fail",
end_timestamp=datetime.utcnow(),
model_path=None,
exception="exception occurred while training model.")
model_training = ModelTraining.objects(bot="tests", status="Fail")
assert model_training.__len__() == 1
assert model_training.first().id == training_status_inprogress_id
assert model_training.first().bot == "tests"
assert model_training.first().user == "testUser"
assert model_training.first().status == "Fail"
assert model_training.first().model_path is None
assert model_training.first(
).exception == "exception occurred while training model."
assert ModelTraining.objects(bot="tests",
status="Inprogress").__len__() == 0
def test_set_training_status_Done(self,
test_set_training_status_inprogress):
assert test_set_training_status_inprogress.__len__() == 1
assert test_set_training_status_inprogress.first().bot == "tests"
assert test_set_training_status_inprogress.first().user == "testUser"
assert test_set_training_status_inprogress.first(
).status == "Inprogress"
training_status_inprogress_id = test_set_training_status_inprogress.first(
).id
ModelProcessor.set_training_status(bot="tests",
user="******",
status="Done",
end_timestamp=datetime.utcnow(),
model_path="model_path")
model_training = ModelTraining.objects(bot="tests", status="Done")
assert model_training.count() == 2
assert model_training[1].id == training_status_inprogress_id
assert model_training[1].bot == "tests"
assert model_training[1].user == "testUser"
assert model_training[1].status == "Done"
assert model_training[1].model_path == "model_path"
assert ModelTraining.objects(bot="tests",
status="Inprogress").__len__() == 0
def test_set_training_status_inprogress(self):
ModelProcessor.set_training_status("tests", "testUser", "Inprogress",
datetime.utcnow())
model_training = ModelTraining.objects(bot="tests",
status="Inprogress")
return model_training
async def train(
background_tasks: BackgroundTasks,
current_user: User = Depends(auth.get_current_user),
):
""" This is used to train the chatbot """
ModelProcessor.is_training_inprogress(current_user.get_bot())
ModelProcessor.is_daily_training_limit_exceeded(current_user.get_bot())
background_tasks.add_task(start_training, current_user.get_bot(),
current_user.get_user())
return {"message": "Model training started."}
def test_is_daily_training_limit_exceeded_exception(self, monkeypatch):
monkeypatch.setitem(Utility.environment,
"MODEL_TRAINING_LIMIT_PER_DAY", 3)
with pytest.raises(AppException) as exp:
assert ModelProcessor.is_daily_training_limit_exceeded("tests")
assert str(exp.value) == "Daily model training limit exceeded."
def test_is_training_inprogress_True(self,
test_set_training_status_inprogress):
assert test_set_training_status_inprogress.__len__() == 1
assert test_set_training_status_inprogress.first().bot == "tests"
assert test_set_training_status_inprogress.first().user == "testUser"
assert test_set_training_status_inprogress.first(
).status == "Inprogress"
actual_response = ModelProcessor.is_training_inprogress("tests", False)
assert actual_response is True
async def get_model_training_history(
current_user: User = Depends(auth.get_current_user), ):
training_history = list(
ModelProcessor.get_training_history(current_user.get_bot()))
return {"data": {"training_history": training_history}}
def test_get_training_history(self):
actual_response = ModelProcessor.get_training_history("tests")
assert actual_response
def test_is_daily_training_limit_exceeded_True(self, monkeypatch):
monkeypatch.setitem(Utility.environment,
"MODEL_TRAINING_LIMIT_PER_DAY", 3)
actual_response = ModelProcessor.is_daily_training_limit_exceeded(
"tests", False)
assert actual_response is True
def test_is_training_inprogress_exception(
self, test_set_training_status_inprogress):
with pytest.raises(AppException) as exp:
assert ModelProcessor.is_training_inprogress("tests")
assert str(exp.value) == "Previous model training in progress."
def test_is_training_inprogress_False(self):
actual_response = ModelProcessor.is_training_inprogress("tests")
assert actual_response is False