Update app.py

app.py

@@ -5,6 +5,7 @@ import time
 
 import gradio as gr
 from transformers import AutoTokenizer
+import psycopg2
 
 
 DESCRIPTION = """
@@ -47,6 +48,28 @@ MAX_SEC = 30
 
 tokenizer = AutoTokenizer.from_pretrained("MediaTek-Research/Breeze-7B-Instruct-v0_1")
 
+def insert_to_db(prompt, response, temperature, top_p):
+    #Establishing the connection
+    conn = psycopg2.connect(
+       database=os.environ.get("DB"), user=os.environ.get("USER"), password=os.environ.get("DB_PASS"), host=os.environ.get("DB_HOST"), port= '5432'
+    )
+    #Setting auto commit false
+    conn.autocommit = True
+    
+    #Creating a cursor object using the cursor() method
+    cursor = conn.cursor()
+
+    # Preparing SQL queries to INSERT a record into the database.
+    cursor.execute(f"INSERT INTO EMPLOYEE(prompt, response, temperature, top_p) VALUES ('{prompt}', '{response}', {temperature}, {top_p})")
+
+    # Commit your changes in the database
+    conn.commit()
+
+    # Closing the connection
+    conn.close()
+
+
+
 def refusal_condition(query):
     # 不要再問這些問題啦！
 
@@ -187,8 +210,9 @@ with gr.Blocks() as demo:
             if refusal_condition(history[-1][1]):
                 history[-1][1] = history[-1][1] + '\n\n**[免責聲明: Breeze-7B-Instruct 和 Breeze-7B-Instruct-64k 並未針對問答進行安全保護，因此語言模型的任何回應不代表 MediaTek Research 立場。]**'
                 yield history
-    
+
         print('== Record ==\nQuery: {query}\nResponse: {response}'.format(query=repr(message), response=repr(history[-1][1])))
+        insert_to_db(message, history[-1][1], float(temperature), float(top_p))
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         fn=bot,