Update app.py

app.py

@@ -2,10 +2,10 @@ import gradio as gr
 import random
 import time
 
-from langchain.llms import OpenAI, OpenAIChat
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Pinecone
+from langchain.chains import LLMChain
 from langchain.chains.retrieval_qa.base import RetrievalQA
 from langchain.chains.question_answering import load_qa_chain
 import pinecone
@@ -13,7 +13,8 @@ import pinecone
 import os
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-OPENAI_KEY = ""
+#OPENAI_API_KEY = ""
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
 OPENAI_TEMP  = 0
 PINECONE_KEY = os.environ["PINECONE_KEY"]
 PINECONE_ENV = "asia-northeast1-gcp"
@@ -28,11 +29,8 @@ LLM_HISTORY_LEN = 3
 
 BUTTON_MIN_WIDTH = 150
 
-MODEL_STATUS = "Wait for API Key to Initialize."
 
-MODEL_LOADED = "Model Loaded"
-
-MODEL_WARNING = "Please paste your OpenAI API Key from openai.com to initialize this application!"
+MODEL_WARNING = "Please paste your OpenAI API Key from openai.com and press 'Enter' to initialize this application!"
 
 
 webui_title = """
@@ -46,35 +44,31 @@ Please insert your question and click 'Submit'
 """
 
 
-def init_model(openai_key):
+def init_model(api_key):
     try:
-        
-        os.environ["OPENAI_API_KEY"] = openai_key
-        
-        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+        if api_key and api_key.startswith("sk-") and len(api_key) > 50:
+            
+            embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-        pinecone.init(api_key     = PINECONE_KEY,
-                      environment = PINECONE_ENV)
+            pinecone.init(api_key     = PINECONE_KEY,
+                          environment = PINECONE_ENV)
 
-        llm = OpenAI(temperature=OPENAI_TEMP,
-                     model_name="gpt-3.5-turbo-0301")
-        
-        # ChatOpenAI(temperature    = OPENAI_TEMP, openai_api_key = openai_key)
-        
+            #llm = OpenAI(temperature=OPENAI_TEMP, model_name="gpt-3.5-turbo-0301")
 
-        global db
-        db = Pinecone.from_existing_index(index_name = PINECONE_INDEX,
-                                          embedding  = embeddings)
-        global chain
-        chain = load_qa_chain(llm, chain_type="stuff")
-        
-        global MODEL_STATUS
-        MODEL_STATUS = MODEL_LOADED
+            llm = ChatOpenAI(temperature = OPENAI_TEMP,
+                             openai_api_key = api_key)
+
+            chain = load_qa_chain(llm, chain_type="stuff")
+
+            db = Pinecone.from_existing_index(index_name = PINECONE_INDEX,
+                                              embedding  = embeddings)
 
-        return openai_key, ""
+            return api_key, chain, db, None
+        else:
+            return None,None,None,None
     except Exception as e:
         print(e)
-        return "",""
+        return None,None,None,None
 
 def get_chat_history(inputs) -> str:
     res = []
@@ -82,29 +76,64 @@ def get_chat_history(inputs) -> str:
         res.append(f"Human: {human}\nAI: {ai}")
     return "\n".join(res)
 
-css = """.bigbox {
-    min-height:200px;
-}"""
+def user(user_message, history):
+    return "", history+[[user_message, None]]
 
-with gr.Blocks(css=css) as demo:
+def bot(box_message, ref_message, chain, db, top_k):
+
+    # bot_message = random.choice(["Yes", "No"])
+    # 0 is user question, 1 is bot response
+    question = box_message[-1][0]
+    history  = box_message[:-1]
     
