Building Enterprise-Level Data Analysis Agent: Architecture Design and Implementation

from typing import Dict, List, Union from abc import ABC, abstractmethod class DataConnector(ABC): """Data source connector base class Provides unified interface specifications for different types of data sources: - Databases (MySQL, PostgreSQL, etc.) - Data warehouses (Snowflake, Redshift, etc.) - File systems (CSV, Excel, etc.) - API interfaces """ @abstractmethod async def connect(self) -> bool: """Establish connection with data source Returns: bool: Whether connection is successful """ pass @abstractmethod async def fetch_data(self, query: str) -> pd.DataFrame: """Fetch data from data source Args: query: Data query statement/parameters Returns: pd.DataFrame: Query result dataframe """ pass class DataProcessor: def __init__(self): # Store instances of various data source connectors  self.connectors: Dict[str, DataConnector] = {} # Preprocessing step pipeline  self.preprocessing_pipeline = [] async def process_data( self, source: str, # Data source identifier  query: str, # Query statement  preprocessing_steps: List[Dict] = None # Preprocessing step configuration  ) -> pd.DataFrame: """Data processing main function Complete data processing workflow includes: 1. Get raw data from specified data source 2. Execute configured preprocessing steps 3. Return processed dataframe Args: source: Data source identifier query: Query statement preprocessing_steps: Preprocessing step configuration list Returns: pd.DataFrame: Processed dataframe """ # Get raw data  raw_data = await self.connectors[source].fetch_data(query) # Apply preprocessing steps  processed_data = raw_data for step in (preprocessing_steps or []): processed_data = await self._apply_preprocessing( processed_data, step ) return processed_data async def _apply_preprocessing( self, data: pd.DataFrame, step: Dict ) -> pd.DataFrame: """Apply single preprocessing step Supported preprocessing types: - missing_value: Missing value handling - outlier: Outlier handling - normalization: Data standardization - encoding: Feature encoding Args: data: Input dataframe step: Preprocessing step configuration Returns: pd.DataFrame: Processed dataframe """ step_type = step["type"] params = step["params"] if step_type == "missing_value": return await self._handle_missing_values(data, **params) elif step_type == "outlier": return await self._handle_outliers(data, **params) # ... other preprocessing types  return data 

class DataTransformer: def __init__(self, llm_service): self.llm = llm_service # LLM service for intelligent data transformation  self.transformation_cache = {} # Cache commonly used transformation results  async def transform_data( self, data: pd.DataFrame, transformation_rules: List[Dict] ) -> pd.DataFrame: """Data transformation main function Execute data transformations according to rule list order: 1. Data type conversion 2. Feature engineering 3. Data aggregation Args: data: Input dataframe transformation_rules: Transformation rule configuration list Returns: pd.DataFrame: Transformed dataframe """ transformed_data = data.copy() for rule in transformation_rules: transformed_data = await self._apply_transformation( transformed_data, rule ) return transformed_data async def _apply_transformation( self, data: pd.DataFrame, rule: Dict ) -> pd.DataFrame: """Apply single transformation rule Supported transformation types: - type_conversion: Data type conversion - feature_engineering: Feature engineering - aggregation: Data aggregation Args: data: Input dataframe rule: Transformation rule configuration Returns: pd.DataFrame: Transformed dataframe """ rule_type = rule["type"] if rule_type == "type_conversion": return await self._convert_types(data, rule["params"]) elif rule_type == "feature_engineering": return await self._engineer_features(data, rule["params"]) elif rule_type == "aggregation": return await self._aggregate_data(data, rule["params"]) return data 

from typing import Dict, List, Optional from dataclasses import dataclass @dataclass class TableSchema: """Table schema definition""" name: str columns: List[Dict[str, str]] # Column names and data types  primary_key: List[str] foreign_keys: Dict[str, str] # Foreign key relationships  class SQLGenerator: def __init__(self, llm_service, schema_manager): self.llm = llm_service self.schema_manager = schema_manager self.query_templates = self._load_query_templates() async def generate_sql( self, user_intent: str, context: Dict = None ) -> str: """Generate SQL based on user intent Args: user_intent: User query intent context: Context information (like time range, filter conditions, etc.) Returns: str: Generated SQL statement """ # 1. Parse user intent  parsed_intent = await self._parse_intent(user_intent) # 2. Identify relevant tables and fields  relevant_tables = await self._identify_tables(parsed_intent) # 3. Construct SQL statement  sql = await self._construct_sql(parsed_intent, relevant_tables, context) # 4. SQL optimization  optimized_sql = await self._optimize_sql(sql) return optimized_sql async def _parse_intent(self, user_intent: str) -> Dict: """Parse user intent Use LLM to convert natural language into structured query intent: - Query type (aggregation/detail/statistics etc.) - Target metrics - Dimension fields - Filter conditions - Sorting requirements """ prompt = f""" Convert the following data analysis requirement into structured format: {user_intent} Please provide: 1. Query type 2. Required metrics 3. Analysis dimensions 4. Filter conditions 5. Sorting rules """ response = await self.llm.generate(prompt) return self._parse_llm_response(response) 

