Update app.py

app.py

@@ -1,56 +1,78 @@
 import streamlit as st
 import openai
+import os
 
-# Set up the OpenAI API key
-openai.api_key = st.secrets['API_KEY']
-
-# Define the models in order of preference
-MODELS = ["gpt-4", "gpt-3.5-turbo", "gpt-3.5-turbo-16k"]
-
-# Initialize session state to store conversation history
-if "conversation" not in st.session_state:
-    st.session_state.conversation = []
+# Function to get the API key from Streamlit secrets
+def get_api_key():
+    try:
+        return st.secrets["API_KEY"]
+    except KeyError:
+        st.error("API_KEY not found in Streamlit secrets. Please add it.")
+        return None
 
-# Function to get a response from the OpenAI API
-def get_chat_response(prompt, model_index=0):
+# Function to interact with the OpenAI API
+def generate_response(prompt, model_name, api_key):
+    openai.api_key = api_key
     try:
-        response = openai.ChatCompletion.create(
-            model=MODELS[model_index],
+        completion = openai.ChatCompletion.create(
+            model=model_name,
             messages=[{"role": "user", "content": prompt}]
         )
-        return response.choices[0].message['content']
-    except openai.error.RateLimitError:
-        if model_index + 1 < len(MODELS):
-            st.warning(f"Quota exceeded for {MODELS[model_index]}. Switching to {MODELS[model_index + 1]}.")
-            return get_chat_response(prompt, model_index + 1)
-        else:
-            st.error("All models have exceeded their quota. Please try again later.")
-            return None
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
+        return completion.choices[0].message.content
+    except openai.APIError as e:  # General API Error
+        st.error(f"OpenAI API Error with {model_name}: {e}")
         return None
+    except openai.RateLimitError as e: # Rate Limit Error
+        st.error(f"OpenAI Rate Limit Error with {model_name}: {e}")
+        return None
+    except openai.AuthenticationError as e: # Authentication Error
+        st.error(f"OpenAI Authentication Error with {model_name}: {e}")
+        return None
+    except Exception as e:  # Catch any other exception
+        st.error(f"An unexpected error occurred with {model_name}: {e}")
+        return None
+
+# Main Streamlit app
+def main():
+    st.title("Chatbot with Model Switching")
 
-# Streamlit UI
-st.title("ChatGPT with Model Fallback")
-
-# Display previous conversation
-st.write("### Conversation History")
-for message in st.session_state.conversation:
-    st.write(f"**{message['role']}**: {message['content']}")
-
-# User input
-user_input = st.text_input("You: ", "")
-
-if user_input:
-    # Add user input to conversation history
-    st.session_state.conversation.append({"role": "User", "content": user_input})
-    
-    # Get response from the model
-    response = get_chat_response(user_input)
-    
-    if response:
-        # Add model response to conversation history
-        st.session_state.conversation.append({"role": "Assistant", "content": response})
+    # Initialize conversation history in session state
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display previous messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Get user input
+    prompt = st.chat_input("Say something")
+
+    if prompt:
+        # Add user message to the state
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        # Define model priority
+        models = ["gpt-4", "gpt-3.5-turbo"]  # Add more models as needed
+
+        # Get API key
+        api_key = get_api_key()
+        if not api_key:
+            return
+
+        response = None
+        for model in models:
+            response = generate_response(prompt, model, api_key)
+            if response:
+                break # If a response is generated, break out of loop.
         
-        # Display the response
-        st.write(f"**Assistant**: {response}")
+        if response:
+            # Add bot message to state
+            st.session_state.messages.append({"role": "assistant", "content": response})
+            with st.chat_message("assistant"):
+                st.markdown(response)
+
+if __name__ == "__main__":
+    main()