Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import random
 import time
 
+from langchain import PromptTemplate
+from langchain.llms import OpenAI
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Pinecone
@@ -15,7 +17,7 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 #OPENAI_API_KEY = ""
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
-OPENAI_TEMP  = 0
+OPENAI_TEMP  = 1
 
 PINECONE_KEY = os.environ.get("PINECONE_KEY", "")
 PINECONE_ENV = os.environ.get("PINECONE_ENV", "asia-northeast1-gcp")
@@ -28,7 +30,7 @@ TOP_K_DEFAULT = 10
 TOP_K_MAX = 25
 
 
-BUTTON_MIN_WIDTH = 180
+BUTTON_MIN_WIDTH = 205
 
 STATUS_NOK = "404-MODEL UNREADY-critical"
 STATUS_OK  = "200-MODEL LOADED-9cf"
@@ -57,10 +59,19 @@ MODEL_WARNING = f"Please paste your OpenAI API Key from \
 [openai.com](https://platform.openai.com/account/api-keys) and then **{KEY_INIT}**"
 
 
-TAB_1 = "3GPP Chatbot"
+TAB_1 = "Chatbot"
 
 FAVICON = './icon.svg'
 
+LLM_LIST = ["gpt-3.5-turbo", "text-davinci-003"]
+
+
+DOC_1 = '3GPP'
+DOC_2 = 'HTTP2'
+
+DOC_SUPPORTED = [DOC_1, DOC_2]
+DOC_DEFAULT = [DOC_1]
+
 webui_title = """
 # OpenAI Chatbot Based on Vector Database
 ## Example of 3GPP
@@ -91,20 +102,41 @@ def init_model(api_key, emb_name, db_api_key, db_env, db_index):
 
             #llm = OpenAI(temperature=OPENAI_TEMP, model_name="gpt-3.5-turbo-0301")
 
-            llm = ChatOpenAI(temperature = OPENAI_TEMP,
+            
+            llm_dict = {}
+            for llm_name in LLM_LIST:
+                if llm_name == "gpt-3.5-turbo":
+                    llm_dict[llm_name] = ChatOpenAI(model_name=llm_name,
+                                                    temperature = OPENAI_TEMP,
+                                                    openai_api_key = api_key)
+                else:
+                    llm_dict[llm_name] = OpenAI(model_name=llm_name,
+                                                temperature = OPENAI_TEMP,
+                                                openai_api_key = api_key)
+
+            '''
+            ChatOpenAI(model_name="gpt-3.5-turbo",
+                             temperature = OPENAI_TEMP,
                              openai_api_key = api_key)
+            chain_1 = load_qa_chain(llm, chain_type="stuff")
+            
+            #LLMChain(llm=llm, prompt=condense_question_prompt)
 
-            chain = load_qa_chain(llm, chain_type="stuff")
-
+            chain_2 = LLMChain(llm = llm,
+                             prompt = PromptTemplate(template='{question}',
+                                                     input_variables=['question']),
+                             output_key = 'output_text')
+                             '''
+            
             db = Pinecone.from_existing_index(index_name = db_index,
                                               embedding  = embeddings)
 
-            return api_key, MODEL_DONE, chain, db, None
+            return api_key, MODEL_DONE, llm_dict, None, db, None
         else:
-            return None,MODEL_NULL,None,None,None
+            return None,MODEL_NULL,None,None,None,None
     except Exception as e:
         print(e)
-        return None,MODEL_NULL,None,None,None
+        return None,MODEL_NULL,None,None,None,None
 
 
 def get_chat_history(inputs) -> str:
@@ -131,14 +163,16 @@ def doc_similarity(query, db, top_k):
 def user(user_message, history):
     return "", history+[[user_message, None]]
 
-def bot(box_message, ref_message, chain, db, top_k):
+def bot(box_message, ref_message,
+        llm_dropdown, llm_dict, doc_list,
+        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):
+    if (not llm_dict) or (not doc_check) or (not db):
         box_message[-1][1] = MODEL_WARNING
         return box_message, "", ""
 
@@ -149,17 +183,30 @@ def bot(box_message, ref_message, chain, db, top_k):
         details = f"Q:  {question}\nR: {ref_message}"
         
         
-    docs = doc_similarity(ref_message, db, top_k)
+    llm = llm_dict[llm_dropdown]
     
-    delta_top_k = top_k - len(docs)
+    print(llm)
+    print(doc_list)
     
-    if delta_top_k > 0:
-        docs = doc_similarity(ref_message, db, top_k+delta_top_k)
+    if DOC_1 in doc_list:
+        chain = load_qa_chain(llm, chain_type="stuff")
+        docs = doc_similarity(ref_message, db, top_k)
+        delta_top_k = top_k - len(docs)
+
+        if delta_top_k > 0:
+            docs = doc_similarity(ref_message, db, top_k+delta_top_k)
+            
+    else:
+        chain = LLMChain(llm = llm,
+                         prompt = PromptTemplate(template='{question}',
+                                                input_variables=['question']),
+                         output_key = 'output_text')
+        docs = []
 
     all_output = chain({"input_documents": docs,
                         "question": question,
                         "chat_history": get_chat_history(history)})
-
+    
     bot_message = all_output['output_text']
 
 
@@ -171,7 +218,7 @@ def bot(box_message, ref_message, chain, db, top_k):
     #print(source)
 
     box_message[-1][1] = bot_message
-    return box_message, "", [[details, source]]
+    return box_message, "", [[details, bot_message + source]]
 
 #----------------------------------------------------------------------------------------------------------
 #----------------------------------------------------------------------------------------------------------
@@ -180,10 +227,11 @@ with gr.Blocks(
     title = TAB_1,
     theme = "Base",
     css = """.bigbox {
-    min-height:200px;
+    min-height:250px;
 }
 """) as demo:
-    llm_chain = gr.State()
+    llm = gr.State()
+    chain_2 = gr.State() # not inuse
     vector_db = gr.State()
     gr.Markdown(webui_title)
     gr.HTML(dup_link)
@@ -208,13 +256,26 @@ with gr.Blocks(
         with gr.Row():
             with gr.Column(scale=10):
                 chatbot = gr.Chatbot(elem_classes="bigbox")
-
+            #with gr.Column(scale=1):
+            with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH):
+                doc_check = gr.CheckboxGroup(choices = DOC_SUPPORTED,
+                                             value   = DOC_DEFAULT,
+                                             label   = "Reference Docs",
+                                             interactive=True)
+                llm_dropdown = gr.Dropdown(LLM_LIST,
+                                           value=LLM_LIST[0],
+                                           multiselect=False,
+                                           interactive=True,
+                                           label="LLM Selection",
+                                           )
         with gr.Row():
             with gr.Column(scale=10):
                 query = gr.Textbox(label="Question:",
                                    lines=2)
                 ref = gr.Textbox(label="Reference(optional):")
+
             with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH):
+
                 clear = gr.Button(KEY_CLEAR)
                 submit = gr.Button(KEY_SUBMIT,variant="primary")
                 
@@ -238,35 +299,38 @@ with gr.Blocks(
                 lines=1,
                 interactive=True,
                 type='email')
-        with gr.Row():
-            db_api_textbox = gr.Textbox(
-                label = "Pinecone API Key",
-                # show_label = False,
-                value = PINECONE_KEY,
-                placeholder = "Paste Your Pinecone API Key (xx-xx-xx-xx-xx) and Hit ENTER",
-                lines=1,
-                interactive=True,
-                type='password')
-        with gr.Row():
-            db_env_textbox = gr.Textbox(
-                label = "Pinecone Environment",
-                # show_label = False,
-                value = PINECONE_ENV,
-                placeholder = "Paste Your Pinecone Environment (xx-xx-xx) and Hit ENTER",
-                lines=1,
-                interactive=True,
-                type='email')
-            db_index_textbox = gr.Textbox(
-                label = "Pinecone Index",
-                # show_label = False,
-                value = PINECONE_INDEX,
-                placeholder = "Paste Your Pinecone Index (xxxx) and Hit ENTER",
-                lines=1,
-                interactive=True,
-                type='email')
-
-    init_input = [llm_api_textbox, emb_textbox, db_api_textbox, db_env_textbox, db_index_textbox]
-    init_output = [llm_api_textbox, model_statusbox, llm_chain, vector_db, chatbot]
+        with gr.Accordion("Pinecone Database for "+DOC_1):
+            with gr.Row():
+                db_api_textbox = gr.Textbox(
+                    label = "Pinecone API Key",
+                    # show_label = False,
+                    value = PINECONE_KEY,
+                    placeholder = "Paste Your Pinecone API Key (xx-xx-xx-xx-xx) and Hit ENTER",
+                    lines=1,
+                    interactive=True,
+                    type='password')
+            with gr.Row():
+                db_env_textbox = gr.Textbox(
+                    label = "Pinecone Environment",
+                    # show_label = False,
+                    value = PINECONE_ENV,
+                    placeholder = "Paste Your Pinecone Environment (xx-xx-xx) and Hit ENTER",
+                    lines=1,
+                    interactive=True,
+                    type='email')
+                db_index_textbox = gr.Textbox(
+                    label = "Pinecone Index",
+                    # show_label = False,
+                    value = PINECONE_INDEX,
+                    placeholder = "Paste Your Pinecone Index (xxxx) and Hit ENTER",
+                    lines=1,
+                    interactive=True,
+                    type='email')
+
+    init_input  = [llm_api_textbox, emb_textbox, db_api_textbox, db_env_textbox, db_index_textbox]
+    init_output = [llm_api_textbox, model_statusbox,
+                   llm, chain_2,
+                   vector_db, chatbot]
                 
     llm_api_textbox.submit(init_model, init_input, init_output)
     init.click(init_model, init_input, init_output)
@@ -276,7 +340,9 @@ with gr.Blocks(
                  [query, chatbot],
                  queue=False).then(
         bot,
-        [chatbot, ref, llm_chain, vector_db, top_k],
+        [chatbot, ref,
+         llm_dropdown, llm, doc_check,
+         vector_db, top_k],
         [chatbot, ref, detail_panel]
     )