Spaces:
Sleeping
Sleeping
| # data_registry.py | |
| import pandas as pd | |
| import numpy as np | |
| from typing import Dict, Any, List, Optional | |
| import os | |
| class DataRegistry: | |
| def __init__(self): | |
| self.data = {} | |
| self.metadata = {} | |
| self.healthcare_metadata = {} | |
| def add_path(self, path: str) -> bool: | |
| """Add a data file to the registry with healthcare-specific handling.""" | |
| try: | |
| file_name = os.path.basename(path) | |
| if file_name.endswith('.csv'): | |
| df = pd.read_csv(path) | |
| # Standardize column names | |
| df.columns = [col.strip().lower().replace(' ', '_').replace('-', '_') for col in df.columns] | |
| self.data[file_name] = df | |
| # Basic metadata | |
| self.metadata[file_name] = { | |
| 'type': 'csv', | |
| 'columns': list(df.columns), | |
| 'shape': df.shape, | |
| 'sample': df.head(3).to_dict('records') | |
| } | |
| # Healthcare-specific metadata extraction | |
| self._extract_healthcare_metadata(file_name, df) | |
| return True | |
| return False | |
| except Exception as e: | |
| print(f"Error adding {path}: {e}") | |
| return False | |
| def _extract_healthcare_metadata(self, file_name: str, df: pd.DataFrame): | |
| """Extract healthcare-specific metadata from the dataframe.""" | |
| healthcare_meta = {} | |
| # Check for healthcare facility data | |
| if any(col in df.columns for col in ['facility_name', 'facility_type', 'odhf_facility_type']): | |
| healthcare_meta['data_type'] = 'healthcare_facilities' | |
| if 'facility_type' in df.columns: | |
| healthcare_meta['facility_types'] = df['facility_type'].value_counts().to_dict() | |
| if 'city' in df.columns: | |
| healthcare_meta['cities'] = df['city'].value_counts().head(10).to_dict() | |
| # Check for bed capacity data | |
| if any(col in df.columns for col in ['beds_current', 'beds_prev', 'bed_count']): | |
| healthcare_meta['data_type'] = 'bed_capacity' | |
| if 'zone' in df.columns: | |
| healthcare_meta['zones'] = df['zone'].unique().tolist() | |
| if 'teaching_status' in df.columns: | |
| healthcare_meta['teaching_status_counts'] = df['teaching_status'].value_counts().to_dict() | |
| # Calculate derived metrics | |
| if 'beds_current' in df.columns and 'beds_prev' in df.columns: | |
| df['bed_change'] = df['beds_current'] - df['beds_prev'] | |
| df['percent_change'] = (df['bed_change'] / df['beds_prev']) * 100 | |
| healthcare_meta['has_derived_metrics'] = True | |
| # Check for patient data (with privacy warning) | |
| if any(col in df.columns for col in ['patient_id', 'patient_name', 'mrn']): | |
| healthcare_meta['data_type'] = 'patient_data' | |
| healthcare_meta['privacy_warning'] = "This file contains patient identifiers. Ensure proper handling." | |
| if healthcare_meta: | |
| self.healthcare_metadata[file_name] = healthcare_meta | |
| def get_healthcare_metadata(self, name: str) -> Dict[str, Any]: | |
| """Get healthcare-specific metadata for a file.""" | |
| return self.healthcare_metadata.get(name, {}) | |
| def get_data_type(self, name: str) -> str: | |
| """Get the healthcare data type of a file.""" | |
| meta = self.get_healthcare_metadata(name) | |
| return meta.get('data_type', 'unknown') | |
| def names(self): | |
| return list(self.data.keys()) | |
| def get(self, name): | |
| return self.data.get(name) | |
| def summarize_for_prompt(self) -> str: | |
| """Generate a summary of all data for prompt inclusion.""" | |
| if not self.data: | |
| return "No data files registered." | |
| summary_parts = [] | |
| for file_name in self.names(): | |
| meta = self.metadata.get(file_name, {}) | |
| health_meta = self.get_healthcare_metadata(file_name) | |
| summary_parts.append(f"File: {file_name}") | |
| summary_parts.append(f"Type: {meta.get('type', 'unknown')}") | |
| summary_parts.append(f"Columns: {', '.join(meta.get('columns', []))}") | |
| summary_parts.append(f"Shape: {meta.get('shape', 'unknown')}") | |
| if health_meta: | |
| summary_parts.append("Healthcare Context:") | |
| for key, value in health_meta.items(): | |
| if key != 'privacy_warning': # Don't include warnings in prompt | |
| summary_parts.append(f" {key}: {value}") | |
| summary_parts.append("") | |
| return "\n".join(summary_parts) | |
| def clear(self): | |
| self.data.clear() | |
| self.metadata.clear() | |
| self.healthcare_metadata.clear() |