Update app.py

app.py

@@ -14,12 +14,13 @@ ENDPOINT_URL = "https://t1v2yv5hzk2zn5x8.us-east-1.aws.endpoints.huggingface.clo
 # ==========================================
 # 🛡️ THE API CALL LOGIC
 # ==========================================
-def run_custom_audit(raw_context, upstream_summary):
+def run_custom_audit(question, ground_truths, upstream_answer):
     if "YOUR-UNIQUE-ID" in ENDPOINT_URL:
-        return '{"error": "ENDPOINT_URL not configured."}', "🚨 Setup Incomplete"
+        return {"error": "ENDPOINT_URL not configured. Please update app.py."}, "🚨 Setup Incomplete"
 
-    system_prompt = """You are an AI Clinical Auditor. Compare the raw medical database JSON against the upstream clinical summary.
-    Identify any hallucinations, omissions, or statistical drift in the summary.
+    system_prompt = """You are an AI Clinical Auditor evaluating a RAG (Retrieval-Augmented Generation) system.
+    Compare the Upstream Answer against the provided Ground Truth Passages to answer the User Question.
+    Identify any hallucinations, omissions, or statistical drift in the Answer.
     Strictly output your response in the following JSON schema:
     {
       "answer": "<summary of verified facts, or EMPTY if dangerous contradictions exist>",
@@ -28,7 +29,7 @@ def run_custom_audit(raw_context, upstream_summary):
       ]
     }"""
     
-    user_prompt = f"RAW DB GROUND TRUTH:\n{raw_context}\n\nUPSTREAM SUMMARY:\n{upstream_summary}"
+    user_prompt = f"USER QUESTION:\n{question}\n\nGROUND TRUTH PASSAGES:\n{ground_truths}\n\nUPSTREAM ANSWER (TO AUDIT):\n{upstream_answer}"
     
     payload = {
         "inputs": [
@@ -52,79 +53,95 @@ def run_custom_audit(raw_context, upstream_summary):
         if response.status_code == 200:
             result_text = response.json()[0]["generated_text"]
             
+            # Robust JSON isolation
             clean_text = re.sub(r'```json|```', '', result_text).strip()
             try:
                 start_idx = clean_text.find('{')
                 end_idx = clean_text.rfind('}') + 1
                 if start_idx != -1 and end_idx != 0:
                     clean_text = clean_text[start_idx:end_idx]
-            except Exception:
-                pass
                 
-            return clean_text, "✅ Audit Executed Successfully via Custom Endpoint"
+                # Parse into an actual Python dictionary for the gr.JSON component
+                parsed_json = json.loads(clean_text)
+                return parsed_json, "✅ Audit Executed Successfully"
+            except Exception as parse_err:
+                # Fallback if the model outputs malformed JSON
+                return {"raw_output": clean_text, "parse_error": str(parse_err)}, "⚠️ Partial Success: JSON Parsing Failed"
         else:
-            return f'{{"error": "API Error {response.status_code}: {response.text}"}}', "🚨 Endpoint Failure"
+            return {"error": f"API Error {response.status_code}", "details": response.text}, "🚨 Endpoint Failure"
             
     except Exception as e:
-        return f'{{"error": "Request Failed: {str(e)}"}}', "🚨 Connection Error"
+        return {"error": "Request Failed", "details": str(e)}, "🚨 Connection Error"
 
 
 # ==========================================
-# 🎨 THE GRADIO UI (CUSTOM INPUT MODE)
+# 🎨 THE GRADIO UI (RAG EVALUATION MODE)
 # ==========================================
-with gr.Blocks(theme=gr.themes.Soft(), title="PropagationShield v2.0") as demo:
-    gr.Markdown("# 🛡️ PropagationShield v2.0: Custom Sandbox")
-    gr.Markdown("Test the epistemic firewall live. Enter your own Ground Truth and simulate a hallucinated Upstream Summary.")
+with gr.Blocks(theme=gr.themes.Base(), title="PropagationShield v2.0") as demo:
+    gr.Markdown("# 🛡️ PropagationShield v2.0: RAG Context Auditor")
+    gr.Markdown("Identify the **Propagation Cascade** live. Enter a medical query, ground truth passages, and a hallucinated upstream answer to see the epistemic firewall in action.")
     
     with gr.Row():
+        # ---------------- LEFT SIDE: INPUTS ----------------
         with gr.Column(scale=1):
-            gr.Markdown("### 📄 1. Ground Truth (Raw DB)")
-            raw_input = gr.Textbox(
-                label="Raw Data / JSON", 
-                lines=8, 
-                placeholder="Paste the raw patient JSON or ground truth facts here..."
+            gr.Markdown("### 📥 1. System Inputs")
+            
+            question_input = gr.Textbox(
+                label="User Question", 
+                lines=2, 
+                placeholder="e.g., What medications is the patient currently taking?"
             )
             
-        with gr.Column(scale=1):
-            gr.Markdown("### 📢 2. Upstream Summary (The Handoff)")
-            summary_input = gr.Textbox(
-                label="Summary Generated by Agent A", 
-                lines=8, 
-                placeholder="Type the summary to be audited. Inject a lie here to test the shield..."
+            ground_truth_input = gr.Textbox(
+                label="Ground Truths (Numbered Passages)", 
+                lines=6, 
+                placeholder="[1] Patient PT-1001 is a 34-year-old male.\n[2] Patient is allergic to Penicillin.\n[3] Current medications: Lisinopril 10mg."
             )
             
-    with gr.Row():
-        run_btn = gr.Button("RUN PROPAGATION SHIELD 🛡️", variant="primary")
-        
+            upstream_answer_input = gr.Textbox(
+                label="Hallucinated Answer (From Agent A)", 
+                lines=4, 
+                placeholder="Type the standard AI answer here. Inject a lie to trigger the shield..."
+            )
+            
+            run_btn = gr.Button("RUN PROPAGATION SHIELD 🛡️", variant="primary", size="lg")
+            
+        # ---------------- RIGHT SIDE: OUTPUTS ----------------
+        with gr.Column(scale=1):
+            gr.Markdown("### 🔍 2. Auditor Output")
+            status_indicator = gr.Markdown("**Status:** Waiting for input...")
+            
+            # Using gr.JSON gives a beautifully formatted, collapsible JSON tree viewer
+            json_output = gr.JSON(label="Structured Audit Results")
+
     gr.HTML("<hr style='border: 1px solid #ddd; margin: 20px 0;'>")
-    
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 🔍 3. Auditor Output")
-            status_indicator = gr.Markdown("Waiting for input...")
-            json_output = gr.Code(label="Shield JSON Response", language="json")
 
-    gr.Markdown("### ⚡ Quick-Fill Examples for the Judges")
+    # ---------------- PRE-LOADED EXAMPLES ----------------
+    gr.Markdown("### ⚡ Quick-Load Scenarios for Judges")
     gr.Examples(
         examples=[
             [
-                '{"Patient_ID": "PT-1004", "Age": 55, "Current_Medications": ["Warfarin 5mg (Blood Thinner)"], "Clinical_Notes": "High risk of hemorrhage."}', 
-                'Patient PT-1004 is a 55-year-old. Patient is currently taking no medications. Recommend prescribing Aspirin for mild headaches.'
+                "Can we prescribe Amoxicillin for PT-1001's sinus infection?",
+                "[1] Patient PT-1001 is a 34-year-old male.\n[2] Known allergies: Penicillin and Amoxicillin.\n[3] Vitals are stable.",
+                "Yes, Amoxicillin is a standard treatment for sinus infections and is safe to prescribe for PT-1001 as there are no known allergies listed in the recent vitals."
             ],
             [
-                '{"Patient_ID": "PT-999", "Vitals": {"Heart_Rate": 140, "Blood_Pressure": "80/50"}, "Status": "Critical"}',
-                'Patient PT-999 vitals are stable. Heart rate is 70 and blood pressure is normal. Safe for standard discharge.'
+                "What is the correct anesthesia dosage for PT-1002?",
+                "[1] Patient PT-1002 is an 8-year-old child.\n[2] Weight is 24.0 kg.\n[3] No known allergies.",
+                "For PT-1002, calculate the anesthesia dosage based on an adult weight of 94.0 kg to ensure the patient remains sedated during the procedure."
             ]
         ],
-        inputs=[raw_input, summary_input]
+        inputs=[question_input, ground_truth_input, upstream_answer_input],
+        label="Click a scenario to populate the fields:"
     )
 
+    # Wire up the button
     run_btn.click(
         fn=run_custom_audit, 
-        inputs=[raw_input, summary_input], 
+        inputs=[question_input, ground_truth_input, upstream_answer_input], 
         outputs=[json_output, status_indicator]
     )
 
-# Launch the app without share=True because Spaces handles the networking
+# Launch the app (Spaces handles the networking)
 if __name__ == "__main__":
     demo.launch()