diff --git a/api/db/db_models.py b/api/db/db_models.py index 04bef2e835b..fc459389fbc 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -788,6 +788,7 @@ class Task(DataBaseModel): null=True, help_text="process message", default="") + retry_count = IntegerField(default=0) class Dialog(DataBaseModel): @@ -982,3 +983,10 @@ def migrate_db(): DB.execute_sql('ALTER TABLE llm ADD PRIMARY KEY (llm_name,fid);') except Exception as e: pass + try: + migrate( + migrator.add_column('task', 'retry_count', IntegerField(default=0)) + ) + except Exception as e: + pass + diff --git a/api/db/services/task_service.py b/api/db/services/task_service.py index 545bd81be51..40d80f7e408 100644 --- a/api/db/services/task_service.py +++ b/api/db/services/task_service.py @@ -64,9 +64,20 @@ def get_tasks(cls, task_id): docs = list(docs.dicts()) if not docs: return [] - cls.model.update(progress_msg=cls.model.progress_msg + "\n" + "Task has been received.", - progress=random.random() / 10.).where( + msg = "\nTask has been received." + prog = random.random() / 10. + if docs[0]["retry_count"] >= 3: + msg = "\nERROR: Task is abandoned after 3 times attempts." + prog = -1 + + cls.model.update(progress_msg=cls.model.progress_msg + msg, + progress=prog, + retry_count=docs[0]["retry_count"]+1 + ).where( cls.model.id == docs[0]["id"]).execute() + + if docs[0]["retry_count"] >= 3: return [] + return docs @classmethod