Update app.py

app.py

@@ -19,14 +19,19 @@ import shutil
 OPENAI_API_KEY = ''
 PINECONE_API_KEY = ''
 PINECONE_API_ENV = ''
-pinecone_index_name = ''
-chroma_collection_name = ''
-persist_directory = ''
-docsearch_ready = False
-directory_name = 'tmp_docs'
 langchain.verbose = False
 
 
+@st.cache_data()
+def init():
+    pinecone_index_name = ''
+    chroma_collection_name = ''
+    persist_directory = ''
+    docsearch_ready = False
+    directory_name = 'tmp_docs'
+    return pinecone_index_name, chroma_collection_name, persist_directory, docsearch_ready, directory_name
+
+
 @st.cache_data()
 def save_file(files):
     # Remove existing files in the directory
@@ -122,7 +127,7 @@ def setup_docsearch(use_pinecone, pinecone_index_name, embeddings, chroma_collec
     return docsearch, n_texts
 
 
-def get_response(query, chat_history):
+def get_response(query, chat_history, CRqa):
     result = CRqa({"question": query, "chat_history": chat_history})
     return result['answer'], result['source_documents']
 
@@ -137,83 +142,94 @@ def setup_em_llm(OPENAI_API_KEY, temperature):
     return embeddings, llm
 
 
-# Get user input of whether to use Pinecone or not
-col1, col2, col3 = st.columns([1, 1, 1])
-# create the radio buttons and text input fields
-with col1:
-    r_pinecone = st.radio('Use Pinecone?', ('Yes', 'No'))
-    r_ingest = st.radio(
-        'Ingest file(s)?', ('Yes', 'No'))
-with col2:
-    OPENAI_API_KEY = st.text_input(
-        "OpenAI API key:", type="password")
-    temperature = st.slider('Temperature', 0.0, 1.0, 0.1)
-    k_sources = st.slider('# source(s) to print out', 0, 20, 2)
-with col3:
-    if OPENAI_API_KEY:
-        embeddings, llm = setup_em_llm(OPENAI_API_KEY, temperature)
-        if r_pinecone.lower() == 'yes':
-            use_pinecone = True
-            PINECONE_API_KEY = st.text_input(
-                "Pinecone API key:", type="password")
-            PINECONE_API_ENV = st.text_input(
-                "Pinecone API env:", type="password")
-            pinecone_index_name = st.text_input('Pinecone index:')
-            pinecone.init(api_key=PINECONE_API_KEY,
-                          environment=PINECONE_API_ENV)
-        else:
-            use_pinecone = False
-            chroma_collection_name = st.text_input(
-                '''Chroma collection name of 3-63 characters:''')
-            persist_directory = "./vectorstore"
+pinecone_index_name, chroma_collection_name, persist_directory, docsearch_ready, directory_name = init()
 
-if pinecone_index_name or chroma_collection_name:
+
+def main(pinecone_index_name, chroma_collection_name, persist_directory, docsearch_ready, directory_name):
+    docsearch_ready = False
     chat_history = []
-    if r_ingest.lower() == 'yes':
-        files = st.file_uploader('Upload Files', accept_multiple_files=True)
-        if files:
-            save_file(files)
-            all_texts, n_texts = load_files()
-            docsearch = ingest(all_texts, use_pinecone, embeddings, pinecone_index_name,
-                               chroma_collection_name, persist_directory)
+    # Get user input of whether to use Pinecone or not
+    col1, col2, col3 = st.columns([1, 1, 1])
+    # create the radio buttons and text input fields
+    with col1:
+        r_pinecone = st.radio('Use Pinecone?', ('Yes', 'No'))
+        r_ingest = st.radio(
+            'Ingest file(s)?', ('Yes', 'No'))
+    with col2:
+        OPENAI_API_KEY = st.text_input(
+            "OpenAI API key:", type="password")
+        temperature = st.slider('Temperature', 0.0, 1.0, 0.1)
+        k_sources = st.slider('# source(s) to print out', 0, 20, 2)
+    with col3:
+        if OPENAI_API_KEY:
+            embeddings, llm = setup_em_llm(OPENAI_API_KEY, temperature)
+            if r_pinecone.lower() == 'yes':
+                use_pinecone = True
+                PINECONE_API_KEY = st.text_input(
+                    "Pinecone API key:", type="password")
+                PINECONE_API_ENV = st.text_input(
+                    "Pinecone API env:", type="password")
+                pinecone_index_name = st.text_input('Pinecone index:')
+                pinecone.init(api_key=PINECONE_API_KEY,
+                              environment=PINECONE_API_ENV)
+            else:
+                use_pinecone = False
+                chroma_collection_name = st.text_input(
+                    '''Chroma collection name of 3-63 characters:''')
+                persist_directory = "./vectorstore"
+
+    if pinecone_index_name or chroma_collection_name:
+        if r_ingest.lower() == 'yes':
+            files = st.file_uploader(
+                'Upload Files', accept_multiple_files=True)
+            if files:
+                save_file(files)
+                all_texts, n_texts = load_files()
+                docsearch = ingest(all_texts, use_pinecone, embeddings, pinecone_index_name,
+                                   chroma_collection_name, persist_directory)
+                docsearch_ready = True
+        else:
+            st.write(
+                'No data is to be ingested. Make sure the Pinecone index or Chroma collection name you provided contains data.')
+            docsearch, n_texts = setup_docsearch(use_pinecone, pinecone_index_name,
+                                                 embeddings, chroma_collection_name, persist_directory)
             docsearch_ready = True
-    else:
-        st.write(
-            'No data is to be ingested. Make sure the Pinecone index or Chroma collection name you provided contains data.')
-        docsearch, n_texts = setup_docsearch(use_pinecone, pinecone_index_name,
-                                             embeddings, chroma_collection_name, persist_directory)
-        docsearch_ready = True
-if docsearch_ready:
-    # number of sources (split-documents when ingesting files); default is 4
-    k = min([20, n_texts])
-    retriever = setup_retriever(docsearch, k)
-    CRqa = ConversationalRetrievalChain.from_llm(
-        llm, retriever=retriever, return_source_documents=True)
-
-    st.title('Chatbot')
-    # Get user input
-    query = st.text_area('Enter your question:', height=10,
-                         placeholder='Summarize the context.')
-    if query:
-        # Generate a reply based on the user input and chat history
-        reply, source = get_response(query, chat_history)
-        # Update the chat history with the user input and system response
-        chat_history.append(('User', query))
-        chat_history.append(('Bot', reply))
-        chat_history_str = '\n'.join(
-            [f'{x[0]}: {x[1]}' for x in chat_history])
-        st.text_area('Chat record:', value=chat_history_str, height=250)
-        # Display sources
-        for i, source_i in enumerate(source):
-            if i < k_sources:
-                if len(source_i.page_content) > 400:
-                    page_content = source_i.page_content[:400]
-                else:
-                    page_content = source_i.page_content
-                if source_i.metadata:
-                    metadata_source = source_i.metadata['source']
-                    st.write(
-                        f"**_Source {i+1}:_** {metadata_source}: {page_content}")
-                    st.write(source_i.metadata)
-                else:
-                    st.write(f"**_Source {i+1}:_** {page_content}")
+    if docsearch_ready:
+        # number of sources (split-documents when ingesting files); default is 4
+        k = min([20, n_texts])
+        retriever = setup_retriever(docsearch, k)
+        CRqa = ConversationalRetrievalChain.from_llm(
+            llm, retriever=retriever, return_source_documents=True)
+
+        st.title('Chatbot')
+        # Get user input
+        query = st.text_area('Enter your question:', height=10,
+                             placeholder='Summarize the context.')
+        if query:
+            # Generate a reply based on the user input and chat history
+            reply, source = get_response(query, chat_history, CRqa)
+            # Update the chat history with the user input and system response
+            chat_history.append(('User', query))
+            chat_history.append(('Bot', reply))
+            chat_history_str = '\n'.join(
+                [f'{x[0]}: {x[1]}' for x in chat_history])
+            st.text_area('Chat record:', value=chat_history_str, height=250)
+            # Display sources
+            for i, source_i in enumerate(source):
+                if i < k_sources:
+                    if len(source_i.page_content) > 400:
+                        page_content = source_i.page_content[:400]
+                    else:
+                        page_content = source_i.page_content
+                    if source_i.metadata:
+                        metadata_source = source_i.metadata['source']
+                        st.write(
+                            f"**_Source {i+1}:_** {metadata_source}: {page_content}")
+                        st.write(source_i.metadata)
+                    else:
+                        st.write(f"**_Source {i+1}:_** {page_content}")
+
+
+if __name__ == '__main__':
+    main(pinecone_index_name, chroma_collection_name, persist_directory,
+         docsearch_ready, directory_name)