Update app.py

app.py

@@ -6,9 +6,13 @@ Pipeline:
  - run VLM (remote gradio / chat_fn) -> JSON feature vector + raw text
  - run LLM (remote gradio /chat) -> structured risk JSON (per requested schema)
  - continue rest of processing and store results
+
 Notes:
  - Add gradio_client==1.13.2 (or another compatible 1.x) to requirements.txt
  - If VLM/LLM Spaces are private, set HF_TOKEN in the environment for authentication.
+ - This version includes a robust regex-based extractor that finds the outermost {...} block
+   in the LLM output, extracts numeric values for the required keys, and always returns
+   numeric defaults (no NaN) so frontends will not receive null/None for numeric fields.
 """
 
 import io
@@ -18,6 +22,7 @@ import json
 import asyncio
 import logging
 import traceback
+import re
 from typing import Dict, Any, Optional, Tuple
 from datetime import datetime
 
@@ -132,6 +137,127 @@ def estimate_eye_openness_from_detection(confidence: float) -> float:
     except Exception:
         return 0.0
 
+# -----------------------
+# Regex-based robust extractor
+# -----------------------
+def extract_json_via_regex(raw_text: str) -> Dict[str, Any]:
+    """
+    1) Finds the outermost { ... } block in raw_text.
+    2) Extracts numeric values after the listed keys using regex, tolerating:
+       - quotes, spaces, percent signs, percent numbers like "55%", strings like "0.12", integers, or numbers in quotes.
+    3) Returns a dict with numeric fields GUARANTEED to be floats (no None/NaN), and string fields for summary/recommendation.
+    """
+    # Find the first {...} block (outermost approximation)
+    match = re.search(r"\{[\s\S]*\}", raw_text)
+    if not match:
+        raise ValueError("No JSON-like block found in LLM output")
+
+    block = match.group(0)
+
+    def find_number_for_key(key: str) -> Optional[float]:
+        """
+        Returns a float in range 0..1 for probabilities, and raw numeric for other keys depending on usage.
+        This helper returns None if not found; caller will replace with defaults (0.0).
+        """
+        # Try multiple patterns to be robust
+        # Pattern captures numbers possibly with % and optional quotes, e.g. "45%", '0.12', 0.5, " 87 "
+        patterns = [
+            rf'"{key}"\s*:\s*["\']?\s*([-+]?\d+(\.\d+)?)\s*%?\s*["\']?',  # "key": "45%" or "key": 0.45
+            rf"'{key}'\s*:\s*['\"]?\s*([-+]?\d+(\.\d+)?)\s*%?\s*['\"]?",
+            rf'\b{key}\b\s*:\s*["\']?\s*([-+]?\d+(\.\d+)?)\s*%?\s*["\']?',  # key: 45%
+            rf'"{key}"\s*:\s*["\']([^"\']+)["\']',  # capture quoted text (for non-numeric attempts)
+            rf"'{key}'\s*:\s*['\"]([^'\"]+)['\"]"
+        ]
+        for pat in patterns:
+            m = re.search(pat, block, flags=re.IGNORECASE)
+            if not m:
+                continue
+            g = m.group(1)
+            if g is None:
+                continue
+            s = str(g).strip()
+            # Remove percent sign if present
+            s = s.replace("%", "").strip()
+            # Try to coerce to float
+            try:
+                val = float(s)
+                return val
+            except Exception:
+                # not numeric
+                return None
+        return None
+
+    def find_text_for_key(key: str) -> str:
+        # capture "key": "some text" allowing single/double quotes and also unquoted until comma/}
+        m = re.search(rf'"{key}"\s*:\s*"([^"]*)"', block, flags=re.IGNORECASE)
+        if m:
+            return m.group(1).strip()
+        m = re.search(rf"'{key}'\s*:\s*'([^']*)'", block, flags=re.IGNORECASE)
+        if m:
+            return m.group(1).strip()
+        # fallback: key: some text (unquoted) up to comma or }
+        m = re.search(rf'\b{key}\b\s*:\s*([^\n,}}]+)', block, flags=re.IGNORECASE)
+        if m:
+            return m.group(1).strip().strip('",')
+        return ""
+
+    # Extract raw numeric candidates
+    raw_risk = find_number_for_key("risk_score")
+    raw_jaundice = find_number_for_key("jaundice_probability")
+    raw_anemia = find_number_for_key("anemia_probability")
+    raw_hydration = find_number_for_key("hydration_issue_probability")
+    raw_neuro = find_number_for_key("neurological_issue_probability")
+    raw_conf = find_number_for_key("confidence")
+
+    # Normalize:
+    # - For probabilities: if value > 1 and <=100 => treat as percent -> divide by 100. If <=1 treat as fraction.
+    def normalize_prob(v: Optional[float]) -> float:
+        if v is None:
+            return 0.0
+        if v > 1.0 and v <= 100.0:
+            return max(0.0, min(1.0, v / 100.0))
+        # if v is large >100, clamp to 1.0
+        if v > 100.0:
+            return 1.0
+        # otherwise assume already 0..1
+        return max(0.0, min(1.0, v))
+
+    jaundice_probability = normalize_prob(raw_jaundice)
+    anemia_probability = normalize_prob(raw_anemia)
+    hydration_issue_probability = normalize_prob(raw_hydration)
+    neurological_issue_probability = normalize_prob(raw_neuro)
+    confidence = normalize_prob(raw_conf)
+
+    # risk_score: return in 0..100
+    def normalize_risk(v: Optional[float]) -> float:
+        if v is None:
+            return 0.0
+        if v <= 1.0:
+            # fraction given -> scale to 0..100
+            return round(max(0.0, min(100.0, v * 100.0)), 2)
+        # if between 1 and 100, assume it's already 0..100
+        if v > 1.0 and v <= 100.0:
+            return round(max(0.0, min(100.0, v)), 2)
+        # clamp anything insane
+        return round(max(0.0, min(100.0, v if v < float('inf') else 100.0)), 2)
+
+    risk_score = normalize_risk(raw_risk)
+
+    summary = find_text_for_key("summary")
+    recommendation = find_text_for_key("recommendation")
+
+    out = {
+        "risk_score": risk_score,
+        "jaundice_probability": round(jaundice_probability, 4),
+        "anemia_probability": round(anemia_probability, 4),
+        "hydration_issue_probability": round(hydration_issue_probability, 4),
+        "neurological_issue_probability": round(neurological_issue_probability, 4),
+        "confidence": round(confidence, 4),
+        "summary": summary,
+        "recommendation": recommendation
+    }
+    return out
+
 # -----------------------
 # Gradio / VLM helper (returns parsed dict OR None, plus raw text)
 # -----------------------
@@ -208,7 +334,7 @@ def run_vlm_and_get_features(face_path: str, eye_path: str, prompt: Optional[str
     return parsed_features, text_out
 
 # -----------------------
-# Gradio / LLM helper (always prompts with VLM output + strict instruction)
+# Gradio / LLM helper (uses regex extractor on LLM output)
 # -----------------------
 def run_llm_on_vlm(vlm_features_or_raw: Any,
                    max_new_tokens: int = 1024,
@@ -219,11 +345,12 @@ def run_llm_on_vlm(vlm_features_or_raw: Any,
                    developer_prompt: Optional[str] = None) -> Dict[str, Any]:
     """
     Call the remote LLM Space's /chat endpoint.
