Update scenario_planner.py

scenario_planner.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, List
 from pydantic import ValidationError
 import json
 from schema import ScenarioPlan, TaskPlan
@@ -21,9 +21,7 @@ If you cannot infer a field/column name, leave it as-is; the executor will attem
 Output STRICT JSON only with keys: tasks, notes. DO NOT include explanations outside JSON."""
 
 def build_parser_prompt(scenario_text: str, dataset_catalog: Dict[str, list]) -> str:
-    cat_lines = []
-    for k, cols in dataset_catalog.items():
-        cat_lines.append(f"- {k}: {', '.join(cols)}")
+    cat_lines = [f"- {k}: {', '.join(cols)}" for k, cols in dataset_catalog.items()]
     catalog_str = "\n".join(cat_lines) if cat_lines else "- (no files uploaded yet)"
     return f"""{_PARSER_PROMPT}
 
@@ -34,17 +32,90 @@ def build_parser_prompt(scenario_text: str, dataset_catalog: Dict[str, list]) ->
 {scenario_text}
 """
 
+def _llm_parse(prompt: str) -> str | None:
+    return cohere_chat(prompt) or open_fallback_chat(prompt)
+
+def _safe_json_slice(raw: str) -> str | None:
+    raw = (raw or "").strip()
+    i, j = raw.find("{"), raw.rfind("}")
+    if i == -1 or j == -1: return None
+    return raw[i:j+1]
+
+def _make_safety_net_plan(scenario_text: str, dataset_catalog: Dict[str, list]) -> ScenarioPlan:
+    csvs = list(dataset_catalog.keys())
+    wt = next((k for k in csvs if "wait" in k.lower()), (csvs[0] if csvs else None))
+    fac = next((k for k in csvs if "facility" in k.lower() or "hospital" in k.lower()), wt)
+
+    tasks: List[TaskPlan] = []
+    text = scenario_text.lower()
+    wants_map = "map" in text or "geographic distribution" in text or "geographic" in text
+    wants_reco = "recommend" in text or "allocation" in text or "plan" in text
+
+    tasks.append(TaskPlan(
+        title="Top facilities by average surgery wait",
+        data_key=wt, format="table",
+        group_by=["Facility","Zone"],
+        agg=["avg(Surgery_Median)","avg(Surgery_90th)","count(*)"],
+        sort_by="avg_Surgery_Median", sort_dir="desc", top=5,
+        fields=["Facility","Zone","avg_Surgery_Median","avg_Surgery_90th","count"],
+        number_format={"avg_Surgery_Median":"0", "avg_Surgery_90th":"0"}
+    ))
+
+    tasks.append(TaskPlan(
+        title="Top specialties by average surgery wait",
+        data_key=wt, format="table",
+        group_by=["Specialty"],
+        agg=["avg(Surgery_Median)","avg(Consult_Median)","count(*)"],
+        sort_by="avg_Surgery_Median", sort_dir="desc", top=5,
+        fields=["Specialty","avg_Surgery_Median","avg_Consult_Median","count"],
+        number_format={"avg_Surgery_Median":"0", "avg_Consult_Median":"0"}
+    ))
+
+    tasks.append(TaskPlan(
+        title="Zone-level surgery wait comparison",
+        data_key=wt, format="table",
+        group_by=["Zone"], agg=["avg(Surgery_Median)","count(*)"],
+        sort_by="avg_Surgery_Median", sort_dir="desc",
+        fields=["Zone","avg_Surgery_Median","count"],
+        number_format={"avg_Surgery_Median":"0"}
+    ))
+
+    if wants_map and fac:
+        tasks.append(TaskPlan(
+            title="Geographic distribution of high-wait facilities",
+            data_key=wt, format="map",
+            group_by=["Facility","Zone"], agg=["avg(Surgery_Median)"],
+            sort_by="avg_Surgery_Median", sort_dir="desc", top=5,
+            joins=[{"right_key": fac, "left_on": "Facility", "right_on": "facility_name", "how": "left"}],
+            fields=["Facility","Zone","city","latitude","longitude","avg_Surgery_Median"],
+            number_format={"avg_Surgery_Median":"0"}
+        ))
+
+    if wants_reco:
+        tasks.append(TaskPlan(
+            title="Recommendations (inputs for narrative)",
+            data_key=wt, format="narrative",
+            group_by=["Facility","Zone"], agg=["avg(Surgery_Median)","count(*)"],
+            sort_by="avg_Surgery_Median", sort_dir="desc", top=10,
+            fields=["Facility","Zone","avg_Surgery_Median","count"],
+            number_format={"avg_Surgery_Median":"0"}
+        ))
+
+    return ScenarioPlan(tasks=tasks, notes="Safety-net plan (LLM planner failed).")
+
 def parse_to_plan(scenario_text: str, dataset_catalog: Dict[str, list]) -> ScenarioPlan:
     prompt = build_parser_prompt(scenario_text, dataset_catalog)
-    raw = cohere_chat(prompt) or open_fallback_chat(prompt)
+    raw = _llm_parse(prompt)
     if not raw:
-        return ScenarioPlan(tasks=[TaskPlan(title="Scenario Summary", format="narrative")], notes="LLM unavailable; minimal plan.")
-    raw = raw.strip()
-    start = raw.find("{"); end = raw.rfind("}")
-    if start != -1 and end != -1:
-        raw = raw[start:end+1]
+        return _make_safety_net_plan(scenario_text, dataset_catalog)
+    js = _safe_json_slice(raw)
+    if not js:
+        return _make_safety_net_plan(scenario_text, dataset_catalog)
     try:
-        obj = json.loads(raw)
-        return ScenarioPlan(**obj)
+        obj = json.loads(js)
+        plan = ScenarioPlan(**obj)
+        if not plan.tasks or (len(plan.tasks) == 1 and plan.tasks[0].format == "narrative"):
+            return _make_safety_net_plan(scenario_text, dataset_catalog)
+        return plan
     except (json.JSONDecodeError, ValidationError):
-        return ScenarioPlan(tasks=[TaskPlan(title="Scenario Summary", format="narrative")], notes="Plan validation failed, fallback.")
+        return _make_safety_net_plan(scenario_text, dataset_catalog)