Update rag-chat.py

rag-chat.py

@@ -1,3 +1,4 @@
+# import all necessary packages
 import os
 
 from langchain.document_loaders import DirectoryLoader
@@ -14,7 +15,32 @@ from langchain.memory import ChatMessageHistory, ConversationBufferMemory
 
 import chainlit as cl
 
-# llm
+from langchain.prompts.chat import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+)
+
+# define prompt template 
+system_template = """Use the following pieces of context to answer the users question.
+If you don't know the answer, just say that you don't know, 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.
+And if the user greets with greetings like Hi, hello, How are you, etc reply accordingly as well.
+Example of your response should be:
+The answer is foo
+SOURCES: xyz
+Begin!
+----------------
+{summaries}"""
+messages = [
+    SystemMessagePromptTemplate.from_template(system_template),
+    HumanMessagePromptTemplate.from_template("{question}"),
+]
+prompt = ChatPromptTemplate.from_messages(messages)
+chain_type_kwargs = {"prompt": prompt}
+
+# define the llm
 model_id = "tiiuae/falcon-7b-instruct"
 conv_model = HuggingFaceHub(
     huggingfacehub_api_token=os.environ['HF_API_TOKEN'], 
@@ -22,7 +48,7 @@ conv_model = HuggingFaceHub(
     model_kwargs={"temperature":0.8,"max_length": 1000}
     )
 
-# chroma
+# set up vector db with chroma
 data_path = "data/html"
 embed_model = "all-MiniLM-L6-v2" # Chroma defaults to "sentence-transformers/all-MiniLM-L6-v2"
 
@@ -90,26 +116,32 @@ def prepare_documents(documents):
         i += 1
     return documents
 
+# define a function to execute when a chat starts
 @cl.on_chat_start
 async def on_chat_start():
-    # Instantiate the chain for that user session
+    # instantiate the chain for that user session
     embedding_func = SentenceTransformerEmbeddings(model_name=embed_model)
-
+    
+    # display a message indicating document loading
     msg = cl.Message(
         content="Loading and processing documents. This may take a while...",
         disable_human_feedback=True)
     await msg.send()
 
+    # load and prepare documents for processing
     documents = load_documents(data_path)
     documents = prepare_documents(documents)
 
+    # create a document search object asynchronously
     docsearch = await cl.make_async(Chroma.from_documents)(
         documents,
         embedding_func
     )
     
+    # initialize ChatMessageHistory object to store message history
     message_history = ChatMessageHistory()
      
+    # initialize ConversationBufferMemory object to store conversation history
     memory = ConversationBufferMemory(
         memory_key="chat_history",
         output_key="answer",
@@ -117,6 +149,7 @@ async def on_chat_start():
         return_messages=True,
     )
 
+    # create a ConversationalRetrievalChain object
     chain = ConversationalRetrievalChain.from_llm(
         conv_model,
         chain_type="stuff",
@@ -124,36 +157,44 @@ async def on_chat_start():
         memory=memory,
         return_source_documents=True,
     )
+    
+    # indicate readiness for questions
     msg.content = "Ready. You can now ask questions!"
-
     await msg.update()
+    
+    # store the chain in the user's session
     cl.user_session.set("chain", chain)
 
-
+# define a function to handle messages
 @cl.on_message
 async def main(message):
+    # retrieve the chain object from the user's session
     chain = cl.user_session.get("chain")  # type: ConversationalRetrievalChain
     cb = cl.AsyncLangchainCallbackHandler()
 
+    # call the chain to process the incoming message
     res = await chain.acall(message.content, callbacks=[cb])
     
+    # retrieve the answer and source documents from the chain's response
     answer = res["answer"]
     source_documents = res["source_documents"]
 
-    text_elements = []
-
-    source_names = set()  # Use a set to store unique source names
+    text_elements = []  # list to store text elements
+    source_names = set()  # set to store unique source names
 
+    # iterate through source documents and extract relevant information
     for idx, source_doc in enumerate(source_documents):
         source_name = source_doc.metadata["source"]
         text_elements.append(
                 cl.Text(content=source_doc.page_content, 
                         name=source_name))
-        source_names.add(source_name)  # Add the source name to the set
+        source_names.add(source_name)  # add the source name to the set
 
+    # append sources information to the answer if available
     if source_names:
             answer += f"\nSources: {', '.join(source_names)}"
     else:
             answer += "\nNo sources found"
 
+    # send the answer along with any extracted text elements
     await cl.Message(content=answer, elements=text_elements).send()