Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-# import base64 # No longer needed for Gemini vision part
 import gradio as gr
 import pandas as pd
 from groq import Groq
@@ -9,45 +8,48 @@ import datetime
 import re
 import google.generativeai as genai # Added for Gemini
 from google.generativeai import types # Added for Gemini
+import traceback # For detailed error logging
 
 # Initialize Groq client (for chat/assessment)
-groq_client = Groq(
-    api_key=os.environ.get("GROQ_API_KEY")
-)
+# Ensure GROQ_API_KEY is set in your environment variables
+try:
+    groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+except Exception as e:
+    print(f"Error initializing Groq client: {e}")
+    groq_client = None # Set to None to handle initialization errors gracefully
 
 # 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 (Unchanged)
+# Process patient history file
 def process_patient_history(file):
     if file is None:
         return ""
 
     try:
         # Check file extension
-        file_ext = os.path.splitext(file.name)[1].lower()
+        file_path = file.name # Gradio File object usually has a .name attribute with the path
+        file_ext = os.path.splitext(file_path)[1].lower()
 
         if file_ext == '.txt':
             # Read text file
-            # 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:
+            with open(file_path, '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)
+                df = pd.read_csv(file_path)
             else:
-                df = pd.read_excel(file.name)
+                # Specify engine for newer xlsx files if needed
+                try:
+                    df = pd.read_excel(file_path)
+                except ImportError:
+                     return "Error: `openpyxl` needed for .xlsx files. Install with `pip install openpyxl`"
+                except Exception as e_excel:
+                     return f"Error reading Excel file: {e_excel}"
+
 
             # Convert dataframe to formatted string
             formatted_data = "PATIENT INFORMATION:\n\n"
@@ -56,16 +58,22 @@ def process_patient_history(file):
                 for column in df.columns:
                     # Handle potential missing values gracefully
                     value = df.iloc[0].get(column, 'N/A')
-                    formatted_data += f"{column}: {value}\n"
+                    # Convert value to string to avoid potential type issues
+                    formatted_data += f"{column}: {str(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."
+            return f"Unsupported file format ({file_ext}). Please upload a .txt, .csv, or .xlsx file."
 
+    except AttributeError:
+         return "Error: Could not get file path from Gradio File object. Ensure a file was uploaded."
+    except FileNotFoundError:
+         return f"Error: File not found at path: {file_path}"
     except Exception as e:
+        print(f"Error processing patient history file:\n{traceback.format_exc()}")
         return f"Error processing patient history file: {str(e)}"
 
 
@@ -80,11 +88,22 @@ def analyze_ecg_image(image, vision_model="gemini-2.0-flash-exp"):
     # Ensure image is PIL Image
     if not isinstance(image, Image.Image):
         try:
-            image = Image.open(image)
+            # If 'image' is a path (from older Gradio versions or specific setups)
+            if isinstance(image, str) and os.path.exists(image):
+                 image = Image.open(image)
+            # If 'image' is file-like object from Gradio upload
+            elif hasattr(image, 'name'):
+                 image = Image.open(image.name)
+            else:
+                 # Assume it might be bytes or needs loading differently
+                 # This part might need adjustment depending on how Gradio passes the image
+                 return f"<strong style='color:red'>Unrecognized image input format: {type(image)}</strong>"
         except Exception as e:
+             print(f"Error opening image:\n{traceback.format_exc()}")
              return f"<strong style='color:red'>Error opening image: {str(e)}</strong>"
 
     # --- Gemini Specific Part ---
+    gemini_client = None # Initialize to None
     try:
         # Get Gemini API Key
         gemini_api_key = os.environ.get("GEMINI_API_KEY")
@@ -92,18 +111,22 @@ def analyze_ecg_image(image, vision_model="gemini-2.0-flash-exp"):
             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)
+        # Use Client() for explicit initialization per request (safer for concurrent use)
         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
+        img_format = "JPEG"
         if image.mode in ('RGBA', 'P'):
             image = image.convert('RGB')
