Update my_memory_logic.py

my_memory_logic.py

@@ -1,47 +1,64 @@
 # my_memory_logic.py
 import os
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import LLMChain
-from langchain.prompts.chat import (
-    ChatPromptTemplate,
-    SystemMessagePromptTemplate,
-    MessagesPlaceholder,
-    HumanMessagePromptTemplate,
-)
-# Import ChatGroq from the langchain_groq package
-from langchain_groq import ChatGroq
-
-# 1) Memory object for storing conversation messages
-memory = ConversationBufferMemory(return_messages=True)
-
-# 2) Restatement system prompt for question rewriting
-restatement_system_prompt = """
-You have a conversation history plus a new user question. 
-Your ONLY job is to rewrite the user's latest question so it 
-makes sense on its own. 
-- Do NOT repeat or quote large sections of the conversation. 
-- Do NOT provide any answer or summary. 
-- Just produce a short, standalone question as your final output.
-"""
-
-
-# 3) Build the ChatPromptTemplate
-restatement_prompt = ChatPromptTemplate.from_messages([
-    SystemMessagePromptTemplate.from_template(restatement_system_prompt),
-    MessagesPlaceholder(variable_name="chat_history"),
-    HumanMessagePromptTemplate.from_template("{input}")
-])
-
-# 4) Initialize the ChatGroq LLM
-#    Ensure you have your GROQ_API_KEY set in the environment
-restatement_llm = ChatGroq(
-    model="llama3-70b-8192",
-    # model="mixtral-8x7b-32768"# or whichever model
-    groq_api_key=os.environ["GROQ_API_KEY"]
-)
 
-# 5) Create the LLMChain for restatement
-restatement_chain = LLMChain(
-    llm=restatement_llm,
-    prompt=restatement_prompt
+# We'll import the session-based classes from langchain_core if you have them installed:
+# If not, you'll need to install the correct package versions or adapt to your environment.
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_community.chat_message_histories import ChatMessageHistory
+from langchain_core.runnables.history import RunnableWithMessageHistory
+
+# We'll assume you have a `rag_chain` from your pipeline code or can import it.
+# For example:
+# from pipeline import rag_chain
+
+# For demonstration, let's just define a dummy "rag_chain" that returns "answer".
+# In your real code, import your actual chain. 
+class DummyRagChain:
+    def invoke(self, inputs):
+        # returns a dictionary with "answer"
+        return {"answer": f"Dummy answer to '{inputs['input']}'."}
+
+rag_chain = DummyRagChain()
+
+###############################################################################
+# 1) We'll keep an in-memory store of session_id -> ChatMessageHistory
+###############################################################################
+store = {}  # { "abc123": ChatMessageHistory(...) }
+
+def get_session_history(session_id: str) -> BaseChatMessageHistory:
+    """
+    Retrieve or create a ChatMessageHistory object for the given session_id.
+    """
+    if session_id not in store:
+        store[session_id] = ChatMessageHistory()
+    return store[session_id]
+
+###############################################################################
+# 2) Create the RunnableWithMessageHistory (conversational chain)
+###############################################################################
+# If your snippet references `rag_chain`, combine it with get_session_history.
+conversational_rag_chain = RunnableWithMessageHistory(
+    rag_chain,               # your main chain (RAG or pipeline)
+    get_session_history,     # function to fetch chat history for a session
+    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
+###############################################################################
+def run_with_session_memory(user_query: str, session_id: str) -> str:
+    """
+    A convenience wrapper that calls our `conversational_rag_chain` 
+    with a specific session_id. This returns the final 'answer'.
+    """
+    response = conversational_rag_chain.invoke(
+        {"input": user_query},
+        config={
+            "configurable": {
+                "session_id": session_id
+            }
+        }
+    )
+    return response["answer"]