Update app.py

app.py

@@ -87,110 +87,29 @@ def get_remaining_queries(user_id):
     count = rate_limits.get(user_id, {}).get(today, 0)
     return MAX_REQUESTS_PER_DAY - count
 
-# Add custom CSS
-def local_css():
-    st.markdown("""
-    <style>
-    .main {
-        background-color: #f8f9fa;
-    }
-    .stApp {
-        max-width: 900px;
-        margin: 0 auto;
-    }
-    .chat-message {
-        padding: 1.5rem;
-        border-radius: 0.5rem;
-        margin-bottom: 1rem;
-        display: flex;
-        flex-direction: column;
-    }
-    .chat-message.user {
-        background-color: #e6f3ff;
-        border-left: 5px solid #2b71d9;
-    }
-    .chat-message.bot {
-        background-color: #f0f7f0;
-        border-left: 5px solid #28a745;
-    }
-    .chat-message .message-content {
-        display: flex;
-        margin-top: 0.5rem;
-    }
-    .chat-message .avatar {
-        min-width: 30px;
-        margin-right: 1rem;
-    }
-    .chat-message .avatar img {
-        max-width: 30px;
-        max-height: 30px;
-    }
-    .chat-message .content {
-        width: 100%;
-    }
-    .stTextInput {
-        padding: 0.5rem 1rem;
-        border-radius: 0.5rem;
-        border: 1px solid #ced4da;
-    }
-    .stButton>button {
-        background-color: #28a745;
-        color: white;
-        border-radius: 0.5rem;
-        border: none;
-        padding: 0.5rem 1rem;
-    }
-    h1 {
-        color: #2b71d9;
-        font-weight: 700;
-    }
-    .highlight {
-        background-color: #f8f9fa;
-        padding: 1.5rem;
-        border-radius: 0.5rem;
-        border-left: 5px solid #dc3545;
-        margin: 1rem 0;
-    }
-    .stProgress > div > div > div {
-        background-color: #28a745;
-    }
-    .sidebar .sidebar-content {
-        background-color: #f8f9fa;
-    }
-    .stAlert {
-        border-radius: 0.5rem;
-    }
-    .progress-container {
-        width: 100%;
-        background-color: #e9ecef;
-        border-radius: 20px;
-        margin-top: 10px;
-        margin-bottom: 20px;
-    }
-    .progress-bar {
-        height: 10px;
-        border-radius: 20px;
-        background-color: #28a745;
-    }
-    .queries-text {
-        color: #6c757d;
-        font-size: 0.9rem;
-        text-align: center;
-        margin-top: 5px;
-    }
-    </style>
-    """, unsafe_allow_html=True)
-
 # Set up page configuration
 st.set_page_config(
-    page_title="MedPrep AI Assistant",
+    page_title="USMLE Step 1 AI",
     page_icon="🩺",
-    layout="wide",
-    initial_sidebar_state="expanded"
+    layout="centered"
 )
 
-# Load custom CSS
-local_css()
+# Custom CSS to style images and improve layout
+st.markdown("""
+<style>
+.header-container {
+    display: flex;
+    align-items: center;
+    gap: 20px;
+}
+.header-image {
+    width: 80px;
+}
+.sidebar-image {
+    margin-bottom: 20px;
+}
+</style>
+""", unsafe_allow_html=True)
 
 # Initialize session state for chat history
 if 'messages' not in st.session_state:
@@ -199,40 +118,16 @@ if 'messages' not in st.session_state:
 # Initialize rate limiting
 init_rate_limiting()
 
-# Sidebar content
+# Add image to sidebar
 with st.sidebar:
-    # Using a more reliable image URL - or you can upload an image to your Space
-    st.image("https://img.freepik.com/free-vector/online-doctor-concept_23-2148522199.jpg", width=200)
-    st.title("MedPrep AI")
-    
-    st.markdown("### About")
-    st.markdown("""
-    This AI assistant is powered by First Aid USMLE Step 1 content and helps medical students prepare for their exams.
-    
-    Ask questions about:
-    - Pathology concepts
-    - Pharmacology
-    - Biochemistry
-    - Anatomy
-    - And more!
-    """)
+    st.image("https://img.freepik.com/free-vector/medical-education-abstract-concept-vector-illustration_107173-26967.jpg", width=200, caption="First Aid USMLE", use_column_width=True)
+    st.markdown("#### About")
+    st.write("This AI assistant is powered by First Aid USMLE Step 1 content and helps medical students prepare for their exams.")
     
-    # Display remaining queries with progress bar
+    # Display remaining queries
     user_id = get_user_id()
     remaining_queries = get_remaining_queries(user_id)
-    used_queries = MAX_REQUESTS_PER_DAY - remaining_queries
-    percentage = (used_queries / MAX_REQUESTS_PER_DAY) * 100
-    
-    st.markdown("### Daily Query Limit")
-    st.markdown(
-        f"""
-        <div class="progress-container">
-            <div class="progress-bar" style="width: {percentage}%;"></div>
-        </div>
-        <p class="queries-text">{remaining_queries} of {MAX_REQUESTS_PER_DAY} queries remaining today</p>
-        """,
-        unsafe_allow_html=True
-    )
+    st.write(f"Remaining queries today: {remaining_queries}/{MAX_REQUESTS_PER_DAY}")
     
     st.markdown("---")
     st.markdown("### Tips for Better Results")
@@ -258,123 +153,101 @@ os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 @st.cache_resource
 def initialize_rag_chain():
     try:
