Update app.py

app.py

@@ -1,203 +1,325 @@
+import random
+from datetime import datetime
+
 import streamlit as st
 from openai import OpenAI
 from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
-from datetime import datetime
-import random
 
+
+# -----------------------------
+# Page config
+# -----------------------------
 st.set_page_config(
     page_title="Bot",
     page_icon="🤖",
     initial_sidebar_state="collapsed",
     layout="wide",
     menu_items={
-        'Report a bug': "mailto:yk408@cam.ac.uk",
-        'About': "Bot"
-    }
+        "Report a bug": "mailto:yk408@cam.ac.uk",
+        "About": "Bot",
+    },
 )
+
 st.markdown(
-    """ <style>
-            div[role="radiogroup"] >  :first-child{
-                display: none !important;
-            }
-        </style>
-        """,
-    unsafe_allow_html=True
+    """
+    <style>
+        div[role="radiogroup"] > :first-child {
+            display: none !important;
+        }
+    </style>
+    """,
+    unsafe_allow_html=True,
 )
 
-### Setting up the session state 
 
-if 'inserted' not in st.session_state:
-    # web app state
-    st.session_state.max_messages = 50
+# -----------------------------
+# Constants
+# -----------------------------
+MAX_MESSAGES_DEFAULT = 50
+SUBMIT_AFTER_USER_TURNS = 5
+
+GREETINGS = {
+    "hi", "hello", "hey", "heya", "hiya", "yo", "howdy", "sup",
+    "good morning", "good afternoon", "good evening", "good day",
+    "what's up", "whats up", "how do you do", "greetings",
+    "salutations", "hi there", "hello there", "hey there",
+    "how's it going", "hows it going", "how are you", "how are you doing"
+}
+
+CONDITIONS = {
+    "1": {
+        "label": "true control",
+        "base_url": None,
+        "model": None,
+        "api_secret": None,
+    },
+    "2": {
+        "label": "base",
+        "base_url": "https://openrouter.ai/api/v1",
+        "model": "meta-llama/llama-3.3-70b-instruct",
+        "api_secret": "OPENROUTER_API_KEY",
+    },
+    "3": {
+        "label": "bridging",
+        "base_url": "https://tinker.thinkingmachines.dev/services/tinker-prod/oai/api/v1",
+        "model": "tinker://808e4f02-e847-54ae-bc75-f14ee885ce5a:train:0/sampler_weights/final_sampler",
+        "api_secret": "TINKER_API_KEY",
+    },
+    "4": {
+        "label": "gpt",
+        "base_url": "https://openrouter.ai/api/v1",
+        "model": "openai/gpt-5.4",
+        "api_secret": "OPENROUTER_API_KEY",
+    },
+}
+
+
+# -----------------------------
+# Session state helpers
+# -----------------------------
+def init_session_state() -> None:
+    if "initialized" in st.session_state:
+        return
+
+    user_id = str(random.randint(100000, 999999))
+
+    st.session_state.initialized = True
+    st.session_state.inserted = 0
+    st.session_state.max_messages = MAX_MESSAGES_DEFAULT
     st.session_state.messages = []
-    st.session_state.API = ""
-    st.session_state.client = ""
-    
-    st.session_state.user_data = {}
-    st.session_state.user_data["BASE_URL"] = ""
-    st.session_state.user_data["MODEL_PATH"] = ""
-    st.session_state.user_data["url_id"] = True
-    st.session_state.user_data["user_id"] = str(random.randint(100000, 999999))
-    st.session_state.user_data["start_time"] = datetime.now()
-    st.session_state.inserted = 1
-
-    
-    st.session_state.convo_start_time = ''
+    st.session_state.client = None
     st.session_state.setup = False
+    st.session_state.convo_start_time = None
 
