Changed to conversation style

app.py

@@ -8,6 +8,7 @@ import spaces
 # --- Initialize the Model Pipeline ---
 print("Loading MedGemma model...")
 try:
+    # We use the same pipeline, but our interaction with it will be different.
     pipe = pipeline(
         "image-text-to-text",
         model="google/medgemma-4b-it",
@@ -21,115 +22,118 @@ except Exception as e:
     model_loaded = False
     print(f"Error loading model: {e}")
 
-# --- Core Analysis Function (Final Robust Version) ---
+# --- Core CONVERSATIONAL Logic ---
 @spaces.GPU()
-def analyze_symptoms(symptom_image: Image.Image, symptoms_text: str):
+def handle_conversation_turn(user_input: str, user_image: Image.Image, history: list):
     """
-    Analyzes user's symptoms with separate, robust logic for image and text-only inputs.
+    Manages a single turn of the conversation, maintaining history.
     """
     if not model_loaded:
-        return "Error: The AI model could not be loaded. Please check the Space logs."
-
-    symptoms_text = symptoms_text.strip() if symptoms_text else ""
-    if symptom_image is None and not symptoms_text:
-        return "Please describe your symptoms or upload an image for analysis."
+        # Append an error message to the chatbot history
+        history.append((user_input, "Error: The AI model is not loaded. Please contact the administrator."))
+        return history, None
 
     try:
+        # --- 1. Define the AI's persona and instructions ---
+        # This is the most critical part for controlling the conversational flow.
         system_prompt = (
-            "You are an expert, empathetic AI medical assistant. Analyze the potential "
-            "medical condition based on the following information. Provide a list of "
-            "possible conditions, your reasoning, and a clear, actionable next-steps plan. "
-            "Start your analysis by describing the user-provided information."
+            "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. "
+            "Your first step is ALWAYS to ask relevant follow-up questions. Ask only one or two focused questions per turn. "
+            "If the user provides an image, acknowledge it by describing what you see in the image first, then ask your questions. "
+            "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,
-        }
-
-        result = ""
-
-        # --- THE FIX: Create two different paths for the logic ---
-
-        if symptom_image:
-            # --- PATH 1: Image is present. Use the proven 'messages' format. ---
-            print("Image detected. Using multimodal 'messages' format...")
-            user_content = []
-            
-            # Only add text content if it actually exists.
-            if symptoms_text:
-                user_content.append({"type": "text", "text": symptoms_text})
-            
-            user_content.append({"type": "image", "image": symptom_image})
-            
-            messages = [
-                {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
-                {"role": "user", "content": user_content}
-            ]
-            
-            output = pipe(text=messages, **generation_args)
-            result = output[0]["generated_text"][-1]["content"]
-
-        else:
-            # --- PATH 2: No image. Use a simple, robust prompt string for text-only. ---
-            print("No image detected. Using robust 'text-only' format...")
-            
-            # Manually construct the prompt to ensure correct formatting.
-            prompt = (
-                f"<start_of_turn>system\n{system_prompt}<start_of_turn>user\n"
-                f"{symptoms_text}<start_of_turn>model\n"
-            )
-
-            # A simple string call returns a different output format.
-            output = pipe(prompt, **generation_args)
-
-            # The full generated text includes the prompt, so we must split it off.
-            full_text = output[0]["generated_text"]
-            result = full_text.split("<start_of_turn>model\n")[-1]
+        # --- 2. Format the conversation for the model ---
+        # The history needs to be converted into the format the model expects.
+        messages = [{"role": "system", "content": [{"type": "text", "text": system_prompt}]}]
+        
+        # Add past interactions from the history
+        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}]})
+
+        # Add the LATEST user input, including the image if provided
+        latest_user_content = []
+        if user_input:
+            latest_user_content.append({"type": "text", "text": user_input})
+        if user_image:
+            latest_user_content.append({"type": "image", "image": user_image})
+        
+        messages.append({"role": "user", "content": latest_user_content})
+        
+        # --- 3. Call the pipeline ---
+        generation_args = {"max_new_tokens": 1024, "do_sample": True, "temperature": 0.7}
+        
+        output = pipe(text=messages, **generation_args)
+        ai_response = output[0]["generated_text"][-1]["content"]
 
-        disclaimer = "\n\n***Disclaimer: I am an AI assistant and not a medical professional. This is not a diagnosis. Please consult a doctor for any health concerns.***"
+        # --- 4. Update the history ---
+        # We store the user's text input and the AI's response. The image is "consumed" in the turn.
+        history.append((user_input, ai_response))
         
