Update app.py

app.py

@@ -1,16 +1,31 @@
 from langchain.prompts.prompt import PromptTemplate
 from langchain.llms import OpenAIChat
-from langchain.chains import ChatVectorDBChain
+from langchain.chains import ConversationalRetrievalChain
 from langchain.embeddings import HuggingFaceEmbeddings, HuggingFaceInstructEmbeddings
-from langchain.callbacks.base import CallbackManager
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.callbacks import StdOutCallbackHandler
+from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT
 from langchain.vectorstores import FAISS
+from langchain.memory import ConversationBufferMemory
 import os
 from typing import Optional, Tuple
 import gradio as gr
 import pickle
 from threading import Lock
 
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    AIMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+)
+from langchain.schema import (
+    AIMessage,
+    HumanMessage,
+    SystemMessage
+)
+
+from langchain.prompts import PromptTemplate
+
 prefix_messages = [{"role": "system", "content": "You are a helpful assistant that is very good at answering questions about investments using the information given."}]
 
 model_options = {'all-mpnet-base-v2': "sentence-transformers/all-mpnet-base-v2",
@@ -18,6 +33,38 @@ model_options = {'all-mpnet-base-v2': "sentence-transformers/all-mpnet-base-v2",
 
 model_options_list = list(model_options.keys())
 
+memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True, output_key='answer')
+
+def load_prompt():
+
+    system_template="""Use only the following pieces of context that has been scraped from a website to answer the users question accurately.
+    Do not use any information not provided in the website context.
+    If you don't know the answer, just say 'There is no relevant answer in the Investor Website', 
+    don't try to make up an answer.
+    
+    ALWAYS return a "SOURCES" part in your answer.
+    The "SOURCES" part should be a reference to the source of the document from which you got your answer.
+    
+    Remember, do not reference any information not given in the context.
+    If the answer is not available in the given context just say 'There is no relevant answer in the website content'
+    
+    Follow the below format when answering:
+    
+    Question: {question}
+    SOURCES: [xyz]
+      
+    Begin!
+    ----------------
+    {context}"""
+    
+    messages = [
+        SystemMessagePromptTemplate.from_template(system_template),
+        HumanMessagePromptTemplate.from_template("{question}")
+    ]
+    prompt = ChatPromptTemplate.from_messages(messages)
+        
+    return prompt
+    
 def load_vectorstore(model):
     '''load embeddings and vectorstore'''
 
@@ -64,15 +111,24 @@ QA_PROMPT = PromptTemplate(template=template, input_variables=["question", "cont
 
 
 def get_chain(vectorstore):
-    llm = OpenAIChat(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0,\
-                    prefix_messages=prefix_messages,model_name='gpt-4-0613')
-    qa_chain = ChatVectorDBChain.from_llm(
-        llm,
-        vectorstore,
-        qa_prompt=QA_PROMPT,
-        condense_question_prompt=CONDENSE_QUESTION_PROMPT,
-    )
-    return qa_chain
+    llm = OpenAIChat(streaming=True, 
+                     callbacks=[StdOutCallbackHandler()]), 
+                     verbose=True, 
+                     temperature=0,
+                     model_name='gpt-4-0613')
+
+    question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)
+    doc_chain = load_qa_chain(llm=llm,chain_type="stuff",prompt=load_prompt())
+
+    chain = ConversationalRetrievalChain(retriever=vectorstore.as_retriever(search_kwags={"k": 3}), 
+                                     question_generator=question_generator, 
+                                     combine_docs_chain=doc_chain, 
+                                     memory=memory, 
+                                     return_source_documents=True, 
+                                     get_chat_history=lambda h :h)
+
+    
+    return chain
 
 def load_chain():
     chain = get_chain(vectorstore)
@@ -93,7 +149,7 @@ class ChatWrapper:
             # Set OpenAI key
             # chain = get_chain(vectorstore)
             # Run chain and append input.
-            output = chain({"question": inp, "chat_history": history})["answer"]
+            output = chain({"question": inp})["answer"]
             history.append((inp, output))
         except Exception as e:
             raise e