Update my_memory_logic.py

my_memory_logic.py

@@ -5,39 +5,47 @@ import os
 # Import the PipelineRunnable from pipeline.py
 from pipeline import pipeline_runnable
 
-from langchain_core.chat_history import BaseChatMessageHistory
+from langchain.schema import BaseChatMessageHistory
 from langchain_community.chat_message_histories import ChatMessageHistory
-from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain.runnables.history import RunnableWithMessageHistory
 
 ###############################################################################
 # 1) In-memory store: session_id -> ChatMessageHistory
 ###############################################################################
-store = {}  # e.g. { "abc123": ChatMessageHistory() }
+store = {}  # e.g., { "abc123": ChatMessageHistory(...) }
 
 def get_session_history(session_id: str) -> BaseChatMessageHistory:
+    """
+    Retrieve or create a ChatMessageHistory object for the given session_id.
+    This ensures each session_id has its own conversation history.
+    """
     if session_id not in store:
         store[session_id] = ChatMessageHistory()
     return store[session_id]
 
 ###############################################################################
-# 2) RunnableWithMessageHistory referencing pipeline_runnable
+# 2) Create the RunnableWithMessageHistory (session-based chain)
 ###############################################################################
+# This wraps your `pipeline_runnable` so it automatically reads/writes
+# conversation history from get_session_history for each session.
 conversational_rag_chain = RunnableWithMessageHistory(
-    pipeline_runnable,   # The Runnable from pipeline.py
-    get_session_history,
-    input_messages_key="input",
-    history_messages_key="chat_history",
-    output_messages_key="answer"
+    pipeline_runnable,               # the Runnable from pipeline.py
+    get_session_history,             # fetches or creates ChatMessageHistory by session_id
+    input_messages_key="input",      # key in the dict for user's new query
+    history_messages_key="chat_history",  # key for existing chat logs
+    output_messages_key="answer"      # key for final output
 )
 
 ###############################################################################
-# 3) Convenience function to run a query with session-based memory
+# 3) A convenience function to run a query with session-based memory
 ###############################################################################
 def run_with_session_memory(user_query: str, session_id: str) -> str:
     """
-    Calls our `conversational_rag_chain` with session_id,
-    returns the final 'answer' from pipeline_runnable.
+    A helper that calls our `conversational_rag_chain` 
+    with a given session_id. Returns the final 'answer'.
     """
+    # We invoke the chain with the user query; 
+    # the chain automatically updates the session’s chat history.
     response = conversational_rag_chain.invoke(
         {"input": user_query},
         config={