Update app.py

app.py

@@ -19,7 +19,14 @@ from privacy import redact_text, safety_filter, refusal_reply
 from data_registry import DataRegistry
 from upload_ingest import extract_text_from_files
 from healthcare_analysis import HealthcareAnalyzer
-from response_formatter import ResponseFormatter
+
+# ---- NEW: scenario-first engine (keeps general chat intact)
+from scenario_engine import ScenarioEngine
+# (Optional) keep old formatter if you want a fallback:
+try:
+    from response_formatter import ResponseFormatter  # noqa: F401
+except Exception:
+    ResponseFormatter = None  # type: ignore
 
 # ---------- Writable caches (HF Spaces-safe) ----------
 HOME = pathlib.Path.home()
@@ -62,6 +69,9 @@ COHERE_API_KEY = os.getenv("COHERE_API_KEY")
 USE_HOSTED_COHERE = bool(COHERE_API_KEY and _HAS_COHERE)
 MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", MODEL_SETTINGS.get("max_new_tokens", 2048)))
 
+# ---- NEW: feature flag to toggle engines without code edits
+USE_SCENARIO_ENGINE = os.getenv("USE_SCENARIO_ENGINE", "1") not in ("0", "false", "False")
+
 # ---------- Helper Functions ----------
 def find_column(df, patterns):
     """Find the first column in df that matches any of the patterns."""
@@ -188,18 +198,23 @@ def handle_healthcare_scenario(scenario_text, data_registry, history):
         # Initialize analyzer
         analyzer = HealthcareAnalyzer(data_registry)
         
-        # Perform comprehensive analysis
+        # Perform comprehensive analysis (returns dict of datasets/results)
         results = analyzer.comprehensive_analysis(scenario_text)
         
-        # Format response
-        formatter = ResponseFormatter()
-        response = formatter.format_healthcare_response(scenario_text, results)
+        # ---- NEW: Scenario-first exact-output engine
+        if USE_SCENARIO_ENGINE:
+            response = ScenarioEngine.render(scenario_text, results)
+        else:
+            # Optional fallback to legacy formatter if desired
+            if ResponseFormatter is None:
+                raise RuntimeError("ResponseFormatter not available and USE_SCENARIO_ENGINE is disabled.")
+            formatter = ResponseFormatter()
+            response = formatter.format_healthcare_response(scenario_text, results)
         
         return response
     except Exception as e:
         log_event("healthcare_scenario_error", None, {"error": str(e)})
         # Log the full traceback for better debugging
-        import traceback
         tb_str = traceback.format_exc()
         log_event("healthcare_scenario_traceback", None, {"traceback": tb_str})
         return f"Error analyzing healthcare scenario: {str(e)}\n\nTechnical details:\n{tb_str}"
@@ -352,11 +367,11 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths, awaiting_answe
 
         # Determine the mode: healthcare scenario or general conversation
         if is_healthcare_scenario(safe_in, uploaded_files_paths):
-            # Healthcare scenario mode
+            # Healthcare scenario mode (ScenarioEngine enforces exact asks)
             response = handle_healthcare_scenario(safe_in, data_registry, history)
             return history + [(user_msg, response)], False
         else:
-            # General conversation mode with enhanced handling
+            # General conversation mode with enhanced handling (unchanged)
             if USE_HOSTED_COHERE:
                 out = cohere_chat(safe_in, history)
                 if out:
@@ -449,7 +464,13 @@ with gr.Blocks(theme=theme, css=custom_css, analytics_enabled=False) as demo:
                     elem_classes="hero-box"
                 )
                 hero_send = gr.Button("➤", scale=0, elem_id="hero-send")
-            gr.Markdown('<div class="hint">I can help with general questions or analyze healthcare scenarios when you upload data files and describe your analysis needs.</div>')
+            # ---- NEW: hint that directive-driven scenarios are supported
+            gr.Markdown(
+                '<div class="hint">I can chat normally or run directive-based analyses. '
+                'In scenarios, add directives like <code>format:</code>, <code>data_key:</code>, '
+                '<code>filter:</code>, <code>group_by:</code>, <code>agg:</code>, <code>pivot:</code>, '
+                '<code>sort_by:</code>, <code>top:</code>, <code>fields:</code>, <code>chart:</code> to control the output exactly.</div>'
+            )
 
     # --- MAIN APP (hidden until first message) ---
     with gr.Column(elem_id="chat-container", visible=False) as app_wrap:
@@ -545,4 +566,4 @@ with gr.Blocks(theme=theme, css=custom_css, analytics_enabled=False) as demo:
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", "7860"))
-    demo.launch(server_name="0.0.0.0", server_port=port, show_api=False, max_threads=40)
+    demo.launch(server_name="0.0.0.0", server_port=port, show_api=False, max_threads=40)