Update chatbot.py

chatbot.py

@@ -6,6 +6,7 @@ from pymongo import MongoClient
 from langchain.prompts import ChatPromptTemplate
 from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
 from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
 
 from llm_provider import llm
 from vectorstore_manager import get_user_retriever
@@ -49,6 +50,7 @@ db = client[DB_NAME]
 sessions_collection = db[SESSIONS_COLLECTION]
 chains_collection = db[CHAINS_COLLECTION]
 
+
 # === Core Functions ===
 
 def create_new_chat(user_id: str) -> str:
@@ -78,7 +80,7 @@ def create_new_chat(user_id: str) -> str:
     # If the user has no chain/vectorstore registered yet, register it
     if chains_collection.count_documents({"user_id": user_id}, limit=1) == 0:
         # This also creates the vectorstore on disk via vectorstore_manager.ingest_report
-        # you should call ingest_report first elsewhere before chat
+        # You should call ingest_report first elsewhere before chat
         chains_collection.insert_one({
             "user_id": user_id,
             "vectorstore_path": f"user_vectorstores/{user_id}_faiss"
@@ -86,38 +88,47 @@ def create_new_chat(user_id: str) -> str:
 
     return chat_id
 
+
 def get_chain_for_user(user_id: str, chat_id: str) -> ConversationalRetrievalChain:
     """
     Reconstructs (or creates) the user's ConversationalRetrievalChain
     using their vectorstore and the chat-specific memory object.
     """
-    # Load chat history memory
-    chat_history = MongoDBChatMessageHistory(
+    # Step 1: Load raw MongoDB-backed chat history
+    mongo_history = MongoDBChatMessageHistory(
         session_id=chat_id,
         connection_string=MONGO_URI,
         database_name=DB_NAME,
         collection_name=HISTORY_COLLECTION,
     )
 
-    # Look up vectorstore path
+    # Step 2: Wrap it in a ConversationBufferMemory so that LangChain accepts it
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        chat_memory=mongo_history,
+        return_messages=True
+    )
+
+    # Step 3: Look up vectorstore path for this user
     chain_doc = chains_collection.find_one({"user_id": user_id})
     if not chain_doc:
         raise ValueError(f"No vectorstore registered for user {user_id}")
 
-    # Initialize retriever from vectorstore
+    # Step 4: Initialize retriever from vectorstore
     retriever = get_user_retriever(user_id)
 
-    # Create and return the chain
+    # Step 5: Create and return the chain with a valid Memory instance
     return ConversationalRetrievalChain.from_llm(
         llm=llm,
         retriever=retriever,
         return_source_documents=True,
         chain_type="stuff",
         combine_docs_chain_kwargs={"prompt": user_prompt},
-        memory=chat_history,
+        memory=memory,
         verbose=False,
     )
 
+
 def summarize_messages(chat_history: MongoDBChatMessageHistory) -> bool:
     """
     If the chat history grows too long, summarize it to keep the memory concise.
@@ -138,6 +149,7 @@ def summarize_messages(chat_history: MongoDBChatMessageHistory) -> bool:
     chat_history.add_ai_message(summary.content)
     return True
 
+
 def stream_chat_response(user_id: str, chat_id: str, query: str):
     """
     Given a user_id, chat_id, and a query string, streams back the AI response
@@ -145,17 +157,20 @@ def stream_chat_response(user_id: str, chat_id: str, query: str):
     """
     # Ensure the chain and memory are set up
     chain = get_chain_for_user(user_id, chat_id)
-    chat_history = chain.memory  # the MongoDBChatMessageHistory instance
+
+    # Since we used ConversationBufferMemory, the underlying MongoDBChatMessageHistory is accessible at:
+    chat_memory_wrapper = chain.memory  # type: ConversationBufferMemory
+    mongo_history = chat_memory_wrapper.chat_memory  # type: MongoDBChatMessageHistory
 
     # Optionally summarize if too many messages
-    summarize_messages(chat_history)
+    summarize_messages(mongo_history)
 
     # Add the user message to history
-    chat_history.add_user_message(query)
+    mongo_history.add_user_message(query)
 
     # Stream the response
     response_accum = ""
-    for chunk in chain.stream({"question": query, "chat_history": chat_history.messages}):
+    for chunk in chain.stream({"question": query, "chat_history": mongo_history.messages}):
         if "answer" in chunk:
             print(chunk["answer"], end="", flush=True)
             response_accum += chunk["answer"]
@@ -165,4 +180,4 @@ def stream_chat_response(user_id: str, chat_id: str, query: str):
 
     # Persist the AI's final message
     if response_accum:
-        chat_history.add_ai_message(response_accum)
+        mongo_history.add_ai_message(response_accum)