from __future__ import annotations 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"], # 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"] } 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 _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.""" 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} if remaining_patterns: # Generate dynamic patterns based on scenario concept_patterns = _generate_dynamic_patterns(scenario_terms, remaining_patterns) # 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 = [] # 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 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)}") # 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.)") if not questions: return "**Analysis Ready**: Data structure understood. Proceeding with analysis..." return "**Clarification Questions**\n\n" + "\n".join(f"{i+1}. {q}" for i, q in enumerate(questions))