James Li

Posted on Nov 19

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

#aiagent #llm #ai #architecture

Introduction

Data Analysis Agent is a crucial component in modern enterprise data stacks, capable of automating data analysis processes and providing intelligent data insights. This article will detail how to build an enterprise-level data analysis agent system.

1. Data Processing Toolchain Design

The data processing toolchain is the fundamental infrastructure of the entire analysis system, determining the system's capability and efficiency in handling data. An excellent toolchain design should have:

Good scalability: Ability to easily add new data sources and processing methods
High configurability: Adjust processing logic through configuration rather than code modification
Stable fault tolerance: Gracefully handle various exceptions
Comprehensive monitoring mechanism: Full monitoring of the processing workflow

1.1 Data Access Layer Design

The data access layer is responsible for interacting with various data sources, securely and efficiently introducing raw data into the system. Here's the core implementation code:

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

💡 Best Practices

Implement automatic retry and failover for data source connectors

Set maximum retry attempts and intervals

Implement graceful degradation strategies

Add circuit breaker to prevent cascading failures

Use connection pools to manage database connections

Pre-create connection pools for better performance

Automatically manage connection lifecycles

Implement connection health checks

Make data preprocessing steps configurable

Define processing workflows through configuration files

Support dynamic loading of new processors

Provide dependency management for processing steps

Add data quality check mechanisms

Data integrity validation

Data type checks

Business rule validation

Anomaly data flagging

1.2 Data Cleaning and Transformation

Data cleaning and transformation is one of the most important aspects of data analysis, directly affecting the quality of subsequent analysis. Here's the core implementation:

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

💡 Data Transformation Best Practices

Type Conversion

Automatically identify and correct data types

Handle special formats (like datetime)

Keep backup of original data

Feature Engineering

Use LLM to assist feature creation

Automated feature selection

Feature importance evaluation

Data Aggregation

Multi-dimensional aggregation support

Flexible aggregation function configuration

Result correctness validation

2. SQL Generation and Optimization

In the Data Analysis Agent, SQL generation and optimization is the key link connecting user intent and data queries. We need to build an intelligent SQL generator that can convert natural language into efficient SQL queries.

2.1 Intelligent SQL Generator

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)

2.2 SQL Optimization Engine

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)

💡 SQL Optimization Best Practices

Index Optimization

Automatically identify required indexes

Evaluate index usage

Regular cleanup of invalid indexes

Query Rewriting

Optimize JOIN order

Simplify complex subqueries

Use temp tables for large data processing

Performance Monitoring

Log slow queries

Analyze execution plans

Monitor resource usage

3. Visualization Integration Solution

Data visualization is a crucial output form of data analysis, requiring automatic selection of appropriate visualization schemes based on data characteristics and analysis purposes.

3.1 Intelligent Chart Recommendation

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

3.2 Visualization Rendering Engine

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)

4. Analysis Pipeline Orchestration

Analysis pipeline orchestration is crucial for organizing various analysis steps into a complete workflow. We need to build a flexible and reliable orchestration system.

4.1 Workflow Engine

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

4.2 Task Orchestration Configuration

@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

5. Result Validation Mechanism

The result validation mechanism ensures the accuracy and reliability of analysis results, including data quality checks, result consistency validation, and anomaly detection.

5.1 Validation Framework

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)

5.2 Specific Validator Implementations

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

💡 Validation Best Practices

Data Quality Validation

Set thresholds for key metrics

Monitor data trend changes

Record anomalous data samples

Result Consistency Validation

Compare with historical results

Cross-validation

Business rule validation

Anomaly Detection

Statistical methods for anomaly detection

Time series trend analysis

Multi-dimensional cross-validation

With this, we have completed the design and implementation of a comprehensive enterprise-level data analysis Agent system. The system features:

Modular design with clear component responsibilities
Extensible architecture supporting new functionality
Robust error handling and validation mechanisms
Flexible configuration and scheduling capabilities
Comprehensive monitoring and logging

In practical applications, customization and optimization based on specific business scenarios will be needed.

DEV Community