adding survey logic

chatbot.py

@@ -5,7 +5,7 @@ import psycopg2
 from urllib.parse import urlparse
 
 # Initialize OpenAI client
-client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Database URL
 DATABASE_URL = os.getenv('DATABASE_URL')
@@ -50,6 +50,12 @@ def create_table():
             transcript TEXT NOT NULL,
             timestamp TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
         );
+        CREATE TABLE IF NOT EXISTS survey_responses (
+            id SERIAL PRIMARY KEY,
+            user_id INTEGER REFERENCES users(id),
+            responses JSONB,
+            timestamp TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
+        );
         '''
         cur.execute(create_table_query)
         conn.commit()
@@ -103,9 +109,42 @@ def authenticate(username, password):
             return user_id[0]
     return None
 
+# Store survey responses
+def store_survey_responses(user_id, responses):
+    conn = connect_db()
+    if conn:
+        cur = conn.cursor()
+        insert_query = '''
+        INSERT INTO survey_responses (user_id, responses)
+        VALUES (%s, %s);
+        '''
+        cur.execute(insert_query, (user_id, responses))
+        conn.commit()
+        cur.close()
+        conn.close()
+
 # Initialize the table
 create_table()
 
+# Survey questions with different input types
+survey_questions = [
+    {"question": "On a scale of 1 to 5, how are you feeling today?", "type": "scale", "options": ["1", "2", "3", "4", "5"]},
+    {"question": "How often do you feel stressed?", "type": "multiple_choice", "options": ["Never", "Rarely", "Sometimes", "Often", "Always"]},
+    {"question": "Do you have trouble sleeping? (True/False)", "type": "true_false"}
+]
+
+# Function to create survey inputs dynamically
+def create_survey_inputs():
+    inputs = []
+    for q in survey_questions:
+        if q["type"] == "scale":
+            inputs.append(gr.Radio(label=q["question"], choices=q["options"]))
+        elif q["type"] == "multiple_choice":
+            inputs.append(gr.Radio(label=q["question"], choices=q["options"]))
+        elif q["type"] == "true_false":
+            inputs.append(gr.Radio(label=q["question"], choices=["True", "False"]))
+    return inputs
+
 # Initial system messages for the chatbot
 initial_messages = [
     {"role": "system", "content": "You are an attachment and close relationship research surveyor"},
@@ -121,11 +160,11 @@ def chatbot(input, state):
     
     if input:
         messages.append({"role": "user", "content": input})
-        response = client.chat.completions.create(
+        response = openai.ChatCompletion.create(
             model="gpt-3.5-turbo",
             messages=messages
         )
-        reply = response.choices[0].message.content
+        reply = response.choices[0].message['content']
         messages.append({"role": "assistant", "content": reply})
         
         conversation = ""
@@ -147,6 +186,9 @@ with gr.Blocks() as demo:
     chat_output = gr.Textbox(label="Conversation", visible=False)
     state = gr.State([None, initial_messages.copy()])
 
+    survey_inputs = create_survey_inputs()
+    submit_survey = gr.Button("Submit Survey")
+
     def login(username, password):
         user_id = authenticate(username, password)
         if user_id:
@@ -161,8 +203,15 @@ with gr.Blocks() as demo:
         else:
             return gr.update(visible=False), gr.update(visible=False), [None, initial_messages.copy()], "Registration failed, try a different username."
 
+    def submit_survey_fn(responses, state):
+        user_id = state[0]
+        survey_responses = {survey_questions[i]["question"]: responses[i] for i in range(len(responses))}
+        store_survey_responses(user_id, survey_responses)
+        return gr.update(visible=True), gr.update(visible=True), state
+
     login_button.click(login, inputs=[username, password], outputs=[chat_input, chat_output, state, auth_message])
     register_button.click(register, inputs=[username, password], outputs=[chat_input, chat_output, state, auth_message])
+    submit_survey.click(submit_survey_fn, inputs=survey_inputs + [state], outputs=[chat_input, chat_output, state])
 
     chat_interface = gr.Interface(
         fn=chatbot,