Skip to content

Commit

Permalink
add retry count to task (#2152)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?


### Type of change

- [x] New Feature (non-breaking change which adds functionality)
  • Loading branch information
KevinHuSh authored Aug 29, 2024
1 parent 06abef6 commit 212bb8e
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
8 changes: 8 additions & 0 deletions api/db/db_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -788,6 +788,7 @@ class Task(DataBaseModel):
null=True,
help_text="process message",
default="")
retry_count = IntegerField(default=0)


class Dialog(DataBaseModel):
Expand Down Expand Up @@ -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

15 changes: 13 additions & 2 deletions api/db/services/task_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 212bb8e

Please sign in to comment.