From c529d066ef61f642d0c65a1d63b9a1d3e9a81615 Mon Sep 17 00:00:00 2001
From: Avram Tudor <tudor.potecaru@8x8.com>
Date: Fri, 11 Oct 2024 11:27:39 +0300
Subject: [PATCH] feat: add priority flag (#106)

* jobs submitted with priority high will be added to the start of the redis queue

Co-authored-by: Avram Tudor <tudor.avram@8x8.com>
---
 run.sh                                    |  4 +++-
 skynet/modules/ttt/summaries/jobs.py      |  8 ++++++--
 skynet/modules/ttt/summaries/v1/models.py | 15 ++++++++++++++-
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/run.sh b/run.sh
index 30c11ee..0fffd92 100755
--- a/run.sh
+++ b/run.sh
@@ -3,7 +3,9 @@
 if nvcc --version
 then
   export CUDA_VISIBLE_DEVICES=0
+  export LLAMA_N_CTX=44000
+else
+  export LLAMA_N_CTX=8182
 fi
 
-export LLAMA_N_CTX=44000
 poetry run python -m uvicorn skynet.main:app --reload
diff --git a/skynet/modules/ttt/summaries/jobs.py b/skynet/modules/ttt/summaries/jobs.py
index e1ddac4..46bdc8e 100644
--- a/skynet/modules/ttt/summaries/jobs.py
+++ b/skynet/modules/ttt/summaries/jobs.py
@@ -19,7 +19,7 @@
 
 from .persistence import db
 from .processor import process, process_azure, process_open_ai
-from .v1.models import DocumentMetadata, DocumentPayload, Job, JobId, JobStatus, JobType, Processors
+from .v1.models import DocumentMetadata, DocumentPayload, Job, JobId, JobStatus, JobType, Priority, Processors
 
 log = get_logger(__name__)
 
@@ -101,7 +101,11 @@ async def create_job(job_type: JobType, payload: DocumentPayload, metadata: Docu
 
     log.info(f"Created job {job.id}.")
 
-    await db.rpush(PENDING_JOBS_KEY, job_id)
+    if payload.priority == Priority.HIGH:
+        await db.lpush(PENDING_JOBS_KEY, job_id)
+    else:
+        await db.rpush(PENDING_JOBS_KEY, job_id)
+
     await update_summary_queue_metric()
 
     return JobId(id=job_id)
diff --git a/skynet/modules/ttt/summaries/v1/models.py b/skynet/modules/ttt/summaries/v1/models.py
index a8fe89b..eb8064a 100644
--- a/skynet/modules/ttt/summaries/v1/models.py
+++ b/skynet/modules/ttt/summaries/v1/models.py
@@ -10,14 +10,27 @@ class HintType(Enum):
     TEXT = 'text'
 
 
+class Priority(Enum):
+    NORMAL = 'normal'
+    HIGH = 'high'
+
+
 class DocumentPayload(BaseModel):
     text: str
     hint: HintType = HintType.TEXT
+    priority: Priority = Priority.NORMAL
     prompt: str | None = None
 
     model_config = {
         'json_schema_extra': {
-            'examples': [{'text': 'Your text here', 'hint': 'text', 'prompt': 'Summarize the following text {text}'}]
+            'examples': [
+                {
+                    'text': 'Your text here',
+                    'hint': 'text',
+                    'priority': 'normal',
+                    'prompt': 'Summarize the following text {text}',
+                }
+            ]
         }
     }