UPDATE: chat history retention

functions.py

@@ -5,6 +5,9 @@ from langchain_qdrant import QdrantVectorStore
 from langchain_core.prompts.chat import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain.retrievers import ParentDocumentRetriever
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain.memory import ChatMessageHistory
+from langchain_core.chat_history import BaseChatMessageHistory
 from langchain.storage import InMemoryStore
 from langchain.docstore.document import Document
 from langchain_huggingface import HuggingFaceEmbeddings
@@ -35,16 +38,24 @@ prompt = """
     3. **Exclusive Reliance on Training Data**: Answer user queries exclusively based on the provided training data. If a query is not covered by the training data, use the fallback response.
     4. **Restrictive Role Focus**: Do not answer questions or perform tasks unrelated to your role and training data.
 DO NOT ADD ANYTHING BY YOURSELF OR ANSWER ON YOUR OWN!
-Based on the context answer the following question.
+Based on the context answer the following question. Remember that you need to frame a meaningful answer in under 512 words.
 Context:
 =====================================
 {context}
 =====================================
+Question:
+=====================================
 {question}
-NOTE: generate responses WITHOUT prepending phrases like "Response:", "Output:", or "Answer:", etc
+
+Also, below I am providing you the previous question you were asked and the output you generated. It's just for your reference so that you know the topic you have been talking about and nothing else:
+=====================================
+{chatHistory}
+=====================================
+NOTE: generate responses WITHOUT prepending phrases like "Response:", "Output:", or "Answer:", etc. Also do not let the user know that you are answering from any extracted context or something.
 """
 prompt = ChatPromptTemplate.from_template(prompt)
 store = InMemoryStore()
+chatHistoryStore = dict()
 
 
 def createUser(username: str, password: str) -> None:
@@ -146,6 +157,25 @@ def format_docs(docs: str):
     else: pass
     return context
 
+
+def get_session_history(session_id: str) -> BaseChatMessageHistory:
+    if session_id not in store:
+        store[session_id] = ChatMessageHistory()
+    return store[session_id]
+
+
+def trimMessages(chain_input):
+    for storeName in chatHistoryStore:
+        messages = chatHistoryStore[storeName].messages
+        if len(messages) <= 2:
+            pass
+        else:
+            chatHistoryStore[storeName].clear()
+            for message in messages[-2: ]: 
+                chatHistoryStore[storeName].add_message(message)
+    return True
+
+
 def answerQuery(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192") -> str:
     global prompt 
     global client
@@ -167,11 +197,18 @@ def answerQuery(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192")
         base_compressor=compressor, base_retriever=retriever
     )
     chain = (
-        {"context": retriever | RunnableLambda(format_docs), "question": RunnablePassthrough()}
+        {"context": retriever | RunnableLambda(format_docs), "question": RunnablePassthrough(), "chatHistory": RunnablePassthrough()}
         | prompt
         | ChatGroq(model = llmModel, temperature = 0.75, max_tokens = 512)
         | StrOutputParser()
         )
+    chain = RunnableWithMessageHistory(
+        chain,
+        get_session_history,
+        input_messages_key = "question",
+        history_messages_key = "chatHistory"
+    )
+    chain = RunnablePassthrough.assign(messages_trimmed = trimMessages) | chain
     return {
         "output": chain.invoke(query)
     }