Update app.py

app.py

@@ -91,9 +91,57 @@ def get_remaining_queries(user_id):
 st.set_page_config(
     page_title="USMLE Step 1 AI",
     page_icon="🩺",
-    layout="centered"
+    layout="centered",
+    initial_sidebar_state="expanded"
 )
 
+# Apply custom CSS for better visual appearance
+st.markdown("""
+<style>
+    .main-header {
+        font-size: 2.5rem !important;
+        margin-bottom: 1rem !important;
+        color: #2c3e50;
+    }
+    .sub-header {
+        font-size: 1.2rem !important;
+        color: #34495e;
+        margin-bottom: 2rem !important;
+    }
+    .stAlert {
+        padding: 15px !important;
+        border-radius: 8px !important;
+    }
+    .usage-metric {
+        padding: 10px;
+        background-color: #f8f9fa;
+        border-radius: 8px;
+        margin-bottom: 15px;
+        border-left: 5px solid #4CAF50;
+    }
+    .usage-metric-warning {
+        border-left: 5px solid #FFC107;
+    }
+    .usage-metric-danger {
+        border-left: 5px solid #F44336;
+    }
+    .footer-text {
+        font-size: 0.85rem !important;
+        color: #7f8c8d;
+    }
+    .stChatMessage div[data-testid="stChatMessageContent"] {
+        border-radius: 15px !important;
+        padding: 15px !important;
+    }
+    .user-message {
+        background-color: #f1f8ff !important;
+    }
+    .assistant-message {
+        background-color: #f9f9f9 !important;
+    }
+</style>
+""", unsafe_allow_html=True)
+
 # Initialize session state for chat history
 if 'messages' not in st.session_state:
     st.session_state.messages = []
@@ -101,17 +149,57 @@ if 'messages' not in st.session_state:
 # Initialize rate limiting
 init_rate_limiting()
 
-# Display remaining queries
-user_id = get_user_id()
-remaining_queries = get_remaining_queries(user_id)
-st.sidebar.write(f"Remaining queries today: {remaining_queries}/{MAX_REQUESTS_PER_DAY}")
+# Sidebar content
+with st.sidebar:
+    st.image("https://img.icons8.com/color/96/000000/caduceus.png", width=80)
+    st.markdown("### USMLE Step 1 Assistant")
+    st.markdown("---")
+    
+    # Display remaining queries with visual indicator
+    user_id = get_user_id()
+    remaining_queries = get_remaining_queries(user_id)
+    
+    # Determine styling based on remaining queries
+    usage_class = "usage-metric"
+    if remaining_queries <= 2:
+        usage_class += " usage-metric-danger"
+    elif remaining_queries <= 3:
+        usage_class += " usage-metric-warning"
+    
+    st.markdown(f"""
+    <div class="{usage_class}">
+        <strong>Daily Usage</strong><br>
+        {remaining_queries}/{MAX_REQUESTS_PER_DAY} queries remaining
+    </div>
+    """, unsafe_allow_html=True)
+    
+    # Help section in sidebar
+    with st.expander("ℹ️ How to use"):
+        st.markdown("""
+        1. Type your USMLE Step 1 question in the chat input
+        2. The AI will search First Aid content and respond
+        3. You have 5 queries per day
+        
+        **Best for:**
+        - Fact checking First Aid content
+        - Understanding complex topics
+        - Quick reference during study
+        """)
+    
+    with st.expander("🔍 Example Questions"):
+        st.markdown("""
+        - "Explain the Krebs cycle"
+        - "What are the symptoms of Parkinson's disease?"
+        - "Differentiate between type 1 and type 2 diabetes"
+        - "What antibiotics are used for MRSA?"
+        """)
 
 # Check for API keys
 PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 
 if not PINECONE_API_KEY or not OPENAI_API_KEY:
-    st.error("Missing API keys. Please set PINECONE_API_KEY and OPENAI_API_KEY environment variables.")
+    st.error("⚠️ Missing API keys. Please set PINECONE_API_KEY and OPENAI_API_KEY environment variables.")
     st.stop()
 
 os.environ["PINECONE_API_KEY"] = PINECONE_API_KEY
@@ -121,19 +209,30 @@ os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 @st.cache_resource
 def initialize_rag_chain():
     try:
-        st.sidebar.write("Loading embeddings model...")
+        progress_text = st.sidebar.empty()
+        progress_bar = st.sidebar.progress(0)
+        
+        # Step 1: Load embeddings
+        progress_text.text("Loading embeddings model... (1/4)")
         embeddings = download_hugging_face_embeddings()
+        progress_bar.progress(25)
         
-        st.sidebar.write("Connecting to Pinecone...")
+        # Step 2: Connect to Pinecone
+        progress_text.text("Connecting to Pinecone database... (2/4)")
         index_name = "medprep"
         docsearch = Pinecone.from_existing_index(
             index_name=index_name,
             embedding=embeddings
         )
+        progress_bar.progress(50)
         