-        image.save(buffered, format="JPEG")
+        # Handle potential transparency issues for PNG -> JPEG conversion
+        elif image.mode == 'LA':
+             image = image.convert('RGB') # Convert Luminance Alpha to RGB
+        # Check format if needed
+        # if image.format == 'PNG': img_format = 'PNG' # Gemini might prefer JPEG? Test this.
+
+        image.save(buffered, format=img_format)
         image_bytes = buffered.getvalue()
 
         # Get current timestamp (computer time)
@@ -113,52 +136,61 @@ def analyze_ecg_image(image, vision_model="gemini-2.0-flash-exp"):
         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)
+        1. Heart rate (bpm)
+        2. PR interval (ms)
+        3. QRS duration (ms)
+        4. QT/QTc interval (ms)
+        5. P wave morphology (e.g., upright in lead II, biphasic, absent)
+        6. ST segment changes (e.g., elevation, depression, location)
+        7. T wave morphology (e.g., upright, inverted, peaked, flattened)
+        8. Rhythm classification (e.g., Sinus Rhythm, Atrial Fibrillation, etc.)
+        9. Axis deviation (if determinable)
+        10. Specific patterns (e.g., Bundle Branch Block, LVH criteria, WPW)
 
-        Report exact numerical values where visible. Format your response using HTML list elements for better readability.
+        Report exact numerical values where visible. If a range is shown, report the range. 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 certain measurements aren't clearly visible or determinable from the provided image quality, explicitly state that they cannot be determined.
 
-        If you notice any abnormalities or concerning patterns, highlight them clearly but avoid making definitive diagnoses.
+        If you notice any abnormalities or concerning patterns, highlight them clearly but avoid making definitive diagnoses. State observations neutrally.
 
-        Format your response like this:
+        Format your response strictly like this example:
         <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 -->
+          <li><strong>Analysis Time:</strong> {timestamp}</li>
+          <li><strong>Heart Rate:</strong> [value] bpm (or 'Not determinable')</li>
+          <li><strong>Rhythm:</strong> [description] (or 'Not determinable')</li>
+          <li><strong>PR Interval:</strong> [value] ms (or 'Not determinable')</li>
+          <li><strong>QRS Duration:</strong> [value] ms (or 'Not determinable')</li>
+          <li><strong>QT/QTc Interval:</strong> [value]/[value] ms (or 'Not determinable')</li>
+          <li><strong>Axis:</strong> [description] (or 'Not determinable')</li>
+          <li><strong>P Waves:</strong> [description]</li>
+          <li><strong>ST Segment:</strong> [description]</li>
+          <li><strong>T Waves:</strong> [description]</li>
+          <!-- Add other relevant findings as list items -->
         </ul>
 
-        <h3>Additional Observations</h3>
+        <h3>Key Observations / Potential Abnormalities</h3>
         <ul>
-          <li>Observation 1</li>
-          <li>Observation 2</li>
+          <li>[Observation 1, e.g., Possible ST elevation in leads V1-V3]</li>
+          <li>[Observation 2, e.g., Inverted T waves in lead III]</li>
+          <!-- List significant observations, use <span style="color:red"> for critical findings -->
         </ul>
 
