Update schema_mapper.py

schema_mapper.py

@@ -152,8 +152,127 @@ class MappingResult:
     missing: List[str] = field(default_factory=list)
     discovered: Dict[str, str] = field(default_factory=dict)  # Discovered column types
 
+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:
-    """Dynamically map concepts based on scenario content and available data."""
+    """Enhanced mapping that extracts explicit information from scenario text."""
     result = MappingResult()
     
     if not registry.names():
@@ -163,9 +282,6 @@ def map_concepts(scenario_text: str, registry: DataRegistry) -> MappingResult:
     # Extract key terms from scenario
     scenario_terms = _extract_key_terms_from_scenario(scenario_text)
     
-    # Generate dynamic patterns based on scenario
-    concept_patterns = _generate_dynamic_patterns(scenario_terms, UNIVERSAL_CONCEPT_PATTERNS)
-    
     # Collect all available columns
     all_columns = []
     for table in registry.iter_tables():
@@ -176,87 +292,120 @@ def map_concepts(scenario_text: str, registry: DataRegistry) -> MappingResult:
         for col in table.df.columns:
             all_columns.append((table.name, str(col)))
 
-    # Map concepts to columns
-    for concept, patterns in concept_patterns.items():
-        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)
+    # 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}
+    
+    if remaining_patterns:
+        # Generate dynamic patterns based on scenario
+        concept_patterns = _generate_dynamic_patterns(scenario_terms, remaining_patterns)
         
-        if not scores or scores[0][1] == 0:
-            result.missing.append(concept)
-            continue
+        # 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)
+            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
+    
+    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 = []
         
-        top_score = scores[0][1]
+        # 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)}?")
         
-        # Find all columns with similar high scores (potential ambiguity)
-        threshold = max(50, top_score - 20)
-        high_scoring = [pair for pair, score in scores if score >= threshold]
+        if not critical_questions:
+            return "**Proceeding with Analysis**: Scenario and data mappings are clear. Analyzing now..."
         
-        if len(high_scoring) == 1:
-            tbl, col = high_scoring[0]
-            result.resolved[concept] = (tbl, col)
-        else:
-            # Multiple good matches - mark as ambiguous
-            result.ambiguous[concept] = high_scoring[:5]  # Limit to top 5
+        return "**Quick Clarification**\n\n" + "\n".join(critical_questions)
     
-    return result
-
-def build_phase1_questions(scenario_text: str, registry: DataRegistry, mapping: MappingResult, max_questions: int = 6) -> str:
-    """Build clarifying questions based on scenario and data gaps."""
+    # Fallback to standard question generation for less comprehensive scenarios
     questions = []
     scenario_lower = scenario_text.lower() if scenario_text else ""
     
-    # Data structure questions
-    if not mapping.resolved and not mapping.ambiguous:
-        questions.append("**Data Structure**: I don't see clear patterns in your data. Could you describe what each column represents?")
-        return "\n".join(questions)
-    
     # Ambiguous mappings - ask for clarification
     important_concepts = ['facility', 'organization', 'department', 'specialty', 'region']
     for concept in important_concepts:
-        if concept in mapping.ambiguous:
-            options = [f"{tbl}.{col}" for tbl, col in mapping.ambiguous[concept][:4]]
-            questions.append(f"**Entity Identification**: Which column represents the main {concept.replace('_', ' ')}? Options: {', '.join(options)}")
-            if len(questions) >= max_questions:
-                break
-    
-    # Metric clarification
-    metric_concepts = ['wait_time', 'cost', 'score', 'performance', 'quality']
-    for concept in metric_concepts:
-        if concept in mapping.ambiguous:
+        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"**Metric Clarification**: Which column best represents {concept.replace('_', ' ')}? Options: {', '.join(options)}")
-            if len(questions) >= max_questions:
-                break
+            questions.append(f"**Entity**: Which column represents {concept.replace('_', ' ')}? Options: {', '.join(options)}")
     
     # Missing critical data
-    if not any(concept in mapping.resolved for concept in ['facility', 'organization', 'department']):
-        questions.append("**Grouping Variable**: What should I group the analysis by? (e.g., facilities, departments, regions)")
-    
-    if not any(concept in mapping.resolved for concept in ['wait_time', 'cost', 'score', 'performance']):
-        questions.append("**Key Metric**: What is the main metric you want to analyze? (e.g., performance scores, wait times, costs)")
-    
-    # Scenario-specific questions
-    if any(term in scenario_lower for term in ['resource', 'allocation', 'priority']):
-        questions.append("**Resource Allocation**: What factors should guide resource prioritization? (e.g., volume, urgency, equity)")
-    
-    if any(term in scenario_lower for term in ['comparison', 'benchmark', 'performance']):
-        questions.append("**Comparison Criteria**: How should different entities be compared? What constitutes good vs. poor performance?")
-    
-    if any(term in scenario_lower for term in ['recommendation', 'decision', 'strategy']):
-        questions.append("**Decision Context**: What constraints or preferences should influence the recommendations? (e.g., budget limits, operational requirements)")
-    
-    # Limit questions and format
-    questions = questions[:max_questions]
+    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.)")
     
     if not questions:
-        return "**Data Analysis Ready**: Your data appears well-structured. Please provide any additional context about your analysis goals."
-    
-    formatted_questions = ["**Clarification Questions**", ""]
-    for i, q in enumerate(questions, 1):
-        formatted_questions.append(f"{i}. {q}")
+        return "**Analysis Ready**: Data structure understood. Proceeding with analysis..."
     
-    return "\n".join(formatted_questions)
+    return "**Clarification Questions**\n\n" + "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))