Update my_memory_logic.py

my_memory_logic.py

@@ -2,52 +2,55 @@
 
 import os
 
-# Import your actual RAG chain (or pipeline) from pipeline.py
-# We'll assume `rag_chain` is exposed by pipeline.py
-from pipeline import run_with_chain_context  
-# We'll import the session-based classes from langchain_core
-# If they're in different modules, adjust accordingly.
+# Import the "run_with_chain_context" function from pipeline.py
+# This function must accept a dict with { "input": ..., "chat_history": ... }
+# and return a dict with { "answer": ... }.
+from pipeline import run_with_chain_context
+
+# For session-based chat history
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_community.chat_message_histories import ChatMessageHistory
 from langchain_core.runnables.history import RunnableWithMessageHistory
 
 ###############################################################################
-# 1) We'll keep an in-memory store mapping session_id -> ChatMessageHistory
+# 1) In-Memory Store: session_id -> 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.
+    Retrieve (or create) a ChatMessageHistory for the given session_id.
+    This ensures each session_id has its own conversation transcripts.
     """
     if session_id not in store:
         store[session_id] = ChatMessageHistory()
     return store[session_id]
 
 ###############################################################################
-# 2) Create the RunnableWithMessageHistory (session-based chain)
+# 2) Build a RunnableWithMessageHistory that wraps "run_with_chain_context"
 ###############################################################################
-# This wraps your `rag_chain` so it automatically reads/writes
-# conversation history from get_session_history for each session.
+# "run_with_chain_context" must be a function returning a dict,
+# e.g. { "answer": "... final string ..." }
+# input_messages_key -> "input"
+# history_messages_key -> "chat_history"
+# output_messages_key -> "answer"
+
 conversational_rag_chain = RunnableWithMessageHistory(
-    run_with_chain_context,               # the main chain 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
+    run_with_chain_context,    # from pipeline.py
+    get_session_history,
+    input_messages_key="input",
+    history_messages_key="chat_history",
+    output_messages_key="answer"
 )
 
 ###############################################################################
-# 3) A convenience function to run a query with session-based memory
+# 3) A convenience function that calls our chain with session-based memory
 ###############################################################################
 def run_with_session_memory(user_query: str, session_id: str) -> str:
     """
-    A helper that calls our `conversational_rag_chain` 
-    with a given session_id. Returns the final 'answer'.
+    Calls the 'conversational_rag_chain' with a given session_id and user_query.
+    This returns the final 'answer' from run_with_chain_context.
     """
-    # 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={