class SQLOptimizer: def __init__(self, db_engine): self.db_engine = db_engine self.optimization_rules = self._load_optimization_rules() async def optimize_sql(self, sql: str) -> str: """Main SQL optimization function Optimization strategies include: 1. Index optimization 2. Join optimization 3. Subquery optimization 4. Aggregation optimization """ # 1. Parse SQL  parsed_sql = self._parse_sql(sql) # 2. Get execution plan  execution_plan = await self._get_execution_plan(sql) # 3. Apply optimization rules  optimizations = [] for rule in self.optimization_rules: if rule.should_apply(parsed_sql, execution_plan): optimization = await rule.apply(parsed_sql) optimizations.append(optimization) # 4. Rewrite SQL  optimized_sql = self._rewrite_sql(parsed_sql, optimizations) return optimized_sql async def _get_execution_plan(self, sql: str) -> Dict: """Get SQL execution plan""" explain_sql = f"EXPLAIN ANALYZE {sql}" return await self.db_engine.execute(explain_sql) 

class ChartRecommender: def __init__(self, llm_service): self.llm = llm_service self.chart_templates = self._load_chart_templates() async def recommend_chart( self, data: pd.DataFrame, analysis_goal: str ) -> Dict: """Recommend suitable chart type Args: data: Data to visualize analysis_goal: Analysis objective Returns: Dict: Chart configuration """ # 1. Analyze data characteristics  data_profile = await self._analyze_data(data) # 2. Match chart type  chart_type = await self._match_chart_type( data_profile, analysis_goal ) # 3. Generate chart configuration  chart_config = await self._generate_chart_config( chart_type, data, analysis_goal ) return chart_config 

class VisualizationEngine: def __init__(self): self.renderers = { 'plotly': PlotlyRenderer(), 'echarts': EChartsRenderer(), 'matplotlib': MatplotlibRenderer() } async def render_chart( self, data: pd.DataFrame, chart_config: Dict, renderer: str = 'plotly' ) -> str: """Render chart Args: data: Data chart_config: Chart configuration renderer: Renderer type Returns: str: Rendered chart (HTML or image URL) """ renderer = self.renderers.get(renderer) if not renderer: raise ValueError(f"Unsupported renderer: {renderer}") return await renderer.render(data, chart_config) 

