Create chat-app

src/chat-app

@@ -0,0 +1,117 @@
+import requests
+import json
+import streamlit as st
+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
+st.set_page_config(
+    page_title="Chat with Einstein LLMs!",
+    page_icon=":brain:",  # Favicon emoji
+    layout="wide",  # Page layout option
+)
+
+# 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)