Update app.py

app.py

@@ -1,15 +1,17 @@
 import os
+import time
 from typing import Optional, Tuple
 
 import gradio as gr
 from query_data import get_chain
 from threading import Lock
 import pinecone
-from langchain.vectorstores import Chroma, Pinecone
+from langchain.vectorstores import Pinecone
 from langchain.embeddings.openai import OpenAIEmbeddings
 
+from dotenv import load_dotenv
 
-embeddings = OpenAIEmbeddings()
+load_dotenv()
 
 PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
 PINECONE_API_ENV = os.getenv("PINECONE_API_ENV")
@@ -22,6 +24,12 @@ pinecone.init(
 )
 
 index_name = PINECONE_INDEX
+
+embeddings = OpenAIEmbeddings()
+
+# with open("posts.pkl", "rb") as f:
+#     vectorstore = pickle.load(f)
+
 vectorstore = Pinecone.from_existing_index(index_name=index_name, embedding=embeddings)
 
 chain = get_chain(vectorstore)
@@ -39,18 +47,35 @@ class ChatWrapper:
         self.lock.acquire()
         try:
             history = history or []
+
             # If chain is None, that is because no API key was provided.
             # if chain is None:
-            #    history.append((inp, "Please paste your OpenAI key to use"))
+            #     history.append((inp, "Please paste your OpenAI key to use"))
             #     return history, history
             # Set OpenAI key
 
             import openai
             openai.api_key = os.getenv("OPENAI_API_KEY")
-            # Run chain and append input.
-            output = chain({"question": inp, "chat_history": history})["answer"]
+            start_time = time.time()
+            chain_obj = chain({"question": inp, "chat_history": history})
+            print('=======time===== : ' + str(time.time() - start_time))
+
+            output = chain_obj["answer"]
+
+            sources = chain_obj['source_documents']
+            
+            source_link = []
+            sources_str = 'Sources:'
+            for source in sources:
+                if source.metadata['source'] not in source_link:
+                    source_link.append(source.metadata['source'])
+                    sources_str += '\n' + source.metadata['source']
+            
+            if 'Unfortunately, our current pool of insights does not have an answer to this.' not in output:
+                output += '\n\n' + sources_str
+
             history.append((inp, output))
-            print(history)
+            # print(chain_obj)
         except Exception as e:
             raise e
         finally:
@@ -104,4 +129,4 @@ with block:
     # )
 
 # block.launch(debug=True)
-block.launch(debug=True, auth=('admin', 'Twimbit@2019'), auth_message='enter username password to proceed further')
+# block.launch(debug=True, auth=('admin', 'Twimbit@2019'), auth_message='enter username password to proceed further')