from enum import Enum from typing import Dict, List, Callable from dataclasses import dataclass class TaskStatus(Enum): PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" @dataclass class AnalysisTask: """Analysis task definition""" id: str name: str type: str params: Dict dependencies: List[str] status: TaskStatus = TaskStatus.PENDING result: Dict = None class WorkflowEngine: def __init__(self): self.tasks: Dict[str, AnalysisTask] = {} self.task_handlers: Dict[str, Callable] = {} self.execution_history = [] async def register_task_handler( self, task_type: str, handler: Callable ): """Register task handler""" self.task_handlers[task_type] = handler async def create_workflow( self, tasks: List[AnalysisTask] ) -> str: """Create analysis workflow Args: tasks: List of tasks Returns: str: Workflow ID """ workflow_id = self._generate_workflow_id() # Validate task dependencies  if not self._validate_dependencies(tasks): raise ValueError("Invalid task dependencies") # Register tasks  for task in tasks: self.tasks[task.id] = task return workflow_id async def execute_workflow(self, workflow_id: str): """Execute workflow 1. Build task execution graph 2. Execute independent tasks in parallel 3. Execute subsequent tasks according to dependencies 4. Handle task failures and retries """ execution_graph = self._build_execution_graph() try: # Get executable tasks  ready_tasks = self._get_ready_tasks(execution_graph) while ready_tasks: # Execute tasks in parallel  results = await asyncio.gather( *[self._execute_task(task) for task in ready_tasks], return_exceptions=True ) # Update task status  for task, result in zip(ready_tasks, results): if isinstance(result, Exception): await self._handle_task_failure(task, result) else: await self._handle_task_success(task, result) # Get next batch of executable tasks  ready_tasks = self._get_ready_tasks(execution_graph) except Exception as e: await self._handle_workflow_failure(workflow_id, e) raise async def _execute_task(self, task: AnalysisTask): """Execute single task""" handler = self.task_handlers.get(task.type) if not handler: raise ValueError(f"No handler for task type: {task.type}") task.status = TaskStatus.RUNNING try: result = await handler(**task.params) task.result = result task.status = TaskStatus.COMPLETED return result except Exception as e: task.status = TaskStatus.FAILED raise 

@dataclass class WorkflowConfig: """Workflow configuration""" name: str description: str tasks: List[Dict] schedule: Optional[str] = None # cron expression  retry_policy: Dict = None class WorkflowBuilder: def __init__(self, engine: WorkflowEngine): self.engine = engine async def build_from_config( self, config: WorkflowConfig ) -> str: """Build workflow from configuration Example configuration: { "name": "Sales Data Analysis", "description": "Daily sales data analysis workflow", "tasks": [ { "id": "data_fetch", "type": "sql", "params": { "query": "SELECT * FROM sales" } }, { "id": "data_process", "type": "transform", "dependencies": ["data_fetch"], "params": { "operations": [...] } }, { "id": "visualization", "type": "chart", "dependencies": ["data_process"], "params": { "chart_type": "line", "metrics": [...] } } ], "schedule": "0 0 * * *", "retry_policy": { "max_attempts": 3, "delay": 300 } } """ tasks = [] for task_config in config.tasks: task = AnalysisTask( id=task_config["id"], name=task_config.get("name", task_config["id"]), type=task_config["type"], params=task_config["params"], dependencies=task_config.get("dependencies", []) ) tasks.append(task) workflow_id = await self.engine.create_workflow(tasks) # Set scheduling policy  if config.schedule: await self._setup_schedule(workflow_id, config.schedule) return workflow_id 

from abc import ABC, abstractmethod from typing import Any, List class Validator(ABC): """Validator base class""" @abstractmethod async def validate(self, data: Any) -> bool: pass @abstractmethod async def get_validation_report(self) -> Dict: pass class ResultValidator: def __init__(self): self.validators: List[Validator] = [] self.validation_history = [] async def add_validator(self, validator: Validator): """Add validator""" self.validators.append(validator) async def validate_result( self, result: Any, context: Dict = None ) -> bool: """Validate analysis results Execute all registered validators: 1. Data quality validation 2. Business rule validation 3. Statistical significance tests 4. Anomaly detection """ validation_results = [] for validator in self.validators: try: is_valid = await validator.validate(result) validation_results.append({ 'validator': validator.__class__.__name__, 'is_valid': is_valid, 'report': await validator.get_validation_report() }) except Exception as e: validation_results.append({ 'validator': validator.__class__.__name__, 'is_valid': False, 'error': str(e) }) # Record validation history  self.validation_history.append({ 'timestamp': datetime.now(), 'context': context, 'results': validation_results }) # Return True only if all validations pass  return all(r['is_valid'] for r in validation_results) 

class DataQualityValidator(Validator): """Data quality validator""" def __init__(self, rules: List[Dict]): self.rules = rules self.validation_results = [] async def validate(self, data: pd.DataFrame) -> bool: """Validate data quality Check items include: 1. Null value ratio 2. Anomaly detection 3. Data type consistency 4. Value range check """ for rule in self.rules: result = await self._check_rule(data, rule) self.validation_results.append(result) return all(r['passed'] for r in self.validation_results) async def get_validation_report(self) -> Dict: return { 'total_rules': len(self.rules), 'passed_rules': sum(1 for r in self.validation_results if r['passed']), 'results': self.validation_results } class StatisticalValidator(Validator): """Statistical validator""" def __init__(self, confidence_level: float = 0.95): self.confidence_level = confidence_level self.test_results = [] async def validate(self, data: Any) -> bool: """Statistical validation Including: 1. Significance tests 2. Confidence interval calculation 3. Sample representativeness tests 4. Distribution tests """ # Implement statistical testing logic  pass