Update app.py

app.py

@@ -10,13 +10,13 @@ def reset_conversation():
     Resets Conversation
     '''
     st.session_state.messages = []
+    st.session_state.message_count = 0
+    st.session_state.ask_intervention = False
+    st.session_state.diagnosis = None
     return None
 
-# Define model links for the Addiction Recovery and Mental Health models
-model_links = {
-    "Biblical Religious Advisor": "NousResearch/Nous-Hermes-2-Yi-34B",
-    "Mental health AI": "NousResearch/Nous-Hermes-2-Yi-34B"
-}
+# Define model link for the Addiction Recovery and Mental Health models
+model_link = "NousResearch/Nous-Hermes-2-Yi-34B"
 
 # Define pre-instructions for each model
 model_pre_instructions = {
@@ -25,7 +25,7 @@ model_pre_instructions = {
 }
 
 # Function to interact with the selected model via the Together API
-def interact_with_together_api(messages, model_link):
+def interact_with_together_api(messages, model_link, diagnostic=False):
     all_messages = []
 
     # Add pre-instructions to the message history if it's the first interaction with this model
@@ -35,12 +35,18 @@ def interact_with_together_api(messages, model_link):
         all_messages.append({"role": "system", "content": f"Switched to model: {selected_model}"})
 
     # Append user and assistant messages
-    for human, assistant in messages:
-        all_messages.append({"role": "user", "content": human})
-        all_messages.append({"role": "assistant", "content": assistant})
-
-    # Add the latest user message
-    all_messages.append({"role": "user", "content": messages[-1][1]})
+    for message in messages:
+        if message[0] == "user":
+            all_messages.append({"role": "user", "content": message[1]})
+        else:
+            all_messages.append({"role": "assistant", "content": message[1]})
+
+    if diagnostic:
+        diagnostic_prompt = "Analyze the following conversation and predict the mental issue the user might be suffering from:"
+        all_messages.append({"role": "user", "content": diagnostic_prompt})
+        for message in messages:
+            if message[0] == "user":
+                all_messages.append({"role": "user", "content": message[1]})
 
     url = "https://api.together.xyz/v1/chat/completions"
     payload = {
@@ -60,17 +66,22 @@ def interact_with_together_api(messages, model_link):
         "Authorization": f"Bearer {TOGETHER_API_KEY}",
     }
 
-    response = requests.post(url, json=payload, headers=headers)
-    response.raise_for_status()  # Ensure HTTP request was successful
-
-    # Extract response from JSON
-    response_data = response.json()
-    assistant_response = response_data["choices"][0]["message"]["content"]
-
+    try:
+        response = requests.post(url, json=payload, headers=headers)
+        response.raise_for_status()  # Ensure HTTP request was successful
+        response_data = response.json()
+        assistant_response = response_data["choices"][0]["message"]["content"]
+    except requests.exceptions.RequestException as e:
+        assistant_response = "Sorry, I couldn't connect to the server. Please try again later."
+    
     return assistant_response
 
+# Function to diagnose mental health issue
+def diagnose_mental_health(messages):
+    return interact_with_together_api(messages, model_link, diagnostic=True)
+
 # Create sidebar with model selection dropdown and reset button
-selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
+selected_model = st.sidebar.selectbox("Select Model", list(model_pre_instructions.keys()))
 st.sidebar.button('Reset Chat', on_click=reset_conversation)
 
 # Add cautionary message about testing phase at the bottom of the sidebar
@@ -80,11 +91,18 @@ st.sidebar.markdown("**Note**: This model is still in the beta phase. Responses
 st.sidebar.image("https://assets.isu.pub/document-structure/221118065013-a6029cf3d563afaf9b946bb9497d45d4/v1/2841525b232adaef7bd0efe1da81a4c5.jpeg", width=200)
 st.sidebar.write("A product proudly developed by Kisii University")
 
+# Display diagnosis in the sidebar
+if st.session_state.message_count >= 4:
+    if "diagnosis" not in st.session_state or st.session_state.diagnosis is None:
+        st.session_state.diagnosis = diagnose_mental_health(st.session_state.messages)
+    st.sidebar.markdown(f"### Diagnosis:\n**{st.session_state.diagnosis}**")
+
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
     st.session_state.message_count = 0
     st.session_state.ask_intervention = False
+    st.session_state.diagnosis = None
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
@@ -112,17 +130,25 @@ if prompt := st.chat_input(f"Hi, I'm {selected_model}, ask me a question"):
     for keyword in intervention_keywords:
         if keyword in prompt.lower():
             # Intervention logic here
-            st.markdown("<span style='color:red;'>I have a feeling you may need to talk to a therapist. If you agree with me please contact +254793609747; Name: Davis. If you dont then keep talking to me as we figure this out.</span>", unsafe_allow_html=True)
+            st.markdown("<span style='color:red;'>I have a feeling you may need to talk to a therapist. If you agree with me please contact +254793609747; Name: Davis. If you don't, then keep talking to me as we figure this out.</span>", unsafe_allow_html=True)
             break  # Exit loop once intervention is triggered
 
     # Interact with the selected model
-    assistant_response = interact_with_together_api(st.session_state.messages, model_links[selected_model])
+    placeholder = st.empty()
+    with placeholder:
+        st.markdown("AI is typing...")
+    
+    assistant_response = interact_with_together_api(st.session_state.messages, model_link)
+    placeholder.empty()
 
     # Display assistant response in chat message container
-    with st.empty():
-        st.markdown("AI is typing...")
-        st.empty()
+    with st.chat_message("assistant"):
         st.markdown(assistant_response)
 
     # Add assistant response to chat history
     st.session_state.messages.append(("assistant", assistant_response))
+
+    # Update diagnosis if necessary
+    if st.session_state.message_count >= 4:
+        st.session_state.diagnosis = diagnose_mental_health(st.session_state.messages)
+        st.sidebar.markdown(f"### Diagnosis:\n**{st.session_state.diagnosis}**")