Moving endpoint status to backend

app.py

@@ -1,14 +1,12 @@
 import logging
-import os
 from pathlib import Path
 from time import perf_counter
 
 import gradio as gr
 from jinja2 import Environment, FileSystemLoader
-import requests
 from transformers import AutoTokenizer
 
-from backend.query_llm import generate
+from backend.query_llm import check_endpoint_status, generate
 from backend.semantic_search import retriever
 
 proj_dir = Path(__file__).parent
@@ -27,27 +25,6 @@ template_html = env.get_template('template_html.j2')
 tokenizer = AutoTokenizer.from_pretrained('derek-thomas/jais-13b-chat-hf')
 
 
-def check_endpoint_status():
-    # Replace with the actual API URL and headers
-    api_url = os.getenv("ENDPOINT_URL")
-    headers = {
-        'accept': 'application/json',
-        'Authorization': f'Bearer {os.getenv("BEARER")}'
-        }
-
-    try:
-        response = requests.get(api_url, headers=headers)
-        response.raise_for_status()  # will throw an exception for non-200 status
-        data = response.json()
-
-        # Extracting the status information
-        status = data.get('status', {}).get('state', 'No status found')
-        message = data.get('status', {}).get('message', 'No message found')
-
-        return f"Status: {status}\nMessage: {message}"
-    except requests.exceptions.RequestException as e:
-        return f"Failed to get status: {str(e)}"
-
 def add_text(history, text):
     history = [] if history is None else history
     history = history + [(text, None)]

backend/query_llm.py

@@ -1,3 +1,5 @@
+import os
+
 import requests
 from os import getenv
 
@@ -26,4 +28,26 @@ def call_jais(payload):
 def generate(prompt: str):
     payload = {'inputs': '', 'prompt':prompt}
     response = call_jais(payload)
-    return response
+    return response
+
+
+def check_endpoint_status():
+    # Replace with the actual API URL and headers
+    api_url = os.getenv("ENDPOINT_URL")
+    headers = {
+        'accept': 'application/json',
+        'Authorization': f'Bearer {os.getenv("BEARER")}'
+        }
+
+    try:
+        response = requests.get(api_url, headers=headers)
+        response.raise_for_status()  # will throw an exception for non-200 status
+        data = response.json()
+
+        # Extracting the status information
+        status = data.get('status', {}).get('state', 'No status found')
+        message = data.get('status', {}).get('message', 'No message found')
+
+        return f"Status: {status}\nMessage: {message}"
+    except requests.exceptions.RequestException as e:
+        return f"Failed to get status: {str(e)}"