Update src/st-chat-app.py

src/st-chat-app.py

@@ -5,113 +5,149 @@ import os
 from dotenv import load_dotenv
 load_dotenv()
 
-client_id = os.getenv("CLIENT_ID")
-client_secret = os.getenv("CLIENT_SECRET")
-base_url = os.getenv("BASE_URL")
-
-def get_access_token():
-    url = base_url+"/services/oauth2/token"
-    payload = {
-    "grant_type": "client_credentials",
-    "client_id": client_id,
-    "client_secret": client_secret
-    }
-    response = requests.post(url, data=payload)
-    # Add error handling for response
-    if response.status_code != 200:
-        st.error(f"Error fetching access token: {response.status_code} - {response.text}")
-        return None
-    data = response.json()
-    access_token = data.get('access_token', 'Token not found')
-    return access_token
-
-# Add model selection dictionary
-MODEL_OPTIONS = {
-    "GPT4Omni": "sfdc_ai__DefaultOpenAIGPT4Omni",
-    "Gemini": "sfdc_ai__DefaultVertexAIGemini20Flash001",
-    "Claude": "sfdc_ai__DefaultBedrockAnthropicClaude37Sonnet"
-}
-
-# Configure Streamlit page settings
+# Page config should be the first Streamlit command
 st.set_page_config(
     page_title="Chat with Einstein LLMs!",
-    page_icon=":brain:",  # Favicon emoji
-    layout="wide",  # Page layout option
+    page_icon=":brain:",
+    layout="wide",
 )
 
