Spaces:
Sleeping
Sleeping
File size: 5,594 Bytes
5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 5add07b bc79018 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import requests
import json
import streamlit as st
import os
from dotenv import load_dotenv
load_dotenv()
# Page config should be the first Streamlit command
st.set_page_config(
page_title="Chat with Einstein LLMs!",
page_icon=":brain:",
layout="wide",
)
# 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"
)
# if not st.session_state.authenticated:
# st.error("😕 Password incorrect")
# return False
# return True
# 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")
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"
}
# 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]
# Update the page title to reflect selected model
st.subheader(f"🤖 Chat with {selected_model_name}")
# 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
}
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?"}]
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)
# 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() |