Update app.py

app.py

@@ -1,31 +1,27 @@
 import streamlit as st
 import google.generativeai as genai
 
+# -------------------
+# API Key Setup
+# -------------------
+gemini_api_key = st.secrets.get("GEN_API_KEY", "")
+
 # -------------------
 # Page Configuration
 # -------------------
 st.set_page_config(page_title="Your AI Buddy", layout="wide")
-
-st.title("💬 Your AI Buddy")
-st.markdown("Ask me anything! 📘🤖\n\n"
-            "💡 *Examples:*\n"
-            "- Summarize a news article\n"
-            "- Explain a math problem\n"
-            "- Give me study tips\n")
+st.title("💡 Need answers? Just type below!")
 
 # -------------------
-# API Key Setup
+# Gemini Setup
 # -------------------
-gemini_api_key = st.secrets.get("GEN_API_KEY", "")
 if not gemini_api_key:
     st.error("⚠️ Please set your 'GEN_API_KEY' in Streamlit secrets.")
     st.stop()
 
 genai.configure(api_key=gemini_api_key)
 
-# -------------------
-# Model Selection
-# -------------------
+# Fetch available Gemini models
 available_models = [
     m.name for m in genai.list_models()
     if "generateContent" in m.supported_generation_methods
@@ -35,60 +31,101 @@ if not available_models:
     st.error("⚠️ No Gemini models available for your API key.")
     st.stop()
 
+# Reset session if old model is invalid
 if "model" in st.session_state and st.session_state["model"] not in available_models:
     del st.session_state["model"]
 
-model = st.sidebar.selectbox("Choose a Gemini Model", available_models, index=0)
+model = st.sidebar.selectbox("Model", available_models, index=0)
 
-# -------------------
-# Initialize Chat Session
-# -------------------
+# Initialize Gemini chat if needed
 if "gemini_chat" not in st.session_state or st.session_state.get("model") != model:
     st.session_state.model = model
-    gemini_model = genai.GenerativeModel(model)
-    st.session_state.gemini_chat = gemini_model.start_chat(history=[])
+    try:
+        gemini_model = genai.GenerativeModel(model)
+        st.session_state.gemini_chat = gemini_model.start_chat(history=[])
+    except Exception as e:
+        st.error(f"⚠️ Could not initialize Gemini model: {e}")
+        st.stop()
+
+# -------------------
+# Sidebar Options
+# -------------------
+system_prompt = st.sidebar.text_area(
+    "System Prompt",
+    "You are a helpful AI assistant. Provide concise and accurate answers."
+)
+
+if "messages" not in st.session_state:
     st.session_state.messages = []
 
 # Reset conversation button
-if st.sidebar.button("🔄 Reset Conversation"):
+if st.sidebar.button("Reset Conversation"):
     st.session_state.messages = []
     gemini_model = genai.GenerativeModel(model)
     st.session_state.gemini_chat = gemini_model.start_chat(history=[])
     st.rerun()
 
 # -------------------
-# System Prompt
+# Custom Chat Bubble Styling
 # -------------------
-system_prompt = st.sidebar.text_area(
-    "System Prompt (AI Behavior)",
-    "You are a helpful AI assistant. Provide concise and accurate answers."
-)
+st.markdown("""
+<style>
+/* General chat bubble styling */
+.stChatMessage > div {
+    border-radius: 15px;
+    padding: 10px 14px;
+    margin: 6px 0;
+    max-width: 80%;
+    font-size: 15px;
+}
+
+/* User messages (right side) */
+.stChatMessage[data-testid="stChatMessage-user"] > div {
+    background-color: #d1e9ff;
+    color: #0a2f5c;
+    border: 1px solid #a3cfff;
+    margin-left: auto;
+    text-align: right;
+}
+
+/* Assistant messages (left side) */
+.stChatMessage[data-testid="stChatMessage-assistant"] > div {
+    background-color: #f1f1f1;
+    color: #1a1a1a;
+    border: 1px solid #ddd;
+    margin-right: auto;
+    text-align: left;
+}
+
+/* Avatars */
+.stChatMessage [data-testid="stChatAvatar"] {
+    border-radius: 50%;
+    width: 32px;
+    height: 32px;
+    font-size: 18px;
+    background-color: #ffffff;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.2);
+}
+</style>
+""", unsafe_allow_html=True)
 
 # -------------------
-# Chat History Display
+# Display Chat Messages
 # -------------------
-if st.session_state.messages:
-    for i, msg in enumerate(st.session_state.messages):
-        if msg["role"] == "user":
-            message(msg["content"], is_user=True, key=f"user_{i}")
-        else:
-            message(msg["content"], key=f"bot_{i}")
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
 
 # -------------------
 # User Input
 # -------------------
-user_input = st.chat_input("Type your question here...")
-if user_input:
-    # Save user message
+if user_input := st.chat_input("Type your message..."):
+    # Show user message
+    st.chat_message("user").markdown(user_input)
     st.session_state.messages.append({"role": "user", "content": user_input})
 
     try:
-        with st.spinner("🤖 Thinking..."):
-            full_input = f"{system_prompt}\n\nUser: {user_input}"
-            resp = st.session_state.gemini_chat.send_message(full_input)
-            bot_text = resp.text
-    except Exception as e:
-        bot_text = f"⚠️ Gemini could not respond right now. Please try again. ({e})"
-
-    st.session_state.messages.append({"role": "assistant", "content": bot_text})
-    st.rerun()
+        with st.spinner("Gemini is thinking..."):