Add progress tracking to transcription tasks with new models and service updates

Author: swordbreaker

src/transcribo_backend/models/progress.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class ProgressResponse(BaseModel):
+    progress: float
+    currentTime: float
+    duration: float

src/transcribo_backend/models/task_status.py

@@ -16,6 +16,7 @@ class TaskStatus(BaseModel):
     status: TaskStatusEnum = Field(default=TaskStatusEnum.IN_PROGRESS)
     created_at: datetime | None = None
     executed_at: datetime | None = None
+    progress: float | None = None
 
     class Config:
         use_enum_values = True

src/transcribo_backend/services/whisper_service.py

@@ -1,9 +1,11 @@
 import json
+import uuid
 from typing import Any
 
 import aiohttp
 from fastapi import APIRouter
 
+from transcribo_backend.models.progress import ProgressResponse
 from transcribo_backend.models.response_format import ResponseFormat
 from transcribo_backend.models.task_status import TaskStatus
 from transcribo_backend.models.transcription_response import TranscriptionResponse
@@ -17,6 +19,9 @@
 BENTOML_API_URL = f"{settings.whisper_api}/audio/transcriptions"
 
 
+taskId_to_progressId: dict[str, str] = {}
+
+
 async def transcribe_get_task_status(task_id: str) -> TaskStatus:
     """
     Checks the status of an ongoing transcription task.
@@ -28,11 +33,19 @@ async def transcribe_get_task_status(task_id: str) -> TaskStatus:
         TaskStatus: The current status of the task
     """
     url = f"{settings.whisper_api}/audio/transcriptions/task/status?task_id={task_id}"
+    progress_url = f"{settings.whisper_api}/progress/{taskId_to_progressId[task_id]}"
 
     # Get the status of the transcription task
-    async with aiohttp.ClientSession() as session, session.get(url) as response:
+    async with (
+        aiohttp.ClientSession() as session,
+        session.get(url) as response,
+        session.get(progress_url) as progress_response,
+    ):
         response.raise_for_status()
-        return TaskStatus(**await response.json())
+        progress_response.raise_for_status()
+
+        progress = ProgressResponse(**await progress_response.json())
+        return TaskStatus(**await response.json(), progress=progress.progress)
 
 
 async def transcribe_get_task_result(task_id: str) -> TranscriptionResponse:
@@ -51,6 +64,15 @@ async def transcribe_get_task_result(task_id: str) -> TranscriptionResponse:
     async with aiohttp.ClientSession() as session, session.get(url) as response:
         response.raise_for_status()
         result_data = await response.json()
+
+        taskId_to_progressId.pop(task_id, None)
+
+        transcription = TranscriptionResponse(**result_data)
+        for segment in transcription.segments:
+            segment.text = segment.text.strip()
+            segment.speaker = segment.speaker or "Unknown"
+            segment.speaker = segment.speaker.strip().capitalize()
+
         return TranscriptionResponse(**result_data)
 
 
@@ -135,6 +157,9 @@ async def transcribe_submit_task(
     form_data.add_field("file", audio_data, filename="audio.wav")
     form_data.add_field("model", model)
 
+    progress_id = uuid.uuid4().hex
+    form_data.add_field("progress_id", progress_id)
+
     if language:
         form_data.add_field("language", language)
     if prompt:
@@ -166,4 +191,6 @@ async def transcribe_submit_task(
     # Send the request
     async with aiohttp.ClientSession() as session, session.post(url, data=form_data) as response:
         response.raise_for_status()
-        return TaskStatus(**await response.json())
+        status = TaskStatus(**await response.json())
+        taskId_to_progressId[status.task_id] = progress_id
+        return status