app.py

import gradio as gr
from transformers import pipeline
from PIL import Image
import torch
import os
import spaces

# --- Initialize the Model Pipeline (No changes) ---
print("Loading MedGemma model...")
try:
    pipe = pipeline(
        "image-text-to-text",
        model="google/medgemma-4b-it",
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=os.environ.get("HF_TOKEN")
    )
    model_loaded = True
    print("Model loaded successfully!")
except Exception as e:
    model_loaded = False
    print(f"Error loading model: {e}")

# --- Core CONVERSATIONAL Logic (No changes) ---
@spaces.GPU()
def handle_conversation_turn(user_input: str, user_image: Image.Image, history: list):
    """
    Manages a single conversation turn with state management.
    """
    if not model_loaded:
        history.append((user_input, "Error: The AI model is not loaded."))
        return history, history, None

    try:
        system_prompt = (
            "You are an expert, empathetic AI medical assistant conducting a virtual consultation. "
            "Your primary goal is to ask clarifying questions to understand the user's symptoms thoroughly. "
            "Do NOT provide a diagnosis or a list of possibilities right away. Ask only one or two focused questions per turn. "
            "If the user provides an image, your first step is to analyze it from an expert perspective. Briefly describe the key findings from the image. "
            "Then, use this analysis to ask relevant follow-up questions about the user's symptoms or medical history to better understand the context. "
            "For example, after seeing a rash, you might say, 'I see a reddish rash with well-defined borders on the forearm. To help me understand more, could you tell me when you first noticed this? Is it itchy, painful, or does it have any other sensation?'"
            "After several turns of asking questions, when you feel you have gathered enough information, you must FIRST state that you are ready to provide a summary. "
            "THEN, in the SAME response, provide a list of possible conditions, your reasoning, and a clear, actionable next-steps plan."
        )

        generation_args = {"max_new_tokens": 1024, "do_sample": True, "temperature": 0.7}
        ai_response = ""
        
        if user_image:
            print("Image detected. Using multimodal 'messages' format...")
            messages = [{"role": "system", "content": [{"type": "text", "text": system_prompt}]}]
            for user_msg, assistant_msg in history:
                messages.append({"role": "user", "content": [{"type": "text", "text": user_msg}]})
                messages.append({"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]})
            latest_user_content = [{"type": "text", "text": user_input}, {"type": "image", "image": user_image}]
            messages.append({"role": "user", "content": latest_user_content})
            output = pipe(text=messages, **generation_args)
            ai_response = output[0]["generated_text"][-1]["content"]
        else:
            print("No image detected. Using robust 'text-only' format.")
            prompt_parts = [f"<start_of_turn>system\n{system_prompt}"]
            for user_msg, assistant_msg in history:
                prompt_parts.append(f"<start_of_turn>user\n{user_msg}")
                prompt_parts.append(f"<start_of_turn>model\n{assistant_msg}")
            prompt_parts.append(f"<start_of_turn>user\n{user_input}")
            prompt_parts.append("<start_of_turn>model")
            prompt = "\n".join(prompt_parts)
            output = pipe(prompt, **generation_args)
            full_text = output[0]["generated_text"]
            ai_response = full_text.split("<start_of_turn>model")[-1].strip()

        history.append((user_input, ai_response))
        return history, history, None, "" # Return empty string to clear textbox

    except Exception as e:
        history.append((user_input, f"An error occurred: {str(e)}"))
        print(f"An exception occurred during conversation turn: {type(e).__name__}: {e}")
        return history, history, None, user_input # Return user_input to not clear on error

# --- UI MODIFICATION: Custom CSS for the sticky footer ---
css = """
/* Make the main app container fill the screen height and add padding at the bottom */
div.gradio-container {
    height: 100vh !important;
    padding-bottom: 120px !important; /* Space for the sticky footer */
}
/* Style the input bar to be fixed at the bottom of the screen */
#sticky-footer {
    position: fixed !important;
    bottom: 0px !important;
    left: 0px !important;
    width: 100% !important;
    background: white !important;
    padding: 10px !important;
    border-top: 1px solid #E5E7EB !important; /* Light gray border */
    z-index: 1000;
}
/* Center the clear button */
#clear-button-container {
    display: flex;
    justify-content: center;
    padding: 10px;
}
"""

# --- Gradio Interface with New Chat Layout ---
with gr.Blocks(theme=gr.themes.Soft(), title="AI Doctor Consultation", css=css) as demo:
    conversation_history = gr.State([])

    gr.HTML('<div style="text-align: center; padding: 10px;"><h1>🩺 AI Symptom Consultation</h1></div>')
    
    # The chatbot display is now the main element at the top.
    chatbot_display = gr.Chatbot(height=550, label="Consultation", show_copy_button=True)
    
    # The input controls are wrapped in a Column with a specific ID for CSS styling.
    with gr.Column(elem_id="sticky-footer"):
        with gr.Row(vertical_align="center"):
            image_input = gr.Image(label="Image", type="pil", height=100, show_label=False, container=False)
            user_textbox = gr.Textbox(
                label="Your Message", 
                placeholder="Type your message here...", 
                show_label=False,
                scale=4,
                container=False, # Important for alignment
            )
            send_button = gr.Button("Send", variant="primary", scale=1)

    # The clear button is placed just above the chatbot for easy access.
    with gr.Row(elem_id="clear-button-container"):
        clear_button = gr.Button("🗑️ Start New Consultation")

    def submit_message(user_input, user_image, history):
        # This wrapper calls the main logic.
        updated_history, new_state, cleared_image, cleared_text = handle_conversation_turn(user_input, user_image, history)
        return updated_history, new_state, cleared_image, cleared_text

    # --- Event Handlers ---
    # Handle both button click and Enter key press
    send_button.click(
        fn=submit_message,
        inputs=[user_textbox, image_input, conversation_history],
        outputs=[chatbot_display, conversation_history, image_input, user_textbox]
    )
    user_textbox.submit(
        fn=submit_message,
        inputs=[user_textbox, image_input, conversation_history],
        outputs=[chatbot_display, conversation_history, image_input, user_textbox]
    )
    
    clear_button.click(
        lambda: ([], [], None, ""), 
        outputs=[chatbot_display, conversation_history, image_input, user_textbox]
    )

if __name__ == "__main__":
    print("Starting Gradio interface...")
    demo.launch(debug=True)