-        <h3>Conclusion</h3>
-        <p>[Your conclusion text]</p>
+        <h3>Impression</h3>
+        <p>[Provide a brief summary impression based *only* on the visible findings, e.g., "Sinus tachycardia with possible anterior ST changes." Avoid definitive diagnosis.]</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
+        - Use EXACTLY the HTML structure shown (<h3>, <ul>, <li>, <strong>, <p>).
+        - Do NOT use markdown like asterisks (**) or hashtags (#). Use HTML tags for formatting.
+        - For potentially urgent findings in the 'Key Observations' list, wrap the description in <span style="color:red">text</span>.
         """
 
         # Generate content using Gemini
         response = gemini_client.models.generate_content(
             model=vision_model,
             contents=[vision_prompt,
-                      types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")]
-            # Add safety_settings if needed:
+                      types.Part.from_bytes(data=image_bytes, mime_type=f"image/{img_format.lower()}")]
+            # Add safety_settings if needed (BLOCK_NONE allows more content but use carefully):
             # safety_settings=[
             #     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
             #     {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
@@ -169,174 +201,141 @@ def analyze_ecg_image(image, vision_model="gemini-2.0-flash-exp"):
 
         # 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>"
+             feedback = getattr(response, 'prompt_feedback', None)
+             block_reason = getattr(feedback, 'block_reason', None) if feedback else None
+             if block_reason:
+                 return f"<strong style='color:red'>Analysis blocked by safety filter: {block_reason}. Consider adjusting safety settings or image content if appropriate.</strong>"
              else:
-                 return "<strong style='color:red'>No content generated by the model. The request might have been blocked or failed.</strong>"
+                 # Log the full response for debugging if possible
+                 print(f"Gemini Response Error: No candidates returned. Full response: {response}")
+                 return "<strong style='color:red'>No content generated by the model. The request might have been blocked, timed out, or failed unexpectedly. Check logs.</strong>"
 
+        # Accessing response text - check structure based on library version
+        try:
+            # Preferred way for recent versions
+            ecg_analysis = response.text
+        except ValueError:
+            # Fallback for potential older structures or errors
+             if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
+                 ecg_analysis = "".join(part.text for part in response.candidates[0].content.parts if hasattr(part, 'text'))
+             else:
+                 print(f"Gemini Response Error: Could not extract text. Response structure: {response}")
+                 return "<strong style='color:red'>Error processing model response. Unexpected format. Check logs.</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)
+        # Basic post-processing (Gemini should follow HTML instructions, but just in case)
+        # Remove potential markdown that might slip through
+        ecg_analysis = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', ecg_analysis)
+        ecg_analysis = re.sub(r'^\s*#+\s+(.*?)\s*$', r'<h3>\1</h3>', ecg_analysis, flags=re.MULTILINE)
+        ecg_analysis = re.sub(r'^\s*[\*-]\s+(.*?)\s*$', r'<li>\1</li>', ecg_analysis, flags=re.MULTILINE) # Convert markdown lists
 
-        # 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.strip())
-                    if match:
-                        key, value = match.groups()
-                        formatted_html += f"  <li><strong>{key.strip()}:</strong> {value.strip()}</li>\n"
-                    else:
-                        formatted_html += f"  <li>{line.strip()}</li>\n"
-
-            formatted_html += "</ul>"
-            ecg_analysis = formatted_html
+        # Simple check if the response looks somewhat like the requested HTML structure
+        if not ("<h3>" in ecg_analysis and "<ul>" in ecg_analysis):
+             print(f"Warning: Gemini response might not be in the expected HTML format:\n{ecg_analysis[:500]}...")
+             # Optionally wrap in basic tags if completely unformatted
+             # ecg_analysis = f"<h3>ECG Analysis (Raw Output)</h3><p>{ecg_analysis.replace('\n', '<br>')}</p>"
 
         return ecg_analysis
 
     except Exception as e:
         # 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)}"
+        print(f"Error during Gemini ECG analysis:\n{traceback.format_exc()}")
+        error_type = type(e).__name__
+        return f"<strong style='color:red'>Error analyzing ECG image with Gemini ({error_type}):</strong> {str(e)}"
+    finally:
+        # Clean up client resource if necessary, though Client() might not require explicit closing
+        pass
 
 
-# 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
+# Generate medical assessment based on ECG readings and patient history (Uses Groq)
+def generate_assessment(ecg_analysis, patient_history=None, chat_model="llama3-70b-8192"): # Using a common Groq model
+    # Check Groq client initialization
+    if groq_client is None:
+        return "<strong style='color:red'>Groq client not initialized. Check API Key and startup logs.</strong>"
+
+    # Use a known available and capable Groq model
+    chat_model = "llama3-70b-8192" # Or "mixtral-8x7b-32768" or another available one
 
     if not ecg_analysis or ecg_analysis.startswith("<strong style='color:red'>"):
-        return "<strong style='color:red'>Please analyze an ECG image first.</strong>"
+        return "<strong style='color:red'>Cannot generate assessment. Please analyze a valid 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.
-
-ECG ANALYSIS:
-{ecg_analysis}
-
-PATIENT HISTORY:
-{patient_history}
-
-TIMESTAMP: {timestamp}
-
-Format your assessment using HTML elements for readability:
-
-<h3>Summary of Findings</h3>
-<ul>
-  <li>Finding 1</li>
-  <li>Finding 2</li>
-  <!-- etc. -->
-</ul>
-
-<h3>Key Abnormalities</h3>
-<ul>
-  <li>Abnormality 1</li>
-  <li>Abnormality 2 (if any)</li>
-  <!-- etc. -->
-</ul>
-
-<h3>Potential Clinical Implications</h3>
-<ul>
-  <li>Implication 1</li>
-  <li>Implication 2</li>
-  <!-- etc. -->
-</ul>
+    # Clean up potential HTML issues in the input ECG analysis for the prompt
+    # This helps prevent confusing the assessment model
+    clean_ecg_analysis = re.sub('<[^>]+>', '', ecg_analysis) # Strip HTML tags for the prompt context
 
-<h3>Recommendation</h3>
-<ul>
-  <li>Recommendation 1 (include urgency level)</li>
-  <li>Recommendation 2 (if any)</li>
-  <!-- etc. -->
-</ul>
-
-<h3>Differential Considerations</h3>
-<ul>
-  <li>Differential 1</li>
-  <li>Differential 2</li>
-  <!-- etc. -->
-</ul>
+    # Construct prompt based on available information
+    prompt_parts = [
+        "You are a highly trained cardiologist assistant AI. Your task is to synthesize information from an ECG analysis and patient history (if provided) into a clinical assessment for a reviewing physician.",
+        "Focus on integrating the findings and suggesting potential implications and recommendations.",
+        "Format your response strictly using the specified HTML structure.",
+        "\nECG ANALYSIS SUMMARY (Provided):\n" + clean_ecg_analysis, # Use cleaned text
+    ]
 
-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
-"""
+    if patient_history and patient_history.strip():
+         prompt_parts.append("\nPATIENT HISTORY (Provided):\n" + patient_history)
     else:
-        prompt = f"""You are a highly trained cardiologist assistant. Based on the ECG analysis below, provide a comprehensive assessment of the patient's cardiac status. Indicate clearly if there are any concerning findings that require immediate medical attention.
+         prompt_parts.append("\nPATIENT HISTORY: Not provided.")
 
-ECG ANALYSIS:
-{ecg_analysis}
+    prompt_parts.append(f"\nASSESSMENT TIMESTAMP: {timestamp}")
 
-TIMESTAMP: {timestamp}
+    prompt_parts.append("""
+Format your assessment using ONLY the following HTML structure:
 
-Format your assessment using HTML elements for readability:
-
-<h3>Summary of Findings</h3>
+<h3>Summary of Integrated Findings</h3>
 <ul>
-  <li>Finding 1</li>
-  <li>Finding 2</li>
+  <li>[Combine key ECG findings with relevant patient history points, e.g., "ECG shows sinus tachycardia in the context of reported palpitations."]</li>
+  <li>[Finding 2]</li>
   <!-- etc. -->
 </ul>
 
-<h3>Key Abnormalities</h3>
+<h3>Key Abnormalities and Concerns</h3>
 <ul>
-  <li>Abnormality 1</li>
-  <li>Abnormality 2 (if any)</li>
-  <!-- etc. -->
+  <li>[List specific significant abnormalities from the ECG, potentially contextualized by history, e.g., "ST elevation in anterior leads concerning for possible ischemia."]</li>
+  <li>[Abnormality 2 (if any)]</li>
+  <!-- Use <span style="color:red"> for urgent/critical concerns -->
 </ul>
 
 <h3>Potential Clinical Implications</h3>
 <ul>
-  <li>Implication 1</li>
-  <li>Implication 2</li>
+  <li>[Suggest possible underlying conditions or risks based on findings, e.g., "Findings could be consistent with acute coronary syndrome."]</li>
+  <li>[Implication 2]</li>
   <!-- etc. -->
 </ul>
 
-<h3>Recommendation</h3>
+<h3>Recommendations for Physician Review</h3>
 <ul>
-  <li>Recommendation 1 (include urgency level)</li>
-  <li>Recommendation 2 (if any)</li>
+  <li>[Suggest next steps or urgency, e.g., "<span style="color:red">Urgent clinical correlation and comparison with previous ECGs recommended.</span>"]</li>
+  <li>[Recommendation 2 (e.g., Consider further cardiac workup like troponins, echocardiogram if clinically indicated.)]</li>
   <!-- etc. -->
 </ul>
 
-<h3>Differential Considerations</h3>
+<h3>Differential Considerations (Optional)</h3>
 <ul>
-  <li>Differential 1</li>
-  <li>Differential 2</li>
+  <li>[List possible alternative explanations for the findings, if applicable.]</li>
+  <li>[Differential 2]</li>
   <!-- etc. -->
 </ul>
 
-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
-"""
+Important Instructions:
+- Adhere strictly to the HTML format (<h3>, <ul>, <li>, <strong>, <p>, <span style="color:red">).
+- Do NOT use markdown formatting (**, #, - ).
+- Base your assessment ONLY on the provided ECG analysis and patient history.
+- Do NOT make a definitive diagnosis. Phrase conclusions as possibilities or suggestions for the physician.
+- If the ECG analysis indicated 'Not determinable' for key parameters, acknowledge this limitation.
+""")
+    prompt = "\n".join(prompt_parts)
 
     try:
         assessment_completion = groq_client.chat.completions.create(
             messages=[
                 {
                     "role": "system",
-                    "content": "You are a medical AI assistant specialized in cardiology. Provide accurate, clinically relevant interpretations of ECG data. If there are concerning findings that might indicate a medical emergency, clearly highlight them. Avoid definitive diagnoses but provide reasoned medical assessments based on the data provided."
+                    "content": "You are a medical AI assistant specialized in cardiology. Generate a structured clinical assessment based on the provided ECG and patient data, formatted in HTML for physician review. Highlight urgent findings appropriately. Avoid definitive diagnoses."
                 },
                 {
                     "role": "user",
@@ -345,83 +344,47 @@ Important formatting instructions:
             ],
             model=chat_model,
             temperature=0.2,
-            max_tokens=2048, # Note: Groq uses max_tokens
+            max_tokens=2048, # Adjust as needed
+            # top_p=0.9, # Optional parameter tuning
+            # stop=None, # Optional stop sequences
         )
 
         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 (Improved fallback)
-        if not re.search(r'<[^>]+>', assessment_text):
-            sections = [
-                "Summary of Findings",
-                "Key Abnormalities",
-                "Potential Clinical Implications",
-                "Recommendation",
-                "Differential Considerations"
-            ]
-            formatted_html = ""
-            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:
-                 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
+        # Basic post-processing (Groq should follow HTML instructions, but good to have fallbacks)
+        assessment_text = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', assessment_text)
+        assessment_text = re.sub(r'^\s*#+\s+(.*?)\s*$', r'<h3>\1</h3>', assessment_text, flags=re.MULTILINE)
+
+        # *** CORRECTED FALLBACK LOGIC ***
+        # Check if the response seems to contain the expected HTML structure
+        if not ("<h3>" in assessment_text and "<ul>" in assessment_text):
+            print(f"Warning: Groq assessment response might not be in the expected HTML format:\n{assessment_text[:500]}...")
+            # Fallback: Wrap the raw text in a paragraph tag with line breaks
+            processed_text = assessment_text.replace('\n', '<br>')
+            assessment_text = f"<h3>Assessment (Raw Output)</h3><p>{processed_text}</p>"
 
         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)}"
