Update app.py

app.py

@@ -13,6 +13,8 @@ from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain.memory import ConversationBufferMemory, PostgresChatMessageHistory
 
+import psycopg2
+import urllib.parse as up
 
 
 
@@ -29,18 +31,40 @@ def get_llm_chain():
 
 
 @st.cache_resource
-def get_memory():
-    return PostgresChatMessageHistory(connection_string=POSTGRE_URL, session_id=str(datetime.timestamp(datetime.now())))
+def get_db_connection(conn_url, password=None):
+    
+    url = up.urlparse(conn_url)
+
+    conn = psycopg2.connect(
+        database=url.path[1:],
+        user=url.username,
+        password=password if password is not None else url.password,
+        host=url.hostname,
+        port=url.port
+    )
+
+    print("Connection to database succesfull!")
+    return conn
+
+# @st.cache_resource
+# def get_memory():
+#     return PostgresChatMessageHistory(connection_string=POSTGRE_URL, session_id=str(datetime.timestamp(datetime.now())))
+
 
+if 'conn' not in st.session_state:
+    st.session_state.conn = get_db_connection(POSTGRE_URL)
+
+if 'cursor' not in st.session_state:
+    st.session_state.cursor = st.session_state.conn.cursor()
 
 if 'memory' not in st.session_state:
-    # st.session_state['memory'] = ConversationBufferMemory(return_messages=True)
+    st.session_state['memory'] = ConversationBufferMemory(return_messages=True)
     
     # st.session_state.memory = PostgresChatMessageHistory(connection_string=POSTGRE_URL, session_id=str(datetime.timestamp(datetime.now())))
 
-    st.sessio_state.memory = get_memory()
-    # st.session_state.memory.chat_memory.add_ai_message("Hello, My name is Jonathan Jordan. You can call me Jojo. How can I help you today?")
-    st.session_state.memory.add_ai_message("Hello, My name is Jonathan Jordan. You can call me Jojo. How can I help you today?")
+    # st.session_state.memory = get_memory()
+    st.session_state.memory.chat_memory.add_ai_message("Hello, My name is Jonathan Jordan. You can call me Jojo. How can I help you today?")
+    # st.session_state.memory.add_ai_message("Hello, My name is Jonathan Jordan. You can call me Jojo. How can I help you today?")
 
 if 'chain' not in st.session_state:
     # st.session_state['chain'] = custom_chain_with_history(
@@ -55,7 +79,7 @@ if 'chain' not in st.session_state:
 
 st.title("Chat With Me")
 st.subheader("by Jonathan Jordan")
-st.markdown("Note : Your chat will be recorded in our Database, thank you :)")
+st.markdown("Note : This conversation will be recorded in our private Database, thank you :)")
 
 # Initialize chat history
 if "messages" not in st.session_state:
@@ -78,10 +102,20 @@ if prompt := st.chat_input("Ask me anything.."):
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
         st.markdown(response)
-    st.session_state.memory.add_user_message(prompt)
-    st.session_state.memory.add_ai_message(response)
-    # st.session_state.memory.save_context({"question":prompt}, {"output":response})
-    # st.session_state.memory.chat_memory.messages = st.session_state.memory.chat_memory.messages[-15:]
+        
+    # st.session_state.memory.add_user_message(prompt)
+    # st.session_state.memory.add_ai_message(response)
+    st.session_state.memory.save_context({"question":prompt}, {"output":response})
+    st.session_state.memory.chat_memory.messages = st.session_state.memory.chat_memory.messages[-15:]
     # Add assistant response to chat history
     st.session_state.messages.append({"role": "assistant", "content": response})
-    
+
+    # Insert data into the table
+    st.session_state.cursor.execute(
+        "INSERT INTO chat_history (input_text, response_text, created_at) VALUES (%s, %s, %s)",
+        (prompt, response, datetime.now())
+    )
+    
+    # Commit the transaction
+    st.session_state.conn.commit()
+