-if 'p' not in st.query_params or st.query_params['p'] == '':
-    st.query_params['p'] = st.radio("Select a condition for the conversation",
-                ['', '1', '2', '3', '4'],
-                help="1 = true control, 2 = base, 3 = bridging, 4 = gpt")
-
-if 'id' not in st.query_params:
-    st.session_state.user_data["url_id"] = False
-    st.query_params['id'] = st.session_state.user_data["user_id"] 
-
-def setup_messages():
-    # 1 = true control,
-    # 2 = base,
-    # 3 = bridging,
-    # 4 = GPT
-
-    if st.query_params["p"] == "1":
-        pass
-        
-    elif st.query_params["p"] == "2":
-        st.session_state.user_data["BASE_URL"] = "https://openrouter.ai/api/v1"
-        st.session_state.user_data["MODEL_PATH"] = "meta-llama/llama-3.3-70b-instruct"
-        st.session_state.API = "OPENROUTER_API_KEY"
-        
-    elif st.query_params["p"] == "3":
-        st.session_state.user_data["BASE_URL"] = "https://tinker.thinkingmachines.dev/services/tinker-prod/oai/api/v1"
-        st.session_state.user_data["MODEL_PATH"] = "tinker://808e4f02-e847-54ae-bc75-f14ee885ce5a:train:0/sampler_weights/final_sampler"
-        st.session_state.API = "TINKER_API_KEY"
-        
-    elif st.query_params["p"] == "4":
-        st.session_state.user_data["BASE_URL"] = "https://openrouter.ai/api/v1"
-        st.session_state.user_data["MODEL_PATH"] = "openai/gpt-5.4"
-        st.session_state.API = "OPENROUTER_API_KEY"
+    st.session_state.user_data = {
+        "BASE_URL": "",
+        "MODEL_PATH": "",
+        "url_id": True,
+        "user_id": user_id,
+        "start_time": datetime.now(),
+        "random_pid": None,
+        "condition": None,
+    }
 
 
+def reset_conversation_state() -> None:
+    st.session_state.messages = []
+    st.session_state.max_messages = MAX_MESSAGES_DEFAULT
+    st.session_state.convo_start_time = None
+    st.session_state.setup = False
+    st.session_state.client = None
+
+
+# -----------------------------
+# Query param / condition helpers
+# -----------------------------
+def ensure_query_params() -> None:
+    if "p" not in st.query_params or st.query_params["p"] not in CONDITIONS:
+        st.query_params["p"] = st.radio(
+            "Select a condition for the conversation",
+            ["", "1", "2", "3", "4"],
+            help="1 = true control, 2 = base, 3 = bridging, 4 = gpt",
+        )
+
+    if "id" not in st.query_params or not st.query_params["id"]:
+        st.session_state.user_data["url_id"] = False
+        st.query_params["id"] = st.session_state.user_data["user_id"]
+
+
+def setup_conversation() -> None:
+    condition = st.query_params.get("p", "")
+    if condition not in CONDITIONS:
+        return
+
+    config = CONDITIONS[condition]
+
+    st.session_state.user_data["random_pid"] = st.query_params["id"]
+    st.session_state.user_data["condition"] = condition
+    st.session_state.user_data["BASE_URL"] = config["base_url"] or ""
+    st.session_state.user_data["MODEL_PATH"] = config["model"] or ""
     st.session_state.convo_start_time = datetime.now()
-    st.session_state.user_data["random_pid"] = st.query_params['id']
-    st.session_state.user_data["condition"] = st.query_params['p']
 
-    st.session_state.client = OpenAI(base_url=st.session_state.user_data["BASE_URL"], api_key=st.secrets[st.session_state.API])
+    if condition != "1":
+        st.session_state.client = OpenAI(
+            base_url=config["base_url"],
+            api_key=st.secrets[config["api_secret"]],
+        )
+    else:
+        st.session_state.client = None
+
     st.session_state.setup = True
 
-if st.session_state.setup == False:
-    setup_messages()
-    
-if len(st.session_state.messages) == 1 and st.session_state.inserted < 2:
-    st.success("Ask, request, or talk to the chatbot about something you consider **politically polarizing** or something that people from different US political parties might disagree about.", icon='🎯')
 
-### App interface 
-with st.sidebar:
-    st.markdown("# Let's talk!")
-    st.markdown(f"""# **Step 1. Type in the chat box to start a conversation**""")
+# -----------------------------
+# Utility helpers
+# -----------------------------
+def is_greeting_only(text: str) -> bool:
+    clean = text.strip().lower().strip("!,.?")
+    return clean in GREETINGS
 
-    st.success("Ask, request, or talk to the chatbot about something you consider **politically polarizing** or something that people from different US political parties might disagree about.", icon='🎯')
 
-    st.markdown("🚫 Please avoid greetings and start the conversation with a question or a statement about a politically polarizing topic. **Note: the chatbot's knowledge only goes up to late August 2025.**")
+def user_turn_count() -> int:
+    return sum(1 for m in st.session_state.messages if m["role"] == "user")
 
