Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-import base64
+# import base64 # No longer needed for Gemini vision part
 import gradio as gr
 import pandas as pd
 from groq import Groq
@@ -7,178 +7,225 @@ from PIL import Image
 import io
 import datetime
 import re
+import google.generativeai as genai # Added for Gemini
+from google.generativeai import types # Added for Gemini
 
-
-# Initialize Groq client
-client = Groq(
+# Initialize Groq client (for chat/assessment)
+groq_client = Groq(
     api_key=os.environ.get("GROQ_API_KEY")
 )
 
-# Function to encode images to base64
-def encode_image(image):
-    buffered = io.BytesIO()
-    image.save(buffered, format="JPEG")
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+# NOTE: Gemini client initialization will happen inside the analyze_ecg_image function
+# Ensure GEMINI_API_KEY is set in your environment variables
+
+# Function to encode images to base64 - REMOVED as not needed for Gemini bytes input
+# def encode_image(image):
+#     buffered = io.BytesIO()
+#     image.save(buffered, format="JPEG")
+#     return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-# Process patient history file
+# Process patient history file (Unchanged)
 def process_patient_history(file):
     if file is None:
         return ""
-    
+
     try:
         # Check file extension
         file_ext = os.path.splitext(file.name)[1].lower()
-        
+
         if file_ext == '.txt':
             # Read text file
-            content = file.read().decode('utf-8')
+            # Gradio File object might need .name for path, let's assume it provides a file-like object
+            with open(file.name, 'r', encoding='utf-8') as f:
+                 content = f.read()
+            # If file is already an IO object (depends on Gradio version/usage)
+            # content = file.read().decode('utf-8')
             return content
-        
+
         elif file_ext in ['.csv', '.xlsx', '.xls']:
             # Read spreadsheet file
             if file_ext == '.csv':
                 df = pd.read_csv(file.name)
             else:
                 df = pd.read_excel(file.name)
-            
+
             # Convert dataframe to formatted string
             formatted_data = "PATIENT INFORMATION:\n\n"
-            for column in df.columns:
-                formatted_data += f"{column}: {df.iloc[0][column]}\n"
-            
+            if not df.empty:
+                 # Assuming the first row contains the relevant patient data
+                for column in df.columns:
+                    # Handle potential missing values gracefully
+                    value = df.iloc[0].get(column, 'N/A')
+                    formatted_data += f"{column}: {value}\n"
+            else:
+                formatted_data += "Spreadsheet is empty or format is not recognized correctly."
+
             return formatted_data
-        
+
         else:
             return "Unsupported file format. Please upload a .txt, .csv, or .xlsx file."
-    
+
     except Exception as e:
         return f"Error processing patient history file: {str(e)}"
 
-# Extract ECG readings from image using Llama Vision model
-def analyze_ecg_image(image, vision_model="llama-3.2-90b-vision-preview"):
-    # Fixed model - always use llama-3.2-90b-vision-preview
-    vision_model = "llama-3.2-90b-vision-preview"
-    
+
+# Extract ECG readings from image using Gemini Vision model
+def analyze_ecg_image(image, vision_model="gemini-2.0-flash-exp"):
+    # Fixed model
+    vision_model = "gemini-2.0-flash-exp"
+
     if image is None:
         return "<strong style='color:red'>No image provided.</strong>"
-    
-    # Convert to PIL Image if needed
+
+    # Ensure image is PIL Image
     if not isinstance(image, Image.Image):
-        image = Image.open(image)
-    
-    # Encode the image
-    base64_image = encode_image(image)
-    
-    # Get current timestamp
-    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    
-    # Create chat completion with vision model
-    vision_prompt = f"""Analyze this ECG image carefully. You are a cardiologist analyzing an electrocardiogram (ECG).
-    
-    Extract and report all visible parameters, including but not limited to:
-    1. Heart rate
-    2. PR interval
-    3. QRS duration
-    4. QT/QTc interval
-    5. P wave morphology
-    6. ST segment changes
-    7. T wave morphology
-    8. Rhythm classification
-    9. Specific patterns (if any)
-    
-    Report exact numerical values where visible. Format your response using HTML list elements for better readability.
-    
-    If certain measurements aren't visible in the image, indicate that they cannot be determined.
-    
-    If you notice any abnormalities or concerning patterns, highlight them clearly but avoid making definitive diagnoses.
-    
-    Format your response like this:
-    <h3>ECG Report</h3>
-    <ul>
-      <li><strong>Date and Time:</strong> {timestamp}</li>
-      <li><strong>Heart Rate:</strong> [value] bpm</li>
-      <li><strong>PR Interval:</strong> [value] ms</li>
-      <!-- Other measurements -->
-    </ul>
-    
-    <h3>Additional Observations</h3>
-    <ul>
-      <li>Observation 1</li>
-      <li>Observation 2</li>
-    </ul>
-    
-    <h3>Conclusion</h3>
-    <p>[Your conclusion text]</p>
-    
-    Important formatting instructions:
-    - Use HTML elements for structure (<ul>, <li>, <strong>, <h3>, etc.)
-    - Do not use asterisks (**) for emphasis - use proper HTML formatting instead
-    - For any urgent findings, use <span style="color:red"> to highlight them
-    """
-    
+        try:
+            image = Image.open(image)
+        except Exception as e:
+             return f"<strong style='color:red'>Error opening image: {str(e)}</strong>"
+
+    # --- Gemini Specific Part ---
     try:
-        vision_completion = client.chat.completions.create(
-            messages=[
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": vision_prompt},
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": f"data:image/jpeg;base64,{base64_image}",
-                            },
-                        },
-                    ],
-                }
-            ],
+        # Get Gemini API Key
+        gemini_api_key = os.environ.get("GEMINI_API_KEY")
+        if not gemini_api_key:
+            return "<strong style='color:red'>GEMINI_API_KEY environment variable not set.</strong>"
+
+        # Initialize Gemini client
+        # Using configure for simplicity if preferred, or Client() directly
+        # genai.configure(api_key=gemini_api_key)
+        # model = genai.GenerativeModel(model_name=vision_model)
+        gemini_client = genai.Client(api_key=gemini_api_key)
+
+
+        # Convert PIL image to bytes (JPEG format)
+        buffered = io.BytesIO()
+        # Ensure image is in RGB format if it's RGBA or P which might cause issues
+        if image.mode in ('RGBA', 'P'):
+            image = image.convert('RGB')
+        image.save(buffered, format="JPEG")
+        image_bytes = buffered.getvalue()
+
+        # Get current timestamp (computer time)
+        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+        # Create prompt for Gemini
+        vision_prompt = f"""Analyze this ECG image carefully. You are a cardiologist analyzing an electrocardiogram (ECG).
+
+        Extract and report all visible parameters, including but not limited to:
+        1. Heart rate
+        2. PR interval
+        3. QRS duration
+        4. QT/QTc interval
+        5. P wave morphology
+        6. ST segment changes
+        7. T wave morphology
+        8. Rhythm classification
+        9. Specific patterns (if any)
+
+        Report exact numerical values where visible. Format your response using HTML list elements for better readability.
+
+        If certain measurements aren't visible in the image, indicate that they cannot be determined.
+
+        If you notice any abnormalities or concerning patterns, highlight them clearly but avoid making definitive diagnoses.
+
+        Format your response like this:
+        <h3>ECG Report</h3>
+        <ul>
+          <li><strong>Date and Time:</strong> {timestamp}</li>
+          <li><strong>Heart Rate:</strong> [value] bpm</li>
+          <li><strong>PR Interval:</strong> [value] ms</li>
+          <!-- Other measurements -->
+        </ul>
+
+        <h3>Additional Observations</h3>
+        <ul>
+          <li>Observation 1</li>
+          <li>Observation 2</li>
+        </ul>
+
+        <h3>Conclusion</h3>
+        <p>[Your conclusion text]</p>
+
+        Important formatting instructions:
+        - Use HTML elements for structure (<ul>, <li>, <strong>, <h3>, etc.)
+        - Do not use asterisks (**) for emphasis - use proper HTML formatting instead
+        - For any urgent findings, use <span style="color:red"> to highlight them
+        """
+
+        # Generate content using Gemini
+        response = gemini_client.models.generate_content(
             model=vision_model,
-            temperature=0.2,  # Lower temperature for more factual responses
-            max_completion_tokens=1024,
+            contents=[vision_prompt,
+                      types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")]
+            # Add safety_settings if needed:
+            # safety_settings=[
+            #     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
+            #     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
+            #     {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
+            #     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+            # ]
         )
-        
-        ecg_analysis = vision_completion.choices[0].message.content
-        
+
+        # Handle potential blocks or errors in response
+        if not response.candidates:
+             # Check finish_reason if available, e.g., safety settings block
+             finish_reason = getattr(response, 'prompt_feedback', None)
+             if finish_reason and getattr(finish_reason, 'block_reason', None):
+                 return f"<strong style='color:red'>Analysis blocked due to: {finish_reason.block_reason}. Check safety settings or content.</strong>"
+             else:
+                 return "<strong style='color:red'>No content generated by the model. The request might have been blocked or failed.</strong>"
+
+
+        # Assuming the first candidate has the content
+        ecg_analysis = response.text # Or response.candidates[0].content.parts[0].text
+
+        # --- End of Gemini Specific Part ---
+
         # Process the response to convert any remaining ** to HTML tags
         ecg_analysis = re.sub(r'\*\*([^*]+)\*\*', r'<strong>\1</strong>', ecg_analysis)
-        
+
         # Make sure all headers are properly formatted (with complete pattern)
         ecg_analysis = re.sub(r'^(#+)\s+(.+)$', r'<h3>\2</h3>', ecg_analysis, flags=re.MULTILINE)
-        
+
         # If there's no HTML formatting at all, wrap in basic structure
         if not re.search(r'<[^>]+>', ecg_analysis):
             lines = ecg_analysis.split('\n')
             formatted_html = "<h3>ECG Report</h3>\n<ul>\n"
-            
+
             for line in lines:
                 if line.strip():
                     # Try to identify key-value pairs
-                    match = re.match(r'^([^:]+):\s*(.+)$', line)
+                    match = re.match(r'^([^:]+):\s*(.+)$', line.strip())
                     if match:
                         key, value = match.groups()
-                        formatted_html += f"  <li><strong>{key}:</strong> {value}</li>\n"
+                        formatted_html += f"  <li><strong>{key.strip()}:</strong> {value.strip()}</li>\n"
                     else:
-                        formatted_html += f"  <li>{line}</li>\n"
-            
+                        formatted_html += f"  <li>{line.strip()}</li>\n"
+
             formatted_html += "</ul>"
             ecg_analysis = formatted_html
-        
+
         return ecg_analysis
-    
+
     except Exception as e:
-        return f"<strong style='color:red'>Error analyzing ECG image:</strong> {str(e)}"
-
-# Generate medical assessment based on ECG readings and patient history
-def generate_assessment(ecg_analysis, patient_history=None, chat_model="llama-3.3-70b-versatile"):
-    # Fixed model - always use llama-3.3-70b-versatile
-    chat_model = "llama-3.3-70b-versatile"
-    
-    if not ecg_analysis or ecg_analysis.startswith("<strong style='color:red'>Error"):
+        # Catch potential API errors or other issues
+        import traceback
+        print(traceback.format_exc()) # Print full traceback to console for debugging
+        return f"<strong style='color:red'>Error analyzing ECG image with Gemini:</strong> {str(e)}"
+
+
+# Generate medical assessment based on ECG readings and patient history (Unchanged - uses Groq)
+def generate_assessment(ecg_analysis, patient_history=None, chat_model="llama-3.1-70b-versatile"): # Adjusted default model slightly if needed
+    # Fixed model
+    chat_model = "llama-3.1-70b-versatile" # Or keep llama-3.3 if available/preferred
+
+    if not ecg_analysis or ecg_analysis.startswith("<strong style='color:red'>"):
         return "<strong style='color:red'>Please analyze an ECG image first.</strong>"
-    
+
     # Get current timestamp
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    
+
     # Construct prompt based on available information
     if patient_history and patient_history.strip():
         prompt = f"""You are a highly trained cardiologist assistant. Based on the ECG analysis below and the patient's history, provide a comprehensive assessment of the patient's cardiac status. Indicate clearly if there are any concerning findings that require immediate medical attention.
@@ -283,9 +330,9 @@ Important formatting instructions:
 - Do not use asterisks (**) for emphasis - use proper HTML formatting instead
 - For any urgent findings, use <span style="color:red"> to highlight them
 """
-    
+
     try:
-        assessment_completion = client.chat.completions.create(
+        assessment_completion = groq_client.chat.completions.create(
             messages=[
                 {
                     "role": "system",
@@ -297,265 +344,285 @@ Important formatting instructions:
                 }
             ],
             model=chat_model,
-            temperature=0.2,  # Lower temperature for more factual responses
-            max_completion_tokens=2048,
+            temperature=0.2,
+            max_tokens=2048, # Note: Groq uses max_tokens
         )
-        
+
         assessment_text = assessment_completion.choices[0].message.content
-        
+
         # Process the response to convert any remaining ** to HTML tags
         assessment_text = re.sub(r'\*\*([^*]+)\*\*', r'<strong>\1</strong>', assessment_text)
-        
+
         # Make sure all headers are properly formatted (with complete pattern)
         assessment_text = re.sub(r'^(#+)\s+(.+)$', r'<h3>\2</h3>', assessment_text, flags=re.MULTILINE)
-        
-        # If there's no HTML formatting at all, wrap in basic structure
+
+        # If there's no HTML formatting at all, wrap in basic structure (Improved fallback)
         if not re.search(r'<[^>]+>', assessment_text):
             sections = [
-                "Summary of Findings", 
-                "Key Abnormalities", 
-                "Potential Clinical Implications", 
-                "Recommendation", 
+                "Summary of Findings",
+                "Key Abnormalities",
+                "Potential Clinical Implications",
+                "Recommendation",
                 "Differential Considerations"
             ]
-            
-            # Split by sections and format
             formatted_html = ""
-            current_section = "Summary of Findings"  # Default
-            content_lines = []
-            
             lines = assessment_text.split('\n')
-            
+            current_section_content = []
+            current_section_title = ""
+
+            def format_section(title, content_lines):
+                html = f"<h3>{title}</h3>\n<ul>\n"
+                for line in content_lines:
+                    if line.strip():
+                       html += f"  <li>{line.strip()}</li>\n"
+                html += "</ul>\n"
+                return html
+
             for line in lines:
-                # Check if this line is a section header
-                is_section = False
-                for section in sections:
-                    if section.lower() in line.lower():
-                        # If we have collected content for a previous section, output it
-                        if content_lines:
-                            formatted_html += f"<h3>{current_section}</h3>\n<ul>\n"
-                            for content_line in content_lines:
-                                formatted_html += f"  <li>{content_line}</li>\n"
-                            formatted_html += "</ul>\n"
-                            content_lines = []
-                        
-                        current_section = section
-                        is_section = True
-                        break
-                
-                if not is_section and line.strip():
-                    content_lines.append(line.strip())
-            
-            # Don't forget the last section
-            if content_lines:
-                formatted_html += f"<h3>{current_section}</h3>\n<ul>\n"
-                for content_line in content_lines:
-                    formatted_html += f"  <li>{content_line}</li>\n"
-                formatted_html += "</ul>\n"
-            
-            assessment_text = formatted_html if formatted_html else assessment_text
-        
+                 line_stripped = line.strip()
+                 if not line_stripped:
+                     continue
+
+                 is_header = False
+                 for section_title in sections:
+                     # Check if line looks like a header (case-insensitive)
+                     if section_title.lower() in line_stripped.lower() and len(line_stripped) < len(section_title) + 10:
+                         if current_section_title and current_section_content:
+                             formatted_html += format_section(current_section_title, current_section_content)
+                         current_section_title = section_title # Use the canonical title
+                         current_section_content = []
+                         is_header = True
+                         break
+
+                 if not is_header:
+                     # If no section started yet, assume it's summary
+                     if not current_section_title:
+                         current_section_title = "Summary of Findings"
+                     current_section_content.append(line_stripped)
+
+            # Add the last section
+            if current_section_title and current_section_content:
+                 formatted_html += format_section(current_section_title, current_section_content)
+
+            assessment_text = formatted_html if formatted_html else f"<p>{assessment_text.replace('\n', '<br>')}</p>" # Basic wrap if structure fails
+
         return assessment_text
-    
+
     except Exception as e:
+        import traceback
+        print(traceback.format_exc()) # Print full traceback
         return f"<strong style='color:red'>Error generating assessment:</strong> {str(e)}"
 
-# Doctor's chat interaction with the model about the patient
-def doctor_chat(message, chat_history, ecg_analysis, patient_history, assessment, chat_model="llama-3.3-70b-versatile"):
-    # Fixed model - always use llama-3.3-70b-versatile
-    chat_model = "llama-3.3-70b-versatile"
-    
-    if not ecg_analysis or ecg_analysis.startswith("<strong style='color:red'>Error"):
-        return "Please analyze an ECG image first before starting a chat.", chat_history
-    
+
+# Doctor's chat interaction with the model about the patient (Unchanged - uses Groq)
+def doctor_chat(message, chat_history, ecg_analysis, patient_history, assessment, chat_model="llama-3.1-70b-versatile"): # Adjusted default model slightly if needed
+    # Fixed model
+    chat_model = "llama-3.1-70b-versatile" # Or keep llama-3.3 if available/preferred
+
+    # Check if ECG analysis exists and is not an error message
+    if not ecg_analysis or ecg_analysis.startswith("<strong style='color:red'>"):
+        # Prepend error message to history instead of returning it directly
+        chat_history.append((message, "<strong style='color:red'>Cannot start chat. Please analyze a valid ECG image first.</strong>"))
+        return "", chat_history # Clear input, update history
+
     if not message.strip():
-        return "", chat_history
-    
+        return "", chat_history # Ignore empty messages
+
     # Get current timestamp
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    
+
     # Prepare chat context
-    context = f"""ECG ANALYSIS:
+    context = f"""CURRENT TIMESTAMP: {timestamp}
+
+=== BEGIN CONTEXT ===
+PATIENT HISTORY:
+{patient_history if patient_history and patient_history.strip() else "No patient history provided."}
+
+ECG ANALYSIS:
 {ecg_analysis}
 
 MEDICAL ASSESSMENT:
-{assessment}
+{assessment if assessment and not assessment.startswith("<strong style='color:red'>") else "No assessment generated yet or assessment failed."}
+=== END CONTEXT ===
 
-TIMESTAMP: {timestamp}
+Based *only* on the context provided above, answer the doctor's questions concisely. If the information is not in the context, state that.
 """
-    
-    if patient_history and patient_history.strip():
-        context += f"""PATIENT HISTORY:
-{patient_history}
-"""
-    
+
     # Construct full chat history for context
     messages = [
         {
             "role": "system",
-            "content": f"You are a medical AI assistant specialized in cardiology. You are helping a doctor interpret ECG results and patient data. Answer the doctor's questions based on the following information:\n\n{context}"
+            "content": f"You are a medical AI assistant specialized in cardiology, conversing with a doctor. Your knowledge is strictly limited to the patient information provided in the context below. Do not invent information or access external knowledge.\n\n{context}"
         }
     ]
-    
-    # Add chat history to the context (limited to last 10 exchanges to avoid token limits)
-    for entry in chat_history[-10:]:
-        messages.append({"role": "user", "content": entry[0]})
-        messages.append({"role": "assistant", "content": entry[1]})
-    
+
+    # Add chat history to the context (limited to last 5 exchanges to avoid token limits)
+    for user_msg, assistant_msg in chat_history[-5:]:
+        messages.append({"role": "user", "content": user_msg})
+        # Avoid adding error messages from assistant back into context
+        if not assistant_msg.startswith("<strong style='color:red'>"):
+             messages.append({"role": "assistant", "content": assistant_msg})
+
     # Add the current message
     messages.append({"role": "user", "content": message})
-    
+
     try:
-        chat_completion = client.chat.completions.create(
+        chat_completion = groq_client.chat.completions.create(
             messages=messages,
             model=chat_model,
             temperature=0.3,
-            max_completion_tokens=1024,
+            max_tokens=1024, # Note: Groq uses max_tokens
         )
-        
+
         response = chat_completion.choices[0].message.content
-        
+
         # Process any remaining asterisks to HTML tags in the response
         response = re.sub(r'\*\*([^*]+)\*\*', r'<strong>\1</strong>', response)
-        
+
         chat_history.append((message, response))
-        return "", chat_history
-    
+        return "", chat_history # Clear input message box
     except Exception as e:
+        import traceback
+        print(traceback.format_exc()) # Print full traceback
         error_message = f"<strong style='color:red'>Error in chat:</strong> {str(e)}"
         chat_history.append((message, error_message))
-        return "", chat_history
+        return "", chat_history # Clear input message box
 
 # Create Gradio interface
 with gr.Blocks(title="Cardiac ECG Analysis System", theme=gr.themes.Soft()) as app:
-    # Session state to store data
-    ecg_analysis_state = gr.State("")
-    
+    # Session state to store data (Removed - use outputs directly or manage state differently if needed)
+    # ecg_analysis_state = gr.State("") # Not strictly necessary if passing outputs directly
+
     gr.Markdown("# 🫀 Cardiac ECG Analysis System")
     gr.Markdown("Upload an ECG image and optional patient history to get an automated analysis and assessment.")
-    
+
     with gr.Tabs():
         with gr.TabItem("💻 Main Interface"):
             with gr.Row():
                 with gr.Column(scale=1):
                     # Input components
-                    with gr.Group():  # Using Group instead of Box
+                    with gr.Group():
                         gr.Markdown("### 📊 ECG Image")
                         ecg_image = gr.Image(type="pil", label="Upload ECG Image")
-                        # Display fixed model info
-                        gr.Markdown("**Vision Model:** llama-3.2-90b-vision-preview")
+                        # Display UPDATED fixed model info for vision
+                        gr.Markdown("**Vision Model:** gemini-2.0-flash-exp")
                         analyze_button = gr.Button("Analyze ECG Image", variant="primary")
-                    
-                    with gr.Group():  # Using Group instead of Box
+
+                    with gr.Group():
                         gr.Markdown("### 📋 Patient Information")
                         patient_history_text = gr.Textbox(
-                            lines=8, 
-                            label="Patient History (Manual Entry)", 
-                            placeholder="Enter patient's medical history, age, sex, symptoms, medications, etc."
+                            lines=8,
+                            label="Patient History (Manual Entry or Loaded from File)",
+                            placeholder="Enter patient's medical history, age, sex, symptoms, medications, etc. OR upload a file below and click Load."
                         )
                         patient_history_file = gr.File(
                             label="Upload Patient History File (Optional, .txt, .csv, or .xlsx)",
                             file_types=[".txt", ".csv", ".xlsx", ".xls"]
                         )
                         load_history_button = gr.Button("Load Patient History from File")
-                    
-                    with gr.Group():  # Using Group instead of Box
+
+                    with gr.Group():
                         gr.Markdown("### 🧠 Assessment Settings")
-                        # Display fixed model info
-                        gr.Markdown("**Chat Model:** llama-3.3-70b-versatile")
+                        # Display fixed model info for chat/assessment
+                        # Make sure this model name matches the one used in generate_assessment/doctor_chat
+                        gr.Markdown("**Chat/Assessment Model:** llama-3.1-70b-versatile")
                         assess_button = gr.Button("Generate Assessment", variant="primary")
-                
+
                 with gr.Column(scale=1):
                     # Output components
-                    with gr.Group():  # Using Group instead of Box
+                    with gr.Group():
                         gr.Markdown("### 📈 ECG Analysis Results")
                         ecg_analysis_output = gr.HTML(label="ECG Analysis", elem_id="ecg-analysis")
-                    
-                    with gr.Group():  # Using Group instead of Box
+
+                    with gr.Group():
                         gr.Markdown("### 📝 Medical Assessment")
                         assessment_output = gr.HTML(label="Assessment", elem_id="assessment-output")
-            
+
             gr.Markdown("## 👨‍⚕️ Doctor's Consultation")
             gr.Markdown("Ask follow-up questions about the patient's ECG results and medical condition.")
-            
-            with gr.Group():  # Using Group instead of Box
-                chatbot = gr.Chatbot(label="Consultation", height=400)
+
+            with gr.Group():
+                chatbot = gr.Chatbot(label="Consultation", height=400, bubble_full_width=False)
                 with gr.Row():
                     message = gr.Textbox(
-                        lines=2, 
+                        lines=2,
                         label="Doctor's Question",
                         placeholder="Ask a question about this patient's cardiac status...",
-                        scale=4
+                        scale=4,
+                        show_label=False,
+                        container=False, # Makes textbox slimmer vertically
                     )
                     chat_button = gr.Button("Send", scale=1, variant="primary")
-        
+
         with gr.TabItem("ℹ️ Instructions"):
             gr.Markdown("""
             ## How to Use This Application
-            
+
             ### Step 1: Upload and Analyze ECG
-            1. Upload an ECG image using the file uploader in the Main Interface tab
-            2. Click "Analyze ECG Image" to extract readings from the image
-            
+            1.  Upload an ECG image using the file uploader in the "Main Interface" tab.
+            2.  Click **Analyze ECG Image**. The system will use Gemini Vision to interpret the image.
+            3.  Wait for the "ECG Analysis Results" to appear.
+
             ### Step 2: Add Patient Information (Optional)
-            - Enter patient history directly in the text box, OR
-            - Upload a patient history file (.txt, .csv, or .xlsx) and click "Load Patient History from File"
-            
+            *   Enter patient history directly into the "Patient History" text box, OR
+            *   Upload a patient history file (.txt, .csv, or .xlsx) and click **Load Patient History from File**. The content will populate the text box.
+
             ### Step 3: Generate Assessment
-            - Click "Generate Assessment" to get an AI-assisted interpretation
-            
+            *   Once the ECG analysis is complete, click **Generate Assessment**. The system will use a Llama model (via Groq) combining the ECG analysis and patient history (if provided).
+            *   Wait for the "Medical Assessment" to appear.
+
             ### Step 4: Consultation
-            - Use the chatbot interface to ask follow-up questions
-            - The AI will consider the ECG analysis, patient history, and previous assessment in its responses
-            
+            *   Use the chat interface at the bottom to ask follow-up questions.
+            *   Type your question and click **Send** or press Enter.
+            *   The AI (Llama via Groq) will consider the ECG analysis, patient history, and the generated assessment in its responses.
+
             ### Important Notes
-            - This tool is designed to assist healthcare professionals, not replace clinical judgment
-            - Always validate AI-generated medical interpretations with proper medical expertise
-            - Patient data privacy should be maintained according to relevant regulations
+            *   **API Keys:** Ensure `GEMINI_API_KEY` and `GROQ_API_KEY` environment variables are set correctly before running the application.
+            *   **Purpose:** This tool is designed to assist healthcare professionals and is NOT a substitute for professional clinical judgment or diagnosis.
+            *   **Validation:** Always validate AI-generated medical interpretations with qualified medical expertise.
+            *   **Privacy:** Ensure compliance with patient data privacy regulations (e.g., HIPAA) when using this tool. Do not upload identifiable patient information if not permitted.
             """)
-    
-    # Set up event handlers
+
+    # --- Event Handlers ---
+
+    # Chain analysis button click to update output
     analyze_button.click(
-        analyze_ecg_image, 
-        inputs=[ecg_image], 
+        analyze_ecg_image,
+        inputs=[ecg_image],
         outputs=ecg_analysis_output
-    ).then(
-        lambda x: x,  # Pass through function to update state
-        inputs=ecg_analysis_output,
-        outputs=ecg_analysis_state
     )
-    
-    def process_and_update_history(file):
-        if file is None:
-            return "No file uploaded."
-        processed_text = process_patient_history(file)
-        return processed_text
-    
+
+    # Load history file content into the textbox
     load_history_button.click(
-        process_and_update_history,
+        process_patient_history,
         inputs=[patient_history_file],
         outputs=[patient_history_text]
     )
-    
+
+    # Generate assessment based on ECG output and history textbox
     assess_button.click(
-        generate_assessment, 
-        inputs=[ecg_analysis_output, patient_history_text], 
+        generate_assessment,
+        inputs=[ecg_analysis_output, patient_history_text],
         outputs=assessment_output
     )
-    
+
+    # Handle chat interactions
     chat_button.click(
-        doctor_chat, 
-        inputs=[message, chatbot, ecg_analysis_output, patient_history_text, assessment_output], 
-        outputs=[message, chatbot]
+        doctor_chat,
+        inputs=[message, chatbot, ecg_analysis_output, patient_history_text, assessment_output],
+        outputs=[message, chatbot] # Clear message input, update chatbot history
     )
-    
-    # Also trigger chat on Enter key
+
+    # Also trigger chat on Enter key press in the message textbox
     message.submit(
         doctor_chat,
         inputs=[message, chatbot, ecg_analysis_output, patient_history_text, assessment_output],
-        outputs=[message, chatbot]
+        outputs=[message, chatbot] # Clear message input, update chatbot history
     )
 
+
 # Launch the app
 if __name__ == "__main__":
+    # For development, you might want debug=True
+    # app.launch(debug=True)
     app.launch()