added chat memory and fixed sys_message bugs

app.py

@@ -2,7 +2,6 @@
 
 # OpenAI Chat completion
 import os
-from openai import AsyncOpenAI  # importing openai for API usage
 import chainlit as cl  # importing chainlit for our app
 from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
 from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
@@ -18,18 +17,14 @@ from utils.store import index_documents
 from utils.chain import create_chain
 from langchain.vectorstores import Pinecone
 from langchain.chat_models import ChatOpenAI
-from langchain.prompts import ChatPromptTemplate
-from langchain.prompts import PromptTemplate
-from operator import itemgetter
 from langchain.schema.runnable import RunnableSequence
 from langchain.schema import format_document
-from langchain.schema.output_parser import StrOutputParser
-from langchain.prompts.prompt import PromptTemplate
 from pprint import pprint
-from langchain_core.documents.base import Document
 from langchain_core.vectorstores import VectorStoreRetriever
 import langchain
 from langchain.cache import InMemoryCache
+from langchain_core.messages.human import HumanMessage
+from langchain.memory import ConversationBufferMemory
 
 load_dotenv()
 YOUR_API_KEY = os.environ["PINECONE_API_KEY"]
@@ -97,11 +92,16 @@ async def start_chat():
     # log data in WaB (on start)
     os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
 
+    # setup memory
+    
+    memory = ConversationBufferMemory(memory_key="chat_history")
+
     tools = {
         "arxiv_client": arxiv_client,
         "index": index,
         "embedder": embedder,
-        "llm": llm
+        "llm": llm,
+        "memory": memory
     }
     cl.user_session.set("tools", tools)
     cl.user_session.set("settings", settings)
@@ -111,18 +111,23 @@ async def start_chat():
 @cl.on_message  # marks a function that should be run each time the chatbot receives a message from a user
 async def main(message: cl.Message):
     settings = cl.user_session.get("settings")
-    tools = cl.user_session.get("tools")
+    tools: dict =  cl.user_session.get("tools")
     first_run = cl.user_session.get("first_run")
+    retrieval_augmented_qa_chain = cl.user_session.get("chain", None)
+    memory: ConversationBufferMemory = cl.user_session.get("memory")
 
+    sys_message = cl.Message(content="")
+    await sys_message.send() # renders a loader
+    
     if not first_run:
 
         arxiv_client: arxiv.Client  = tools['arxiv_client']
         index: pinecone.GRPCIndex = tools['index']
         embedder: CacheBackedEmbeddings = tools['embedder']
         llm: ChatOpenAI = tools['llm']
+        memory: ConversationBufferMemory = tools['memory']
 
         # using query search for ArXiv documents (on message)
-        
         search = arxiv.Search(
             query = message.content,
             max_results = 10,
@@ -130,18 +135,10 @@ async def main(message: cl.Message):
         )
         paper_urls = []
         
-        sys_message = cl.Message(content="")
-        await sys_message.send() # renders a loader
+
         for result in arxiv_client.results(search):
             paper_urls.append(result.pdf_url)
-        sys_message.content = """
-        I found some papers, let me study them real quick to help 
-        you learn, don't worry it'll be a few seconds 😉"""
-        await sys_message.update()
-        await sys_message.send()
 
-        sys_message = cl.Message(content="")
-        await sys_message.send() # renders a loader
         # load them and split them (on message)
         docs = []
         for paper_url in paper_urls:
@@ -159,9 +156,6 @@ async def main(message: cl.Message):
 
         # create an index using pinecone (on message)
         index_documents(docs, text_splitter, embedder, index)
-        sys_message.content = "Done studying :)"
-        await sys_message.update()
-        await sys_message.send()
 
         text_field = "source_document"
         index = pinecone.Index(INDEX_NAME)
@@ -174,74 +168,24 @@ async def main(message: cl.Message):
 
         # create the chain (on message)
         retrieval_augmented_qa_chain: RunnableSequence = create_chain(retriever=retriever, llm=llm)
+        cl.user_session.set("chain", retrieval_augmented_qa_chain)
         
-        # message.content = await cl.AskUserMessage(
-        #     content="Ask away"
-        # ).send()
-
+        sys_message.content = """
+        I found some papers and studied them 😉 \n"""
+        await sys_message.update()
+    
     # run
-    msg = cl.Message(content="")
-    for chunk in retrieval_augmented_qa_chain.stream({"question": f"{message.content}"}):
+    for chunk in retrieval_augmented_qa_chain.stream({"question": f"{message.content}", "chat_history": memory.buffer_as_messages}):
         pprint(chunk)
         if res:= chunk.get('response'):
