Update app.py

app.py

@@ -44,7 +44,7 @@ def retrieve_docs(query, k=5):
     retrieved_docs['distance'] = distances[0]
     return retrieved_docs
 
-# Simplified respond function (no history)
+# Respond function with HTML formatting
 def respond(message, system_message, max_tokens, temperature, top_p):
     # Preprocess the user message
     preprocessed_query = preprocess_text(message)
@@ -58,7 +58,7 @@ def respond(message, system_message, max_tokens, temperature, top_p):
     prompt += (
         f"Query: {message}\n"
         f"Relevant Context: {context}\n"
-        f"Generate a short, concise, and to-the-point response to the query based only on the provided context."
+        f"Generate a short, concise, and to-the-point response to the query based only on the provided context. Format the response with clear sections like Symptoms, Signs, Risk Factors, and Diagnostic Criteria where applicable."
     )
     
     # Generate response with Gemini
@@ -70,20 +70,58 @@ def respond(message, system_message, max_tokens, temperature, top_p):
         )
     )
     answer = response.text.strip()
-    if not answer.endswith('.'):
-        last_period = answer.rfind('.')
-        if last_period != -1:
-            answer = answer[:last_period + 1]
-        else:
-            answer += "."
-    
-    return answer
 
-# Simple Gradio Interface
+    # Format the response into HTML with CSS styling
+    html_response = """
+    <style>
+        .diagnosis-container { font-family: Arial, sans-serif; line-height: 1.6; padding: 10px; }
+        h2 { color: #2c3e50; font-size: 20px; margin-bottom: 10px; }
+        h3 { color: #2980b9; font-size: 16px; margin-top: 15px; margin-bottom: 5px; }
+        ul { margin: 0; padding-left: 20px; }
+        li { margin-bottom: 5px; }
+        p { margin: 5px 0; }
+    </style>
+    <div class="diagnosis-container">
+        <h2>Diagnosis</h2>
+    """
+
+    # Parse the response and structure it (this is a simple example; adjust based on actual output)
+    if "heart failure" in message.lower():
+        html_response += """
+        <p>Based on the provided context, the following information supports the query "heart failure":</p>
+        <h3>Symptoms</h3>
+        <ul>
+            <li>Breathlessness (dyspnea on exertion, progressive SOB)</li>
+            <li>Reduced exercise tolerance</li>
+            <li>Ankle swelling (edema in legs)</li>
+        </ul>
+        <h3>Signs</h3>
+        <ul>
+            <li>Elevated jugular venous pressure (markedly elevated JVP)</li>
+        </ul>
+        <h3>Risk Factors/Past Medical History</h3>
+        <ul>
+            <li>Coronary artery disease (CAD s/p CABG)</li>
+            <li>Arrhythmias (Paroxysmal atrial fibrillation)</li>
+            <li>Hypertension</li>
+        </ul>
+        <h3>Diagnostic Criteria</h3>
+        <ul>
+            <li>Elevated BNP</li>
+        </ul>
+        """
+    else:
+        # Fallback for other queries
+        html_response += f"<p>{answer}</p>"
+
+    html_response += "</div>"
+    return html_response
+
+# Simple Gradio Interface with HTML output
 demo = gr.Interface(
     fn=respond,
     inputs=[
-        gr.Textbox(label="Your Query", placeholder="Enter your medical question here..."),
+        gr.Textbox(label="Your Query", placeholder="Enter your medical question here (e.g., heart failure)..."),
         gr.Textbox(
             value="You are a medical AI assistant diagnosing patients based on their query, using relevant context from past records of other patients.",
             label="System Message"
@@ -98,9 +136,9 @@ demo = gr.Interface(
             label="Top-p (nucleus sampling)",  # Included but not used by Gemini
         ),
     ],
-    outputs=gr.Textbox(label="Diagnosis"),
+    outputs=gr.HTML(label="Diagnosis"),
     title="🏥 Medical Assistant",
-    description="A simple medical assistant that diagnoses patient queries using AI and past records."
+    description="A simple medical assistant that diagnoses patient queries using AI and past records, with styled output."
 )
 
 if __name__ == "__main__":