-        return result.strip() + disclaimer
+        # We return the updated history for the chatbot display and None to clear the image box.
+        return history, None
 
     except Exception as e:
-        print(f"An exception occurred during analysis: {type(e).__name__}: {e}")
-        return f"An error occurred during analysis. Please check the logs for details: {str(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, None
+
+# --- Gradio Interface for Conversational Flow ---
+with gr.Blocks(theme=gr.themes.Soft(), title="AI Doctor Consultation") as demo:
+    # We use gr.State to hold the conversation history as a list of tuples.
+    conversation_history = gr.State([])
 
-# --- Gradio Interface (No changes needed) ---
-with gr.Blocks(theme=gr.themes.Soft(), title="AI Symptom Analyzer") as demo:
     gr.HTML("""
         <div style="text-align: center; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 2rem; border-radius: 10px; margin-bottom: 2rem;">
-            <h1>🩺 AI Symptom Analyzer</h1>
-            <p>Advanced symptom analysis powered by Google's MedGemma AI</p>
+            <h1>🩺 AI Symptom Consultation</h1>
+            <p>A conversational AI to help you understand your symptoms, powered by Google's MedGemma</p>
         </div>
     """)
     gr.HTML("""
         <div style="background-color: #fff3cd; border: 1px solid #ffeaa7; border-radius: 8px; padding: 1rem; margin: 1rem 0; color: #856404;">
-            <strong>⚠️ Medical Disclaimer:</strong> This AI tool is for informational purposes only and is not a substitute for professional medical diagnosis or treatment.
+            <strong>⚠️ Medical Disclaimer:</strong> This is not a diagnosis. This AI is for informational purposes and is not a substitute for professional medical advice.
         </div>
     """)
-    
-    with gr.Row(equal_height=True):
-        with gr.Column(scale=1):
-            gr.Markdown("### 1. Describe Your Symptoms")
-            symptoms_input = gr.Textbox(
-                label="Symptoms",
-                placeholder="e.g., 'I have a rash on my arm that is red and itchy...'", lines=5)
-            gr.Markdown("### 2. Upload an Image (Optional)")
-            image_input = gr.Image(label="Symptom Image", type="pil", height=300)
-            with gr.Row():
-                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
-                analyze_btn = gr.Button("🔍 Analyze Symptoms", variant="primary", size="lg")
-                
-        with gr.Column(scale=1):
-            gr.Markdown("### 📊 Analysis Report")
-            output_text = gr.Textbox(
-                label="AI Analysis", lines=25, show_copy_button=True, placeholder="Analysis results will appear here...")
 
-    # Event handlers
-    analyze_btn.click(fn=analyze_symptoms, inputs=[image_input, symptoms_input], outputs=output_text)
-    clear_btn.click(fn=lambda: (None, "", ""), outputs=[image_input, symptoms_input, output_text])
+    # The chatbot component will display the conversation history.
+    chatbot_display = gr.Chatbot(height=500, label="Consultation")
+    
+    with gr.Row():
+        # We need an image box that can be cleared after each turn.
+        image_input = gr.Image(label="Upload Symptom Image (Optional)", type="pil", height=150)
+        
+        with gr.Column(scale=4):
+            # The textbox for the user to type their message.
+            user_textbox = gr.Textbox(
+                label="Your Message",
+                placeholder="Describe your primary symptom to begin...",
+                lines=4,
+            )
+            send_button = gr.Button("Send Message", variant="primary")
+
+    def submit_message(user_input, user_image, history):
+        # This wrapper calls the main logic and then clears the user's input fields.
+        updated_history, cleared_image = handle_conversation_turn(user_input, user_image, history)
+        return updated_history, conversation_history, "", cleared_image
+
+    # The submit action
+    send_button.click(
+        fn=handle_conversation_turn,
+        inputs=[user_textbox, image_input, conversation_history],
+        outputs=[chatbot_display, image_input]  # Update the chatbot and clear the image
+    ).then(
+        # Clear the user's text box after the message is sent.
+        lambda: "",
+        outputs=user_textbox
+    )
+    
+    # Add a clear button for convenience
+    clear_button = gr.Button("🗑️ Start New Consultation")
+    clear_button.click(lambda: ([], [], None, ""), outputs=[chatbot_display, conversation_history, image_input, user_textbox])
 
 if __name__ == "__main__":
     print("Starting Gradio interface...")