[#145] Create health endpoint in assistant

Author: wayangalihpratama

assistant/assistant.py

@@ -4,6 +4,7 @@
 import logging
 import chromadb
 import nltk
+import threading
 
 from time import sleep
 from openai import OpenAI
@@ -15,6 +16,7 @@
 from typing import Optional
 from db import connect_to_sqlite, get_stable_prompt
 from nltk.corpus import words, stopwords
+from health_check_handler import run
 
 
 logger = logging.getLogger(__name__)
@@ -418,4 +420,9 @@ async def consumer():
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
     main()
+
+    # Start the HTTP server in a separate thread
+    server_thread = threading.Thread(target=run, kwargs={"port": 9001})
+    server_thread.start()
+
     asyncio.run(consumer())

assistant/health_check_handler.py

@@ -0,0 +1,24 @@
+import logging
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+
+logger = logging.getLogger(__name__)
+
+
+class HealthCheckHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path == "/health":
+            self.send_response(200)
+            self.send_header("Content-type", "application/json")
+            self.end_headers()
+            self.wfile.write(b'{"status": "ok"}')
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+
+def run(server_class=HTTPServer, handler_class=HealthCheckHandler, port=9001):
+    server_address = ("", port)
+    httpd = server_class(server_address, handler_class)
+    logger.info(f"Starting httpd server on port {port}")
+    httpd.serve_forever()