Update schema_mapper.py

schema_mapper.py

@@ -1,411 +1,128 @@
-from __future__ import annotations
+# schema_mapper.py
+from typing import Dict, List, Any, Set
 import re
-from dataclasses import dataclass, field
-from typing import Dict, List, Any, Tuple, Optional, Set
-import pandas as pd
-from data_registry import DataRegistry
 
-# Generic concept patterns that work across domains
-UNIVERSAL_CONCEPT_PATTERNS = {
-    # Entity/grouping concepts
-    "facility": [r"\bfacilit(y|ies)\b", r"\bhospital\b", r"\bsite\b", r"\bcentre\b", r"\bcenter\b", r"\blocation\b", r"\bprovider\b"],
-    "organization": [r"\borganization\b", r"\bcompany\b", r"\bbusiness\b", r"\bfirm\b", r"\bentity\b"],
-    "department": [r"\bdepartment\b", r"\bdivision\b", r"\bunit\b", r"\bsection\b"],
-    "specialty": [r"\bspecialt(y|ies)\b", r"\bservice\b", r"\btype\b", r"\bcategory\b", r"\bkind\b"],
-    "region": [r"\bzone\b", r"\bregion\b", r"\barea\b", r"\bdistrict\b", r"\bterritory\b"],
-    
-    # Time-based metrics
-    "wait_time": [r"\bwait", r"\bdelay", r"\btime", r"\bduration", r"\blength"],
-    "wait_median": [r"\bmedian\b.*\bwait", r"\bP50\b", r"\bwait.*\bmedian", r"median.*time"],
-    "wait_p90": [r"\bp90\b", r"\b90(th)?\s*percentile\b", r"\bwait.*p90", r"90.*wait"],
-    "response_time": [r"\bresponse\b.*\btime\b", r"\bprocessing\b.*\btime\b"],
-    
-    # Performance metrics
-    "score": [r"\bscore\b", r"\brating\b", r"\bindex\b", r"\brank\b"],
-    "efficiency": [r"\befficiency\b", r"\bthroughput\b", r"\bproductivity\b"],
-    "quality": [r"\bquality\b", r"\bperformance\b", r"\boutcome\b"],
-    "satisfaction": [r"\bsatisfaction\b", r"\bfeedback\b", r"\brating\b"],
+# Healthcare terminology mappings
+HEALTHCARE_CONCEPTS = {
+    # Facility types
+    "hospital": ["hospital", "medical center", "health centre", "clinic"],
+    "nursing_facility": ["nursing home", "long-term care", "residential care", "care facility"],
+    "ambulatory_care": ["ambulatory", "outpatient", "clinic", "surgery center"],
     
     # Capacity metrics
-    "capacity": [r"\bcapacity\b", r"\bvolume\b", r"\bsize\b", r"\blimit\b"],
-    "utilization": [r"\butilization\b", r"\boccupancy\b", r"\busage\b"],
-    "availability": [r"\bavailab\w+", r"\bopen\b", r"\bfree\b"],
-    
-    # Cost/financial metrics
-    "cost": [r"\bcost\b", r"\bprice\b", r"\bexpense\b", r"\bfee\b", r"\bcharge\b"],
-    "budget": [r"\bbudget\b", r"\bfunding\b", r"\ballocation\b"],
-    "revenue": [r"\brevenue\b", r"\bincome\b", r"\bearnings\b"],
-    
-    # Count/volume metrics
-    "count": [r"\bcount\b", r"\bnumber\b", r"\bquantity\b", r"\btotal\b"],
-    "rate": [r"\brate\b", r"\bratio\b", r"\bpercent\b", r"\bfrequency\b"],
-    "volume": [r"\bvolume\b", r"\bamount\b", r"\bquantity\b"]
+    "bed_capacity": ["beds", "capacity", "bed count", "staffed beds"],
+    "occupancy_rate": ["occupancy", "utilization", "bed occupancy"],
+    
+    # Geographic terms
+    "zone": ["zone", "region", "area", "district"],
+    "province": ["province", "state", "territory"],
+    
+    # Time periods
+    "fiscal_year": ["fiscal year", "fy", "financial year"],
+    "current_period": ["current", "2023-24", "present", "latest"],
+    "previous_period": ["previous", "2022-23", "past", "last"],
+    
+    # Healthcare operations
+    "patient_flow": ["patient flow", "throughput", "patient movement"],
+    "resource_allocation": ["resource allocation", "staffing", "resource distribution"],
+    "surge_capacity": ["surge", "overflow", "emergency capacity"],
 }
 
-def _extract_key_terms_from_scenario(scenario_text: str) -> Set[str]:
-    """Extract important terms from scenario text to guide concept detection."""
-    if not scenario_text:
-        return set()
-    
-    # Extract meaningful words, filtering out common stop words
-    stop_words = {
-        'the', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by',
-        'is', 'are', 'was', 'were', 'be', 'been', 'have', 'has', 'had', 'do', 'does', 'did',
-        'a', 'an', 'this', 'that', 'these', 'those', 'i', 'you', 'he', 'she', 'it', 'we', 'they'
-    }
-    
-    words = re.findall(r'\b[a-zA-Z]{3,}\b', scenario_text.lower())
-    key_terms = {word for word in words if word not in stop_words}
-    
-    return key_terms
-
-def _generate_dynamic_patterns(scenario_terms: Set[str], existing_patterns: Dict[str, List[str]]) -> Dict[str, List[str]]:
-    """Generate additional concept patterns based on scenario content."""
-    dynamic_patterns = existing_patterns.copy()
-    
-    # Add scenario-specific terms as potential concepts
-    for term in scenario_terms:
-        if len(term) >= 4:  # Only meaningful terms
-            # Check if term relates to existing concepts
-            term_pattern = rf"\b{re.escape(term)}\b"
-            
-            # Add as potential entity if it sounds like one
-            if any(indicator in term for indicator in ['hospital', 'clinic', 'school', 'department', 'facility']):
-                if 'facility' not in dynamic_patterns:
-                    dynamic_patterns['facility'] = []
-                dynamic_patterns['facility'].append(term_pattern)
-            
-            # Add as potential metric if it sounds like one
-            elif any(indicator in term for indicator in ['time', 'score', 'rate', 'cost', 'wait']):
-                concept_key = f"metric_{term}"
-                dynamic_patterns[concept_key] = [term_pattern]
-    
-    return dynamic_patterns
-
-def _score_column_match(col_name: str, patterns: List[str], scenario_terms: Set[str] = None) -> int:
-    """Score how well a column matches concept patterns."""
-    col_lower = col_name.lower()
-    score = 0
-    
-    # Pattern matching
-    for i, pattern in enumerate(patterns):
-        if re.search(pattern, col_lower):
-            score += 100 - (i * 10)  # Higher score for earlier patterns
-            break
-    
-    # Boost score if column name contains scenario-relevant terms
-    if scenario_terms:
-        for term in scenario_terms:
-            if term in col_lower:
-                score += 25
-    
-    return score
-
-def _detect_column_types(df: pd.DataFrame) -> Dict[str, str]:
-    """Detect the likely type/purpose of each column."""
-    column_types = {}
-    
-    for col in df.columns:
-        col_lower = col.lower()
-        
-        # Detect numeric columns that could be converted
-        sample = df[col].dropna().head(50)
-        numeric_convertible = False
-        if len(sample) > 0:
-            try:
-                numeric_sample = pd.to_numeric(sample, errors='coerce')
-                if numeric_sample.notna().sum() > len(sample) * 0.7:
-                    numeric_convertible = True
-            except:
-                pass
-        
-        # Categorize columns
-        if numeric_convertible:
-            if any(term in col_lower for term in ['id', 'number', 'code', 'index']):
-                column_types[col] = 'identifier'
-            elif any(term in col_lower for term in ['time', 'date', 'duration', 'wait', 'delay']):
-                column_types[col] = 'time_metric'
-            elif any(term in col_lower for term in ['cost', 'price', 'budget', 'fee', 'expense']):
-                column_types[col] = 'cost_metric'
-            elif any(term in col_lower for term in ['count', 'number', 'quantity', 'volume']):
-                column_types[col] = 'count_metric'
-            elif any(term in col_lower for term in ['rate', 'ratio', 'percent', 'score']):
-                column_types[col] = 'performance_metric'
-            else:
-                column_types[col] = 'numeric_metric'
-        else:
-            # String/categorical columns
-            unique_ratio = df[col].nunique() / len(df)
-            if unique_ratio < 0.1:
-                column_types[col] = 'category'
-            elif unique_ratio < 0.5:
-                column_types[col] = 'grouping'
-            else:
-                column_types[col] = 'text'
-    
-    return column_types
-
-@dataclass
 class MappingResult:
-    resolved: Dict[str, Tuple[str, str]] = field(default_factory=dict)
-    ambiguous: Dict[str, List[Tuple[str, str]]] = field(default_factory=dict)
-    missing: List[str] = field(default_factory=list)
-    discovered: Dict[str, str] = field(default_factory=dict)  # Discovered column types
+    def __init__(self):
+        self.resolved = {}  # Successfully mapped concepts
+        self.ambiguous = {}  # Concepts with multiple possible mappings
+        self.missing = set()  # Concepts that couldn't be mapped
 
-def _extract_explicit_mappings_from_scenario(scenario_text: str, available_columns: List[Tuple[str, str]]) -> Dict[str, Tuple[str, str]]:
-    """Extract explicit column mappings from scenario text."""
-    explicit_mappings = {}
-    
-    if not scenario_text:
-        return explicit_mappings
-    
-    scenario_lower = scenario_text.lower()
-    
-    # Create a lookup of available columns (case-insensitive)
-    column_lookup = {}
-    for table_name, col_name in available_columns:
-        column_lookup[col_name.lower()] = (table_name, col_name)
-    
-    # Pattern 1: Direct column descriptions like "Surgery_Median column contains..."
-    column_desc_patterns = [
-        r'(\w+)\s+column\s+(?:contains|reports|shows|includes|represents)',
-        r'column\s+(\w+)\s+(?:contains|reports|shows|includes|represents)',
-        r'(\w+)\s+(?:contains|reports|shows|includes|represents)'
-    ]
-    
-    for pattern in column_desc_patterns:
-        matches = re.findall(pattern, scenario_text, re.IGNORECASE)
-        for match in matches:
-            col_name = match.lower()
-            if col_name in column_lookup:
-                # Determine the concept based on context around the column name
-                context = scenario_text[max(0, scenario_text.lower().find(col_name)-50):scenario_text.lower().find(col_name)+100].lower()
-                
-                if any(term in context for term in ['wait', 'time', 'delay', 'duration']):
-                    if 'median' in col_name:
-                        explicit_mappings['wait_median'] = column_lookup[col_name]
-                    elif '90' in col_name or 'percentile' in col_name:
-                        explicit_mappings['wait_p90'] = column_lookup[col_name]
-                    else:
-                        explicit_mappings['wait_time'] = column_lookup[col_name]
-                
-                elif any(term in context for term in ['facility', 'hospital', 'clinic', 'site']):
-                    explicit_mappings['facility'] = column_lookup[col_name]
-                
-                elif any(term in context for term in ['specialty', 'service', 'department']):
-                    explicit_mappings['specialty'] = column_lookup[col_name]
-                
-                elif any(term in context for term in ['zone', 'region', 'area', 'district']):
-                    explicit_mappings['region'] = column_lookup[col_name]
-    
-    # Pattern 2: Task-based column identification like "calculate average for each facility"
-    task_patterns = [
-        (r'(?:for each|by)\s+(\w+)', ['facility', 'specialty', 'region']),
-        (r'(?:identify|rank|list)\s+(\w+)', ['facility', 'specialty', 'region']),
-        (r'average\s+(\w+)\s+(?:wait|time)', ['wait_median', 'wait_time']),
-        (r'median\s+(\w+)', ['wait_median']),
-        (r'90th\s+percentile\s+(\w+)', ['wait_p90'])
-    ]
-    
-    for pattern, concepts in task_patterns:
-        matches = re.findall(pattern, scenario_lower)
-        for match in matches:
-            match_lower = match.lower()
-            if match_lower in column_lookup:
-                for concept in concepts:
-                    if concept not in explicit_mappings:
-                        explicit_mappings[concept] = column_lookup[match_lower]
-                        break
-    
-    # Pattern 3: Direct column name matches from scenario
-    explicit_columns = re.findall(r'\b([A-Za-z_][A-Za-z0-9_]*)\b', scenario_text)
-    for col_candidate in explicit_columns:
-        col_lower = col_candidate.lower()
-        if col_lower in column_lookup:
-            # Smart concept assignment based on column name patterns
-            if not any(concept in explicit_mappings for concept in ['facility', 'organization', 'department']):
-                if re.search(r'facility|hospital|clinic|site|provider', col_lower):
-                    explicit_mappings['facility'] = column_lookup[col_lower]
-            
-            if not any(concept in explicit_mappings for concept in ['specialty', 'service']):
-                if re.search(r'specialty|service|department|type', col_lower):
-                    explicit_mappings['specialty'] = column_lookup[col_lower]
-            
-            if not any(concept in explicit_mappings for concept in ['region', 'zone']):
-                if re.search(r'zone|region|area|district', col_lower):
-                    explicit_mappings['region'] = column_lookup[col_lower]
-            
-            if not any(concept in explicit_mappings for concept in ['wait_median', 'wait_time']):
-                if re.search(r'.*median.*', col_lower) and re.search(r'wait|time|surgery|consult', col_lower):
-                    explicit_mappings['wait_median'] = column_lookup[col_lower]
-            
-            if not any(concept in explicit_mappings for concept in ['wait_p90']):
-                if re.search(r'.*(90|percentile).*', col_lower) and re.search(r'wait|time|surgery|consult', col_lower):
-                    explicit_mappings['wait_p90'] = column_lookup[col_lower]
-    
-    return explicit_mappings
-
-def _extract_explicit_tasks_from_scenario(scenario_text: str) -> List[str]:
-    """Extract explicit task requirements from scenario text."""
-    tasks = []
-    
-    if not scenario_text:
-        return tasks
-    
-    scenario_lower = scenario_text.lower()
-    
-    # Task extraction patterns
-    task_patterns = [
-        r'(?:your tasks?(?:\s+are)?[:\s]+)([^.]*?)(?:\.|$)',
-        r'(?:you (?:should|need to|are to|must)[:\s]+)([^.]*?)(?:\.|$)',
-        r'(?:tasks?[:\s]+)([^.]*?)(?:\.|deliverables|$)',
-        r'(?:\d+\.?\s*)([^.]*?)(?:\.|$)'  # Numbered tasks
-    ]
-    
-    for pattern in task_patterns:
-        matches = re.findall(pattern, scenario_text, re.IGNORECASE | re.DOTALL)
-        for match in matches:
-            task = match.strip()
-            if len(task) > 10 and any(verb in task.lower() for verb in ['identify', 'calculate', 'analyze', 'compare', 'assess', 'determine', 'rank', 'list']):
-                tasks.append(task)
-    
-    return tasks
-
-def map_concepts(scenario_text: str, registry: DataRegistry) -> MappingResult:
-    """Enhanced mapping that extracts explicit information from scenario text."""
+def map_concepts(scenario_text: str, data_registry) -> MappingResult:
+    """Map healthcare concepts from scenario text to data registry."""
     result = MappingResult()
     
-    if not registry.names():
-        result.missing = list(UNIVERSAL_CONCEPT_PATTERNS.keys())
-        return result
-
     # Extract key terms from scenario
-    scenario_terms = _extract_key_terms_from_scenario(scenario_text)
-    
-    # Collect all available columns
-    all_columns = []
-    for table in registry.iter_tables():
-        # Detect column types for this table
-        column_types = _detect_column_types(table.df)
-        result.discovered.update({f"{table.name}.{col}": col_type for col, col_type in column_types.items()})
-        
-        for col in table.df.columns:
-            all_columns.append((table.name, str(col)))
-
-    # STEP 1: Extract explicit mappings from scenario text
-    explicit_mappings = _extract_explicit_mappings_from_scenario(scenario_text, all_columns)
-    
-    # STEP 2: Use explicit mappings first
-    for concept, (table_name, col_name) in explicit_mappings.items():
-        result.resolved[concept] = (table_name, col_name)
-
-    # STEP 3: For unmapped concepts, use pattern matching with scenario context
-    remaining_patterns = {k: v for k, v in UNIVERSAL_CONCEPT_PATTERNS.items() if k not in result.resolved}
+    scenario_lower = scenario_text.lower()
     
-    if remaining_patterns:
-        # Generate dynamic patterns based on scenario
-        concept_patterns = _generate_dynamic_patterns(scenario_terms, remaining_patterns)
+    # Check for healthcare concepts
+    for concept, synonyms in HEALTHCARE_CONCEPTS.items():
+        # Check if any synonym appears in the scenario
+        found_synonyms = [syn for syn in synonyms if syn in scenario_lower]
         
-        # Map remaining concepts to columns
-        for concept, patterns in concept_patterns.items():
-            if concept in result.resolved:
-                continue  # Skip already resolved
-                
-            scores = [
-                ((tbl, col), _score_column_match(col, patterns, scenario_terms)) 
-                for (tbl, col) in all_columns
-            ]
-            
-            scores.sort(key=lambda x: x[1], reverse=True)
-            
-            if not scores or scores[0][1] == 0:
-                result.missing.append(concept)
-                continue
-            
-            top_score = scores[0][1]
-            
-            # Find all columns with similar high scores (potential ambiguity)
-            threshold = max(70, top_score - 15)  # Higher threshold for explicit scenarios
-            high_scoring = [pair for pair, score in scores if score >= threshold]
-            
-            if len(high_scoring) == 1:
-                tbl, col = high_scoring[0]
-                result.resolved[concept] = (tbl, col)
+        if found_synonyms:
+            # Try to map to data registry
+            mapped_to = _map_to_data_registry(concept, data_registry)
+            if mapped_to:
+                result.resolved[concept] = mapped_to
             else:
-                # Check if scenario text makes disambiguation obvious
-                disambiguated = False
-                for (tbl, col), score in scores[:3]:  # Check top 3
-                    col_mentioned = col.lower() in scenario_text.lower()
-                    if col_mentioned and score >= threshold:
-                        result.resolved[concept] = (tbl, col)
-                        disambiguated = True
-                        break
-                
-                if not disambiguated:
-                    result.ambiguous[concept] = high_scoring[:3]  # Limit to top 3
+                result.missing.add(concept)
+    
+    # Additional mapping for specific healthcare patterns
+    # Check for facility distribution patterns
+    if any(phrase in scenario_lower for phrase in ["facility distribution", "facility count", "number of facilities"]):
+        if any("facility" in name.lower() for name in data_registry.names()):
+            result.resolved["facility_distribution"] = next(
+                (name for name in data_registry.names() if "facility" in name.lower()), None
+            )
+        else:
+            result.missing.add("facility_distribution")
+    
+    # Check for bed capacity patterns
+    if any(phrase in scenario_lower for phrase in ["bed capacity", "bed count", "staffed beds"]):
+        if any("bed" in name.lower() for name in data_registry.names()):
+            result.resolved["bed_capacity"] = next(
+                (name for name in data_registry.names() if "bed" in name.lower()), None
+            )
+        else:
+            result.missing.add("bed_capacity")
+    
+    # Check for long-term care patterns
+    if any(phrase in scenario_lower for phrase in ["long-term care", "ltc", "nursing capacity"]):
+        result.resolved["long_term_care"] = "facility_distribution"  # Usually in facility data
     
     return result
 
-def build_phase1_questions(scenario_text: str, registry: DataRegistry, mapping: MappingResult, max_questions: int = 4) -> str:
-    """Build minimal clarifying questions, only when truly necessary."""
-    
-    # Extract explicit tasks from scenario
-    explicit_tasks = _extract_explicit_tasks_from_scenario(scenario_text)
-    
-    # Check if scenario provides comprehensive instructions
-    has_detailed_tasks = len(explicit_tasks) >= 3
-    has_data_descriptions = any(term in scenario_text.lower() for term in [
-        'column', 'dataset', 'file', 'csv', 'records', 'contains', 'includes'
-    ])
-    
-    # If scenario is comprehensive, minimize questions
-    if has_detailed_tasks and has_data_descriptions:
-        # Only ask about truly ambiguous mappings where scenario doesn't clarify
-        critical_questions = []
-        
-        # Only ask about ambiguities that can't be resolved from context
-        for concept, options in mapping.ambiguous.items():
-            if len(options) > 1:
-                # Check if scenario text clearly indicates which column to use
-                scenario_lower = scenario_text.lower()
-                clear_preference = None
-                
-                for table_name, col_name in options:
-                    if col_name.lower() in scenario_lower:
-                        mentions = scenario_lower.count(col_name.lower())
-                        if mentions > 0:
-                            clear_preference = f"{table_name}.{col_name}"
-                            break
-                
-                if not clear_preference and len(critical_questions) < max_questions:
-                    option_strs = [f"{tbl}.{col}" for tbl, col in options[:3]]
-                    critical_questions.append(f"**Column Clarification**: For {concept.replace('_', ' ')}, use: {', '.join(option_strs)}?")
-        
-        if not critical_questions:
-            return "**Proceeding with Analysis**: Scenario and data mappings are clear. Analyzing now..."
-        
-        return "**Quick Clarification**\n\n" + "\n".join(critical_questions)
-    
-    # Fallback to standard question generation for less comprehensive scenarios
+def _map_to_data_registry(concept: str, data_registry) -> Any:
+    """Helper to map a concept to the data registry."""
+    file_names = data_registry.names()
+    
+    if concept in ["hospital", "facility_distribution", "long_term_care"]:
+        return next((name for name in file_names if "facility" in name.lower() or "health" in name.lower()), None)
+    elif concept == "bed_capacity":
+        return next((name for name in file_names if "bed" in name.lower()), None)
+    elif concept == "zone":
+        # Check if any dataframe has a 'zone' column
+        for name in file_names:
+            df = data_registry.get(name)
+            if df is not None and 'zone' in df.columns:
+                return name
+        return None
+    
+    return None
+
+def build_phase1_questions(scenario_text: str, registry, mapping: MappingResult) -> str:
+    """Build clarifying questions based on mapping results."""
     questions = []
-    scenario_lower = scenario_text.lower() if scenario_text else ""
     
-    # Ambiguous mappings - ask for clarification
-    important_concepts = ['facility', 'organization', 'department', 'specialty', 'region']
-    for concept in important_concepts:
-        if concept in mapping.ambiguous and len(questions) < max_questions:
-            options = [f"{tbl}.{col}" for tbl, col in mapping.ambiguous[concept][:3]]
-            questions.append(f"**Entity**: Which column represents {concept.replace('_', ' ')}? Options: {', '.join(options)}")
+    # If we have good mapping, we might not need questions
+    if len(mapping.resolved) > len(mapping.missing) and len(mapping.ambiguous) == 0:
+        return "**Data Analysis Ready**: Your data appears well-structured. Please provide any additional context about your analysis goals."
+    
+    # Questions for missing concepts
+    if mapping.missing:
+        questions.append("### Missing Information")
+        for concept in mapping.missing:
+            if concept == "facility_distribution":
+                questions.append("- Do you have data about healthcare facilities and their distribution?")
+            elif concept == "bed_capacity":
+                questions.append("- Do you have data about hospital bed capacity and changes over time?")
+            else:
+                questions.append(f"- Can you provide more information about {concept}?")
     
-    # Missing critical data
-    if len(questions) < max_questions:
-        if not any(concept in mapping.resolved for concept in ['facility', 'organization', 'department']):
-            questions.append("**Grouping**: What entities should be analyzed? (facilities, departments, regions, etc.)")
-        
-        if not any(concept in mapping.resolved for concept in ['wait_time', 'wait_median', 'score', 'performance']):
-            questions.append("**Metric**: What is the primary metric to analyze? (wait times, scores, costs, etc.)")
+    # Questions for ambiguous concepts
+    if mapping.ambiguous:
+        questions.append("### Clarification Needed")
+        for concept, options in mapping.ambiguous.items():
+            questions.append(f"- For '{concept}', did you mean: {', '.join(options)}?")
     
     if not questions:
-        return "**Analysis Ready**: Data structure understood. Proceeding with analysis..."
+        return "**Data Analysis Ready**: Your data appears well-structured. Please provide any additional context about your analysis goals."
     
-    return "**Clarification Questions**\n\n" + "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))
+    return "\n".join(questions)