Update app.py

app.py

@@ -12,50 +12,66 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
+# --- Introductory Explanation ---
+st.markdown(
+    """
+    ### Welcome to the AI Clinical Support Demonstrator
+    This application demonstrates how Generative AI (Google Gemini) can be guided to assist with analyzing clinical information.
+    - **Agentic Text Analysis:** Simulates a structured reasoning process on clinical text to identify key findings, suggest considerations, and note information gaps.
+    - **Medical Image Analysis:** Provides descriptive observations of potential anomalies in medical images.
+
+    **Crucially, this tool is for demonstration purposes ONLY. It does NOT provide medical advice or diagnosis.**
+    """
+)
+st.markdown("---") # Visual separator
+
 # --- Configuration and Initialization ---
 
-# Securely load API key
-# Prioritize Streamlit secrets, fall back to environment variable for flexibility
+# Securely load API key (Secrets > Env Var)
 GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY", os.environ.get("GEMINI_API_KEY"))
 
-# Configure Gemini Client (only if key is found)
+# Configure Gemini Client
 genai_client_configured = False
 if GEMINI_API_KEY:
     try:
         genai.configure(api_key=GEMINI_API_KEY)
         genai_client_configured = True
     except Exception as e:
-        # Use st.exception() for a more detailed error in the app if needed during debugging
-        # st.exception(e)
-        st.error(f"Fatal Error: Failed to configure Google Generative AI. Check API Key. Details: {e}")
-        st.stop() # Stop execution if configuration fails
+        st.error(f"Fatal Error: Failed to configure Google Generative AI. Check API Key. Details: {e}", icon="🚨")
+        st.stop()
 else:
-    st.error("⚠️ Gemini API Key not found. Please configure `GEMINI_API_KEY` in Streamlit secrets or environment variables.")
-    st.stop() # Stop execution if no API key
+    st.error("⚠️ Gemini API Key not found. Please configure `GEMINI_API_KEY` in Streamlit secrets or environment variables.", icon="🔑")
+    st.stop()
+
+# Initialize models
+TEXT_MODEL_NAME = 'gemini-1.5-pro-latest' # For agentic text reasoning
+VISION_MODEL_NAME = 'gemini-1.5-flash'   # For image analysis
 
-# Initialize models (Consider more powerful model for agentic reasoning if needed)
-TEXT_MODEL_NAME = 'gemini-1.5-pro-latest' # Model for agentic text reasoning
-VISION_MODEL_NAME = 'gemini-1.5-flash' # Model for image analysis
+# Use session state to initialize models only once
+if 'models_initialized' not in st.session_state:
+    st.session_state.models_initialized = False
+    st.session_state.text_model = None
+    st.session_state.vision_model = None
 
-if genai_client_configured:
+if genai_client_configured and not st.session_state.models_initialized:
     try:
-        text_model = genai.GenerativeModel(TEXT_MODEL_NAME)
-        vision_model = genai.GenerativeModel(VISION_MODEL_NAME)
-        # Optional: Display success message (won't violate the rule as set_page_config was first)
-        # st.sidebar.success(f"Models Initialized ({TEXT_MODEL_NAME}, {VISION_MODEL_NAME})", icon="✅")
+        st.session_state.text_model = genai.GenerativeModel(TEXT_MODEL_NAME)
+        st.session_state.vision_model = genai.GenerativeModel(VISION_MODEL_NAME)
+        st.session_state.models_initialized = True
+        # Display success message subtly in sidebar perhaps, or remove if too noisy
+        # st.sidebar.success("AI Models Ready.", icon="✅")
     except Exception as e:
-        st.error(f"Fatal Error: Failed to initialize Gemini models. Text: {TEXT_MODEL_NAME}, Vision: {VISION_MODEL_NAME}. Details: {e}")
+        st.error(f"Fatal Error: Failed to initialize Gemini models. Text: {TEXT_MODEL_NAME}, Vision: {VISION_MODEL_NAME}. Details: {e}", icon="💥")
         st.stop()
