app.py

# -*- coding: utf-8 -*-
import streamlit as st
import google.generativeai as genai
import os
from PIL import Image
import io # Needed for handling image bytes
from typing import Optional, Tuple, Any # For type hinting

# --- Page Configuration (MUST BE THE FIRST STREAMLIT COMMAND) ---
st.set_page_config(
    page_title="AI Clinical Support Demonstrator",
    layout="wide",
    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.
    - **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 (Secrets > Env Var)
GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY", os.environ.get("GEMINI_API_KEY"))

# 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:
        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.", icon="🔑")
    st.stop()

# Initialize models using Session State for persistence across reruns
TEXT_MODEL_NAME = 'gemini-1.5-pro-latest' # For agentic text reasoning
VISION_MODEL_NAME = 'gemini-1.5-flash'   # For image analysis (or 'gemini-pro-vision')

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 and not st.session_state.models_initialized:
    try:
        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
    except Exception as e:
        st.error(f"Fatal Error: Failed to initialize Gemini models. Text: {TEXT_MODEL_NAME}, Vision: {VISION_MODEL_NAME}. Details: {e}", icon="💥")
        st.stop()
elif not genai_client_configured:
    st.error("AI Models could not be initialized due to configuration issues.", icon="🚫")
    st.stop()


# --- Core AI Interaction Functions ---

# AGENTIC prompt for Text Analysis (Remains the same as previous good version)
AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE = """
**Simulated Clinical Reasoning Agent Task:**

**Role:** AI assistant simulating an agentic clinical reasoning process to support a healthcare professional by structuring information, generating possibilities, and suggesting investigation pathways based *strictly* on the provided text. **This is NOT a diagnosis.**

**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.  **## 1. Information Extraction & Structuring:**
    *   Key Demographics (Age, Sex if provided).
    *   Primary Symptoms/Signs.
    *   Relevant Medical History.
    *   Pertinent Negatives (if mentioned).

2.  **## 2. Differential Considerations Generation:**
    *   Based *only* on Step 1, list **potential differential considerations** (possible conditions).
    *   **Use cautious language:** "could be consistent with," "warrants consideration," "less likely but possible." **AVOID definitive statements.**
    *   Briefly justify each consideration based on findings.

3.  **## 3. Information Gap Analysis:**
    *   Identify critical missing information (e.g., lab results, imaging, exam specifics, duration/onset).

4.  **## 4. Suggested Next Steps for Investigation (for Clinician):**
    *   Propose logical next steps a **healthcare professional might consider**.
    *   Categorize (e.g., Further History, Exam Points, Labs, Imaging).
    *   Frame as *suggestions* (e.g., "Consider ordering...", "Assessment of X may be informative").

5.  **## 5. Mandatory Disclaimer:** Conclude with: "This AI-generated analysis is for informational support only. It is **NOT** a diagnosis and cannot replace the judgment of a qualified healthcare professional."

**Input Clinical Information:**
---
{text_input}
---

**Agentic Analysis:**
"""

# REFINED prompt for Image Analysis - More explicit instructions based on good example
IMAGE_ANALYSIS_PROMPT_TEMPLATE = """
**Medical Image Analysis Request:**

**Context:** Analyze the provided medical image objectively based *only* on visual information. User may provide additional context or questions.

**Output Format:** Structure your response precisely using the following Markdown headings. Be factual and descriptive.

**Task:**

1.  **## 1. Visible Structures:**
    *   Identify the likely imaging modality and view (e.g., PA Chest Radiograph, Axial CT slice of the abdomen).
    *   Briefly list the main anatomical structures clearly visible (e.g., ribs, heart silhouette, lung fields, diaphragm).

2.  **## 2. Identify Potential Anomalies / Key Findings:**
    *   Carefully examine the image for any areas that *appear* abnormal or deviate significantly from typical presentation.
    *   **Use extremely cautious, descriptive language.** Describe *what* you see (e.g., "area of increased opacity," "region of lucency," "asymmetry observed in X," "potential contour abnormality," "patchy distribution").
    *   **Specify location accurately** using standard anatomical terms (e.g., "right lower lung zone," "left hilum," "hepatic flexure region").
    *   **Crucially, AVOID interpretive or diagnostic terms** (DO NOT use words like "pneumonia," "tumor," "fracture," "infection," "inflammation"). Stick strictly to visual observation.
    *   If relevant and clearly discernible, mention the **absence** of certain major expected abnormalities (pertinent negatives, e.g., "No obvious large pneumothorax identified," "Bowel gas pattern appears unremarkable in visualized areas").
    *   Compare sides if applicable and relevant differences are seen (e.g., "Left lung field demonstrates greater transparency compared to the right").

3.  **## 3. Correlate with User Prompt (if provided):**
    *   Address specific user questions based *strictly* on the visual information identifiable in the image.
    *   If the image cannot visually answer the question (e.g., requires clinical context, different view), state that clearly.
    *   If no user prompt was provided, state "N/A".

4.  **## 4. Limitations of this AI Analysis:**
    *   **Explicitly list the following limitations inherent to this analysis:**
        *   Dependency on the **quality, resolution, and potential artifacts** of the provided image.
        *   Analysis is restricted to the **single view/slice(s)** provided; other areas are not assessed.
        *   **Complete lack of clinical context:** Patient history, symptoms, physical exam findings, and laboratory results are unknown and not considered.
        *   **Absence of prior imaging studies:** Comparisons over time are not possible, which is often crucial for interpretation.
        *   The AI functions purely on **visual pattern recognition**; it does not perform clinical reasoning or differential diagnosis.

5.  **## 5. Mandatory Disclaimer:**
    *   State clearly: This is an AI-generated visual analysis intended for informational and demonstration purposes **ONLY**.
    *   It is **NOT** a radiological interpretation or medical diagnosis.
    *   It **CANNOT** substitute for a comprehensive evaluation and interpretation by a qualified radiologist or physician integrating full clinical information.
    *   Any potential observations noted herein **MUST** be correlated with clinical findings and reviewed/confirmed by qualified healthcare professionals.

**User's Additional Context/Question (if any):**
---
{user_prompt}
---

**Image Analysis:**
"""

# --- Backend Functions ---

def run_agentic_text_analysis(text_input: str) -> Tuple[Optional[str], Optional[str]]:
    """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. Please refresh or check configuration."

    try:
        prompt = AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE.format(text_input=text_input)
        response = st.session_state.text_model.generate_content(prompt) # Use model from session state

        # Handle response variations
        if response.parts:
            return response.text, None
        elif response.prompt_feedback.block_reason:
             return None, f"Analysis blocked by safety filters: {response.prompt_feedback.block_reason.name}. Please review or revise input."
        else:
             candidate = response.candidates[0] if response.candidates else None
             if candidate and candidate.finish_reason != "STOP":
                 # If stopped for reasons other than normal completion (e.g., length, safety)
                 return None, f"Analysis stopped prematurely. Reason: {candidate.finish_reason.name}. Input might be too long or triggered other limits."
             else:
                 # General case for empty or unexpected response
                 return None, "Received an empty or unexpected response from the AI model for text analysis."

    except Exception as e:
        print(f"ERROR in run_agentic_text_analysis: {e}") # Log for server/dev console
        st.error("An error occurred during text analysis.", icon="🚨") # User-facing generic error
        return None, "An internal error occurred during text analysis. Please try again later or contact support if the issue persists."

def analyze_medical_image(image_file: Any, user_prompt: str = "") -> Tuple[Optional[str], Optional[str]]:
    """Sends a medical image to the configured vision model for analysis using the refined prompt."""
    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. Please refresh or check configuration."

    try:
        try:
            # Open and prepare the image
            image = Image.open(image_file)
            if image.mode != 'RGB':
                 image = image.convert('RGB') # Ensure RGB format
        except Exception as img_e:
             return None, f"Error opening or processing the uploaded image file: {img_e}. Please ensure it's a valid image."

        # Prepare the prompt using the refined template
        prompt_text = IMAGE_ANALYSIS_PROMPT_TEMPLATE.format(user_prompt=user_prompt if user_prompt else "N/A")
        model_input = [prompt_text, image] # Combine text prompt and image data

        response = st.session_state.vision_model.generate_content(model_input) # Use model from session state

        # Handle response variations
        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}. This might relate to sensitive content policies regarding medical images."
        else:
             candidate = response.candidates[0] if response.candidates else None
             if candidate and candidate.finish_reason != "STOP":
                 return None, f"Image analysis stopped prematurely. Reason: {candidate.finish_reason.name}. Input might be too complex or triggered other limits."
             else:
                 return None, "Received an empty or unexpected response from the AI model for image analysis."

    except Exception as e:
        print(f"ERROR in analyze_medical_image: {e}") # Log for server/dev console
        st.error("An error occurred during image analysis.", icon="🖼️") # User-facing generic error
        return None, "An internal error occurred during image analysis. Please try again later or contact support if the issue persists."


# --- Streamlit User Interface ---

def main():
    # Page title and model info
    st.title("🤖 AI Clinical Support Demonstrator")
    st.caption(f"Utilizing: Text Model ({TEXT_MODEL_NAME}), Vision Model ({VISION_MODEL_NAME})")

    # --- CRITICAL DISCLAIMER ---
    st.warning(
        """
        **🔴 IMPORTANT SAFETY & USE DISCLAIMER 🔴**
        *   This tool **DEMONSTRATES** AI capabilities. It **DOES NOT** provide medical advice or diagnosis.
        *   **Agentic Text Analysis:** Simulates reasoning on text input. Output is illustrative, not diagnostic.
        *   **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). 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:",
        ("Agentic Text Analysis", "Medical Image Analysis"),
        key="input_method_radio",
        help="Choose 'Agentic Text Analysis' for reasoning simulation on clinical text, or 'Medical Image Analysis' for visual observations on images."
    )
    st.sidebar.markdown("---") # Visual separator

    # --- Main Area Layout (Input and Output Columns) ---
    col1, col2 = st.columns(2)

    analysis_result = None # Reset results variables for this run
    error_message = None
    output_header = "Analysis Results" # Default header
    analyze_button_key = None # Initialize key variable

    # --- Column 1: Input Area ---
    with col1:
        st.header("Input Data")

        # 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 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_key = "analyze_text_button" # Set key for this branch
            analyze_button_label = "▶️ Run Agentic Text Analysis"

            if st.button(analyze_button_label, key=analyze_button_key, type="primary"):
                if text_input:
                    with st.spinner("🧠 Simulating agentic reasoning... Please wait."):
                        analysis_result, error_message = run_agentic_text_analysis(text_input)
                    output_header = "Simulated Agentic Analysis Output"
                else:
                    st.warning("Please enter clinical text to analyze.", icon="☝️")

        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(
                "Choose an image file:",
                type=["png", "jpg", "jpeg"],
                key="image_uploader"
            )
            user_image_prompt = st.text_input(
                "Optional: Add context or specific question for image analysis:",
                placeholder="Example: 'Describe findings in the lung fields' or 'Any visible fractures?'",
                key="image_prompt_input"
            )
            analyze_button_key = "analyze_image_button" # Set key for this branch
            analyze_button_label = "🖼️ Analyze Medical Image"

            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:
                    with st.spinner("👁️ Analyzing image... Please wait."):
                        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="☝️")

    # --- Column 2: Output Area ---
    with col2:
        st.header(output_header)

        # Determine if analysis was attempted in this run using the correct button key
        button_pressed = st.session_state.get(analyze_button_key, False) if analyze_button_key else False

        if button_pressed:
            # Only display results if the corresponding button was pressed AND generated output/error
            if error_message:
                 st.error(f"Analysis Failed: {error_message}", icon="❌")
            elif analysis_result:
                 # Use st.markdown to render potential formatting (like headings) from AI
                 st.markdown(analysis_result, unsafe_allow_html=False) # Keep unsafe_allow_html=False for security
            # No explicit 'else' needed here; if button pressed but no result/error, likely handled by input validation warning
        else:
            # Default placeholder message if no analysis attempted yet in this run
            st.info("Analysis results will appear here after providing input and clicking the corresponding analysis button.")


    # --- Sidebar Explanations ---
    st.sidebar.markdown("---")
    st.sidebar.header("About The Prompts")
    with st.sidebar.expander("View Agentic Text Prompt Structure", icon="📄"):
        # Show only the instructive part of the prompt template
        st.code(AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE.split('---')[0] + "...", language='markdown')
        st.caption("Guides the AI through structured reasoning steps for text.")
    with st.sidebar.expander("View Image Analysis Prompt Structure", icon="🖼️"):
        # Show only the instructive part of the prompt template
        st.code(IMAGE_ANALYSIS_PROMPT_TEMPLATE.split('---')[0] + "...", language='markdown')
        st.caption("Guides the AI to provide cautious, descriptive visual observations for 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.",
        icon = "⚕️"
        )

# --- Main Execution Guard ---
if __name__ == "__main__":
    # Check if models are initialized before running the main UI components
    if st.session_state.models_initialized:
        main()
    else:
        # Errors during initialization should have been shown and stopped execution.
        # This path might be reached if initialization is still pending in an async setup (not the case here)
        # or if there was a non-fatal init issue not caught by st.stop().
        st.info("Waiting for AI model initialization or resolving configuration issues...")