+        # Step 3: Set up retriever
+        progress_text.text("Setting up retrieval system... (3/4)")
         retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+        progress_bar.progress(75)
         
-        st.sidebar.write("Initializing OpenAI...")
+        # Step 4: Initialize LLM and chain
+        progress_text.text("Initializing language model... (4/4)")
         llm = OpenAI(temperature=0.4, max_tokens=500)
         
         prompt = ChatPromptTemplate.from_messages([
@@ -143,66 +242,92 @@ def initialize_rag_chain():
         
         question_answer_chain = create_stuff_documents_chain(llm, prompt)
         rag_chain = create_retrieval_chain(retriever, question_answer_chain)
+        progress_bar.progress(100)
+        
+        # Clean up progress indicators
+        progress_text.empty()
+        progress_bar.empty()
         
         st.sidebar.success("✅ System initialized successfully!")
         return rag_chain
     except Exception as e:
-        st.sidebar.error(f"Error initializing system: {str(e)}")
+        st.sidebar.error(f"⚠️ Error initializing system: {str(e)}")
         import traceback
         st.sidebar.text(traceback.format_exc())
         return None
 
-# Main app title
-st.title("First AID USMLE Step 1")
-st.write("Ask me any question from First AID USMLE Step 1 book, and I'll try to help!")
+# Main app content
+st.markdown('<h1 class="main-header">First Aid USMLE Step 1 Assistant</h1>', unsafe_allow_html=True)
+st.markdown('<p class="sub-header">Ask me any question from First Aid USMLE Step 1 book, and I\'ll try to help!</p>', unsafe_allow_html=True)
 
 # Initialize the RAG chain
 rag_chain = initialize_rag_chain()
 
 if rag_chain is None:
-    st.error("Failed to initialize the system. Please check the sidebar for error details.")
+    st.error("⚠️ Failed to initialize the system. Please check the sidebar for error details.")
     st.stop()
 
-# Display chat history
-for message in st.session_state.messages:
+# Display chat history with improved styling
+for i, message in enumerate(st.session_state.messages):
+    message_class = "user-message" if message["role"] == "user" else "assistant-message"
     with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+        st.markdown(f'<div class="{message_class}">{message["content"]}</div>', unsafe_allow_html=True)
 
 # Get user input
-if prompt := st.chat_input("Ask a question..."):
+if prompt := st.chat_input("Ask a USMLE Step 1 question..."):
     # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
     
     # Display user message
     with st.chat_message("user"):
-        st.markdown(prompt)
+        st.markdown(f'<div class="user-message">{prompt}</div>', unsafe_allow_html=True)
     
     # Check rate limit
     user_id = get_user_id()
     allowed, count = check_rate_limit(user_id)
     
     if not allowed:
-        response = f"⚠️ Daily limit reached. You've used {count} queries today. Please try again tomorrow."
+        response = f"⚠️ **Daily limit reached**\n\nYou've used {count} queries today. Please try again tomorrow."
     else:
         # Process the query with the RAG chain
         with st.chat_message("assistant"):
-            with st.spinner("Thinking..."):
+            message_placeholder = st.empty()
+            
+            with st.spinner("Searching First Aid content..."):
                 try:
                     result = rag_chain.invoke({"input": prompt})
                     response = result.get("answer", "Sorry, I couldn't find an answer to that.")
+                    
+                    # Format the remaining queries notification
                     remaining = MAX_REQUESTS_PER_DAY - count
-                    response += f"\n\n\n_You have {remaining} queries remaining today._"
+                    if remaining <= 1:
+                        usage_note = f"⚠️ **{remaining} query remaining today**"
+                    else:
+                        usage_note = f"ℹ️ {remaining} queries remaining today"
+                    
+                    # Add a separator and the usage note
+                    response += f"\n\n---\n\n{usage_note}"
+                    
                 except Exception as e:
-                    response = f"Error processing your request: {str(e)}"
+                    response = f"⚠️ **Error processing your request**\n\n{str(e)}"
             
-            st.markdown(response)
+            message_placeholder.markdown(f'<div class="assistant-message">{response}</div>', unsafe_allow_html=True)
     
     # Add assistant response to chat history
     st.session_state.messages.append({"role": "assistant", "content": response})
 
-# Footer
+# Footer with improved styling
 st.markdown("---")
-# Footer
-st.markdown("---")
-st.markdown("""*This AI assistant uses retrieval augmented generation to provide information from First Aid USMLE Step 1 content.
-        It's designed to help with studying, but should not replace professional medical advice.*""")
+st.markdown("""
+<div class="footer-text">
+<p><strong>About this assistant</strong></p>
+<p>This AI assistant uses retrieval augmented generation to provide information from First Aid USMLE Step 1 content.
+It's designed to help with studying, but should not replace professional medical advice.</p>
+<p>© 2025 USMLE Step 1 Assistant</p>
+</div>
+""", unsafe_allow_html=True)
+
+# Add a reset button at the bottom
+if st.button("Clear Conversation"):
+    st.session_state.messages = []
+    st.experimental_rerun()