Update chatbot.py

chatbot.py

@@ -1,10 +1,74 @@
 import openai
 import gradio as gr
 import os
+import psycopg2
+from urllib.parse import urlparse
 
 # Initialize OpenAI client
 client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
+# Database URL
+DATABASE_URL = os.getenv('DATABASE_URL')
+
+# Parse the database URL
+result = urlparse(DATABASE_URL)
+
+username = result.username
+password = result.password
+database = result.path[1:]
+hostname = result.hostname
+port = result.port
+
+# Connect to the database
+def connect_db():
+    try:
+        conn = psycopg2.connect(
+            dbname=database,
+            user=username,
+            password=password,
+            host=hostname,
+            port=port
+        )
+        return conn
+    except Exception as e:
+        print(f"Error connecting to the database: {e}")
+        return None
+
+# Create table if not exists
+def create_table():
+    conn = connect_db()
+    if conn:
+        cur = conn.cursor()
+        create_table_query = '''
+        CREATE TABLE IF NOT EXISTS chat_transcripts (
+            id SERIAL PRIMARY KEY,
+            user_id VARCHAR(255) NOT NULL,
+            transcript TEXT NOT NULL,
+            timestamp TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
+        );
+        '''
+        cur.execute(create_table_query)
+        conn.commit()
+        cur.close()
+        conn.close()
+
+# Store chat transcript
+def store_transcript(user_id, transcript):
+    conn = connect_db()
+    if conn:
+        cur = conn.cursor()
+        insert_query = '''
+        INSERT INTO chat_transcripts (user_id, transcript)
+        VALUES (%s, %s);
+        '''
+        cur.execute(insert_query, (user_id, transcript))
+        conn.commit()
+        cur.close()
+        conn.close()
+
+# Initialize the table
+create_table()
+
 # Initial system messages for the chatbot
 initial_messages = [
     {"role": "system", "content": "You are an attachment and close relationship research surveyor"},
@@ -14,7 +78,7 @@ Do you usually discuss your problems and concerns with your mother or a mother-l
 """},
 ]
 
-def chatbot(input):
+def chatbot(user_id, input):
     if not hasattr(chatbot, "messages"):
         chatbot.messages = initial_messages.copy()
     
@@ -24,7 +88,7 @@ def chatbot(input):
             model="gpt-3.5-turbo",
             messages=chatbot.messages
         )
-        reply = response.choices[0].message.content
+        reply = response.choices[0].message["content"]
         chatbot.messages.append({"role": "assistant", "content": reply})
         
         conversation = ""
@@ -32,15 +96,36 @@ def chatbot(input):
             role = "You" if message["role"] == "user" else "AttachmentBot"
             conversation += f"{role}: {message['content']}\n"
 
+        store_transcript(user_id, conversation)
         return conversation
 
-inputs = gr.Textbox(lines=7, label="Chat with AttachmentBot")
-outputs = gr.Textbox(label="Conversation")
+def authenticate(username, password):
+    # Replace with actual authentication logic
+    if username == "user" and password == "pass":
+        return username
+    else:
+        return None
+
+with gr.Blocks() as demo:
+    username = gr.Textbox(label="Username")
+    password = gr.Textbox(label="Password", type="password")
+    login_button = gr.Button("Login")
+    
+    chat_interface = gr.Interface(
+        fn=chatbot,
+        inputs=[gr.Textbox(label="User ID"), gr.Textbox(lines=7, label="Chat with AttachmentBot")],
+        outputs=gr.Textbox(label="Conversation"),
+        title="AttachmentBot",
+        description="Let me survey you about your attachment with certain people in your life. To begin, enter 'start'."
+    )
+    
+    def login(username, password):
+        user_id = authenticate(username, password)
+        if user_id:
+            return gr.update(visible=True), user_id, ""
+        else:
+            return gr.update(visible=False), "", "Invalid credentials"
+
+    login_button.click(login, inputs=[username, password], outputs=[chat_interface, gr.Variable(), gr.Textbox(visible=False)])
 
-gr.Interface(
-    fn=chatbot, 
-    inputs=inputs, 
-    outputs=outputs, 
-    title="AttachmentBot",
-    description="Let me survey you about your attachment with certain people in your life. To begin, enter 'start'."
-).launch()
+demo.launch()