+        print(f"Error during Groq assessment generation:\n{traceback.format_exc()}")
+        error_type = type(e).__name__
+        return f"<strong style='color:red'>Error generating assessment with Groq ({error_type}):</strong> {str(e)}"
+
 
+# Doctor's chat interaction with the model about the patient (Uses Groq)
+def doctor_chat(message, chat_history, ecg_analysis, patient_history, assessment, chat_model="llama3-70b-8192"):
+    # Check Groq client initialization
+    if groq_client is None:
+        # Prepend error message to history instead of returning it directly
+        chat_history.append((message, "<strong style='color:red'>Cannot start chat. Groq client not initialized. Check API Key.</strong>"))
+        return "", chat_history # Clear input, update 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
+    chat_model = "llama3-70b-8192" # Consistent with assessment
 
     # 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
+        # Prepend error message to history
         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
 
@@ -431,36 +394,42 @@ def doctor_chat(message, chat_history, ecg_analysis, patient_history, assessment
     # Get current timestamp
     timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
+    # Clean inputs for context
+    clean_ecg = re.sub('<[^>]+>', '', ecg_analysis)
+    clean_assessment = re.sub('<[^>]+>', '', assessment) if assessment and not assessment.startswith("<strong style='color:red'>") else "Assessment not available or failed."
+    clean_history = patient_history if patient_history and patient_history.strip() else "No patient history provided."
+
     # Prepare chat context
     context = f"""CURRENT TIMESTAMP: {timestamp}
 
-=== BEGIN CONTEXT ===
+=== BEGIN PATIENT CONTEXT ===
 PATIENT HISTORY:
-{patient_history if patient_history and patient_history.strip() else "No patient history provided."}
+{clean_history}
 
-ECG ANALYSIS:
-{ecg_analysis}
+ECG ANALYSIS SUMMARY:
+{clean_ecg}
 
-MEDICAL ASSESSMENT:
-{assessment if assessment and not assessment.startswith("<strong style='color:red'>") else "No assessment generated yet or assessment failed."}
-=== END CONTEXT ===
+GENERATED ASSESSMENT SUMMARY:
+{clean_assessment}
+=== END PATIENT CONTEXT ===
 
-Based *only* on the context provided above, answer the doctor's questions concisely. If the information is not in the context, state that.
+Based *only* on the patient context provided above, answer the doctor's questions concisely and professionally. If the information needed to answer is not in the context, explicitly state that. Do not invent information or access external knowledge.
 """
 
     # Construct full chat history for context
     messages = [
         {
             "role": "system",
-            "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}"
+            "content": f"You are a specialized cardiology AI assistant conversing with a doctor. Your knowledge is strictly limited to the patient information provided in the context below. Answer questions based *only* on this context.\n\n{context}"
         }
     ]
 
-    # Add chat history to the context (limited to last 5 exchanges to avoid token limits)
-    for user_msg, assistant_msg in chat_history[-5:]:
+    # Add chat history to the context (limited number of turns to manage token count)
+    history_limit = 5 # Number of past user/assistant pairs
+    for user_msg, assistant_msg in chat_history[-history_limit:]:
         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'>"):
+        if isinstance(assistant_msg, str) and not assistant_msg.startswith("<strong style='color:red'>"):
              messages.append({"role": "assistant", "content": assistant_msg})
 
     # Add the current message
@@ -470,31 +439,30 @@ Based *only* on the context provided above, answer the doctor's questions concis
         chat_completion = groq_client.chat.completions.create(
             messages=messages,
             model=chat_model,
-            temperature=0.3,
-            max_tokens=1024, # Note: Groq uses max_tokens
+            temperature=0.3, # Slightly higher for more conversational answers, but still factual
+            max_tokens=1024, # Adjust as needed
         )
 
         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)
