Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import pandas as pd
 import matplotlib.pyplot as plt
 import google.generativeai as genai
-import json
+import json, re
 import os
 from datetime import datetime
 
@@ -27,10 +27,12 @@ def process_loan_data(df: pd.DataFrame):
     return ecl_df
 
 def get_gemini_decision(segment, pd_val, lgd_val, ead_val, ecl_val):
-    """Ask Gemini to decide the recommended action for a segment."""
+    """Ask Gemini to decide and sanitize any invalid JSON output."""
     model = genai.GenerativeModel("gemini-1.5-pro")
     system_prompt = """You are a financial risk advisor.
-Return only JSON: {"action":"increase_interest"|"reduce_disbursement"|"maintain","rationale":"string","confidence":float}"""
+Return only valid JSON in this exact schema:
+{"action":"increase_interest"|"reduce_disbursement"|"maintain","rationale":"string","confidence":float}
+No text outside the JSON."""
     user_prompt = f"""
 Segment: {segment}
 PD: {pd_val:.3f}
@@ -42,16 +44,27 @@ Rules:
 - PD > 0.20 and ECL rising ⇒ reduce_disbursement
 - PD < 0.15 ⇒ maintain
 """
+
     try:
         response = model.generate_content(
             [{"role": "system", "parts": [system_prompt]},
              {"role": "user", "parts": [user_prompt]}],
-            generation_config={"temperature": 0.2}
+            generation_config={"temperature": 0.1}
         )
-        result = json.loads(response.text)
+        text = response.text.strip()
+
+        # extract JSON substring if extra text exists
+        match = re.search(r'\{.*\}', text, re.DOTALL)
+        if match:
+            text = match.group(0)
+
+        data = json.loads(text)
+        if not all(k in data for k in ["action", "rationale", "confidence"]):
+            raise ValueError("Missing keys")
+        return data
+
     except Exception:
-        result = {"action": "maintain", "rationale": "Fallback - invalid JSON", "confidence": 0.0}
-    return result
+        return {"action": "maintain", "rationale": "Fallback - cleaned JSON failed", "confidence": 0.0}
 
 # ====== UI ======
 st.title("📊 Expected Credit Loss (ECL) Risk Dashboard")