Update my_memory_logic.py

my_memory_logic.py

@@ -1,21 +1,44 @@
 # my_memory_logic.py
+import os
 from langchain.memory import ConversationBufferMemory
-from langchain.chat_models import ChatOpenAI
 from langchain.chains import LLMChain
-from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
+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 = (
     "Given a chat history and the latest user question "
-    "which might reference context in the chat history, ..."
+    "which might reference context in the chat history, "
+    "formulate a standalone question that can be understood "
+    "without the chat history. Do NOT answer the question, "
+    "just reformulate it if needed; otherwise return it as is."
 )
 
+# 3) Build the ChatPromptTemplate
 restatement_prompt = ChatPromptTemplate.from_messages([
     SystemMessagePromptTemplate.from_template(restatement_system_prompt),
     MessagesPlaceholder(variable_name="chat_history"),
     HumanMessagePromptTemplate.from_template("{input}")
 ])
 
-restatement_llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.0)
-restatement_chain = LLMChain(llm=restatement_llm, prompt=restatement_prompt)
+# 4) Initialize the ChatGroq LLM
+#    Ensure you have your GROQ_API_KEY set in the environment
+restatement_llm = ChatGroq(
+    model="llama3-70b-8192"              # 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
+)