add conversational chat mode

qa/chains.py

@@ -3,4 +3,18 @@ from typing import Callable
 def retrieval_qa(llm: Callable, retriever: Callable) -> Callable: 
     from langchain.chains import RetrievalQA
     qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever)
-    return qa
+    return qa
+
+def conversational_retrieval_qa(llm: Callable, retriever: Callable) -> Callable: 
+    from langchain.memory import ConversationBufferMemory
+    from langchain.chains import ConversationalRetrievalChain
+    memory = ConversationBufferMemory(
+                memory_key="chat_history", 
+                return_messages=True
+    )
+    qa = ConversationalRetrievalChain.from_llm(
+            llm, 
+            retriever=retriever, 
+            memory=memory
+    )    
+    return qa

qa/manager.py

@@ -1,4 +1,4 @@
-from qa.chains import retrieval_qa
+from qa.chains import conversational_retrieval_qa, retrieval_qa
 from qa.loader import youtube_doc_loader
 from qa.model import load_llm
 from qa.split import split_document
@@ -7,7 +7,10 @@ from qa.vector_store import create_vector_store
 class YoutubeQA:
 
     def __init__(self):
-        pass
+        self.CHAT_MODE = 'normal' 
+
+    def change_chat_mode(self, mode: str) -> None:       
+        self.CHAT_MODE = mode
 
     def load_model(self) -> None:
         self.llm = load_llm()
@@ -18,7 +21,15 @@ class YoutubeQA:
         self.retriver = create_vector_store(docs=docs)
      
     def load_retriever(self) -> None:
-        self.retrieval_qa = retrieval_qa(llm=self.llm, retriever=self.retriver)
+        if self.CHAT_MODE == 'normal':
+            self.retrieval_qa = retrieval_qa(llm=self.llm, retriever=self.retriver)
+        elif self.CHAT_MODE == 'conversational':
+            self.retrieval_qa = conversational_retrieval_qa(llm=self.llm, retriever=self.retriver)
+        else:
+            raise ValueError('Chat Mode not implemented')
 
     def run(self, question: str) -> str:
-        return self.retrieval_qa.run(question)
+        if self.CHAT_MODE == 'normal':
+            return self.retrieval_qa.run(question)
+        elif self.CHAT_MODE == 'conversational':
+            return self.retrieval_qa({'question': question})['answer']