Update app.py

app.py

@@ -10,11 +10,16 @@ import torch
 import regex as re2
 
 # Import necessary modules
-from settings import SNAPSHOT_PATH, PERSIST_CONTENT, HEALTHCARE_SETTINGS, MODEL_SETTINGS
+from settings import (
+    SNAPSHOT_PATH, PERSIST_CONTENT, HEALTHCARE_SETTINGS, MODEL_SETTINGS,
+    HEALTHCARE_SYSTEM_PROMPT, GENERAL_CONVERSATION_PROMPT
+)
 from audit_log import log_event, hash_summary
 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
 
 # ---------- Writable caches (HF Spaces-safe) ----------
 HOME = pathlib.Path.home()
@@ -57,22 +62,6 @@ 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)))
 
-# ---------- Generic System Prompt ----------
-SYSTEM_MASTER = """
-SYSTEM ROLE
-You are an AI analytical system that provides data-driven insights for any scenario.
-Absolute rules:
-- Use ONLY information provided in this conversation (scenario text + uploaded files + user answers).
-- Never invent data. If something required is missing after clarifications, write the literal token: INSUFFICIENT_DATA.
-- Provide clear analysis with calculations, evidence, and reasoning.
-- Maintain privacy safeguards (aggregate data; suppress small cohorts <10).
-- Adapt your analysis approach to the specific scenario and data provided.
-Formatting rules for structured analysis:
-- Start with the header: "Structured Analysis"
-- Organize analysis into logical sections based on the scenario requirements
-- End with concrete recommendations and a brief "Provenance" mapping outputs to scenario text, uploaded files, and answers.
-""".strip()
-
 # ---------- Helper Functions ----------
 def find_column(df, patterns):
     """Find the first column in df that matches any of the patterns."""
@@ -193,447 +182,18 @@ def process_healthcare_data(uploaded_files_paths, data_registry):
                 "error": str(e)
             })
 