+    if (not chain) or (not db):
+        box_message[-1][1] = MODEL_WARNING
+        return box_message, "", ""
+
+    if not ref_message:
+        ref_message = question
+        details = f"Q:  {question}"
+    else:
+        details = f"Q:  {question}\nR: {ref_message}"
+
+    docsearch = db.as_retriever(search_kwargs={'k':top_k})
+    docs = docsearch.get_relevant_documents(ref_message)
+
+
+    all_output = chain({"input_documents": docs,
+                        "question": question,
+                        "chat_history": get_chat_history(history)})
+
+    bot_message = all_output['output_text']
+
+
+    source = "".join([f"""<details> <summary>{doc.metadata["source"]}</summary>
+{doc.page_content}
+
+</details>""" for i, doc in enumerate(docs)])
+
+    #print(source)
+
+    box_message[-1][1] = bot_message
+    return box_message, "", [[details, source]]
+
+
+with gr.Blocks(css=""".bigbox {
+    min-height:200px;
+}""") as demo:
+    llm_chain = gr.State()
+    vector_db = gr.State()
     gr.Markdown(webui_title)
     gr.Markdown(init_message)
     
-    if OPENAI_KEY and OPENAI_KEY.startswith("sk-") and len(OPENAI_KEY) > 50:
-        api_textbox_ph = "API Founded in Environment Variable: sk-..." + OPENAI_KEY[-4:]
-        api_textbox_edit = False
-        init_model(OPENAI_KEY)
-    else:
-        api_textbox_ph = "Paste Your OpenAI API Key (sk-...) and Hit ENTER"
-        api_textbox_edit = True
-    
-    api_textbox = gr.Textbox(placeholder = api_textbox_ph,
-                             interactive = api_textbox_edit,
-                            show_label=False, lines=1, type='password')
+    with gr.Row():
+        api_textbox = gr.Textbox(
+            value = OPENAI_API_KEY,
+            placeholder = "Paste Your OpenAI API Key (sk-...) and Hit ENTER",
+            show_label=False, lines=1, type='password')
+        init = gr.Button("Initialize Model").style(full_width=False)
     
-    
-    with gr.Tab("Chatbot"):
+    with gr.Tab("3GPP-Chatbot"):
         with gr.Row():
             with gr.Column(scale=10):
                 chatbot = gr.Chatbot(elem_classes="bigbox")
@@ -138,51 +167,22 @@ with gr.Blocks(css=css) as demo:
         detail_panel = gr.Chatbot(label="Related Docs")
         
                 
-    def user(user_message, history):
-        return "", history+[[user_message, None]]
-
-    def bot(box_message, ref_message, top_k):
-        if MODEL_STATUS != MODEL_LOADED:
-            box_message[-1][1] = MODEL_WARNING
-            return box_message, "", ""
-
-        # bot_message = random.choice(["Yes", "No"])
-        # 0 is user question, 1 is bot response
-        question = box_message[-1][0]
-        history  = box_message[:-1]
-        
-        if not ref_message:
-            ref_message = question
-            details = f"Q:  {question}"
-        else:
-            details = f"Q:  {question}\nR: {ref_message}"
-            
-        #print(question, ref_message)
-        #print(history)
-        #print(get_chat_history(history))
-        
-        docsearch = db.as_retriever(search_kwargs={'k':top_k})
-        docs = docsearch.get_relevant_documents(ref_message)
-        all_output = chain({"input_documents": docs,
-                             "question": question,
-                             "chat_history": get_chat_history(history)})
-        bot_message = all_output['output_text']
-        #print(docs)
-        
-        source = "".join([f"""<details> <summary>{doc.metadata["source"]}</summary>
-{doc.page_content}
-
-</details>""" for i, doc in enumerate(docs)])
-        
-        #print(source)
-
-        box_message[-1][1] = bot_message
-        return box_message, "", [[details, source]]
+    api_textbox.submit(init_model,
+                       api_textbox,
+                       [api_textbox, llm_chain, vector_db, chatbot])
+    init.click(init_model,
+                       api_textbox,
+                       [api_textbox, llm_chain, vector_db, chatbot])
     
-    submit.click(user, [query, chatbot], [query, chatbot], queue=False).then(
-        bot, [chatbot, ref, top_k], [chatbot, ref, detail_panel]
+    submit.click(user,
+                 [query, chatbot],
+                 [query, chatbot],
+                 queue=False).then(
+        bot,
+        [chatbot, ref, llm_chain, vector_db, top_k],
+        [chatbot, ref, detail_panel]
     )
-    api_textbox.submit(init_model, api_textbox, [api_textbox, chatbot])
+    
     clear.click(lambda: (None,None,None), None, [query, ref, chatbot], queue=False)
 
 if __name__ == "__main__":