-else:
-    # This state should not be reached due to earlier st.stop() calls,
-    # but included for completeness.
-    st.error("AI Models could not be initialized due to configuration issues.")
+elif not genai_client_configured:
+    # Should have stopped already, but defensive check
+    st.error("AI Models could not be initialized due to configuration issues.", icon="🚫")
     st.stop()
 
 
 # --- Core AI Interaction Functions ---
 
-# AGENTIC prompt for Text Analysis
+# Refined AGENTIC prompt for Text Analysis with formatting request
 AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE = """
 **Simulated Clinical Reasoning Agent Task:**
 
@@ -63,6 +79,8 @@ AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE = """
 
 **Input Data:** Unstructured clinical information (e.g., symptoms, history, basic findings).
 
+**Output Format:** Please structure your response using Markdown headings for each step (e.g., `## 1. Information Extraction`).
+
 **Simulated Agentic Steps (Perform sequentially):**
 
 1.  **Information Extraction & Structuring:**
@@ -94,17 +112,21 @@ AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE = """
 **Agentic Analysis:**
 """
 
-# Standard prompt for Image Analysis
+# Refined prompt for Image Analysis with formatting request
 IMAGE_ANALYSIS_PROMPT_TEMPLATE = """
 **Medical Image Analysis Request:**
 
 **Context:** Analyze the provided medical image. User may provide additional context or questions.
+
+**Output Format:** Please structure your response using Markdown headings for each step (e.g., `## 1. Visible Structures`).
+
 **Task:**
+
 1.  **Describe Visible Structures:** Briefly describe main anatomical structures.
-2.  **Identify Potential Anomalies:** Point out areas that *appear* abnormal or deviate from typical presentation (e.g., "potential opacity," "altered signal intensity," "possible asymmetry"). Use cautious, descriptive language.
-3.  **Correlate with User Prompt (if provided):** Address specific user questions based *only* on visual information.
-4.  **Limitations:** State that image quality, view, and lack of clinical context limit analysis.
-5.  **Disclaimer:** Explicitly state this is AI visual analysis, not radiological interpretation or diagnosis, requiring review by a qualified professional with clinical context.
+2.  **Identify Potential Anomalies:** Point out areas that *appear* abnormal or deviate from typical presentation (e.g., "potential opacity," "altered signal intensity," "possible asymmetry"). Use cautious, descriptive language. **This is not a definitive finding.**
+3.  **Correlate with User Prompt (if provided):** Address specific user questions based *only* on visual information. State if the image cannot answer the question.
+4.  **Limitations:** State that image quality, view, lack of clinical context, and the AI's nature limit analysis.
+5.  **Mandatory Disclaimer:** Explicitly state this is AI visual analysis, not radiological interpretation or diagnosis, requiring review by a qualified professional with clinical context.
 
 **User's Additional Context/Question (if any):**
 ---
@@ -115,21 +137,15 @@ IMAGE_ANALYSIS_PROMPT_TEMPLATE = """
 """
 
 def run_agentic_text_analysis(text_input: str) -> Tuple[Optional[str], Optional[str]]:
-    """
-    Sends clinical text to the Gemini text model for simulated agentic analysis.
-
-    Args:
-        text_input: The clinical text provided by the user.
-
-    Returns:
-        Tuple: (analysis_text, error_message)
-    """
+    """Sends clinical text to the configured text model for simulated agentic analysis."""
     if not text_input or not text_input.strip():
         return None, "Input text cannot be empty."
+    if not st.session_state.models_initialized or not st.session_state.text_model:
+         return None, "Text analysis model not initialized."
+
     try:
         prompt = AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE.format(text_input=text_input)
-        # Using the designated text model
-        response = text_model.generate_content(prompt)
+        response = st.session_state.text_model.generate_content(prompt)
 
         if response.parts:
             return response.text, None
@@ -143,41 +159,34 @@ def run_agentic_text_analysis(text_input: str) -> Tuple[Optional[str], Optional[
                  return None, "Received an empty or unexpected response from the AI model."
 
     except Exception as e:
-        # Log error for debugging purposes server-side if possible
-        print(f"ERROR in run_agentic_text_analysis: {e}") # Simple print for demo
-        st.error(f"Error during agentic text analysis. Check logs if possible.", icon="🚨") # Keep UI error generic
-        return None, f"An internal error occurred during text analysis." # Generic message to user
+        print(f"ERROR in run_agentic_text_analysis: {e}") # Basic server-side log
+        st.error("An error occurred during text analysis.", icon="🚨") # User-facing error
+        return None, "An internal error occurred during text analysis. Please try again later."
 
 def analyze_medical_image(image_file: Any, user_prompt: str = "") -> Tuple[Optional[str], Optional[str]]:
-    """
-    Sends a medical image (and optional prompt) to the Gemini Vision model for analysis.
-
-    Args:
-        image_file: Uploaded image file object from Streamlit.
-        user_prompt: Optional text context/questions from the user.
-
-    Returns:
-        Tuple: (analysis_text, error_message)
-    """
+    """Sends a medical image to the configured vision model for analysis."""
     if not image_file:
         return None, "Image file cannot be empty."
+    if not st.session_state.models_initialized or not st.session_state.vision_model:
+         return None, "Image analysis model not initialized."
+
     try:
         try:
             image = Image.open(image_file)
+            # Convert to RGB ensure compatibility, common requirement for vision models
             if image.mode != 'RGB':
                  image = image.convert('RGB')
         except Exception as img_e:
-             return None, f"Error opening or processing image file: {img_e}"
+             return None, f"Error opening or processing the uploaded image file: {img_e}"
 
         prompt_text = IMAGE_ANALYSIS_PROMPT_TEMPLATE.format(user_prompt=user_prompt if user_prompt else "N/A")
         model_input = [prompt_text, image]
-        # Using the designated vision model
-        response = vision_model.generate_content(model_input)
+        response = st.session_state.vision_model.generate_content(model_input)
 
         if response.parts:
             return response.text, None
         elif response.prompt_feedback.block_reason:
-            return None, f"Image analysis blocked by safety filters: {response.prompt_feedback.block_reason.name}. May relate to sensitive content policies."
+            return None, f"Image analysis blocked by safety filters: {response.prompt_feedback.block_reason.name}. This may relate to sensitive content policies."
         else:
              candidate = response.candidates[0] if response.candidates else None
              if candidate and candidate.finish_reason != "STOP":
@@ -186,22 +195,20 @@ def analyze_medical_image(image_file: Any, user_prompt: str = "") -> Tuple[Optio
                  return None, "Received an empty or unexpected response from the AI model for image analysis."
 
     except Exception as e:
-        # Log error for debugging purposes server-side if possible
-        print(f"ERROR in analyze_medical_image: {e}") # Simple print for demo
-        st.error(f"Error during image analysis. Check logs if possible.", icon="🖼️") # Keep UI error generic
-        return None, f"An internal error occurred during image analysis." # Generic message to user
+        print(f"ERROR in analyze_medical_image: {e}") # Basic server-side log
+        st.error("An error occurred during image analysis.", icon="🖼️") # User-facing error
+        return None, "An internal error occurred during image analysis. Please try again later."
 
 
 # --- Streamlit User Interface ---
 
 def main():
-    # REMOVED st.set_page_config() FROM HERE
-
+    # Page title and model info
     st.title("🤖 AI Clinical Support Demonstrator")
-    st.caption(f"Agentic Text Analysis ({TEXT_MODEL_NAME}) & Image Analysis ({VISION_MODEL_NAME})")
-    st.markdown("---")
+    st.caption(f"Utilizing: Text Model ({TEXT_MODEL_NAME}), Vision Model ({VISION_MODEL_NAME})")
 
     # --- CRITICAL DISCLAIMER ---
+    # Positioned prominently after title/intro
     st.warning(
         """
         **🔴 IMPORTANT SAFETY & USE DISCLAIMER 🔴**