-    Accepts **either**:
+    Accepts either:
       - a dict (parsed VLM features) -> will be JSON-dumped (backwards compatible)
       - a raw string (the exact VLM text output) -> will be forwarded AS-IS (no extra JSON quoting)
 
-    Enforces a strict instruction to OUTPUT ONLY a single valid JSON object.
+    After the LLM returns, we use a regex-based extractor to pull numeric values and strings,
+    reconstruct a clean JSON dict with numeric defaults (no NaN).
     """
     if not GRADIO_AVAILABLE:
         raise RuntimeError("gradio_client not installed. Add gradio_client to requirements.txt")
@@ -285,80 +412,60 @@ def run_llm_on_vlm(vlm_features_or_raw: Any,
     if not text_out or len(text_out.strip()) == 0:
         raise RuntimeError("LLM returned empty response")
 
-    # Try strict parse first (expecting exactly one JSON object)
-    parsed = None
+    # Use regex-based extraction (robust)
     try:
-        parsed = json.loads(text_out)
-    except Exception:
-        # Forgiving extraction: find first {...} block and parse it
+        parsed = extract_json_via_regex(text_out)
+    except Exception as e:
+        logging.exception("Regex JSON extraction failed")
+        # As a last fallback, attempt naive JSON parsing; if that fails, raise with raw output
         try:
-            s = text_out
-            first = s.find("{")
-            last = s.rfind("}")
-            if first != -1 and last != -1 and last > first:
-                maybe = s[first:last+1]
-                parsed = json.loads(maybe)
-            else:
-                raise ValueError("No JSON object found in LLM output")
-        except Exception as e:
-            logging.exception("Failed to parse JSON from LLM output")
-            raise ValueError(f"Failed to parse JSON from LLM output: {e}\nRaw output: {text_out}")
+            parsed = json.loads(text_out)
+        except Exception:
+            raise ValueError(f"Failed to extract JSON from LLM output: {e}\nRaw Output:\n{text_out}")
 
     if not isinstance(parsed, dict):
         raise ValueError("Parsed LLM output is not a JSON object/dict")
 
-    # Validate and coerce expected probability fields to floats between 0..1 and risk_score 0..100
+    # Final safety clamps (already ensured by extractor, but keep defensive checks)
     def safe_prob(val):
         try:
             v = float(val)
-            if v > 1:
-                # if model returned 0-100 percentage, convert
-                if v <= 100:
-                    v = v / 100.0
             return max(0.0, min(1.0, v))
         except Exception:
-            return None
+            return 0.0
 
-    expected_prob_keys = [
+    for k in [
         "jaundice_probability",
         "anemia_probability",
         "hydration_issue_probability",
-        "neurological_issue_probability",
-    ]
-    for k in expected_prob_keys:
-        if k in parsed:
-            parsed[k] = safe_prob(parsed[k])
-        else:
-            parsed[k] = None
+        "neurological_issue_probability"
+    ]:
+        parsed[k] = safe_prob(parsed.get(k, 0.0))
 
-    # risk_score: coerce to 0..100
-    if "risk_score" in parsed:
-        try:
-            rs = float(parsed["risk_score"])
-            if rs <= 1:
-                rs = rs * 100.0
-            parsed["risk_score"] = round(max(0.0, min(100.0, rs)), 2)
-        except Exception:
-            parsed["risk_score"] = None
-    else:
-        probs = [p for p in (parsed.get(k) for k in expected_prob_keys) if isinstance(p, (int, float))]
-        parsed["risk_score"] = round((sum(probs) / len(probs) * 100.0) if probs else 0.0, 2)
+    # risk_score clamp 0..100
+    try:
+        rs = float(parsed.get("risk_score", 0.0))
+        parsed["risk_score"] = round(max(0.0, min(100.0, rs)), 2)
+    except Exception:
+        parsed["risk_score"] = 0.0
 
-    # Ensure confidence exists and is 0..1
-    if "confidence" in parsed:
-        try:
-            c = float(parsed["confidence"])
-            if c > 1 and c <= 100:
-                c = c / 100.0
-            parsed["confidence"] = max(0.0, min(1.0, c))
-        except Exception:
-            parsed["confidence"] = None
-    else:
-        parsed["confidence"] = None
+    # confidence clamp 0..1
+    parsed["confidence"] = safe_prob(parsed.get("confidence", 0.0))
+
+    # Ensure summary/recommendation are strings
+    parsed["summary"] = str(parsed.get("summary", "") or "").strip()
+    parsed["recommendation"] = str(parsed.get("recommendation", "") or "").strip()
 
-    # summary and recommendation must be strings (neutral wording)
-    parsed["summary"] = str(parsed.get("summary", "")).strip()
-    parsed["recommendation"] = str(parsed.get("recommendation", "")).strip()
+    # Optional: add flags indicating missing values (useful for frontend)
+    for k in [
+        "jaundice_probability",
+        "anemia_probability",
+        "hydration_issue_probability",
+        "neurological_issue_probability",
+        "confidence",
+        "risk_score"
+    ]:
+        parsed[f"{k}_was_missing"] = False  # extractor already returned defaults; mark as False
 
     return parsed