Update chatbot.py

chatbot.py

@@ -40,9 +40,14 @@ def create_table():
     if conn:
         cur = conn.cursor()
         create_table_query = '''
+        CREATE TABLE IF NOT EXISTS users (
+            id SERIAL PRIMARY KEY,
+            username VARCHAR(255) UNIQUE NOT NULL,
+            password VARCHAR(255) NOT NULL
+        );
         CREATE TABLE IF NOT EXISTS chat_transcripts (
             id SERIAL PRIMARY KEY,
-            user_id VARCHAR(255) NOT NULL,
+            user_id INTEGER REFERENCES users(id),
             transcript TEXT NOT NULL,
             timestamp TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
         );
@@ -66,6 +71,39 @@ def store_transcript(user_id, transcript):
         cur.close()
         conn.close()
 
+# Register new user
+def register_user(username, password):
+    conn = connect_db()
+    if conn:
+        cur = conn.cursor()
+        insert_query = '''
+        INSERT INTO users (username, password)
+        VALUES (%s, %s) RETURNING id;
+        '''
+        cur.execute(insert_query, (username, password))
+        user_id = cur.fetchone()[0]
+        conn.commit()
+        cur.close()
+        conn.close()
+        return user_id
+    return None
+
+# Authenticate user
+def authenticate(username, password):
+    conn = connect_db()
+    if conn:
+        cur = conn.cursor()
+        select_query = '''
+        SELECT id FROM users WHERE username = %s AND password = %s;
+        '''
+        cur.execute(select_query, (username, password))
+        user_id = cur.fetchone()
+        cur.close()
+        conn.close()
+        if user_id:
+            return user_id[0]
+    return None
+
 # Initialize the table
 create_table()
 
@@ -99,18 +137,12 @@ def chatbot(user_id, input):
         store_transcript(user_id, conversation)
         return 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")
-    
+    register_button = gr.Button("Register")
+
     chat_interface = gr.Interface(
         fn=chatbot,
         inputs=[gr.Textbox(label="User ID"), gr.Textbox(lines=7, label="Chat with AttachmentBot")],
@@ -118,7 +150,7 @@ with gr.Blocks() as demo:
         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:
@@ -126,6 +158,14 @@ with gr.Blocks() as demo:
         else:
             return gr.update(visible=False), "", "Invalid credentials"
 
+    def register(username, password):
+        user_id = register_user(username, password)
+        if user_id:
+            return gr.update(visible=True), user_id, "Registration successful, you can now login."
+        else:
+            return gr.update(visible=False), "", "Registration failed, try a different username."
+
     login_button.click(login, inputs=[username, password], outputs=[chat_interface, gr.Variable(), gr.Textbox(visible=False)])
+    register_button.click(register, inputs=[username, password], outputs=[chat_interface, gr.Variable(), gr.Textbox(visible=False)])
 
 demo.launch()