add ai diagnosis and medication interaction

cdss.py

@@ -35,6 +35,7 @@ from simulator import (
 )
 from editor import editor_ui, save_rules
 from validator import validator_ui, test_condition, add_rule_to_set
+from diagnosis import diagnosis_ui, generate_diagnosis, check_medication_interaction
 
 
 # --- Build UI ---
@@ -67,6 +68,14 @@ with gr.Blocks(
             add_rule_button,
             add_rule_status,
         ) = validator_ui()
+        (
+            generate_button,
+            diagnosis_output,
+            medication_output,
+            medication_input,
+            check_button,
+            interaction_output,
+        ) = diagnosis_ui()
 
     with gr.Row():
         with gr.Column(scale=2):
@@ -111,6 +120,16 @@ with gr.Blocks(
         historic_box = gr.Textbox(label="Historic Text", lines=12, interactive=False)
 
     # --- Event Handlers ---
+    def update_medication_input(patient_type):
+        if patient_type == "Mother":
+            return gr.update(value="Aspirin, Ibuprofen")
+        elif patient_type == "Gyn":
+            return gr.update(value="Clopidogrel, Omeprazole")
+        elif patient_type == "Neonate":
+            return gr.update(value="Ceftriaxone, Calcium")
+
+    patient_type_radio.change(update_medication_input, inputs=patient_type_radio, outputs=medication_input)
+
     ui_outputs = [
         state,
         scenario_lbl,
@@ -181,6 +200,9 @@ with gr.Blocks(
     gr.Timer(30.0).tick(tick_timer, timer_inputs, ui_outputs)
     gr.Timer(1.0).tick(countdown_tick, [last_tick_ts], [countdown_lbl])
 
+    generate_button.click(generate_diagnosis, inputs=state, outputs=[diagnosis_output, medication_output])
+    check_button.click(check_medication_interaction, inputs=[patient_type_radio, medication_input], outputs=interaction_output)
+
     demo.load(inject_scenario, [gr.State("A0"), cdss_toggle, history_df, historic_text], ui_outputs)
 
 if __name__ == "__main__":

diagnosis.py

@@ -0,0 +1,117 @@
+"""
+CDSS Diagnosis Component
+"""
+
+import gradio as gr
+import os
+import google.generativeai as genai
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
+from models import PatientState, Vitals
+
+GEMINI_MODEL_NAME = "gemini-2.5-flash"
+# --- Gemini setup (simplified) ---
+try:
+    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+    GEMINI_MODEL = genai.GenerativeModel(GEMINI_MODEL_NAME)
+    GEMINI_ERR = None
+except Exception as e:
+    GEMINI_MODEL, GEMINI_ERR = None, f"Gemini import/config error: {e}"
+
+
+def generate_diagnosis(patient_state: dict) -> tuple[str, str]:
+    if not GEMINI_MODEL:
+        return f"[CDSS AI ERROR] {GEMINI_ERR}", ""
+
+    ps = PatientState(**patient_state)
+    ps.vitals = Vitals(**ps.vitals)
+
+    prompt = f"""Generate a diagnosis in a medical record statement format (Subjective, Anamnese, Plan, Objective) and provide medication recommendations for the following patient data:
+
+    - Patient Type: {ps.patient_type}
+    - Vitals: {ps.vitals}
+    - Labs: {ps.labs}
+    - Notes: {ps.notes}
+    """
+
+    try:
+        response = GEMINI_MODEL.generate_content(
+            prompt,
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+            },
+        )
+        print(response)
+        if response.parts:
+            # simple parsing, assuming the response is well-formed
+            diagnosis = response.text.split("Medication Recommendations:")[0]
+            medication = response.text.split("Medication Recommendations:")[1]
+            return diagnosis, medication
+        else:
+            return "No response from AI.", ""
+    except Exception as e:
+        return f"[CDSS AI error] {e}", ""
+
+
+def check_medication_interaction(patient_type: str, medications: str) -> str:
+    if not GEMINI_MODEL:
+        return f"[CDSS AI ERROR] {GEMINI_ERR}"
+
+    prompt = f"""Check for dangerous medication interactions in the following list of medications for a {patient_type} patient: {medications}.
+
+    Provide a clear warning if any dangerous interactions are found."""
+
+    try:
+        response = GEMINI_MODEL.generate_content(
+            prompt,
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+            },
+        )
+        print(response)
+        if response.parts:
+            return response.text
+        else:
+            return "No response from AI."
+    except Exception as e:
+        return f"[CDSS AI error] {e}"
+
+
+def diagnosis_ui():
+    with gr.TabItem("Diagnosis"):
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("## Generate Diagnosis and Medication")
+                generate_button = gr.Button("Generate", variant="primary")
+                diagnosis_output = gr.Textbox(
+                    label="Diagnosis (S.O.A.P)", lines=10, interactive=False
+                )
+                medication_output = gr.Textbox(
+                    label="Medication Recommendations", lines=5, interactive=False
+                )
+
+            with gr.Column():
+                gr.Markdown("## Check Medication Interaction")
+                medication_input = gr.Textbox(
+                    label="Medications (comma-separated)",
+                    lines=3,
+                    value="Aspirin, Ibuprofen",  # Default for Mother
+                )
+                check_button = gr.Button("Check Interaction", variant="secondary")
+                interaction_output = gr.Textbox(
+                    label="Interaction Result", lines=10, interactive=False
+                )
+
+    return (
+        generate_button,
+        diagnosis_output,
+        medication_output,
+        medication_input,
+        check_button,
+        interaction_output,
+    )

simulator.py

@@ -1,6 +1,7 @@
 """
 CDSS Simulator Component
 """
+
 import random
 import time
 from dataclasses import asdict
@@ -13,13 +14,15 @@ import plotly.express as px
 
 from models import Vitals, PatientState
 from rules import rule_based_cdss
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
 
 # --- Gemini setup (simplified) ---
 try:
     import google.generativeai as genai
     import os
+
     genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
-    GEMINI_MODEL = genai.GenerativeModel("gemini-1.5-pro")
+    GEMINI_MODEL = genai.GenerativeModel("gemini-2.5-flash")
     GEMINI_ERR = None
 except Exception as e:
     GEMINI_MODEL, GEMINI_ERR = None, f"Gemini import/config error: {e}"
@@ -27,6 +30,7 @@ except Exception as e:
 
 # --- Data structures & Scenarios (Full list included) ---
 
+
 def scenario_A0_Normal() -> PatientState:
     return PatientState(
         "A0 Normal Case",
@@ -36,6 +40,7 @@ def scenario_A0_Normal() -> PatientState:
         Vitals(110, 70, 80, 16, 36.7, 99),
     )
 
+
 def scenario_A1_PPH() -> PatientState:
     return PatientState(
         "A1 PPH",
@@ -45,6 +50,7 @@ def scenario_A1_PPH() -> PatientState:
         Vitals(90, 60, 120, 24, 36.8, 96),
     )
 
+
 def scenario_A2_Preeclampsia() -> PatientState:
     return PatientState(
         "A2 Preeklampsia",
@@ -54,6 +60,7 @@ def scenario_A2_Preeclampsia() -> PatientState:
         Vitals(165, 105, 98, 20, 36.9, 98),
     )
 
+
 def scenario_A3_MaternalSepsis() -> PatientState:
     return PatientState(
         "A3 Sepsis Maternal",
@@ -63,6 +70,7 @@ def scenario_A3_MaternalSepsis() -> PatientState:
         Vitals(95, 60, 110, 24, 39.0, 96),
     )
 
+
 def scenario_B1_Prematurity() -> PatientState:
     return PatientState(
         "B1 Prematuritas/BBLR",
@@ -72,6 +80,7 @@ def scenario_B1_Prematurity() -> PatientState:
         Vitals(60, 35, 150, 50, 35.0, 90),
     )
 
+
 def scenario_B2_Asphyxia() -> PatientState:
     return PatientState(
         "B2 Asfiksia Perinatal",
@@ -81,6 +90,7 @@ def scenario_B2_Asphyxia() -> PatientState:
         Vitals(55, 30, 80, 10, 36.5, 82),
     )
 
+
 def scenario_B3_NeonatalSepsis() -> PatientState:
     return PatientState(
         "B3 Sepsis Neonatal",
@@ -90,6 +100,7 @@ def scenario_B3_NeonatalSepsis() -> PatientState:
         Vitals(60, 35, 170, 60, 38.5, 93),
     )
 
+
 def scenario_C1_GynSurgComp() -> PatientState:
     return PatientState(
         "C1 Komplikasi Bedah Ginekologis",
@@ -99,6 +110,7 @@ def scenario_C1_GynSurgComp() -> PatientState:
         Vitals(100, 65, 105, 20, 37.8, 98),
     )
 
+
 def scenario_C2_PostOpInfection() -> PatientState:
     return PatientState(
         "C2 Infeksi Pasca-Bedah",
@@ -108,6 +120,7 @@ def scenario_C2_PostOpInfection() -> PatientState:
         Vitals(105, 70, 108, 22, 38.0, 98),
     )
 
+
 def scenario_C3_DelayedGynCancer() -> PatientState:
     return PatientState(
         "C3 Keterlambatan Diagnostik Kanker Ginekologi",
@@ -117,6 +130,7 @@ def scenario_C3_DelayedGynCancer() -> PatientState:
         Vitals(120, 78, 86, 18, 36.8, 99),
     )
 
+
 SCENARIOS = {
     "A0": scenario_A0_Normal,
     "A1": scenario_A1_PPH,
@@ -145,13 +159,27 @@ def drift_vitals(state: PatientState) -> PatientState:
 
 # --- Rule-based fallback (no AI or AI disabled) ---
 
+
 def gemini_cdss(state: PatientState) -> str:
     if not GEMINI_MODEL:
         return f"[CDSS AI ERROR] {GEMINI_ERR}"
     try:
         v = state.vitals
         prompt = f"CDSS for {state.scenario}. Vitals: SBP {v.sbp}/{v.dbp}, HR {v.hr}. Analyze risks, give concise steps in Indonesian."
-        return GEMINI_MODEL.generate_content(prompt).text or "[CDSS AI] No response."
+        response = GEMINI_MODEL.generate_content(
+            prompt,
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+            },
+        )
+        print(response)
+        if response.parts:
+            return response.text or "[CDSS AI] No response."
+        else:
+            return "[CDSS AI] No response due to safety settings."
     except Exception as e:
         return f"[CDSS AI error] {e}"
 
@@ -180,9 +208,11 @@ def create_vital_plot(
     )
     return fig
 
+
 def _row_from_state(ps: PatientState) -> Dict[str, Any]:
     return {"timestamp": datetime.now(), "scenario": ps.scenario, **asdict(ps.vitals)}
 
+
 def prepare_df_for_display(df: pd.DataFrame) -> pd.DataFrame:
     if df is None or df.empty:
         return pd.DataFrame(
@@ -203,6 +233,7 @@ def prepare_df_for_display(df: pd.DataFrame) -> pd.DataFrame:
     df_display["timestamp"] = df_display["timestamp"].dt.strftime("%Y-%m-%d %H:%M:%S")
     return df_display
 
+
 def generate_all_plots(df: pd.DataFrame):
     """Helper to generate all 5 plot figures from a dataframe."""
     df_display = prepare_df_for_display(df)
@@ -257,6 +288,7 @@ def process_and_update(
         spo2_fig,
     )
 
+
 def state_to_panels(state: PatientState) -> Tuple:
     v = state.vitals
     return (
@@ -271,6 +303,7 @@ def state_to_panels(state: PatientState) -> Tuple:
         v.spo2,
     )
 
+
 def inject_scenario(
     tag: str, cdss_on: bool, history_df: pd.DataFrame, historic_text: str
 ):
@@ -283,6 +316,7 @@ def inject_scenario(
         )
     return process_and_update(ps, history_df, historic_text, cdss_on)
 
+
 def manual_edit(
     sbp,
     dbp,
@@ -313,6 +347,7 @@ def manual_edit(
         historic_text += f"\n[{datetime.now().strftime('%H:%M:%S')}] {ps.notes}"
     return process_and_update(ps, history_df, historic_text, cdss_on)
 
+
 def tick_timer(cdss_on, current_state, history_df, historic_text):
     if not current_state:
         return [gr.update()] * 22
@@ -321,6 +356,7 @@ def tick_timer(cdss_on, current_state, history_df, historic_text):
     ps = drift_vitals(ps)
     return process_and_update(ps, history_df, historic_text, cdss_on)
 
+
 def load_csv(file, history_df: pd.DataFrame):
     try:
         if file is not None:
@@ -338,11 +374,13 @@ def load_csv(file, history_df: pd.DataFrame):
     )
     return history_df, df_for_table, bp_fig, hr_fig, rr_fig, temp_fig, spo2_fig
 
+
 def countdown_tick(last_tick_ts: float):
     if not last_tick_ts:
         return "Next update in —"
     return f"Next update in {max(0, 30 - int(time.time() - last_tick_ts))}s"
 
+
 def simulator_ui():
     with gr.TabItem("CDSS Simulator"):
         with gr.Accordion("History, Trends, and Data Loading", open=True):