Update scenario_planner.py

scenario_planner.py

@@ -1,43 +1,50 @@
+from typing import Dict, Any
+from pydantic import ValidationError
 import json
-from schemas import ScenarioPlan
-from settings import HEALTHCARE_SYSTEM_PROMPT
-from llm_router import generate_with_fallback
+from schema import ScenarioPlan, TaskPlan
+from llm_router import cohere_chat, open_fallback_chat
 
-PLAN_INSTRUCTIONS = """
-Return ONLY valid JSON. Schema:
+_PARSER_PROMPT = """You are a Canadian healthcare analysis planner.
+Read the scenario and produce a JSON plan with tasks that exactly match the scenario requests.
+Each task must specify: title, format, data_key (if you can infer), and any needed steps:
+- filter: boolean expression using dataframe column names
+- derive: list of 'col = expression' for calculated columns
+- joins: list of objects {right_key, left_on, right_on, how}
+- group_by: list of columns
+- agg: list of aggregations like 'avg(x)', 'median(y)', 'p90(z)', 'count(*)'
+- pivot: {"index":[...], "columns":"...", "values":"..."}
+- sort_by, sort_dir, top
+- fields: columns to output in order
+- chart + encodings: when a chart is requested (x,y,color,column)
+- number_format: per-column display formats (e.g., {"wait":"0", "rate":"0.0%"})
+If you cannot infer a field/column name, leave it as-is; the executor will attempt to resolve.
+Output STRICT JSON only with keys: tasks, notes. DO NOT include explanations outside JSON."""
 
-{
-  "tasks": [
-    {
-      "title": "string",
-      "data_key": "string|null",
-      "format": "table|list|comparison|map|narrative|chart",
-      "filter": "expr|null",
-      "derive": ["col=expr"]|null,
-      "group_by": ["col"]|null,
-      "agg": ["sum(col)","avg(col)",...]|null,
-      "pivot": {"index":"a","columns":"b","values":"c"}|null,
-      "join": [{"right_key":"ds","left_on":"x","right_on":"y","how":"left"}]|null,
-      "sort_by": "col|null",
-      "sort_dir": "asc|desc",
-      "top": int|null,
-      "fields": ["col"]|null,
-      "chart": "bar|line|area|point|tick|rule"|null,
-      "x": "col|null", "y": "col|null", "color":"col|null", "column":"col|null"
-    }
-  ],
-  "narrative_required": true,
-  "notes": "optional"
-}
-"""
+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)}")
+    catalog_str = "\n".join(cat_lines) if cat_lines else "- (no files uploaded yet)"
+    return f"""{_PARSER_PROMPT}
+
+# Uploaded datasets and their columns
+{catalog_str}
 
-def build_prompt(scenario: str, catalog: dict) -> str:
-    catalog_str = "\n".join([f"- {k}: {', '.join(v)}" for k,v in catalog.items()])
-    return f"{HEALTHCARE_SYSTEM_PROMPT}\n\nDATASETS:\n{catalog_str}\n\n{PLAN_INSTRUCTIONS}\n\nSCENARIO:\n{scenario}\n\nJSON:"
+# Scenario
+{scenario_text}
+"""
 
-def plan_from_llm(scenario: str, catalog: dict) -> ScenarioPlan:
-    prompt = build_prompt(scenario, catalog)
-    raw = generate_with_fallback(prompt)
-    start, end = raw.find("{"), raw.rfind("}")
-    data = json.loads(raw[start:end+1])
-    return ScenarioPlan.model_validate(data)
+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)
+    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]
+    try:
+        obj = json.loads(raw)
+        return ScenarioPlan(**obj)
+    except (json.JSONDecodeError, ValidationError):
+        return ScenarioPlan(tasks=[TaskPlan(title="Scenario Summary", format="narrative")], notes="Plan validation failed, fallback.")