feat: Add support for chat endpoint (#150)

* feat: Add support for chat endpoint

This commit introduces a new endpoint called chat.
It maintains context of the last 5 user messages
and continues to do semantic search to find the
context.

embedchain/embedchain.py

@@ -5,6 +5,7 @@ from chromadb.utils import embedding_functions
 from dotenv import load_dotenv
 from langchain.docstore.document import Document
 from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.memory import ConversationBufferMemory
 
 from embedchain.loaders.youtube_video import YoutubeVideoLoader
 from embedchain.loaders.pdf_file import PdfFileLoader
@@ -26,6 +27,8 @@ load_dotenv()
 ABS_PATH = os.getcwd()
 DB_DIR = os.path.join(ABS_PATH, "db")
 
+memory = ConversationBufferMemory()
+
 
 class EmbedChain:
     def __init__(self, db=None, ef=None):
@@ -218,7 +221,47 @@ class EmbedChain:
         prompt = self.generate_prompt(input_query, context)
         answer = self.get_answer_from_llm(prompt)
         return answer
-    
+
+    def generate_chat_prompt(self, input_query, context, chat_history=''):
+        """
+        Generates a prompt based on the given query, context and chat history
+        for chat interface. This is then passed to an LLM.
+
+        :param input_query: The query to use.
+        :param context: Similar documents to the query used as context.
+        :param chat_history: User and bot conversation that happened before.
+        :return: The prompt
+        """
+        prefix_prompt = f"""You are a chatbot having a conversation with a human. You are given chat history and context. You need to answer the query considering context, chat history and your knowledge base. If you don't know the answer or the answer is neither contained in the context nor in history, then simply say "I don't know"."""
+        chat_history_prompt = f"""\n----\nChat History: {chat_history}\n----"""
+        suffix_prompt = f"""\n####\nContext: {context}\n####\nQuery: {input_query}\nHelpful Answer:"""
+        prompt = prefix_prompt
+        if chat_history:
+            prompt += chat_history_prompt
+        prompt += suffix_prompt
+        return prompt
+
+    def chat(self, input_query):
+        """
+        Queries the vector database on the given input query.
+        Gets relevant doc based on the query and then passes it to an
+        LLM as context to get the answer.
+
+        Maintains last 5 conversations in memory.
+        """
+        context = self.retrieve_from_database(input_query)
+        global memory
+        chat_history = memory.load_memory_variables({})["history"]
+        prompt = self.generate_chat_prompt(
+            input_query,
+            context,
+            chat_history=chat_history,
+        )
+        answer = self.get_answer_from_llm(prompt)
+        memory.chat_memory.add_user_message(input_query)
+        memory.chat_memory.add_ai_message(answer)
+        return answer
+
     def dry_run(self, input_query):
         """
         A dry run does everything except send the resulting prompt to
@@ -292,7 +335,7 @@ class OpenSourceApp(EmbedChain):
 
     def get_llm_model_answer(self, prompt):
         from gpt4all import GPT4All
-    
+
         global gpt4all_model
         if gpt4all_model is None:
             gpt4all_model = GPT4All("orca-mini-3b.ggmlv3.q4_0.bin")