-        with st.spinner("Loading knowledge base..."):
-            embeddings = download_hugging_face_embeddings()
-            
-            index_name = "medprep"
-            docsearch = Pinecone.from_existing_index(
-                index_name=index_name,
-                embedding=embeddings
-            )
-            
-            retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-            
-            llm = OpenAI(model="gpt-3.5-turbo", temperature=0.4, max_tokens=500)
-            
-            prompt = ChatPromptTemplate.from_messages([
-                ("system", system_prompt),
-                ("human", "{input}")
-            ])
-            
-            question_answer_chain = create_stuff_documents_chain(llm, prompt)
-            rag_chain = create_retrieval_chain(retriever, question_answer_chain)
-            
-            return rag_chain
+        st.sidebar.write("Loading embeddings model...")
+        embeddings = download_hugging_face_embeddings()
+        
+        st.sidebar.write("Connecting to Pinecone...")
+        index_name = "medprep"
+        docsearch = Pinecone.from_existing_index(
+            index_name=index_name,
+            embedding=embeddings
+        )
+        
+        retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+        
+        st.sidebar.write("Initializing OpenAI...")
+        llm = OpenAI(model="gpt-3.5-turbo", temperature=0.4, max_tokens=500)
+        
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", system_prompt),
+            ("human", "{input}")
+        ])
+        
+        question_answer_chain = create_stuff_documents_chain(llm, prompt)
+        rag_chain = create_retrieval_chain(retriever, question_answer_chain)
+        
+        st.sidebar.success("✅ System initialized successfully!")
+        return rag_chain
     except Exception as e:
-        st.error(f"Error initializing system: {str(e)}")
+        st.sidebar.error(f"Error initializing system: {str(e)}")
         import traceback
-        st.text(traceback.format_exc())
+        st.sidebar.text(traceback.format_exc())
         return None
 
-# Main app content
-st.title("MedPrep AI Assistant")
-st.markdown("##### Powered by First Aid USMLE Step 1 content")
+# Main app header with logo
+st.markdown("""
+<div class="header-container">
+    <img src="https://med.stanford.edu/content/dam/sm-news/images/2014/09/first-aid-logo.gif" class="header-image">
+    <h1>First AID USMLE Step 1</h1>
+</div>
+""", unsafe_allow_html=True)
 
-# Separator
-st.markdown("---")
+st.write("Ask me any question from First AID USMLE Step 1 book, and I'll try to help!")
 
 # Initialize the RAG chain
 rag_chain = initialize_rag_chain()
 
 if rag_chain is None:
-    st.error("Failed to initialize the system. Please try again later.")
+    st.error("Failed to initialize the system. Please check the sidebar for error details.")
     st.stop()
 
-# Create two columns for user avatar
-def display_message(message, is_user=False):
-    if is_user:
-        avatar = "👤"
-        message_type = "user"
-    else:
-        avatar = "🩺"
-        message_type = "bot"
-    
-    st.markdown(
-        f"""
-        <div class="chat-message {message_type}">
-            <div class="message-content">
-                <div class="avatar">{avatar}</div>
-                <div class="content">{message}</div>
-            </div>
-        </div>
-        """,
-        unsafe_allow_html=True
-    )
-
 # Display chat history
 for message in st.session_state.messages:
-    display_message(message["content"], message["role"] == "user")
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
 # Get user input
-with st.container():
-    prompt = st.text_input("Ask your medical question:", key="user_input", placeholder="e.g., Explain the mechanism of action for ACE inhibitors...")
-
-    if st.button("Submit", key="send"):
-        if prompt:
-            # Add user message to chat history
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            
-            # Display user message
-            display_message(prompt, 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."
-                display_message(response)
-                st.session_state.messages.append({"role": "assistant", "content": response})
-            else:
-                # Process the query with the RAG chain
-                with st.spinner("Analyzing First Aid content..."):
-                    try:
-                        result = rag_chain.invoke({"input": prompt})
-                        response = result.get("answer", "Sorry, I couldn't find an answer to that in First Aid.")
-                        
-                        # Display response without footer
-                        display_message(response)
-                        
-                    except Exception as e:
-                        error_response = f"Error processing your request: {str(e)}"
-                        display_message(error_response)
-                        st.session_state.messages.append({"role": "assistant", "content": error_response})
-                        st.experimental_rerun()
-                
-                # Add assistant response to chat history
-                st.session_state.messages.append({"role": "assistant", "content": response})
+if prompt := st.chat_input("Ask a 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)
+    
+    # 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."
+    else:
+        # Process the query with the RAG chain
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                try:
+                    result = rag_chain.invoke({"input": prompt})
+                    response = result.get("answer", "Sorry, I couldn't find an answer to that.")
+                    remaining = MAX_REQUESTS_PER_DAY - count
+                    response += f"\n\n\n_You have {remaining} queries remaining today._"
+                except Exception as e:
+                    response = f"Error processing your request: {str(e)}"
             
-            # Clear the input box - using a different approach that's more reliable
-            st.session_state.user_input = ""
-            st.experimental_rerun()
+            st.markdown(response)
+    
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})
 
 # Footer
 st.markdown("---")
+# Add a small image in the footer
+cols = st.columns([1, 2, 1])
+with cols[1]:
+    st.image("https://i.pinimg.com/originals/ec/81/1c/ec811c4588d10f8286aad7e3ed33e478.png", width=80)
+
 st.markdown("""
-<div style="text-align: center">
-    <p style="color: #6c757d; font-size: 0.8rem;">
-        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>
+<div style="text-align: center; margin-top: 10px;">
+    <p style="font-size: 0.8rem;">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>
 </div>
 """, unsafe_allow_html=True)