From a57752446ab5a52cdb20f8ddff6850c2db90fed2 Mon Sep 17 00:00:00 2001
From: Jaap Blom <jblom@beeldengeluid.nl>
Date: Wed, 25 Sep 2024 10:28:22 +0200
Subject: [PATCH] added Status; added ping; removed old health_check endpoint

---
 health_check.py | 56 -------------------------------------------------
 whisper_api.py  | 30 ++++++++++++++++++++++----
 2 files changed, 26 insertions(+), 60 deletions(-)
 delete mode 100644 health_check.py

diff --git a/health_check.py b/health_check.py
deleted file mode 100644
index 0fc7c1c..0000000
--- a/health_check.py
+++ /dev/null
@@ -1,56 +0,0 @@
-import logging
-import json
-from typing import Dict, Any
-from flask import Response
-import requests
-
-
-logger = logging.getLogger(f"search.{__name__}")
-
-
-class HealthCheckProxy:
-    def __init__(self, service):
-        self.service = service
-        self.service.add_url_rule("/ping", view_func=self.ping, methods=["GET"])
-        self.service.add_url_rule("/ready", view_func=self.ready, methods=["GET"])
-
-    def ping(self) -> Response:
-        logger.debug("Received ping")
-        return Response("pong", mimetype="text/plain")
-
-    def ready(self) -> Response:
-        logger.debug("Checking availability all external connections")
-
-        connections: dict[str, Dict[str, Any]] = {}
-        logger.info(connections)
-
-        return self.check_connections(connections)
-
-    def check_connections(self, connections: Dict[str, Dict[str, Any]]) -> Response:
-        logger.info("Checking the following connections")
-        logger.debug(connections)
-        overall_status = 200
-        for c in connections.values():
-            logger.debug(c)
-            try:
-                c["statusCode"] = requests.get(
-                    c["connectionUrl"], verify=False, timeout=5
-                ).status_code
-            except Exception as e:
-                logger.exception(
-                    f"Unhandled ready check error for {c['connectionUrl']}"
-                )
-                c["error"] = str(e)
-                c["statusCode"] = 502
-
-            if c["statusCode"] < 200 or c["statusCode"] >= 300:
-                logger.error(
-                    f"Connection {c['connectionUrl']} failed with {c['statusCode']}"
-                )
-                overall_status = 502
-
-        return Response(
-            json.dumps({"status": overall_status}, indent=4, sort_keys=True),
-            overall_status,
-            mimetype="application/json",
-        )
diff --git a/whisper_api.py b/whisper_api.py
index 99b8d63..4e981b4 100644
--- a/whisper_api.py
+++ b/whisper_api.py
@@ -2,7 +2,7 @@
 from uuid import uuid4
 from fastapi import BackgroundTasks, FastAPI, HTTPException, status
 from asr import run
-
+from enum import Enum
 from typing import Optional
 from pydantic import BaseModel
 from base_util import LOG_FORMAT
@@ -11,11 +11,18 @@
 api = FastAPI()
 
 
+class Status(Enum):
+    CREATED = "CREATED"
+    PROCESSING = "PROCESSING"
+    DONE = "DONE"
+    ERROR = "ERROR"
+
+
 class Task(BaseModel):
     input_uri: str
     output_uri: str
+    status: Status = Status.CREATED
     id: str | None = None
-    status: str | None = None
 
 
 all_tasks = [
@@ -39,7 +46,6 @@ def get_task_index(task_id: str) -> int:
     return -1
 
 
-# delete after processing is done or if client calls for it
 def delete_task(task_id) -> bool:
     task_index = get_task_index(task_id)
     if task_index == -1:
@@ -48,13 +54,24 @@ def delete_task(task_id) -> bool:
     return True
 
 
+def update_task(task: Task) -> bool:
+    task_index = get_task_index(task.id)
+    if task_index == -1:
+        return False
+    all_tasks[task_index] = task.dict()
+    return True
+
+
 def try_whisper(task: Task):
     logger.info(f"Trying to call Whisper for task {task.id}")
     try:
         run(task.input_uri, task.output_uri)
+        task.status = Status.DONE
     except Exception:
         logger.exception("Failed to run whisper")
-    logger.info(f"Successfully ran Whisper for task {task.id}")
+        task.status = Status.ERROR
+    update_task(task)
+    logger.info(f"Done running Whisper for task {task.id}")
 
 
 @api.get("/tasks")
@@ -98,6 +115,11 @@ async def remove_task(task_id: str):
     }
 
 
+@api.get("/ping")
+async def ping():
+    return "pong"
+
+
 if __name__ == "__main__":
     import sys
     import uvicorn