Update app.py

app.py

@@ -1,11 +1,6 @@
 import streamlit as st
 import google.generativeai as genai
 
-# -------------------
-# API Key Setup
-# -------------------
-gemini_api_key = st.secrets.get("GEN_API_KEY", "")
-
 # -------------------
 # Page Configuration
 # -------------------
@@ -13,15 +8,73 @@ st.set_page_config(page_title="Your AI Buddy", layout="wide")
 st.title("💡 Need answers? Just type below!")
 
 # -------------------
-# Gemini Setup
+# CSS Styling for Chat Bubbles
+# -------------------
+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:has(.user-message-marker) > div {
+    flex-direction: row-reverse;
+    text-align: right;
+    background-color: #d1e9ff;   /* Light blue bubble */
+    color: #0a2f5c;              /* Dark blue text */
+    border: 1px solid #a3cfff;
+}
+
+/* Assistant messages (left side) */
+.stChatMessage:has(.assistant-message-marker) > div {
+    flex-direction: row;
+    text-align: left;
+    background-color: #f1f1f1;   /* Light gray bubble */
+    color: #1a1a1a;
+    border: 1px solid #ddd;
+}
+
+/* Ensure markdown text follows bubble alignment */
+.stChatMessage:has(.user-message-marker) .stMarkdown {
+    text-align: right;
+}
+.stChatMessage:has(.assistant-message-marker) .stMarkdown {
+    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)
+
 # -------------------
+# API Key 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)
 
-# Fetch available Gemini models
+# -------------------
+# Gemini Setup
+# -------------------
 available_models = [
     m.name for m in genai.list_models()
     if "generateContent" in m.supported_generation_methods
@@ -31,13 +84,12 @@ 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("Model", available_models, index=0)
 
-# Initialize Gemini chat if needed
+# Initialize Gemini chat
 if "gemini_chat" not in st.session_state or st.session_state.get("model") != model:
     st.session_state.model = model
     try:
@@ -72,28 +124,30 @@ if st.sidebar.button("Reset Conversation"):
 # Display Chat Messages
 # -------------------
 for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+    role = msg["role"]
+    avatar = "🙂" if role == "user" else "🤖"
+    with st.chat_message(role, avatar=avatar):
+        # Add hidden marker for CSS styling
+        st.markdown(
+            f"<span class='{role}-message-marker'></span>{msg['content']}",
+            unsafe_allow_html=True
+        )
 
 # -------------------
 # User Input
 # -------------------
 if user_input := st.chat_input("Type your message..."):
-    # Show user message
-    st.chat_message("user").markdown(user_input)
+    # Save and display user input
     st.session_state.messages.append({"role": "user", "content": user_input})
+    st.rerun()
 
     try:
         with st.spinner("Gemini is thinking..."):
-            # Prepend system prompt for context
             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})"
 
-    with st.chat_message("assistant"):
-        st.markdown(bot_text)
-
     st.session_state.messages.append({"role": "assistant", "content": bot_text})
     st.rerun()