@@ -210,13 +217,13 @@ def main():
         *   **Image Analysis:** Provides observations on images. Output is **NOT** a radiological interpretation.
         *   AI analysis lacks full clinical context, may be inaccurate, and **CANNOT** replace professional judgment.
         *   **ALWAYS consult qualified healthcare professionals** for diagnosis and treatment.
-        *   **PRIVACY:** Do **NOT** upload identifiable patient information (PHI) without explicit consent and adherence to all privacy laws (e.g., HIPAA).
+        *   **PRIVACY:** Do **NOT** upload identifiable patient information (PHI) without explicit consent and adherence to all privacy laws (e.g., HIPAA). You are responsible for the data you input.
         """,
         icon="⚠️"
     )
     st.markdown("---")
 
-
+    # --- Sidebar Controls ---
     st.sidebar.header("Analysis Options")
     input_method = st.sidebar.radio(
         "Select Analysis Type:",
@@ -226,26 +233,31 @@ def main():
     )
     st.sidebar.markdown("---") # Visual separator
 
+    # --- Main Area Layout (Input and Output Columns) ---
     col1, col2 = st.columns(2)
 
+    analysis_result = None # Initialize results variables for this run
+    error_message = None
+    output_header = "Analysis Results" # Default header
+
+    # --- Column 1: Input Area ---
     with col1:
         st.header("Input Data")
-        analysis_result = None # Initialize results variables
-        error_message = None
-        output_header = "Analysis Results" # Default header for the output column
 
-        # --- Agentic Text Analysis Input ---
+        # Conditional Input UI based on selection
         if input_method == "Agentic Text Analysis":
             st.subheader("Clinical Text for Agentic Analysis")
+            st.caption("Please ensure data is de-identified before pasting.")
             text_input = st.text_area(
-                "Paste de-identified clinical information (symptoms, history, findings):",
-                height=350, # Slightly larger text area
+                "Paste clinical information:",
+                height=350,
                 placeholder="Example: 68yo male, sudden SOB & pleuritic chest pain post-flight. HR 110, SpO2 92% RA. No known cardiac hx...",
                 key="text_input_area"
             )
-            analyze_button = st.button("▶️ Run Agentic Text Analysis", key="analyze_text_button", type="primary")
+            analyze_button_key = "analyze_text_button"
+            analyze_button_label = "▶️ Run Agentic Text Analysis"
 
-            if analyze_button:
+            if st.button(analyze_button_label, key=analyze_button_key, type="primary"):
                 if text_input:
                     with st.spinner("🧠 Simulating agentic reasoning..."):
                         analysis_result, error_message = run_agentic_text_analysis(text_input)
@@ -253,11 +265,11 @@ def main():
                 else:
                     st.warning("Please enter clinical text to analyze.", icon="☝️")
 
-        # --- Medical Image Analysis Input ---
         elif input_method == "Medical Image Analysis":
             st.subheader("Medical Image for Analysis")
+            st.caption("Upload a de-identified medical image (PNG, JPG, JPEG).")
             image_file = st.file_uploader(
-                "Upload a de-identified medical image (e.g., X-ray, CT slice). Supported: PNG, JPG, JPEG.",
+                "Choose an image file:",
                 type=["png", "jpg", "jpeg"],
                 key="image_uploader"
             )
@@ -266,60 +278,63 @@ def main():
                 placeholder="Example: 'Describe findings in the lung fields' or 'Any visible fractures?'",
                 key="image_prompt_input"
             )
-            analyze_button = st.button("🖼️ Analyze Medical Image", key="analyze_image_button", type="primary")
+            analyze_button_key = "analyze_image_button"
+            analyze_button_label = "🖼️ Analyze Medical Image"
 
-            if analyze_button:
+            if image_file:
+                # Display preview immediately after upload inside the input column
+                st.image(image_file, caption="Uploaded Image Preview", use_column_width=True)
+
+            if st.button(analyze_button_label, key=analyze_button_key, type="primary"):
                 if image_file:
-                    # Using st.image within the column context
-                    st.image(image_file, caption="Uploaded Image Preview", use_column_width=True)
                     with st.spinner("👁️ Analyzing image..."):
                         analysis_result, error_message = analyze_medical_image(image_file, user_image_prompt)
                     output_header = "Medical Image Analysis Output"
                 else:
                     st.warning("Please upload an image file to analyze.", icon="☝️")
 
-    # --- Output Column ---
+    # --- Column 2: Output Area ---
     with col2:
         st.header(output_header)
-        # Determine if analysis was attempted in this run
-        # Use session state keys which are set when buttons are clicked
-        button_pressed = st.session_state.get('analyze_text_button', False) or \
-                         st.session_state.get('analyze_image_button', False)
-
-        if button_pressed:
-            # Check if the button was clicked AND we have results or an error from that action
-            # analysis_result and error_message are set within the button's if block in col1
-            if error_message:
+
+        # Display results or errors if an analysis was attempted
+        # Check if button was pressed and corresponding result/error exists
+        button_pressed = st.session_state.get(analyze_button_key, False) if 'analyze_button_key' in locals() else False
+
+        if button_pressed and (analysis_result or error_message):
+             if error_message:
                  st.error(f"Analysis Failed: {error_message}", icon="❌")
-            elif analysis_result:
-                 st.markdown(analysis_result) # Display the successful result
-            # If button was pressed but neither error nor result is set, maybe input was invalid
-            # The warnings for empty input are handled in col1, so this case might not be needed
-            # else:
-            #    st.info("Provide valid input and click Analyze.") # Fallback?
-
-        else: # No button pressed in this run
+             elif analysis_result:
+                 # Use st.markdown to render potential formatting from AI
+                 st.markdown(analysis_result)
+        elif not button_pressed :
             st.info("Analysis results will appear here after providing input and clicking the corresponding analysis button.")
+        # Add a condition for button pressed but no input provided (already handled by st.warning in col1)
 
 
     # --- Sidebar Explanations ---
     st.sidebar.markdown("---")
     st.sidebar.header("About The Prompts")
-    with st.sidebar.expander("View Agentic Text Prompt Structure"):
+    with st.sidebar.expander("View Agentic Text Prompt Structure", icon="📄"):
         st.markdown(f"```plaintext\n{AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE.split('---')[0]} ... [Input Text] ...\n```")
         st.caption("Guides the AI through structured reasoning steps for text.")
-    with st.sidebar.expander("View Image Analysis Prompt Structure"):
+    with st.sidebar.expander("View Image Analysis Prompt Structure", icon="🖼️"):
         st.markdown(f"```plaintext\n{IMAGE_ANALYSIS_PROMPT_TEMPLATE.split('---')[0]} ... [User Prompt] ...\n```")
         st.caption("Guides the AI to describe visual features and potential anomalies in images.")
 
     st.sidebar.markdown("---")
     st.sidebar.error(
-        "**Ethical Use Reminder:** AI in medicine requires extreme caution. This tool is for demonstration and education, not clinical practice. Verify all information and rely on professional expertise."
+        "**Ethical Use Reminder:** AI in medicine requires extreme caution. This tool is for demonstration and education, not clinical practice. Verify all information and rely on professional expertise.",
+        icon = "⚕️"
         )
 
 # --- Main Execution Guard ---
 if __name__ == "__main__":
-    # Ensure configuration happens before main UI logic runs
-    # NOTE: The config/init code now runs automatically when the script is imported/run,
-    # BEFORE the main() function is called here. This is correct.
-    main()
+    # Check if models are initialized before running the main UI components
+    if st.session_state.models_initialized:
+        main()
+    else:
+        # If models failed to initialize, errors would have been shown already.
+        # We might add a placeholder or further message here if needed.
+        st.info("Waiting for model initialization...") # Or check specific errors if needed
+        # The script might stop earlier if initialization fails critically.