-# Add sidebar with model selection
-with st.sidebar:
-    st.title("Model Settings")
-    selected_model_name = st.selectbox(
-        "Choose AI Model",
-        options=list(MODEL_OPTIONS.keys()),
-        index=0
+# Initialize session state for authentication
+if "authenticated" not in st.session_state:
+    st.session_state.authenticated = False
+
+def check_password():
+    """Returns `True` if the user had the correct password."""
+    password_input = os.getenv("PASSWORD")
+    def password_entered():
+        """Checks whether a password entered by the user is correct."""
+        if st.session_state["password"] == password_input:
+            st.session_state.authenticated = True
+            del st.session_state["password"]  # Remove password from session state for security
+        else:
+            st.session_state.authenticated = False
+            st.error("😕 Password incorrect")
+
+    # Show input for password
+    st.text_input(
+        "Please enter the password to access the Einstein Assistant",
+        type="password",
+        on_change=password_entered,
+        key="password"
     )
-    model = MODEL_OPTIONS[selected_model_name]
+    
+    # if not st.session_state.authenticated:
+    #     st.error("😕 Password incorrect")
+    #     return False
+    # return True
 
-# Update the page title to reflect selected model
-st.subheader(f"🤖 Chat with {selected_model_name}")
+# Show chat interface only if authenticated
+if st.session_state.authenticated:
+    client_id = os.getenv("CLIENT_ID")
+    client_secret = os.getenv("CLIENT_SECRET")
+    base_url = os.getenv("BASE_URL")
 
-# Modify get_gpt_response function to use selected model
-def get_gpt_response(prompt):
-    url = f"https://api.salesforce.com/einstein/platform/v1/models/{model}/chat-generations"
-    access_token = get_access_token()
-    headers = {
-        "Authorization": f"Bearer {access_token}",
-        "Content-Type": "application/json;charset=utf-8",
-        'x-sfdc-app-context': 'EinsteinGPT',
-        'x-client-feature-id': 'ai-platform-models-connected-app'
-    }
-    chat_payload = {
-        "messages": prompt
+    def get_access_token():
+        url = base_url+"/services/oauth2/token"
+        payload = {
+        "grant_type": "client_credentials",
+        "client_id": client_id,
+        "client_secret": client_secret
+        }
+        response = requests.post(url, data=payload)
+        # Add error handling for response
+        if response.status_code != 200:
+            st.error(f"Error fetching access token: {response.status_code} - {response.text}")
+            return None
+        data = response.json()
+        access_token = data.get('access_token', 'Token not found')
+        return access_token
+
+    # Add model selection dictionary
+    MODEL_OPTIONS = {
+        "GPT4-Omni": "sfdc_ai__DefaultOpenAIGPT4Omni",
+        "Gemini": "sfdc_ai__DefaultVertexAIGemini20Flash001",
+        "Claude": "sfdc_ai__DefaultBedrockAnthropicClaude37Sonnet"
     }
 
-    try:
-        response = requests.post(url, headers=headers, data=json.dumps(chat_payload))
-        response.raise_for_status()  # Raise exception for bad status codes
-        data = response.json()
-        return data["generationDetails"]["generations"][0]["content"]
-    except requests.exceptions.RequestException as e:
-        st.error(f"Error calling the API: {str(e)}")
-        return "I apologize, but I encountered an error. Please try again."
-    except (KeyError, IndexError) as e:
-        st.error(f"Error parsing response: {str(e)}")
-        return "I apologize, but I received an invalid response. Please try again."
+    # Add sidebar with model selection
+    with st.sidebar:
+        st.title("Model Settings")
+        selected_model_name = st.selectbox(
+            "Choose AI Model",
+            options=list(MODEL_OPTIONS.keys()),
+            index=0
+        )
+        model = MODEL_OPTIONS[selected_model_name]
 
-# Display the chatbot's title on the page
-if "messages" not in st.session_state:
-    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
+    # Update the page title to reflect selected model
+    st.subheader(f"🤖 Chat with {selected_model_name}")
 
-for msg in st.session_state.messages:
-    st.chat_message(msg["role"], avatar="🤖").write(msg["content"])
-    if "image" in msg:
-        st.image(msg["image"])
+    # Modify get_gpt_response function to use selected model
+    def get_gpt_response(prompt):
+        url = f"https://api.salesforce.com/einstein/platform/v1/models/{model}/chat-generations"
+        access_token = get_access_token()
+        headers = {
+            "Authorization": f"Bearer {access_token}",
+            "Content-Type": "application/json;charset=utf-8",
+            'x-sfdc-app-context': 'EinsteinGPT',
+            'x-client-feature-id': 'ai-platform-models-connected-app'
+        }
+        chat_payload = {
+            "messages": prompt
+        }
 
-prompt = st.chat_input(
-    "Say something and/or attach an image",
-    accept_file=True,
-    file_type=["jpg", "jpeg", "png"],
-)
+        try:
+            response = requests.post(url, headers=headers, data=json.dumps(chat_payload))
+            response.raise_for_status()  # Raise exception for bad status codes
+            data = response.json()
+            return data["generationDetails"]["generations"][0]["content"]
+        except requests.exceptions.RequestException as e:
+            st.error(f"Error calling the API: {str(e)}")
+            return "I apologize, but I encountered an error. Please try again."
+        except (KeyError, IndexError) as e:
+            st.error(f"Error parsing response: {str(e)}")
+            return "I apologize, but I received an invalid response. Please try again."
+
+    # Display the chatbot's title on the page
+    if "messages" not in st.session_state:
+        st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
 
-if prompt:
-    # Handle text input
-    if prompt.text:
-        st.session_state.messages.append({"role": "user", "content": prompt.text})
-        st.chat_message("user").write(prompt.text)
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"], avatar="🤖").write(msg["content"])
+        if "image" in msg:
+            st.image(msg["image"])
+
+    prompt = st.chat_input(
+        "Say something and/or attach an image",
+        accept_file=True,
+        file_type=["jpg", "jpeg", "png"],
+    )
+
+    if prompt:
+        # Handle text input
+        if prompt.text:
+            st.session_state.messages.append({"role": "user", "content": prompt.text})
+            st.chat_message("user").write(prompt.text)
+            
+        # Handle image upload
+        if prompt.get("files"):
+            uploaded_file = prompt["files"][0]
+            st.session_state.messages.append({
+                "role": "user", 
+                "content": "Uploaded an image",
+                "image": uploaded_file
+            })
+            st.chat_message("user").write("Uploaded an image")
+            st.image(uploaded_file)
         
-    # Handle image upload
-    if prompt.get("files"):
-        uploaded_file = prompt["files"][0]
-        st.session_state.messages.append({
-            "role": "user", 
-            "content": "Uploaded an image",
-            "image": uploaded_file
-        })
-        st.chat_message("user").write("Uploaded an image")
-        st.image(uploaded_file)
-    
-    # Get AI response if there's any input
-    if prompt.text or prompt.get("files"):
-        msg = get_gpt_response(st.session_state.messages)
-        st.session_state.messages.append({"role": "assistant", "content": msg})
-        st.chat_message("assistant", avatar="🤖").write(msg)
+        # Get AI response if there's any input
+        if prompt.text or prompt.get("files"):
+            msg = get_gpt_response(st.session_state.messages)
+            st.session_state.messages.append({"role": "assistant", "content": msg})
+            st.chat_message("assistant", avatar="🤖").write(msg)
+else:
+    # Show login page
+    st.title("Welcome to Einstein Assistant")
+    #st.markdown("Please log in to continue")
+    check_password()