From b0b2eca84f7a683b7c7ee9648eb6e7993b770be5 Mon Sep 17 00:00:00 2001
From: Michael Clifford <mcliffor@redhat.com>
Date: Mon, 19 Feb 2024 10:08:17 -0500
Subject: [PATCH] add model service check (#41)

* add model service check

* add model service check to chatbot

* improve checking_model output UI
---
 chatbot-langchain/chatbot_ui.py    | 22 +++++++++++++++++++++-
 chatbot/ai_applications/chat_ui.py | 16 ++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/chatbot-langchain/chatbot_ui.py b/chatbot-langchain/chatbot_ui.py
index 22bacdf1..cc8e696b 100644
--- a/chatbot-langchain/chatbot_ui.py
+++ b/chatbot-langchain/chatbot_ui.py
@@ -3,12 +3,32 @@
 from langchain_community.callbacks import StreamlitCallbackHandler
 from langchain_core.prompts import ChatPromptTemplate
 import streamlit as st
+import requests
+import time
 import os 
 
-
 model_service = os.getenv("MODEL_SERVICE_ENDPOINT",
                           "http://localhost:8001/v1")
 
+@st.cache_resource(show_spinner=False)
+def checking_model_service():
+    start = time.time()
+    print("Checking Model Service Availability...")
+    ready = False
+    while not ready:
+        try:
+            request = requests.get(f'{model_service}/models')
+            if request.status_code == 200:
+                ready = True
+        except:
+            pass
+        time.sleep(1) 
+    print("Model Service Available")
+    print(f"{time.time()-start} seconds")
+
+with st.spinner("Checking Model Service Availability..."):
+    checking_model_service()
+
 st.title("💬 Chatbot")  
 if "messages" not in st.session_state:
     st.session_state["messages"] = [{"role": "assistant", 
diff --git a/chatbot/ai_applications/chat_ui.py b/chatbot/ai_applications/chat_ui.py
index 1374f7ff..d4d6e02a 100644
--- a/chatbot/ai_applications/chat_ui.py
+++ b/chatbot/ai_applications/chat_ui.py
@@ -1,5 +1,7 @@
 import gradio as gr
 from gradio_client import Client
+import requests
+import time
 import os
 
 
@@ -17,8 +19,22 @@ def ask(self, prompt, history):
                 yield r
         yield str(job.outputs()[-1])
 
+def checking_model_service(model_service):
+    print("Waiting for Model Service Availability...")
+    ready = False
+    while not ready:
+        try:
+            request = requests.get(f'{model_service}')
+            if request.status_code == 200:
+                ready = True
+        except:
+            pass
+        time.sleep(1) 
+    print("Model Service Available")
+
 if __name__ == "__main__":
     model_endpoint = os.getenv('MODEL_ENDPOINT', "http://0.0.0.0:7860")
+    checking_model_service(model_endpoint)
     chat = Chat(model_endpoint)
     demo = gr.ChatInterface(chat.ask)
     demo.launch(server_name="0.0.0.0", server_port=8080)
\ No newline at end of file