-    st.markdown(f"""# **Step 2. Use the *Submit Interaction* button to get your chatbot word** 
-    
-    ⚠️ You must respond **at least 5 times** before you will see a *Submit Interaction* button. You can continue before submitting, but **you must Submit Interaction and enter your chatbot word to proceed with the survey**.
 
-    ❗ Do not share any personal information (e.g., name or address). Do not use AI tools to write your responses. If you encounter any technical issues, please let us know. It might sometimes take 30 seconds or more to generate a response, so please be patient.
-    """)
-                
-for message in st.session_state.messages:
-    if message['role']!='system':
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+def can_submit() -> bool:
+    return (
+        user_turn_count() >= SUBMIT_AFTER_USER_TURNS
+        or len(st.session_state.messages) >= st.session_state.max_messages
+    )
 
-if len(st.session_state.messages) >= st.session_state.max_messages:
-    st.info(
-        "You have reached the limit of messages for this conversation. Please end and submit the conversation."
+
+def render_sidebar() -> None:
+    with st.sidebar:
+        st.markdown("# Let's talk!")
+        st.markdown("# **Step 1. Type in the chat box to start a conversation**")
+
+        st.success(
+            "Ask, request, or talk to the chatbot about something you consider "
+            "**politically polarizing** or something that people from different "
+            "US political parties might disagree about.",
+            icon="🎯",
+        )
+
+        st.markdown(
+            "🚫 Please avoid greetings and start the conversation with a question "
+            "or a statement about a politically polarizing topic. "
+            "**Note: the chatbot's knowledge only goes up to late August 2025.**"
+        )
+
+        st.markdown(
+            "# **Step 2. Use the *Submit Interaction* button to get your chatbot word**\n\n"
+            "⚠️ You must respond **at least 5 times** before you will see a *Submit Interaction* button. "
+            "You can continue before submitting, but **you must Submit Interaction and enter your chatbot word "
+            "to proceed with the survey**.\n"
+            "❗ Do not share any personal information (e.g., name or address). "
+            "Do not use AI tools to write your responses. "
+            "If you encounter any technical issues, please let us know. "
+            "It might sometimes take 30 seconds or more to generate a response, so please be patient."
+        )
+
+
+def render_messages() -> None:
+    for message in st.session_state.messages:
+        if message["role"] != "system":
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+
+
+def save_conversation() -> None:
+    payload = dict(st.session_state.user_data)
+    payload["messages"] = st.session_state.messages
+    payload["convo_start_time"] = st.session_state.convo_start_time
+    payload["convo_end_time"] = datetime.now()
+    payload["inserted"] = st.session_state.inserted + 1
+
+    with MongoClient(st.secrets["mongo"], server_api=ServerApi("1")) as mongo_client:
+        db = mongo_client.bridge
+        collection = db.app2
+        collection.insert_one(payload)
+
+
+def get_assistant_response() -> str:
+    completion = st.session_state.client.chat.completions.create(
+        model=st.session_state.user_data["MODEL_PATH"],
+        messages=st.session_state.messages,
+        stream=False,
+    )
+    return completion.choices[0].message.content or ""
+
+
+# -----------------------------
+# App bootstrap
+# -----------------------------
+init_session_state()
+ensure_query_params()
+
+if not st.session_state.setup and st.query_params.get("p") in CONDITIONS:
+    setup_conversation()
+
+render_sidebar()
+
+if st.session_state.setup and not st.session_state.messages:
+    st.success(
+        "Ask, request, or talk to the chatbot about something you consider "
+        "**politically polarizing** or something that people from different US political parties might disagree about.",
+        icon="🎯",
     )
 
-elif st.session_state.inserted > 1:
+render_messages()
+
+
+# -----------------------------
+# Main app flow
+# -----------------------------
+if len(st.session_state.messages) >= st.session_state.max_messages:
+    st.info("You have reached the limit of messages for this conversation. Please end and submit the conversation.")
+
+elif st.session_state.inserted > 0:
     st.markdown("## Copy your WORD!")
-    st.markdown('**Your chatbot WORD is:**')
-    st.markdown(f'## TOMATOES')
-    st.markdown('**Please copy the WORD and enter it into the survey field below.**')
-
-elif prompt := st.chat_input("Type to ask a question or respond..."):   
-    if len(st.session_state.messages) == 1:
-        clean_prompt = prompt.strip().lower()
-        clean_prompt = clean_prompt.strip("!,.?")
-        greetings = {
-            "hi", "hello", "hey", "heya", "hiya", "yo", "howdy", "sup",
-            "good morning", "good afternoon", "good evening", "good day",
-            "what's up", "whats up", "how do you do", "greetings",
-            "salutations", "hi there", "hello there", "hey there",
-            "how's it going", "hows it going", "how are you", "how are you doing"
-        }
+    st.markdown("**Your chatbot WORD is:**")
+    st.markdown("## TOMATOES")
+    st.markdown("**Please copy the WORD and enter it into the survey field below.**")
 
-        if clean_prompt in greetings:
-            st.error(
-                "Please avoid greetings and start the conversation with a question or a statement about a politically polarizing topic.",
-                icon="🚫"
-            )
-            st.stop()
+elif prompt := st.chat_input("Type to ask a question or respond..."):
+    if not st.session_state.setup:
+        st.error("Please select a condition first.")
+        st.stop()
+
+    if not st.session_state.messages and is_greeting_only(prompt):
+        st.error(
+            "Please avoid greetings and start the conversation with a question or a statement about a politically polarizing topic.",
+            icon="🚫",
+        )
+        st.stop()
 
     st.session_state.messages.append({"role": "user", "content": prompt})
+
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    if st.query_params["p"] == "1":
-        rate_limit_message = "Thank you for your question! You have been randomly assigned to a condition without a chatbot. **Please submit your interaction anyway** to get your chatbot word and proceed with the survey. Do not worry, this will not influence your compensation."
-        st.session_state.messages.append({"role": "assistant", "content": rate_limit_message})
+    condition = st.session_state.user_data["condition"]
+
+    if condition == "1":
+        response = (
+            "Thank you for your question. You have been randomly assigned to a condition "
+            "without a chatbot. **Please submit your interaction anyway** to get your chatbot word "
+            "and proceed with the survey. Do not worry, this will not influence your compensation."
+        )
+        st.session_state.messages.append({"role": "assistant", "content": response})
         st.session_state.max_messages = len(st.session_state.messages)
         st.rerun()
 
-    else:
-        with st.chat_message("assistant"):
-            try:
-                with st.spinner("Typing..."):
-                    completion = st.session_state.client.chat.completions.create(
-                        model=st.session_state.user_data["MODEL_PATH"],
-                        messages=[
-                            {"role": m["role"], "content": m["content"]}
-                            for m in st.session_state.messages
-                        ],
-                        stream=False,
-                    )
-        
-                response = completion.choices[0].message.content or ""
-                st.markdown(response)
-        
-                st.session_state.messages.append({"role": "assistant", "content": response})
-
-            except:
-                rate_limit_message = "An error has occurred or you've reached the maximum conversation length. Please submit the conversation."
-                st.session_state.messages.append(
-                    {"role": "assistant", "content": rate_limit_message}
-                )
-                st.session_state.max_messages = len(st.session_state.messages)
-                st.rerun()
-
-if len(st.session_state.messages) > 10 or st.session_state.max_messages == len(st.session_state.messages):
-    columns = st.columns((1,1,1))
-    with columns[2]:
-        if st.button("Submit Interaction",use_container_width=True):
-            keys = ["inserted", "messages", "convo_start_time"]
+    with st.chat_message("assistant"):
+        try:
+            with st.spinner("Typing..."):
+                response = get_assistant_response()
 
-            st.session_state.user_data.update({key: st.session_state[key] for key in keys})
-            
-            st.session_state.user_data["convo_end_time"] = datetime.now()
+            st.markdown(response)
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
-            with MongoClient(st.secrets["mongo"],server_api=ServerApi('1')) as client:
-                    db = client.bridge
-                    collection = db.app2
-                    user_data = st.session_state.user_data
+        except Exception as e:
+            error_message = (
+                "An error has occurred or you've reached the maximum conversation length. "
+                "Please submit the conversation."
+            )
+            st.session_state.messages.append({"role": "assistant", "content": error_message})
+            st.session_state.max_messages = len(st.session_state.messages)
+            st.error(f"Request failed: {e}")
 
-                    collection.insert_one(user_data)  
-                    st.session_state.inserted += 1
-                    done = True
 
-                    setup_messages()
-                    st.rerun()
+# -----------------------------
+# Submit button
+# -----------------------------
+if can_submit() and st.session_state.inserted == 0:
+    cols = st.columns((1, 1, 1))
+    with cols[2]:
+        if st.button("Submit Interaction", use_container_width=True):
+            try:
+                save_conversation()
+                st.session_state.inserted += 1
+                reset_conversation_state()
+                st.rerun()
+            except Exception as e:
+                st.error(f"Failed to save conversation: {e}")