Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py - Complete Dynamic Healthcare Scenario Analysis System
+# app.py - Complete Dual-Mode Healthcare Analysis System
 import os, re, json, traceback, pathlib
 from functools import lru_cache
 from typing import List, Dict, Any, Tuple, Optional
@@ -73,7 +73,35 @@ Formatting rules for structured analysis:
 - End with concrete recommendations and a brief "Provenance" mapping outputs to scenario text, uploaded files, and answers.
 """.strip()
 
-# ---------- Session RAG Class (Simplified) ----------
+# ---------- Helper Functions ----------
+def find_column(df, patterns):
+    """Find the first column in df that matches any of the patterns."""
+    if df is None or df.empty:
+        return None
+    for col in df.columns:
+        if any(pattern.lower() in col.lower() for pattern in patterns):
+            return col
+    return None
+
+def extract_scenario_tasks(scenario_text):
+    """Extract specific tasks from scenario text."""
+    tasks = []
+    lines = scenario_text.split('\n')
+    in_tasks = False
+    for line in lines:
+        line = line.strip()
+        if line.lower().startswith('tasks'):
+            in_tasks = True
+            continue
+        if in_tasks:
+            if line.lower().startswith('operational recommendations') or line.lower().startswith('future integration'):
+                in_tasks = False
+                continue
+            if line and (line.startswith(('1.', '2.', '3.', '4.', '5.')) or line.startswith(('•', '-', '*'))):
+                tasks.append(line)
+    return tasks
+
+# ---------- Session RAG Class ----------
 class SessionRAG:
     def __init__(self):
         self.docs = []
@@ -134,6 +162,22 @@ def is_healthcare_scenario(text: str, uploaded_files_paths) -> bool:
     return (has_healthcare_keywords or has_facility_types or has_healthcare_tasks) and \
            (has_healthcare_files or has_scenario_structure)
 
+def is_general_conversation(text: str, uploaded_files_paths) -> bool:
+    """Determine if this is a general conversation rather than a scenario analysis."""
+    # If there are uploaded files, it's likely a scenario
+    if uploaded_files_paths:
+        return False
+    
+    # Check for scenario indicators
+    scenario_indicators = [
+        "scenario", "analyze", "analysis", "assess", "evaluate", "recommend",
+        "tasks", "background", "situation", "dataset", "data"
+    ]
+    
+    # If no scenario indicators, it's likely general conversation
+    text_lower = text.lower()
+    return not any(indicator in text_lower for indicator in scenario_indicators)
+
 def process_healthcare_data(uploaded_files_paths, data_registry):
     """Process healthcare data files with robust error handling."""
     for file_path in uploaded_files_paths:
@@ -152,28 +196,25 @@ def process_healthcare_data(uploaded_files_paths, data_registry):
 def analyze_facility_distribution(facilities_df):
     """Analyze healthcare facility distribution dynamically."""
     try:
-        # Filter to Alberta if province column exists
-        province_col = facilities_df.columns[facilities_df.columns.str.contains('province', case=False)]
-        if len(province_col) > 0:
-            province_col = province_col[0]
+        # Find province column
+        province_col = find_column(facilities_df, ['province', 'state', 'territory'])
+        if province_col:
             alberta_mask = facilities_df[province_col].str.lower().isin(['alberta', 'ab'])
             ab_facilities = facilities_df[alberta_mask].copy()
         else:
             ab_facilities = facilities_df.copy()
         
         # Find facility type column
-        type_col = facilities_df.columns[facilities_df.columns.str.contains('type', case=False)]
-        if len(type_col) == 0:
+        type_col = find_column(facilities_df, ['type', 'category', 'class', 'facility_type', 'odhf_facility_type'])
+        if not type_col:
             return {"error": "Facility type column not found"}
-        type_col = type_col[0]
         
         # Facility type frequency
         type_counts = ab_facilities[type_col].value_counts().to_dict()
         
-        # Top cities by facility count
-        city_col = facilities_df.columns[facilities_df.columns.str.contains('city', case=False)]
-        if len(city_col) > 0:
-            city_col = city_col[0]
+        # Find city column
+        city_col = find_column(facilities_df, ['city', 'municipality', 'town'])
+        if city_col:
             city_counts = ab_facilities[city_col].value_counts().head(5)
             top_cities = city_counts.index.tolist()
             
@@ -193,8 +234,8 @@ def analyze_facility_distribution(facilities_df):
             "city_breakdown": city_breakdown,
             "columns_used": {
                 "facility_type": type_col,
-                "city": city_col[0] if len(city_col) > 0 else None,
-                "province": province_col[0] if len(province_col) > 0 else None
+                "city": city_col,
+                "province": province_col
             }
         }
     except Exception as e:
@@ -205,14 +246,11 @@ def analyze_bed_capacity(beds_df):
     """Analyze bed capacity dynamically."""
     try:
         # Find required columns
-        current_cols = beds_df.columns[beds_df.columns.str.contains('current|2023|2024', case=False)]
-        prev_cols = beds_df.columns[beds_df.columns.str.contains('prev|2022|previous', case=False)]
-        
-        if len(current_cols) == 0 or len(prev_cols) == 0:
-            return {"error": f"Missing required columns. Found current: {current_cols.tolist()}, prev: {prev_cols.tolist()}"}
+        current_col = find_column(beds_df, ['current', '2023', '2024', 'beds_current', 'staffed_beds', 'capacity'])
+        prev_col = find_column(beds_df, ['prev', 'previous', '2022', 'beds_prev', 'previous_beds'])
         
-        current_col = current_cols[0]
-        prev_col = prev_cols[0]
+        if not current_col or not prev_col:
+            return {"error": f"Missing required columns. Found current: {current_col}, prev: {prev_col}"}
         
         # Ensure derived columns exist
         if 'bed_change' not in beds_df.columns:
@@ -225,18 +263,16 @@ def analyze_bed_capacity(beds_df):
             )
         
         # Filter to Alberta if province column exists
-        province_col = beds_df.columns[beds_df.columns.str.contains('province', case=False)]
-        if len(province_col) > 0:
-            province_col = province_col[0]
+        province_col = find_column(beds_df, ['province', 'state', 'territory'])
+        if province_col:
             alberta_mask = beds_df[province_col].str.lower().isin(['alberta', 'ab'])
             ab_beds = beds_df[alberta_mask].copy()
         else:
             ab_beds = beds_df.copy()
         
         # Calculate zone-level summaries if zone column exists
-        zone_col = beds_df.columns[beds_df.columns.str.contains('zone|region|area', case=False)]
-        if len(zone_col) > 0:
-            zone_col = zone_col[0]
+        zone_col = find_column(beds_df, ['zone', 'region', 'area', 'district'])
+        if zone_col:
             zone_summary = ab_beds.groupby(zone_col).agg({
                 current_col: 'sum',
                 prev_col: 'sum',
@@ -275,8 +311,8 @@ def analyze_bed_capacity(beds_df):
             "columns_used": {
                 "beds_current": current_col,
                 "beds_prev": prev_col,
-                "zone": zone_col[0] if len(zone_col) > 0 else None,
-                "province": province_col[0] if len(province_col) > 0 else None
+                "zone": zone_col,
+                "province": province_col
             }
         }
     except Exception as e:
@@ -287,32 +323,27 @@ def assess_long_term_capacity(facilities_df, beds_df, zone_name):
     """Assess long-term care capacity dynamically."""
     try:
         # Find relevant columns
-        zone_col = facilities_df.columns[facilities_df.columns.str.contains('zone|region|area', case=False)]
-        city_col = facilities_df.columns[facilities_df.columns.str.contains('city|municipality|town', case=False)]
-        type_col = facilities_df.columns[facilities_df.columns.str.contains('type|category|class', case=False)]
+        zone_col = find_column(facilities_df, ['zone', 'region', 'area', 'district'])
+        city_col = find_column(facilities_df, ['city', 'municipality', 'town'])
+        type_col = find_column(facilities_df, ['type', 'category', 'class', 'facility_type', 'odhf_facility_type'])
         
-        if len(type_col) == 0:
+        if not type_col:
             return {"error": "Facility type column not found"}
         
-        type_col = type_col[0]
-        
         # Get facilities in the specified zone
-        if len(zone_col) > 0:
-            zone_col = zone_col[0]
+        if zone_col:
             zone_facilities = facilities_df[facilities_df[zone_col] == zone_name].copy()
         else:
             # If zone column not available, use province
-            province_col = facilities_df.columns[facilities_df.columns.str.contains('province', case=False)]
-            if len(province_col) > 0:
-                province_col = province_col[0]
+            province_col = find_column(facilities_df, ['province', 'state', 'territory'])
+            if province_col:
                 alberta_mask = facilities_df[province_col].str.lower().isin(['alberta', 'ab'])
                 zone_facilities = facilities_df[alberta_mask].copy()
             else:
                 zone_facilities = facilities_df.copy()
         
         # Find major city in zone
-        if len(city_col) > 0:
-            city_col = city_col[0]
+        if city_col:
             city_counts = zone_facilities[city_col].value_counts()
             if len(city_counts) > 0:
                 major_city = city_counts.index[0]
@@ -355,7 +386,8 @@ def generate_operational_recommendations(analysis_results):
     if 'bed_capacity' in analysis_results:
         bed_data = analysis_results['bed_capacity']
         if 'max_percentage_decrease' in bed_data and isinstance(bed_data['max_percentage_decrease'], dict):
-            zone = bed_data['max_percentage_decrease'].get('zone', '')
+            zone_col = bed_data.get('columns_used', {}).get('zone')
+            zone = bed_data['max_percentage_decrease'].get(zone_col, '') if zone_col else ''
             decrease = bed_data['max_percentage_decrease'].get('percent_change', 0)
             if zone and decrease:
                 recommendations.append({
@@ -399,6 +431,9 @@ def format_healthcare_analysis_response(scenario_text, results, recommendations,
     """Format the healthcare analysis response with tables and sections."""
     response = "# Structured Analysis: Healthcare Scenario\n\n"
     
+    # Extract tasks from scenario to ensure we address all requirements
+    tasks = extract_scenario_tasks(scenario_text)
+    
     # Data Preparation Section
     if 'facility_distribution' in results:
         fd = results['facility_distribution']
@@ -442,10 +477,14 @@ def format_healthcare_analysis_response(scenario_text, results, recommendations,
                 response += "| Zone | Beds (Current) | Beds (Previous) | Absolute Change | Percent Change |\n"
                 response += "|------|---------------|-----------------|-----------------|----------------|\n"
                 
+                zone_col = bc.get('columns_used', {}).get('zone')
+                current_col = bc.get('columns_used', {}).get('beds_current')
+                prev_col = bc.get('columns_used', {}).get('beds_prev')
+                
                 for zone_data in bc['zone_summary']:
-                    zone = zone_data.get(bc['columns_used']['zone'], 'N/A') if bc['columns_used'].get('zone') else 'N/A'
-                    current = zone_data.get(bc['columns_used']['beds_current'], 'N/A')
-                    prev = zone_data.get(bc['columns_used']['beds_prev'], 'N/A')
+                    zone = zone_data.get(zone_col, 'N/A') if zone_col else 'N/A'
+                    current = zone_data.get(current_col, 'N/A') if current_col else 'N/A'
+                    prev = zone_data.get(prev_col, 'N/A') if prev_col else 'N/A'
                     change = zone_data.get('bed_change', 'N/A')
                     pct = zone_data.get('percent_change', 'N/A')
                     response += f"| {zone} | {current} | {prev} | {change} | {pct:.1f}% |\n"
@@ -453,19 +492,21 @@ def format_healthcare_analysis_response(scenario_text, results, recommendations,
             
             if 'max_absolute_decrease' in bc and isinstance(bc['max_absolute_decrease'], dict) and \
                'max_percentage_decrease' in bc and isinstance(bc['max_percentage_decrease'], dict):
+                zone_col = bc.get('columns_used', {}).get('zone')
                 abs_dec = bc['max_absolute_decrease']
                 pct_dec = bc['max_percentage_decrease']
-                response += f"**Zone with largest absolute decrease**: {abs_dec.get(bc['columns_used']['zone'], 'N/A')} ({abs_dec.get('bed_change', 'N/A')} beds)\n\n"
-                response += f"**Zone with largest percentage decrease**: {pct_dec.get(bc['columns_used']['zone'], 'N/A')} ({pct_dec.get('percent_change', 'N/A'):.1f}%)\n\n"
+                response += f"**Zone with largest absolute decrease**: {abs_dec.get(zone_col, 'N/A') if zone_col else 'N/A'} ({abs_dec.get('bed_change', 'N/A')} beds)\n\n"
+                response += f"**Zone with largest percentage decrease**: {pct_dec.get(zone_col, 'N/A') if zone_col else 'N/A'} ({pct_dec.get('percent_change', 'N/A'):.1f}%)\n\n"
             
             if 'facilities_with_largest_declines' in bc and bc['facilities_with_largest_declines']:
                 response += "### Facilities with Largest Bed Declines\n\n"
                 response += "| Facility | Zone | Teaching Status | Beds Lost |\n"
                 response += "|----------|------|----------------|-----------|\n"
                 
+                zone_col = bc.get('columns_used', {}).get('zone')
                 for facility in bc['facilities_with_largest_declines']:
                     name = facility.get('facility_name', 'N/A')
-                    zone = facility.get(bc['columns_used']['zone'], 'N/A') if bc['columns_used'].get('zone') else 'N/A'
+                    zone = facility.get(zone_col, 'N/A') if zone_col else 'N/A'
                     teaching = facility.get('teaching_status', 'N/A')
                     change = facility.get('bed_change', 'N/A')
                     response += f"| {name} | {zone} | {teaching} | {change} |\n"
@@ -494,7 +535,7 @@ def format_healthcare_analysis_response(scenario_text, results, recommendations,
                     response += f"- {ftype}: {count}\n"
                 response += "\n"
     
-    # Recommendations Section
+    # Operational Recommendations Section
     response += "## 4. Operational Recommendations\n\n"
     if recommendations:
         for rec in recommendations:
@@ -524,10 +565,14 @@ def format_healthcare_analysis_response(scenario_text, results, recommendations,
     return response
 
 def handle_healthcare_scenario(scenario_text, data_registry, history):
-    """Handle healthcare scenarios dynamically."""
+    """Handle healthcare scenarios dynamically with explicit task following."""
     try:
         results = {}
         
+        # Extract tasks from scenario to ensure we address all requirements
+        tasks = extract_scenario_tasks(scenario_text)
+        print(f"Extracted tasks: {tasks}")
+        
         # Dynamically identify relevant files
         facility_files = data_registry.get_data_by_type('facility_data')
         bed_files = data_registry.get_data_by_type('bed_data')
@@ -549,14 +594,15 @@ def handle_healthcare_scenario(scenario_text, data_registry, history):
             "bed_files": bed_files
         })
         
-        # Perform analyses based on available data
+        # Task 1: Data preparation (facility distribution)
         if facilities_df is not None:
             results['facility_distribution'] = analyze_facility_distribution(facilities_df)
         
+        # Task 2: Bed capacity analysis
         if beds_df is not None:
             results['bed_capacity'] = analyze_bed_capacity(beds_df)
             
-            # Long-term care assessment if we have both data types
+            # Task 3: Long-term care capacity assessment
             if 'bed_capacity' in results and 'max_percentage_decrease' in results['bed_capacity']:
                 zone_col = results['bed_capacity'].get('columns_used', {}).get('zone')
                 if zone_col:
@@ -568,13 +614,13 @@ def handle_healthcare_scenario(scenario_text, data_registry, history):
                             worst_zone
                         )
         
-        # Generate recommendations
+        # Generate operational recommendations (Task 4.1)
         recommendations = generate_operational_recommendations(results)
         
-        # Generate AI integration discussion
+        # Generate AI integration discussion (Task 4.2)
         ai_integration = generate_ai_integration_discussion(results)
         
-        # Format response
+        # Format response ensuring all tasks are addressed
         response = format_healthcare_analysis_response(scenario_text, results, recommendations, ai_integration)
         
         return response
@@ -704,14 +750,14 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths, awaiting_answe
             return history + [(user_msg, ans)], awaiting_answers
 
         if is_identity_query(safe_in, history):
-            ans = "I am an AI analytical system designed to help you analyze healthcare scenarios and make data-driven decisions."
+            ans = "I am an AI analytical system designed to help with both general conversations and healthcare scenario analysis. I can answer your questions and also analyze healthcare data when you upload files and describe a scenario."
             return history + [(user_msg, ans)], awaiting_answers
 
         # Initialize data registry and session RAG
         data_registry = DataRegistry()
         session_rag = SessionRAG()
 
-        # Process uploaded files
+        # Process uploaded files if any
         if uploaded_files_paths:
             process_healthcare_data(uploaded_files_paths, data_registry)
             
@@ -728,16 +774,55 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths, awaiting_answe
                     df = data_registry.get(file_name)
                     session_rag.csv_columns = list(df.columns)
 
-        # Check if this is a healthcare scenario
+        # Determine the mode: healthcare scenario or general conversation
         if is_healthcare_scenario(safe_in, uploaded_files_paths):
-            # Handle healthcare scenario directly
+            # Healthcare scenario mode
             response = handle_healthcare_scenario(safe_in, data_registry, history)
             return history + [(user_msg, response)], False
-
-        # For non-healthcare scenarios, use the original logic
-        # For now, provide a fallback response
-        response = "I can help you analyze this scenario. Please provide more details about what you'd like to analyze."
-        return history + [(user_msg, response)], awaiting_answers
+        else:
+            # General conversation mode
+            # Try Cohere first if available
+            if USE_HOSTED_COHERE:
+                out = cohere_chat(safe_in, history)
+                if out:
+                    out = _sanitize_text(out)
+                    safe_out, blocked_out, reason_out = safety_filter(out, mode="output")
+                    if blocked_out:
+                        safe_out = refusal_reply(reason_out)
+                    log_event("assistant_reply", None, {
+                        **hash_summary("prompt", safe_in if not PERSIST_CONTENT else ""),
+                        **hash_summary("reply", safe_out if not PERSIST_CONTENT else ""),
+                        "mode": "general_cohere",
+                    })
+                    return history + [(user_msg, safe_out)], False
+            
+            # Fall back to local model
+            try:
+                model, tokenizer = load_local_model()
+                inputs = build_inputs(tokenizer, safe_in, history)
+                out = local_generate(model, tokenizer, inputs, max_new_tokens=MAX_NEW_TOKENS)
+                
+                if isinstance(out, str):
+                    for tag in ("Assistant:", "System:", "User:"):
+                        if out.startswith(tag):
+                            out = out[len(tag):].strip()
+                
+                out = _sanitize_text(out or "")
+                safe_out, blocked_out, reason_out = safety_filter(out, mode="output")
+                if blocked_out:
+                    safe_out = refusal_reply(reason_out)
+                
+                log_event("assistant_reply", None, {
+                    **hash_summary("prompt", safe_in if not PERSIST_CONTENT else ""),
+                    **hash_summary("reply", safe_out if not PERSIST_CONTENT else ""),
+                    "mode": "general_local",
+                })
+                
+                return history + [(user_msg, safe_out)], False
+            except Exception as e:
+                err = f"Error generating response: {str(e)}"
+                log_event("model_error", None, {"error": str(e)})
+                return history + [(user_msg, err)], False
 
     except Exception as e:
         err = f"Error: {e}"
@@ -778,16 +863,16 @@ with gr.Blocks(theme=theme, css=custom_css, analytics_enabled=False) as demo:
     # --- HERO (initial screen) ---
     with gr.Column(elem_id="hero-wrap", visible=True) as hero_wrap:
         with gr.Column(elem_id="hero"):
-            gr.HTML("<h2>What healthcare scenario can I help you analyze?</h2>")
+            gr.HTML("<h2>How can I help you today?</h2>")
             with gr.Row(elem_classes="search-row"):
                 hero_msg = gr.Textbox(
-                    placeholder="Describe your healthcare scenario or upload data files for analysis…",
+                    placeholder="Ask me anything or upload healthcare data files for scenario analysis…",
                     show_label=False,
                     lines=1,
                     elem_classes="hero-box"
                 )
                 hero_send = gr.Button("➤", scale=0, elem_id="hero-send")
-            gr.Markdown('<div class="hint">Upload healthcare data files (CSV, Excel, JSON, PDF, etc.) and describe your scenario for comprehensive analysis.</div>')
+            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>')
 
     # --- MAIN APP (hidden until first message) ---
     with gr.Column(elem_id="chat-container", visible=False) as app_wrap:
@@ -802,7 +887,7 @@ with gr.Blocks(theme=theme, css=custom_css, analytics_enabled=False) as demo:
             msg = gr.Textbox(
                 label="",
                 show_label=False,
-                placeholder="Continue the conversation. Provide additional details or answer clarifying questions.",
+                placeholder="Ask me anything or continue your healthcare scenario analysis…",
                 scale=10,
                 elem_id="chat-msg",
                 lines=1,