aiecs 1.5.1__py3-none-any.whl

aiecs/tools/task_tools/chart_tool.py

@@ -0,0 +1,732 @@
+import os
+import json
+import csv
+import tempfile
+import logging
+from typing import Dict, Any, List, Optional, Union, Tuple
+from enum import Enum
+
+from pydantic import BaseModel, Field, field_validator, ConfigDict
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+from aiecs.tools import register_tool
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools.tool_executor import measure_execution_time
+
+# Enums for configuration options
+
+
+class ExportFormat(str, Enum):
+    JSON = "json"
+    CSV = "csv"
+    HTML = "html"
+    EXCEL = "excel"
+    MARKDOWN = "markdown"
+
+
+class VisualizationType(str, Enum):
+    HISTOGRAM = "histogram"
+    BOXPLOT = "boxplot"
+    SCATTER = "scatter"
+    BAR = "bar"
+    LINE = "line"
+    HEATMAP = "heatmap"
+    PAIR = "pair"
+
+
+@register_tool("chart")
+class ChartTool(BaseTool):
+    """Chart and visualization tool: creates charts and exports data in various formats."""
+
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the chart tool"""
+
+        model_config = ConfigDict(env_prefix="CHART_TOOL_")
+
+        export_dir: str = Field(
+            default=os.path.join(tempfile.gettempdir(), "chart_exports"),
+            description="Directory to export files to",
+        )
+        plot_dpi: int = Field(default=100, description="DPI for plot exports")
+        plot_figsize: Tuple[int, int] = Field(
+            default=(10, 6),
+            description="Default figure size (width, height) in inches",
+        )
+        allowed_extensions: List[str] = Field(
+            default=[
+                ".csv",
+                ".xlsx",
+                ".xls",
+                ".json",
+                ".parquet",
+                ".feather",
+                ".sav",
+                ".sas7bdat",
+                ".por",
+            ],
+            description="Allowed file extensions",
+        )
+
+    # Input schemas for operations
+    class ReadDataSchema(BaseModel):
+        """Schema for reading data files"""
+
+        file_path: str = Field(description="Path to the data file")
+        nrows: Optional[int] = Field(default=None, description="Number of rows to read")
+        sheet_name: Optional[Union[str, int]] = Field(
+            default=0, description="Sheet name or index for Excel files"
+        )
+        export_format: Optional[ExportFormat] = Field(
+            default=None, description="Format to export results in"
+        )
+        export_path: Optional[str] = Field(default=None, description="Path to export results to")
+
+        @field_validator("file_path")
+        @classmethod
+        def validate_file_path(cls, v):
+            if not os.path.isfile(v):
+                raise ValueError(f"File not found: {v}")
+            return v
+
+        @field_validator("export_path")
+        @classmethod
+        def validate_export_path(cls, v, info):
+            if v and "export_format" not in info.data:
+                raise ValueError("export_format must be specified when export_path is provided")
+            return v
+
+    class VisualizationSchema(BaseModel):
+        """Schema for data visualization"""
+
+        file_path: str = Field(description="Path to the data file")
+        plot_type: VisualizationType = Field(description="Type of visualization to create")
+        x: Optional[str] = Field(default=None, description="Column to use for x-axis")
+        y: Optional[str] = Field(default=None, description="Column to use for y-axis")
+        hue: Optional[str] = Field(default=None, description="Column to use for color encoding")
+        variables: Optional[List[str]] = Field(
+            default=None,
+            description="List of variables to include in the visualization",
+        )
+        title: Optional[str] = Field(default=None, description="Title for the visualization")
+        figsize: Optional[Tuple[int, int]] = Field(
+            default=None, description="Figure size (width, height) in inches"
+        )
+        output_path: Optional[str] = Field(
+            default=None, description="Path to save the visualization"
+        )
+        dpi: Optional[int] = Field(default=None, description="DPI for the visualization")
+        export_format: Optional[ExportFormat] = Field(
+            default=None, description="Format to export results in"
+        )
+        export_path: Optional[str] = Field(default=None, description="Path to export results to")
+
+        @field_validator("file_path")
+        @classmethod
+        def validate_file_path(cls, v):
+            if not os.path.isfile(v):
+                raise ValueError(f"File not found: {v}")
+            return v
+
+        @field_validator("export_path")
+        @classmethod
+        def validate_export_path(cls, v, info):
+            if v and "export_format" not in info.data:
+                raise ValueError("export_format must be specified when export_path is provided")
+            return v
+
+    class ExportDataSchema(BaseModel):
+        """Schema for exporting data"""
+
+        file_path: str = Field(description="Path to the data file")
+        variables: Optional[List[str]] = Field(
+            default=None,
+            description="List of variables to include in the export",
+        )
+        format: ExportFormat = Field(description="Format to export data in")
+        export_path: Optional[str] = Field(
+            default=None, description="Path to save the exported data"
+        )
+        export_format: Optional[ExportFormat] = Field(
+            default=None, description="Format to export results in"
+        )
+
+        @field_validator("file_path")
+        @classmethod
+        def validate_file_path(cls, v):
+            if not os.path.isfile(v):
+                raise ValueError(f"File not found: {v}")
+            return v
+
+        @field_validator("export_path")
+        @classmethod
+        def validate_export_path(cls, v, info):
+            if v and "export_format" not in info.data:
+                raise ValueError("export_format must be specified when export_path is provided")
+            return v
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the chart tool
+
+        Args:
+            config: Optional configuration for the tool
+        """
+        super().__init__(config)
+
+        # Parse configuration
+        self.config = self.Config(**(config or {}))
+
+        # Create export directory if it doesn't exist
+        os.makedirs(self.config.export_dir, exist_ok=True)
+
+        # Set up logger
+        self.logger = logging.getLogger(__name__)
+
+        # Set default matplotlib style
+        plt.style.use("seaborn-v0_8-whitegrid")
+
+    def _load_data(
+        self,
+        file_path: str,
+        nrows: Optional[int] = None,
+        sheet_name: Optional[Union[str, int]] = 0,
+    ) -> pd.DataFrame:
+        """
+        Load data from various file formats into a pandas DataFrame
+
+        Args:
+            file_path: Path to the data file
+            nrows: Number of rows to read
+            sheet_name: Sheet name or index for Excel files
+
+        Returns:
+            Loaded DataFrame
+        """
+        # Determine file type and read accordingly
+        ext = os.path.splitext(file_path)[1].lower()
+
+        try:
+            if ext == ".sav":
+                import pyreadstat
+
+                df, meta = pyreadstat.read_sav(file_path)
+                return df
+            elif ext == ".sas7bdat":
+                import pyreadstat
+
+                df, meta = pyreadstat.read_sas7bdat(file_path)
+                return df
+            elif ext == ".por":
+                import pyreadstat
+
+                df, meta = pyreadstat.read_por(file_path)
+                return df
+            elif ext == ".csv":
+                return pd.read_csv(file_path, nrows=nrows)
+            elif ext in [".xlsx", ".xls"]:
+                return pd.read_excel(file_path, sheet_name=sheet_name, nrows=nrows)
+            elif ext == ".json":
+                return pd.read_json(file_path)
+            elif ext == ".parquet":
+                return pd.read_parquet(file_path)
+            elif ext == ".feather":
+                return pd.read_feather(file_path)
+            else:
+                raise ValueError(f"Unsupported file format: {ext}")
+        except Exception as e:
+            raise ValueError(f"Error reading file {file_path}: {str(e)}")
+
+    def _export_result(self, result: Dict[str, Any], path: str, format: ExportFormat) -> None:
+        """
+        Export results to the specified format
+
+        Args:
+            result: Result to export
+            path: Path to save the exported result
+            format: Format to export in
+        """
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+
+        try:
+            if format == ExportFormat.JSON:
+                # Convert numpy types to Python native types
+                def json_serialize(obj):
+                    if isinstance(obj, (np.integer, np.int64)):
+                        return int(obj)
+                    elif isinstance(obj, (np.floating, np.float64)):
+                        return float(obj)
+                    elif isinstance(obj, np.ndarray):
+                        return obj.tolist()
+                    elif isinstance(obj, pd.DataFrame):
+                        return obj.to_dict(orient="records")
+                    return str(obj)
+
+                with open(path, "w") as f:
+                    json.dump(result, f, default=json_serialize, indent=2)
+
+            elif format == ExportFormat.CSV:
+                # Find the first dict or DataFrame in the result
+                data_to_export = None
+                for key, value in result.items():
+                    if isinstance(value, dict) and value:
+                        data_to_export = pd.DataFrame(value)
+                        break
+                    elif isinstance(value, pd.DataFrame):
+                        data_to_export = value
+                        break
+
+                if data_to_export is not None:
+                    data_to_export.to_csv(path, index=False)
+                else:
+                    # Fallback: convert the entire result to a flat structure
+                    flat_data = {}
+                    for k, v in result.items():
+                        if not isinstance(v, (dict, list, pd.DataFrame)):
+                            flat_data[k] = v
+
+                    with open(path, "w", newline="") as f:
+                        writer = csv.writer(f)
+                        writer.writerow(flat_data.keys())
+                        writer.writerow(flat_data.values())
+
+            elif format == ExportFormat.HTML:
+                # Convert to HTML table
+                html_content = "<html><body><h1>Chart Results</h1>"
+                for key, value in result.items():
+                    html_content += f"<h2>{key}</h2>"
+                    if isinstance(value, pd.DataFrame):
+                        html_content += value.to_html()
+                    elif isinstance(value, dict):
+                        html_content += (
+                            "<table border='1'><tr><th>Parameter</th><th>Value</th></tr>"
+                        )
+                        for k, v in value.items():
+                            html_content += f"<tr><td>{k}</td><td>{v}</td></tr>"
+                        html_content += "</table>"
+                    else:
+                        html_content += f"<p>{value}</p>"
+                html_content += "</body></html>"
+
+                with open(path, "w") as f:
+                    f.write(html_content)
+
+            elif format == ExportFormat.EXCEL:
+                with pd.ExcelWriter(path) as writer:
+                    for key, value in result.items():
+                        if isinstance(value, pd.DataFrame):
+                            # Excel sheet names limited to 31 chars
+                            value.to_excel(writer, sheet_name=key[:31])
+                        elif isinstance(value, dict):
+                            pd.DataFrame(value, index=[0]).to_excel(writer, sheet_name=key[:31])
+                        else:
+                            pd.DataFrame({key: [value]}).to_excel(writer, sheet_name="Summary")
+
+            elif format == ExportFormat.MARKDOWN:
+                with open(path, "w") as f:
+                    f.write("# Chart Results\n\n")
+                    for key, value in result.items():
+                        f.write(f"## {key}\n\n")
+                        if isinstance(value, pd.DataFrame):
+                            f.write(value.to_markdown())
+                        elif isinstance(value, dict):
+                            f.write("| Parameter | Value |\n|-----------|-------|\n")
+                            for k, v in value.items():
+                                f.write(f"| {k} | {v} |\n")
+                        else:
+                            f.write(f"{value}\n\n")
+
+            return path
+        except Exception as e:
+            raise ValueError(f"Error exporting to {format}: {str(e)}")
+
+    def _create_visualization(
+        self,
+        df: pd.DataFrame,
+        plot_type: VisualizationType,
+        x: Optional[str] = None,
+        y: Optional[str] = None,
+        hue: Optional[str] = None,
+        variables: Optional[List[str]] = None,
+        title: Optional[str] = None,
+        figsize: Optional[Tuple[int, int]] = None,
+        output_path: Optional[str] = None,
+        dpi: Optional[int] = None,
+    ) -> str:
+        """
+        Create a visualization based on the parameters and return the path to the saved image
+
+        Args:
+            df: DataFrame to visualize
+            plot_type: Type of visualization to create
+            x: Column to use for x-axis
+            y: Column to use for y-axis
+            hue: Column to use for color encoding
+            variables: List of variables to include in the visualization
+            title: Title for the visualization
+            figsize: Figure size (width, height) in inches
+            output_path: Path to save the visualization
+            dpi: DPI for the visualization
+
+        Returns:
+            Path to the saved visualization
+        """
+        if not output_path:
+            output_path = os.path.join(self.config.export_dir, f"plot_{os.urandom(4).hex()}.png")
+        elif not os.path.isabs(output_path):
+            output_path = os.path.join(self.config.export_dir, output_path)
+
+        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+
+        try:
+            figsize = figsize or self.config.plot_figsize
+            dpi = dpi or self.config.plot_dpi
+
+            plt.figure(figsize=figsize)
+
+            if plot_type == VisualizationType.HISTOGRAM:
+                if variables:
+                    for var in variables:
+                        sns.histplot(data=df, x=var, kde=True, label=var)
+                    plt.legend()
+                else:
+                    sns.histplot(data=df, x=x, hue=hue)
+
+            elif plot_type == VisualizationType.BOXPLOT:
+                sns.boxplot(data=df, x=x, y=y, hue=hue)
+
+            elif plot_type == VisualizationType.SCATTER:
+                sns.scatterplot(data=df, x=x, y=y, hue=hue)
+
+            elif plot_type == VisualizationType.BAR:
+                sns.barplot(data=df, x=x, y=y, hue=hue)
+
+            elif plot_type == VisualizationType.LINE:
+                sns.lineplot(data=df, x=x, y=y, hue=hue)
+
+            elif plot_type == VisualizationType.HEATMAP:
+                if variables:
+                    corr = df[variables].corr()
+                else:
+                    corr = df.corr()
+                sns.heatmap(corr, annot=True, cmap="coolwarm", fmt=".2f")
+
+            elif plot_type == VisualizationType.PAIR:
+                if variables:
+                    plot_vars = variables + [hue] if hue else variables
+                    sns.pairplot(df[plot_vars], hue=hue)
+                else:
+                    sns.pairplot(df, hue=hue)
+
+            if title:
+                plt.title(title)
+
+            plt.tight_layout()
+            plt.savefig(output_path, dpi=dpi)
+            plt.close()
+
+            return output_path
+        except Exception as e:
+            raise ValueError(f"Error creating visualization: {str(e)}")
+
+    def _validate_variables(self, df: pd.DataFrame, variables: List[str]) -> None:
+        """
+        Validate that variables exist in the DataFrame
+
+        Args:
+            df: DataFrame to check
+            variables: List of variables to validate
+
+        Raises:
+            ValueError: If any variables are not found in the DataFrame
+        """
+        if not variables:
+            return
+
+        available_columns = set(df.columns)
+        missing = [col for col in variables if col not in available_columns]
+        if missing:
+            raise ValueError(
+                f"Variables not found in dataset: {', '.join(missing)}. Available columns: {list(available_columns)}"
+            )
+
+    def _to_json_serializable(
+        self, result: Union[pd.DataFrame, pd.Series, Dict]
+    ) -> Union[List[Dict], Dict]:
+        """
+        Convert result to JSON serializable format
+
+        Args:
+            result: Result to convert
+
+        Returns:
+            JSON serializable result
+        """
+        if isinstance(result, pd.DataFrame):
+            # Handle datetime columns
+            for col in result.select_dtypes(include=["datetime64"]).columns:
+                result[col] = result[col].dt.strftime("%Y-%m-%d %H:%M:%S")
+            return result.to_dict(orient="records")
+        elif isinstance(result, pd.Series):
+            if pd.api.types.is_datetime64_any_dtype(result):
+                result = result.dt.strftime("%Y-%m-%d %H:%M:%S")
+            return result.to_dict()
+        elif isinstance(result, dict):
+            # Handle numpy types and datetime objects
+            def convert_value(v):
+                if isinstance(v, (np.floating, np.integer)):
+                    return float(v)
+                elif isinstance(v, np.bool_):
+                    return bool(v)
+                elif isinstance(v, (pd.Timestamp, np.datetime64)):
+                    return str(v)
+                elif isinstance(v, np.ndarray):
+                    return v.tolist()
+                elif pd.isna(v):
+                    return None
+                return v
+
+            return {k: convert_value(v) for k, v in result.items()}
+        return result
+
+    @measure_execution_time
+    def read_data(
+        self,
+        file_path: str,
+        nrows: Optional[int] = None,
+        sheet_name: Optional[Union[str, int]] = 0,
+        export_format: Optional[ExportFormat] = None,
+        export_path: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Read data from various file formats
+
+        Args:
+            file_path: Path to the data file
+            nrows: Number of rows to read
+            sheet_name: Sheet name or index for Excel files
+            export_format: Format to export results in
+            export_path: Path to export results to
+
+        Returns:
+            Dictionary with data summary
+        """
+        # Validate file path
+        if not os.path.isfile(file_path):
+            raise ValueError(f"File not found: {file_path}")
+
+        # Check file extension
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext not in self.config.allowed_extensions:
+            raise ValueError(
+                f"Extension '{ext}' not allowed. Supported formats: {', '.join(self.config.allowed_extensions)}"
+            )
+
+        # Load data
+        df = self._load_data(file_path, nrows, sheet_name)
+
+        # Create result
+        result = {
+            "variables": df.columns.tolist(),
+            "observations": len(df),
+            "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
+            # MB
+            "memory_usage": df.memory_usage(deep=True).sum() / (1024 * 1024),
+            "preview": df.head(5).to_dict(orient="records"),
+        }
+
+        # Handle export if requested
+        if export_format and export_path:
+            if not os.path.isabs(export_path):
+                export_path = os.path.join(self.config.export_dir, export_path)
+
+            self._export_result(result, export_path, export_format)
+            result["exported_to"] = export_path
+
+        return result
+
+    @measure_execution_time
+    def visualize(
+        self,
+        file_path: str,
+        plot_type: VisualizationType,
+        x: Optional[str] = None,
+        y: Optional[str] = None,
+        hue: Optional[str] = None,
+        variables: Optional[List[str]] = None,
+        title: Optional[str] = None,
+        figsize: Optional[Tuple[int, int]] = None,
+        output_path: Optional[str] = None,
+        dpi: Optional[int] = None,
+        export_format: Optional[ExportFormat] = None,
+        export_path: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Create data visualizations
+
+        Args:
+            file_path: Path to the data file
+            plot_type: Type of visualization to create
+            x: Column to use for x-axis
+            y: Column to use for y-axis
+            hue: Column to use for color encoding
+            variables: List of variables to include in the visualization
+            title: Title for the visualization
+            figsize: Figure size (width, height) in inches
+            output_path: Path to save the visualization
+            dpi: DPI for the visualization
+            export_format: Format to export results in
+            export_path: Path to export results to
+
+        Returns:
+            Dictionary with visualization details
+        """
+        # Validate file path
+        if not os.path.isfile(file_path):
+            raise ValueError(f"File not found: {file_path}")
+
+        # Check file extension
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext not in self.config.allowed_extensions:
+            raise ValueError(
+                f"Extension '{ext}' not allowed. Supported formats: {', '.join(self.config.allowed_extensions)}"
+            )
+
+        # Load data
+        df = self._load_data(file_path)
+
+        # Validate variables
+        vars_to_check = []
+        if variables:
+            vars_to_check.extend(variables)
+        if x:
+            vars_to_check.append(x)
+        if y:
+            vars_to_check.append(y)
+        if hue:
+            vars_to_check.append(hue)
+
+        self._validate_variables(df, vars_to_check)
+
+        # Create visualization
+        output_path = self._create_visualization(
+            df,
+            plot_type,
+            x,
+            y,
+            hue,
+            variables,
+            title,
+            figsize,
+            output_path,
+            dpi,
+        )
+
+        # Create result
+        result = {
+            "plot_type": plot_type,
+            "output_path": output_path,
+            "variables": variables or [x, y, hue],
+            "title": title or f"{plot_type.capitalize()} Plot",
+        }
+
+        # Handle export if requested
+        if export_format and export_path:
+            if not os.path.isabs(export_path):
+                export_path = os.path.join(self.config.export_dir, export_path)
+
+            self._export_result(result, export_path, export_format)
+            result["exported_to"] = export_path
+
+        return result
+
+    @measure_execution_time
+    def export_data(
+        self,
+        file_path: str,
+        format: ExportFormat,
+        variables: Optional[List[str]] = None,
+        export_path: Optional[str] = None,
+        export_format: Optional[ExportFormat] = None,
+    ) -> Dict[str, Any]:
+        """
+        Export data to various formats
+
+        Args:
+            file_path: Path to the data file
+            format: Format to export data in
+            variables: List of variables to include in the export
+            export_path: Path to save the exported data
+            export_format: Format to export results in
+
+        Returns:
+            Dictionary with export details
+        """
+        # Validate file path
+        if not os.path.isfile(file_path):
+            raise ValueError(f"File not found: {file_path}")
+
+        # Check file extension
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext not in self.config.allowed_extensions:
+            raise ValueError(
+                f"Extension '{ext}' not allowed. Supported formats: {', '.join(self.config.allowed_extensions)}"
+            )
+
+        # Load data
+        df = self._load_data(file_path)
+
+        # Validate variables
+        if variables:
+            self._validate_variables(df, variables)
+            df = df[variables]
+
+        # Determine export path
+        if not export_path:
+            ext = "." + format.value
+            if format == ExportFormat.EXCEL:
+                ext = ".xlsx"
+            export_path = os.path.join(self.config.export_dir, f"export_{os.urandom(4).hex()}{ext}")
+        elif not os.path.isabs(export_path):
+            export_path = os.path.join(self.config.export_dir, export_path)
+
+        # Create export directory if it doesn't exist
+        os.makedirs(os.path.dirname(os.path.abspath(export_path)), exist_ok=True)
+
+        # Export data
+        try:
+            if format == ExportFormat.JSON:
+                df.to_json(export_path, orient="records", indent=2)
+            elif format == ExportFormat.CSV:
+                df.to_csv(export_path, index=False)
+            elif format == ExportFormat.HTML:
+                df.to_html(export_path)
+            elif format == ExportFormat.EXCEL:
+                df.to_excel(export_path, index=False)
+            elif format == ExportFormat.MARKDOWN:
+                with open(export_path, "w") as f:
+                    f.write(df.to_markdown())
+        except Exception as e:
+            raise ValueError(f"Error exporting to {format}: {str(e)}")
+
+        # Create result
+        result = {
+            "format": format,
+            "path": export_path,
+            "rows": len(df),
+            "columns": len(df.columns),
+            "variables": df.columns.tolist(),
+        }
+
+        # Handle export if requested
+        if export_format and export_path:
+            if not os.path.isabs(export_path):
+                export_path = os.path.join(self.config.export_dir, export_path)
+
+            self._export_result(result, export_path, export_format)
+            result["exported_to"] = export_path
+
+        return result