feat: save chat history

app.py

@@ -3,6 +3,7 @@ import json
 
 import streamlit as st
 import streamlit.components.v1 as components
+from chat_history import insert_chat_history
 from css import load_css
 from custom_pgvector import CustomPGVector
 from langchain import OpenAI
@@ -77,6 +78,7 @@ def on_click_callback():
             )
         )
         st.session_state.token_count += cb.total_tokens
+        insert_chat_history(conn, human_prompt, llm_response["answer"])
 
 
 load_css()
@@ -167,4 +169,3 @@ with doc_column:
             expander.markdown("**Authors** : " + doc_content["authors"])
             expander.markdown("**Keywords** : " + doc_content["keywords"])
             expander.markdown("**Distance** : " + str(doc_content["distance"]))
-            

chat_history.py

@@ -0,0 +1,17 @@
+import sqlalchemy
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+
+def insert_chat_history(conn: sqlalchemy.engine.Connection, query: str, answer: str):
+    with Session(conn) as conn:
+        conn.execute(
+            text("INSERT INTO chat_history (query, answer) VALUES (:query, :answer);"),
+            [
+                {
+                    "query": query,
+                    "answer": answer,
+                }
+            ],
+        )
+        conn.commit()