Update app.py

app.py

@@ -23,20 +23,44 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 RANKER_URL = os.getenv("RANKER_URL")
 RETRIEVER_URL = os.getenv("RETRIEVER_URL")
 
-# RETRIEVER_IE = get_inference_endpoint(
-#     "fastrag-retriever", namespace="optimum-intel", token=HF_TOKEN
-# )
-# RANKER_IE = get_inference_endpoint(
-#     "fastrag-ranker", namespace="optimum-intel", token=HF_TOKEN
-# )
+RETRIEVER_IE = get_inference_endpoint(
+    "fastrag-retriever", namespace="optimum-intel", token=HF_TOKEN
+)
+RANKER_IE = get_inference_endpoint(
+    "fastrag-ranker", namespace="optimum-intel", token=HF_TOKEN
+)
+
+def check_inference_endpoints():
+    RETRIEVER_IE.update()
+    RANKER_IE.update()
+
+    messages = []
+
+    if RETRIEVER_IE.status in ["initializing", "pending"]:
+        messages += [
+            f"Retriever Inference Endpoint is {RETRIEVER_IE.status}. Please wait a few seconds and try again."
+        ]
+    elif RETRIEVER_IE.status in ["paused", "scaledToZero"]:
+        messages += [
+            f"Retriever Inference Endpoint is {RETRIEVER_IE.status}. Resuming it. Please wait a few seconds and try again."
+        ]
+        RETRIEVER_IE.resume()
 
-# if RETRIEVER_IE.status != "running":
-#     RETRIEVER_IE.resume()
-#     RETRIEVER_IE.wait()
+    if RANKER_IE.status in ["initializing", "pending"]:
+        messages += [
+            f"Ranker Inference Endpoint is {RANKER_IE.status}. Please wait a few seconds and try again."
+        ]
+    elif RANKER_IE.status in ["paused", "scaledToZero"]:
+        messages += [
+            f"Ranker Inference Endpoint is {RANKER_IE.status}. Resuming it. Please wait a few seconds and try again."
+        ]
+        RANKER_IE.resume()
+
+    if len(messages) > 0:
+        return "<br>".join(messages)
+    else:
+        return None
 
-# if RANKER_IE.status != "running":
-#     RANKER_IE.resume()
-#     RANKER_IE.wait()
 
 
 def post(url, payload):
@@ -189,6 +213,14 @@ pipe.add_node(component=ranker, name="Ranker", inputs=["Retriever"])
 
 
 def run(query: str) -> dict:
+    message = check_inference_endpoints()
+
+    if message is not None:
+        return f"""
+        <h2>Service Unavailable</h2>
+        <p>{message}</p>
+        """
+
     pipe_output = pipe.run(query=query)
 
     output = f"""<h2>Top {TOP_K} Documents</h2>"""