Create task from Youtube link & query status of a task using taskId

entries/run.py

@@ -66,7 +66,6 @@ if __name__ == "__main__":
     task_list.append({"id": task_id, "status": "created", "resource_status:": "local"})
     stream = open(local_dir.joinpath("task_queue.yaml"), "w")
     dump(tasks_queue, stream)
-
     task.run_pipeline()
 
 

src/task.py

@@ -1,12 +1,9 @@
-from src.srt_util.srt import SrtScript
-from src.srt_util.srt2ass import srt2ass
+import threading
+import time
+
 import openai
-import stable_whisper
-import torch
-import whisper
 from pytube import YouTube
 from os import getenv
-from enum import Enum
 from pathlib import Path
 from enum import Enum, auto
 import logging
@@ -44,19 +41,30 @@ SRT_Script : SrtScript
     -  (Optional) mp4
 """
 
-class TaskStatus(Enum):
-    CREATED = (auto(), None)
-    INITIALIZING_ASR = (auto(), None)
-    PRE_PROCESSING = (auto(), None)
-    TRANSLATING = (auto(), 0.0)
-    POST_PROCESSING = (auto(), None)
-    OUTPUT_MODULE = (auto(), None)
+class TaskStatus(str, Enum):
+    CREATED = 'CREATED'
+    INITIALIZING_ASR = 'INITIALIZING_ASR'
+    PRE_PROCESSING = 'PRE_PROCESSING'
+    TRANSLATING = 'TRANSLATING'
+    POST_PROCESSING = 'POST_PROCESSING'
+    OUTPUT_MODULE = 'OUTPUT_MODULE'
 
 
 
 class Task:
+    @property
+    def status(self):
+        with self.__status_lock:
+            return self.__status
+
+    @status.setter
+    def status(self, new_status):
+        with self.__status_lock:
+            self.__status = new_status
+
     def __init__(self, task_id, task_local_dir, launch_info):
-        self.status = TaskStatus.CREATED
+        self.__status_lock = threading.Lock()
+        self.__status = TaskStatus.CREATED
         openai.api_key = getenv("OPENAI_API_KEY")
         self.launch_info = launch_info
         self.task_local_dir = task_local_dir
@@ -66,6 +74,7 @@ class Task:
         self.task_id = task_id
         self.progress = NotImplemented
         self.SRT_Script = None
+        self.result = None
 
 
     @staticmethod
@@ -88,26 +97,36 @@ class Task:
     # Module 1 ASR: audio --> SRT_script
     def get_srt_class(self, whisper_model='tiny', method="stable"):
         # Instead of using the script_en variable directly, we'll use script_input
+        self.status = TaskStatus.INITIALIZING_ASR
+        time.sleep(5)
         pass
 
     # Module 2: SRT preprocess: perform preprocess steps   
     def preprocess(self):
+        self.status = TaskStatus.PRE_PROCESSING
+        time.sleep(5)
         pass
     
     def update_translation_progress(self, new_progress):
         if self.progress == TaskStatus.TRANSLATING:
             self.progress = TaskStatus.TRANSLATING.value[0], new_progress
+        time.sleep(5)
 
     # Module 3: perform srt translation
     def translation(self):
+        time.sleep(5)
         pass
     
     # Module 4: perform srt post process steps
     def postprocess(self):
+        self.status = TaskStatus.POST_PROCESSING
+        time.sleep(5)
         pass
 
     # Module 5: output module
     def output_render(self):
+        self.status = TaskStatus.OUTPUT_MODULE
+        return "TODO"
         pass
     
     def run_pipeline(self):
@@ -115,8 +134,7 @@ class Task:
         self.preprocess()
         self.translation()
         self.postprocess()
-        out = self.output_render()
-        return out
+        self.result = self.output_render()
 
 class YoutubeTask(Task):
     def __init__(self, task_id, task_local_dir, launch_info, youtube_url):

src/web/api_specs.yaml

@@ -0,0 +1,79 @@
+openapi: 3.0.3
+info:
+  title: Pigeon AI
+  description: Pigeon AI
+  version: 1.0.0
+servers:
+  - url: 'https'
+paths:
+  /api/task:
+    post:
+      summary: Create a task
+      operationId: createTask
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/youtubeLink'
+      responses:
+        '200':
+          description: OK
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/task'
+  /api/task/{taskId}/status:
+    get:
+      summary: Get task status
+      operationId: getTask
+      parameters:
+        - name: taskId
+          in: path
+          required: true
+          description: task id
+          schema:
+            type: string
+      responses:
+        '200':
+          description: OK
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/taskStatus'
+        '404':
+          description: Not Found
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/error'
+
+components:
+  schemas:
+    youtubeLink:
+      type: object
+      properties:
+        youtubeLink:
+          type: string
+          description: youtube link
+          example: https://www.youtube.com/watch?v=5qap5aO4i9A
+    task:
+      type: object
+      properties:
+        taskId:
+          type: string
+          description: task id generated by uuid
+          example: 7a765280-1a72-47e4-8747-8a38cdbaca91
+    taskStatus:
+      type: object
+      properties:
+        status:
+          type: string
+          description: task status
+          example: PROCESSING
+    error:
+      type: object
+      properties:
+        error:
+          type: string
+          description: error message
+          example: 'Invalid youtube link'

src/web/web.py

@@ -2,27 +2,38 @@ import yaml
 from flask import Flask, request, jsonify
 from concurrent.futures import ThreadPoolExecutor
 from src.task import Task
+from uuid import uuid4
 
 app = Flask(__name__)
 
 # Global thread pool
 executor = ThreadPoolExecutor(max_workers=4)  # Adjust max_workers as per your requirement
 
+# thread safe task map to store task status
+task_map = {}
+
 @app.route('/api/task', methods=['POST'])
-def create_task():
+def create_task_youtube():
+    global task_map
     data = request.get_json()
-
-    if not data or 'youtube_link' not in data:
+    if not data or 'youtubeLink' not in data:
         return jsonify({'error': 'YouTube link not provided'}), 400
-    youtube_link = data['youtube_link']
+    youtube_link = data['youtubeLink']
     launch_config = yaml.load(open("./configs/local_launch.yaml"), Loader=yaml.Loader)
-    task = Task.fromYoutubeLink(youtube_link, 1, launch_config)
-
+    task_id = str(uuid4())
+    task = Task.fromYoutubeLink(youtube_link, task_id, launch_config)
+    task_map[task_id] = task
     # Submit task to thread pool
     executor.submit(task.run)
 
-    return jsonify({'task_id': task.task_id})
+    return jsonify({'taskId': task.task_id})
 
+@app.route('/api/task/<taskId>/status', methods=['GET'])
+def get_task_status(taskId):
+    global task_map
+    if taskId not in task_map:
+        return jsonify({'error': 'Task not found'}), 404
+    return jsonify({'status': task_map[taskId].status})
 
 if __name__ == '__main__':
     app.run(debug=True)