-def analyze_facility_distribution(facilities_df):
-    """Analyze healthcare facility distribution dynamically."""
-    try:
-        # Validate input
-        if facilities_df is None or facilities_df.empty:
-            return {"error": "No facility data provided"}
-        
-        # 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 = find_column(facilities_df, ['type', 'category', 'class', 'facility_type', 'odhf_facility_type'])
-        if not type_col:
-            return {"error": "Facility type column not found"}
-        
-        # Facility type frequency
-        type_counts = ab_facilities[type_col].value_counts().to_dict()
-        
-        # 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()
-            
-            # Breakdown by facility type for top cities
-            city_breakdown = {}
-            for city in top_cities:
-                city_data = ab_facilities[ab_facilities[city_col] == city]
-                city_breakdown[city] = city_data[type_col].value_counts().to_dict()
-        else:
-            top_cities = []
-            city_breakdown = {}
-        
-        return {
-            "total_facilities": len(ab_facilities),
-            "type_distribution": type_counts,
-            "top_cities": top_cities,
-            "city_breakdown": city_breakdown,
-            "columns_used": {
-                "facility_type": type_col,
-                "city": city_col,
-                "province": province_col
-            }
-        }
-    except Exception as e:
-        log_event("facility_analysis_error", None, {"error": str(e)})
-        return {"error": str(e)}
-
-def analyze_bed_capacity(beds_df):
-    """Analyze bed capacity dynamically."""
-    try:
-        # Validate input
-        if beds_df is None or beds_df.empty:
-            return {"error": "No bed data provided"}
-        
-        # Find required columns
-        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'])
-        
-        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:
-            beds_df['bed_change'] = beds_df[current_col] - beds_df[prev_col]
-        
-        if 'percent_change' not in beds_df.columns:
-            beds_df['percent_change'] = beds_df.apply(
-                lambda row: (row['bed_change'] / row[prev_col] * 100) if row[prev_col] != 0 else 0, 
-                axis=1
-            )
-        
-        # Filter to Alberta if province column exists
-        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 = 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',
-                'bed_change': 'sum'
-            }).reset_index()
-            
-            zone_summary['percent_change'] = zone_summary.apply(
-                lambda row: (row['bed_change'] / row[prev_col] * 100) if row[prev_col] != 0 else 0, 
-                axis=1
-            )
-            
-            # Find zones with largest changes
-            if len(zone_summary) > 0:
-                max_abs_decrease_idx = zone_summary['bed_change'].idxmin()
-                max_pct_decrease_idx = zone_summary['percent_change'].idxmin()
-                
-                max_abs_decrease = zone_summary.loc[max_abs_decrease_idx]
-                max_pct_decrease = zone_summary.loc[max_pct_decrease_idx]
-            else:
-                max_abs_decrease = {}
-                max_pct_decrease = {}
-            
-            # Identify facilities with largest declines
-            facilities_decline = ab_beds.sort_values('bed_change').head(5)
-        else:
-            zone_summary = pd.DataFrame()
-            max_abs_decrease = {}
-            max_pct_decrease = {}
-            facilities_decline = pd.DataFrame()
-        
-        return {
-            "zone_summary": zone_summary.to_dict('records') if not zone_summary.empty else [],
-            "max_absolute_decrease": max_abs_decrease.to_dict() if isinstance(max_abs_decrease, pd.Series) else max_abs_decrease,
-            "max_percentage_decrease": max_pct_decrease.to_dict() if isinstance(max_pct_decrease, pd.Series) else max_pct_decrease,
-            "facilities_with_largest_declines": facilities_decline.to_dict('records') if not facilities_decline.empty else [],
-            "columns_used": {
-                "beds_current": current_col,
-                "beds_prev": prev_col,
-                "zone": zone_col,
-                "province": province_col
-            }
-        }
-    except Exception as e:
-        log_event("bed_analysis_error", None, {"error": str(e)})
-        return {"error": str(e)}
-
-def assess_long_term_capacity(facilities_df, beds_df, zone_name):
-    """Assess long-term care capacity dynamically."""
-    try:
-        # Validate inputs
-        if facilities_df is None or facilities_df.empty:
-            return {"error": "No facility data provided"}
-        
-        # Find relevant columns
-        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 not type_col:
-            return {"error": "Facility type column not found"}
-        
-        # Get facilities in the specified zone
-        if zone_col:
-            zone_facilities = facilities_df[facilities_df[zone_col] == zone_name].copy()
-        else:
-            # If zone column not available, use province
-            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 city_col:
-            city_counts = zone_facilities[city_col].value_counts()
-            if len(city_counts) > 0:
-                major_city = city_counts.index[0]
-                city_facilities = zone_facilities[zone_facilities[city_col] == major_city]
-                
-                # Count facility types
-                facility_counts = city_facilities[type_col].value_counts().to_dict()
-                
-                # Calculate ratio of nursing/residential to hospitals
-                hospitals = sum(count for key, count in facility_counts.items() if 'hospital' in key.lower())
-                nursing = sum(count for key, count in facility_counts.items() if any(word in key.lower() for word in ['nursing', 'residential', 'care']))
-                ratio = nursing / hospitals if hospitals > 0 else 0
-                
-                # Assess capacity
-                capacity_assessment = "sufficient" if ratio >= 1.5 else "insufficient"
-                
-                return {
-                    "zone": zone_name,
-                    "major_city": major_city,
-                    "facility_counts": facility_counts,
-                    "nursing_to_hospital_ratio": ratio,
-                    "capacity_assessment": capacity_assessment,
-                    "columns_used": {
-                        "zone": zone_col,
-                        "city": city_col,
-                        "facility_type": type_col
-                    }
-                }
-        
-        return {"error": "Could not determine major city or facility counts"}
-    except Exception as e:
-        log_event("ltc_assessment_error", None, {"error": str(e)})
-        return {"error": str(e)}
-
-def generate_operational_recommendations(analysis_results):
-    """Generate data-driven operational recommendations."""
-    recommendations = []
-    
-    # Recommendation 1: Address bed capacity issues
-    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_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({
-                    "title": f"Restore staffed beds in {zone} Zone",
-                    "description": f"Priority should be given to reopening closed units and hiring staff to address the {decrease:.1f}% decrease in bed capacity.",
-                    "data_source": "Bed capacity analysis"
-                })
-    
-    # Recommendation 2: Expand long-term care capacity
-    if 'long_term_care' in analysis_results:
-        ltc_data = analysis_results['long_term_care']
-        if ltc_data.get('capacity_assessment') == 'insufficient':
-            city = ltc_data.get('major_city', '')
-            if city:
-                recommendations.append({
-                    "title": f"Expand long-term care capacity in {city}",
-                    "description": f"Invest in new long-term care beds or repurpose existing sites to expedite discharge of stabilized patients.",
-                    "data_source": "Long-term care capacity assessment"
-                })
-    
-    # Recommendation 3: Implement surge plans
-    if 'bed_capacity' in analysis_results:
-        recommendations.append({
-            "title": "Implement surge capacity plans",
-            "description": "Develop modular units and activate staffing pools to handle unpredictable spikes in demand.",
-            "data_source": "Bed capacity trends"
-        })
-    
-    return recommendations
-
-def generate_ai_integration_discussion(analysis_results):
-    """Generate discussion on future AI integration for healthcare operations."""
-    return {
-        "title": "Future Integration for Augmented Decision-Making",
-        "description": "Combining facility information with operational data like emergency department wait times and disease surveillance can enable AI-driven resource optimization.",
-        "example": "A model could ingest current ED wait times, hospital occupancy, and community case counts to forecast bed demand by zone and recommend redirecting ambulances to facilities with spare capacity.",
-        "metrics": ["Hospital occupancy rates", "ED wait times", "Disease surveillance data"]
-    }
-
-def format_healthcare_analysis_response(scenario_text, results, recommendations, ai_integration):
-    """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']
-        if 'error' in fd:
-            response += "## 1. Data Preparation\n\n"
-            response += f"Error in facility distribution analysis: {fd['error']}\n\n"
-        else:
-            response += "## 1. Data Preparation\n\n"
-            response += f"Total healthcare facilities: {fd.get('total_facilities', 'N/A')}\n\n"
-            
-            if 'type_distribution' in fd and isinstance(fd['type_distribution'], dict):
-                response += "### Facility Type Distribution\n\n"
-                for ftype, count in fd['type_distribution'].items():
-                    response += f"- {ftype}: {count}\n"
-                response += "\n"
-            
-            if 'city_breakdown' in fd and isinstance(fd['city_breakdown'], dict):
-                response += "### Top Cities by Facility Count\n\n"
-                response += "| City | Hospitals | Nursing/Residential | Ambulatory | Total |\n"
-                response += "|------|-----------|-------------------|------------|-------|\n"
-                
-                for city, breakdown in fd['city_breakdown'].items():
-                    hospitals = breakdown.get('Hospitals', 0)
-                    nursing = breakdown.get('Nursing and residential care facilities', 0)
-                    ambulatory = breakdown.get('Ambulatory health care services', 0)
-                    total = hospitals + nursing + ambulatory
-                    response += f"| {city} | {hospitals} | {nursing} | {ambulatory} | {total} |\n"
-                response += "\n"
-    
-    # Bed Capacity Analysis Section
-    if 'bed_capacity' in results:
-        bc = results['bed_capacity']
-        if 'error' in bc:
-            response += "## 2. Bed Capacity Analysis\n\n"
-            response += f"Error in bed capacity analysis: {bc['error']}\n\n"
-        else:
-            response += "## 2. Bed Capacity Analysis\n\n"
-            
-            if 'zone_summary' in bc and bc['zone_summary']:
-                response += "### Bed Capacity by Zone\n\n"
-                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(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"
-                response += "\n"
-            
-            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(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(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"
-                response += "\n"
-    
-    # Long-term Care Section
-    if 'long_term_care' in results:
-        ltc = results['long_term_care']
-        if 'error' in ltc:
-            response += "## 3. Long-Term Care Capacity Assessment\n\n"
-            response += f"Error in long-term care assessment: {ltc['error']}\n\n"
-        else:
-            response += "## 3. Long-Term Care Capacity Assessment\n\n"
-            
-            zone = ltc.get('zone', 'N/A')
-            city = ltc.get('major_city', 'N/A')
-            ratio = ltc.get('nursing_to_hospital_ratio', 0)
-            assessment = ltc.get('capacity_assessment', 'N/A')
-            
-            response += f"In {zone} Zone, the major city is {city} with a nursing/residential to hospital ratio of {ratio:.2f}.\n\n"
-            response += f"Long-term care capacity appears **{assessment}** in {city}.\n\n"
-            
-            if 'facility_counts' in ltc and isinstance(ltc['facility_counts'], dict):
-                response += "### Facility Counts\n\n"
-                for ftype, count in ltc['facility_counts'].items():
-                    response += f"- {ftype}: {count}\n"
-                response += "\n"
-    
-    # Operational Recommendations Section
-    response += "## 4. Operational Recommendations\n\n"
-    if recommendations:
-        for rec in recommendations:
-            response += f"### {rec['title']}\n\n"
-            response += f"{rec['description']}\n\n"
-            response += f"*Data source: {rec['data_source']}*\n\n"
-    else:
-        response += "No specific recommendations could be generated due to data limitations.\n\n"
-    
-    # AI Integration Section
-    response += "## 5. Future Integration for Augmented AI\n\n"
-    response += f"### {ai_integration['title']}\n\n"
-    response += f"{ai_integration['description']}\n\n"
-    response += f"**Example**: {ai_integration['example']}\n\n"
-    response += "**Key metrics to incorporate**:\n"
-    for metric in ai_integration['metrics']:
-        response += f"- {metric}\n"
-    response += "\n"
-    
-    # Provenance Section
-    response += "## Provenance\n\n"
-    response += "This analysis is based on:\n"
-    response += "- Scenario description provided by the user\n"
-    response += "- Uploaded data files\n"
-    response += "- Calculations performed on the provided data\n"
-    
-    return response
-
 def handle_healthcare_scenario(scenario_text, data_registry, history):
-    """Handle healthcare scenarios dynamically with explicit task following."""
+    """Handle healthcare scenarios with enhanced analysis"""
     try:
-        results = {}
+        # Initialize analyzer
+        analyzer = HealthcareAnalyzer(data_registry)
         
-        # Extract tasks from scenario to ensure we address all requirements
-        tasks = extract_scenario_tasks(scenario_text)
-        print(f"Extracted tasks: {tasks}")
+        # Perform comprehensive analysis
+        results = analyzer.comprehensive_analysis(scenario_text)
         
-        # Dynamically identify relevant files
-        facility_files = data_registry.get_data_by_type('facility_data')
-        bed_files = data_registry.get_data_by_type('bed_data')
-        
-        # Use the first file of each type (can be enhanced to use multiple)
-        facilities_df = None
-        if facility_files:
-            facilities_df = data_registry.get(facility_files[0])
-        
-        beds_df = None
-        if bed_files:
-            beds_df = data_registry.get(bed_files[0])
-        
-        # Log what we found
-        log_event("data_files_found", None, {
-            "facilities": facilities_df is not None,
-            "beds": beds_df is not None,
-            "facility_files": facility_files,
-            "bed_files": bed_files
-        })
-        
-        # 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)
-            
-            # 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:
-                    worst_zone = results['bed_capacity']['max_percentage_decrease'].get(zone_col, '')
-                    if worst_zone and facilities_df is not None:
-                        results['long_term_care'] = assess_long_term_capacity(
-                            facilities_df, 
-                            beds_df, 
-                            worst_zone
-                        )
-        
-        # Generate operational recommendations (Task 4.1)
-        recommendations = generate_operational_recommendations(results)
-        
-        # Generate AI integration discussion (Task 4.2)
-        ai_integration = generate_ai_integration_discussion(results)
-        
-        # Format response ensuring all tasks are addressed
-        response = format_healthcare_analysis_response(scenario_text, results, recommendations, ai_integration)
+        # Format response
+        formatter = ResponseFormatter()
+        response = formatter.format_healthcare_response(scenario_text, results)
         
         return response
     except Exception as e:
@@ -727,8 +287,8 @@ def cohere_chat(message, history):
     except Exception:
         return None
 
-def build_inputs(tokenizer, message, history):
-    msgs = [{"role": "system", "content": SYSTEM_MASTER}]
+def build_inputs(tokenizer, message, history, system_prompt):
+    msgs = [{"role": "system", "content": system_prompt}]
     for u, a in _iter_user_assistant(history):
         if u: msgs.append({"role": "user", "content": u})
         if a: msgs.append({"role": "assistant", "content": a})
@@ -792,8 +352,7 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths, awaiting_answe
             response = handle_healthcare_scenario(safe_in, data_registry, history)
             return history + [(user_msg, response)], False
         else:
-            # General conversation mode
-            # Try Cohere first if available
+            # General conversation mode with enhanced handling
             if USE_HOSTED_COHERE:
                 out = cohere_chat(safe_in, history)
                 if out:
@@ -808,10 +367,12 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths, awaiting_answe
                     })
                     return history + [(user_msg, safe_out)], False
             
-            # Fall back to local model
+            # Enhanced local model generation
             try:
                 model, tokenizer = load_local_model()
-                inputs = build_inputs(tokenizer, safe_in, history)
+                
+                # Use general conversation prompt
+                inputs = build_inputs(tokenizer, safe_in, history, GENERAL_CONVERSATION_PROMPT)
                 out = local_generate(model, tokenizer, inputs, max_new_tokens=MAX_NEW_TOKENS)
                 
                 if isinstance(out, str):