Update app.py

app.py

@@ -1,89 +1,115 @@
-""" Simple Chatbot
-@author: Nigel Gebodh
-@email: nigel.gebodh@gmail.com
-"""
-
 import streamlit as st
 import os
 import requests
-from dotenv import load_dotenv
 import json
 
-load_dotenv()
+entire_assistant_response = ""
 
-def reset_conversation():
-    '''
-    Resets Conversation
-    '''
-    st.session_state.conversation = []
-    st.session_state.messages = []
-    return None
-
-# Add reset button to clear conversation
-st.sidebar.button('Reset Chat', on_click=reset_conversation)
-
-# Model info for Mistral
-model_info = {
-    "Mistral": {
-        'description': """The Mistral model is a **Large Language Model (LLM)** that's able to have question and answer interactions.\n \
-                         \nIt was created by the [**Mistral AI**](https://mistral.ai/news/announcing-mistral-7b/) team as has over  **7 billion parameters.** \n""",
-        'logo': 'https://mistral.ai/images/logo_hubc88c4ece131b91c7cb753f40e9e1cc5_2589_256x0_resize_q97_h2_lanczos_3.webp'
-    },
-    "BibleLearnerAI": {
-        'description': """You're now chatting with **BibleLearnerAI**. This AI is focused on religion, specifically Christianity, and provides relevant Bible verses. When greeted, it responds with a religious greeting and introduces itself. It knows the Bible more than anything else.""",
-        'logo': 'https://your-bible-teacher.com/wp-content/uploads/2019/03/teacher-300x300.png'
+def get_streamed_response(message, history, model):
+    all_message = []
+    
+    for human, assistant in history:
+        all_message.append({"role": "user", "content": human})
+        all_message.append({"role": "assistant", "content": assistant})
+    
+    global entire_assistant_response
+    entire_assistant_response = ""  # Reset the entire assistant response
+
+    all_message.append({"role": "user", "content": message})
+
+    url = "https://api.together.xyz/v1/chat/completions"
+    payload = {
+        "model": model,
+        "temperature": 1.05,
+        "top_p": 0.9,
+        "top_k": 50,
+        "repetition_penalty": 1,
+        "n": 1,
+        "messages": all_message,
+        "stream_tokens": True,
     }
-}
 
-# Create model description for selected model
-selected_model = st.sidebar.selectbox("Select Model", ["Mistral", "BibleLearnerAI"])
+    TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
+    headers = {
+        "accept": "application/json",
+        "content-type": "application/json",
+        "Authorization": f"Bearer {TOGETHER_API_KEY}",
+    }
+
+    response = requests.post(url, json=payload, headers=headers, stream=True)
+    response.raise_for_status()  # Ensure HTTP request was successful
+
+    for line in response.iter_lines():
+        if line:
+            decoded_line = line.decode('utf-8')
+
+            # Check for the completion signal
+            if decoded_line == "data: [DONE]":
+                yield entire_assistant_response  # Yield the entire response at the end
+                break
+
+            try:
+                # Decode and strip any SSE format specific prefix ("data: ")
+                if decoded_line.startswith("data: "):
+                    decoded_line = decoded_line.replace("data: ", "")
+                    chunk_data = json.loads(decoded_line)
+                    content = chunk_data['choices'][0]['delta']['content']
+                    entire_assistant_response += content  # Aggregate content
+                    yield entire_assistant_response
 
-# Create model description for selected model
-st.sidebar.write(f"You're now chatting with **{selected_model}**")
-st.sidebar.markdown(model_info[selected_model]['description'])
-st.sidebar.image(model_info[selected_model]['logo'])
-st.sidebar.markdown("*Generated content may be inaccurate or false.*")
-st.sidebar.markdown("\nLearn how to build this chatbot [here](https://ngebodh.github.io/projects/2024-03-05/).")
-st.sidebar.markdown("\nRun into issues? Try the [back-up](https://huggingface.co/spaces/ngebodh/SimpleChatbot-Backup).")
+            except json.JSONDecodeError:
+                print(f"Invalid JSON received: {decoded_line}")
+                continue
+            except KeyError as e:
+                print(f"KeyError encountered: {e}")
+                continue
 
-# Initialize chat history
+    print(entire_assistant_response)
+    all_message.append({"role": "assistant", "content": entire_assistant_response})
+
+
+# Initialize Streamlit app
+st.title("AI Chatbot")
+
+# Initialize session state if not present
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display chat messages from history on app rerun
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Pre-instructions for addiction recovery
-pre_instructions = "Welcome to the Addiction Recovery AI. I'm here to help you recover from your worst addictions. Feel free to ask me anything related to your recovery journey."
+# Define available models
+models = {
+    "Addiction Recovery": "model_addiction_recovery",
+    "Mental Health": "model_mental_health",
+    "Wellness": "model_wellness"
+}
 
-# Pre-instructions for BibleLearnerAI
-pre_instructions_bible = "Welcome to BibleLearnerAI. I'm here to assist you in learning more about the Bible and Christianity. Feel free to ask me anything related to religion and spirituality."
+# Allow user to select a model
+selected_model = st.selectbox("Select Model", list(models.keys()))
 
-# Add pre-instructions to chat history based on the selected model
-if selected_model == "Mistral":
-    st.session_state.messages.append({"role": "assistant", "content": pre_instructions})
-elif selected_model == "BibleLearnerAI":
-    st.session_state.messages.append({"role": "assistant", "content": pre_instructions_bible})
+# Define models
+model_addiction_recovery = "model_addiction_recovery"
+model_mental_health = "model_mental_health"
+model_wellness = "model_wellness"
 
 # Accept user input
-if prompt := st.chat_input(f"Hi, I'm {selected_model}, ask me a question"):
-    # Display user message in chat message container
-    with st.chat_message("user"):
-        st.markdown(prompt)
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-
-    # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        # Call the appropriate model based on the selected_model
-        if selected_model == "Mistral":
-            # Code to call Mistral model
-            pass
-        elif selected_model == "BibleLearnerAI":
-            stream = get_streamed_response(prompt, [(prompt, pre_instructions_bible)])
-            for response in stream:
-                st.write(response)
-
-    st.session_state.messages.append({"role": "assistant", "content": response})
+if prompt := st.text_input("You:", key="user_input"):
+    # Display user message
+    with st.spinner("AI is typing..."):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        # Call selected model to get response
+        if selected_model == "Addiction Recovery":
+            response_stream = get_streamed_response(prompt, [(m["content"] for m in st.session_state.messages[:-1])], model_addiction_recovery)
+        elif selected_model == "Mental Health":
+            response_stream = get_streamed_response(prompt, [(m["content"] for m in st.session_state.messages[:-1])], model_mental_health)
+        elif selected_model == "Wellness":
+            response_stream = get_streamed_response(prompt, [(m["content"] for m in st.session_state.messages[:-1])], model_wellness)
+            
+        for response in response_stream:
+            st.session_state.messages.append({"role": "assistant", "content": response})
+
+# Display chat history
+for message in st.session_state.messages:
+    if message["role"] == "user":
+        st.text_input("You:", value=message["content"], disabled=True)
+    else:
+        st.text_input("AI:", value=message["content"], disabled=True)