-            await msg.stream_token(res.content)
-    await msg.send()
-    cl.user_session.set("first_run", True)
-    # first_run = True
+            await sys_message.stream_token(res.content)
+    await sys_message.send()
 
-    
-    # client = AsyncOpenAI()
-
-    # print(message.content)
-
-    # results_list = vector_db.search_by_text(query_text=message.content, k=3, return_as_text=True)
-    # if results_list:
-    #     results_string = "\n\n".join(results_list)
-    # else:
-    #     results_string = ""
-
-    # prompt = Prompt(
-    #     provider=ChatOpenAI.id,
-    #     messages=[
-    #         PromptMessage(
-    #             role="system",
-    #             template=system_template,
-    #             formatted=system_template,
-    #         ),
-    #         PromptMessage(
-    #             role="user",
-    #             template=user_template,
-    #             formatted=user_template.format(input=message.content),
-    #         ),
-    #         PromptMessage(
-    #             role="assistant",
-    #             template=assistant_template,
-    #             formatted=assistant_template.format(context=results_string)
-    #         )
-    #     ],
-    #     inputs={
-    #         "input": message.content,
-    #         "context": results_string
-    #         },
-    #     settings=settings,
-    # )
-
-    # print([m.to_openai() for m in prompt.messages])
-
-    # msg = cl.Message(content="")
-
-    # # Call OpenAI
-    # async for stream_resp in await client.chat.completions.create(
-    #     messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
-    # ):
-    #     token = stream_resp.choices[0].delta.content
-    #     if not token:
-    #         token = ""
-    #     await msg.stream_token(token)
-
-    # # Update the prompt object with the completion
-    # prompt.completion = msg.content
-    # msg.prompt = prompt
-
-    # # Send and close the message stream
-    # await msg.send()
+    memory.chat_memory.add_user_message(message.content)
+    memory.chat_memory.add_ai_message(sys_message.content)
+
+    print(memory.buffer_as_str)
+
+
+    cl.user_session.set("memory", memory)
+    cl.user_session.set("first_run", True)

requirements.txt

@@ -6,4 +6,5 @@ python-dotenv==1.0.0
 numpy==1.25.2
 langchain
 pinecone-client[grpc]
-pypdf
+pypdf
+arxiv

utils/chain.py

@@ -2,10 +2,12 @@ from operator import itemgetter
 from langchain_core.vectorstores import VectorStoreRetriever
 from langchain.schema.runnable import RunnableLambda, RunnableParallel, RunnableSequence
 from langchain.chat_models import ChatOpenAI
-from langchain.prompts import PromptTemplate
+from langchain.prompts import PromptTemplate, ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.documents import Document
 from langchain_core.messages.ai import AIMessage
-
+from langchain_core.messages.human import HumanMessage
+from langchain_core.messages.system import SystemMessage
+from langchain_core.messages.function import FunctionMessage
 
 template = """
 You are a helpful assistant, your job is to answer the user's question using the relevant context.
@@ -16,7 +18,21 @@ CONTEXT
 
 User question: {question}
 """
+
 prompt = PromptTemplate.from_template(template=template)
+chat_prompt = ChatPromptTemplate.from_messages([
+    ("system", """
+     You are a helpful assistant, your job is to answer the user's question using the relevant context:
+     =========
+     CONTEXT: 
+     {context}
+     =========
+     """),
+     MessagesPlaceholder(variable_name="chat_history"),
+     ("human", "{question}")
+])
+
+
 
 
 def to_doc(input: AIMessage) -> list[Document]:
@@ -46,7 +62,7 @@ def create_chain(**kwargs) -> RunnableSequence:
 
     docs_chain = (itemgetter("question") | retriever).with_config(config={"run_name": "docs"})
     self_knowledge_chain = (itemgetter("question") | llm | to_doc).with_config(config={"run_name": "self knowledge"})
-    response_chain = (prompt | llm).with_config(config={"run_name": "response"})
+    response_chain = (chat_prompt | llm).with_config(config={"run_name": "response"})
     merge_docs_link = RunnableLambda(merge_docs).with_config(config={"run_name": "merge docs"})
     context_chain = (
         RunnableParallel(
@@ -61,11 +77,13 @@ def create_chain(**kwargs) -> RunnableSequence:
     retrieval_augmented_qa_chain = (
         RunnableParallel({
             "question": itemgetter("question"),
+            "chat_history": itemgetter("chat_history"),
             "context": context_chain
         })
         | RunnableParallel({
             "response": response_chain,
             "context": itemgetter("context"),
+            "chat_history": itemgetter("chat_history")
         })
     )
     return retrieval_augmented_qa_chain