Update app.py

app.py

@@ -1,66 +1,18 @@
 import streamlit as st
+from streamlit_chat import message
 import google.generativeai as genai
 
 # -------------------
 # Page Configuration
 # -------------------
 st.set_page_config(page_title="Your AI Buddy", layout="wide")
-st.title("💡 Need answers? Just type below!")
 
-# -------------------
-# 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)
+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")
 
 # -------------------
 # API Key Setup
@@ -73,7 +25,7 @@ if not gemini_api_key:
 genai.configure(api_key=gemini_api_key)
 
 # -------------------
-# Gemini Setup
+# Model Selection
 # -------------------
 available_models = [
     m.name for m in genai.list_models()
@@ -87,62 +39,52 @@ if not available_models:
 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 "gemini_chat" not in st.session_state or st.session_state.get("model") != model:
-    st.session_state.model = model
-    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()
-
-# -------------------
-# System Prompt
-# -------------------
-system_prompt = st.sidebar.text_area(
-    "System Prompt",
-    "You are a helpful AI assistant. Provide concise and accurate answers."
-)
+model = st.sidebar.selectbox("Choose a Gemini Model", available_models, index=0)
 
 # -------------------
-# Chat History State
+# Initialize Chat Session
 # -------------------
-if "messages" not in st.session_state:
+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=[])
     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()
 
 # -------------------
-# Display Chat Messages
+# System Prompt
 # -------------------
-for msg in st.session_state.messages:
-    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
-        )
+system_prompt = st.sidebar.text_area(
+    "System Prompt (AI Behavior)",
+    "You are a helpful AI assistant. Provide concise and accurate answers."
+)
+
+# -------------------
+# Chat History Display
+# -------------------
+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}")
 
 # -------------------
 # User Input
 # -------------------
-if user_input := st.chat_input("Type your message..."):
-    # Save and display user input
+user_input = st.chat_input("Type your question here...")
+if user_input:
+    # Save user message
     st.session_state.messages.append({"role": "user", "content": user_input})
-    st.rerun()
 
     try:
-        with st.spinner("Gemini is thinking..."):
+        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