+        # Basic post-processing for the chat response (less critical for HTML here)
+        response = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', response) # Convert bold markdown
+        response = response.replace('\n', '<br>') # Ensure line breaks are rendered in HTML chatbot
 
         chat_history.append((message, response))
         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)}"
+        print(f"Error during Groq chat:\n{traceback.format_exc()}")
+        error_type = type(e).__name__
+        error_message = f"<strong style='color:red'>Error in chat ({error_type}):</strong> {str(e)}"
         chat_history.append((message, error_message))
         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 (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.")
+    gr.Markdown("Upload an ECG image and optional patient history for AI-assisted analysis, assessment, and consultation.")
 
     with gr.Tabs():
         with gr.TabItem("💻 Main Interface"):
@@ -502,10 +470,9 @@ with gr.Blocks(title="Cardiac ECG Analysis System", theme=gr.themes.Soft()) as a
                 with gr.Column(scale=1):
                     # Input components
                     with gr.Group():
-                        gr.Markdown("### 📊 ECG Image")
-                        ecg_image = gr.Image(type="pil", label="Upload ECG Image")
-                        # Display UPDATED fixed model info for vision
-                        gr.Markdown("**Vision Model:** gemini-2.0-flash-exp")
+                        gr.Markdown("### 📊 ECG Image Upload")
+                        ecg_image = gr.Image(type="pil", label="Upload ECG Image", height=300)
+                        gr.Markdown("**Vision Model:** `gemini-2.0-flash-exp` (via Google AI)")
                         analyze_button = gr.Button("Analyze ECG Image", variant="primary")
 
                     with gr.Group():
@@ -513,19 +480,18 @@ with gr.Blocks(title="Cardiac ECG Analysis System", theme=gr.themes.Soft()) as a
                         patient_history_text = gr.Textbox(
                             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."
+                            placeholder="Enter relevant patient details (age, sex, symptoms, meds, conditions) OR upload a file (.txt, .csv, .xlsx) and click Load."
                         )
                         patient_history_file = gr.File(
-                            label="Upload Patient History File (Optional, .txt, .csv, or .xlsx)",
+                            label="Upload Patient History File (Optional)",
                             file_types=[".txt", ".csv", ".xlsx", ".xls"]
                         )
                         load_history_button = gr.Button("Load Patient History from File")
 
                     with gr.Group():
-                        gr.Markdown("### 🧠 Assessment Settings")
-                        # 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")
+                        gr.Markdown("### 🧠 Generate Assessment")
+                        # Ensure this model name matches the one used in generate_assessment/doctor_chat
+                        gr.Markdown("**Assessment/Chat Model:** `llama3-70b-8192` (via Groq)")
                         assess_button = gr.Button("Generate Assessment", variant="primary")
 
                 with gr.Column(scale=1):
@@ -538,91 +504,93 @@ with gr.Blocks(title="Cardiac ECG Analysis System", theme=gr.themes.Soft()) as a
                         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.")
+            gr.Markdown("---") # Separator
+            gr.Markdown("## 👨‍⚕️ Doctor's Consultation Chat")
+            gr.Markdown("Ask follow-up questions based on the analysis and assessment above.")
 
             with gr.Group():
-                chatbot = gr.Chatbot(label="Consultation", height=400, bubble_full_width=False)
+                chatbot = gr.Chatbot(
+                    label="Consultation Log",
+                    height=450,
+                    bubble_full_width=False,
+                    show_label=False # Label provided by Markdown above
+                )
                 with gr.Row():
                     message = gr.Textbox(
-                        lines=2,
-                        label="Doctor's Question",
-                        placeholder="Ask a question about this patient's cardiac status...",
+                        label="Your Question", # Added label for clarity
+                        placeholder="Type your question here and press Enter or click Send...",
                         scale=4,
-                        show_label=False,
-                        container=False, # Makes textbox slimmer vertically
+                        show_label=False, # Hide label visually if desired, but keep for accessibility
+                        container=False,
                     )
                     chat_button = gr.Button("Send", scale=1, variant="primary")
 
-        with gr.TabItem("ℹ️ Instructions"):
+        with gr.TabItem("ℹ️ Instructions & Disclaimer"):
             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**. 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 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
-            *   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 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
-            *   **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.
+            1.  **Upload ECG:** Go to the "Main Interface" tab. Upload an ECG image using the designated area.
+            2.  **Analyze ECG:** Click the **Analyze ECG Image** button. Wait for the results to appear in the "ECG Analysis Results" box. This uses the Gemini model.
+            3.  **Add Patient History (Optional):**
+                *   Type relevant details directly into the "Patient History" text box.
+                *   OR, upload a `.txt`, `.csv`, or `.xlsx` file containing patient info and click **Load Patient History from File**. The content will load into the text box.
+            4.  **Generate Assessment:** Click the **Generate Assessment** button. The system will combine the ECG analysis and patient history (if provided) to generate a structured assessment using a Llama model via Groq. Results appear in the "Medical Assessment" box.
+            5.  **Consult:** Use the chat interface at the bottom to ask follow-up questions. Type your question and click **Send** or press Enter. The chat uses the Llama model via Groq and considers the context provided above it.
+
+            ---
+            ## Important Disclaimer
+
+            *   **Not a Medical Device:** This tool is for informational and educational purposes only. It is **NOT** a certified medical device and should **NOT** be used for primary diagnosis, treatment decisions, or emergency situations.
+            *   **AI Limitations:** AI models can make mistakes, misinterpret images, or generate inaccurate information. The outputs are not a substitute for professional medical expertise.
+            *   **Professional Judgment Required:** All outputs must be critically reviewed and verified by a qualified healthcare professional in the context of the patient's overall clinical picture. Do not rely solely on the AI's interpretation.
+            *   **Data Privacy:** Ensure you comply with all applicable data privacy regulations (e.g., HIPAA) when using this tool. Avoid uploading identifiable patient information unless you have explicit consent and are operating within a secure, compliant environment. The developers of this interface are not responsible for data breaches resulting from user uploads.
+            *   **API Keys:** Securely manage your `GEMINI_API_KEY` and `GROQ_API_KEY`. Do not expose them publicly.
+            *   **No Liability:** Use this tool at your own risk. The creators assume no liability for any decisions made based on its output.
             """)
 
     # --- Event Handlers ---
 
-    # Chain analysis button click to update output
+    # Analyze Button: Input ECG Image -> Output ECG Analysis HTML
     analyze_button.click(
-        analyze_ecg_image,
+        fn=analyze_ecg_image,
         inputs=[ecg_image],
         outputs=ecg_analysis_output
     )
 
-    # Load history file content into the textbox
+    # Load History Button: Input File -> Output Text into History Textbox
     load_history_button.click(
-        process_patient_history,
+        fn=process_patient_history,
         inputs=[patient_history_file],
         outputs=[patient_history_text]
     )
 
-    # Generate assessment based on ECG output and history textbox
+    # Assess Button: Input ECG Analysis HTML, History Text -> Output Assessment HTML
     assess_button.click(
-        generate_assessment,
+        fn=generate_assessment,
         inputs=[ecg_analysis_output, patient_history_text],
         outputs=assessment_output
     )
 
-    # Handle chat interactions
+    # Chat Send Button: Input Message, Chat History, Context -> Output Cleared Message, Updated Chat History
     chat_button.click(
-        doctor_chat,
+        fn=doctor_chat,
         inputs=[message, chatbot, ecg_analysis_output, patient_history_text, assessment_output],
-        outputs=[message, chatbot] # Clear message input, update chatbot history
+        outputs=[message, chatbot]
     )
 
-    # Also trigger chat on Enter key press in the message textbox
+    # Chat Textbox Submit (Enter): Input Message, Chat History, Context -> Output Cleared Message, Updated Chat History
     message.submit(
-        doctor_chat,
+        fn=doctor_chat,
         inputs=[message, chatbot, ecg_analysis_output, patient_history_text, assessment_output],
-        outputs=[message, chatbot] # Clear message input, update chatbot history
+        outputs=[message, chatbot]
     )
 
 
 # Launch the app
 if __name__ == "__main__":
-    # For development, you might want debug=True
-    # app.launch(debug=True)
-    app.launch()
+    print("===== Application Startup =====")
+    print(f"Attempting to launch Gradio app at {datetime.datetime.now()}")
+    # Add share=True for a temporary public link (use with caution regarding data privacy)
+    # Add debug=True for more detailed error output during development
+    app.launch()
+    # app.launch(share=True, debug=True)