aiecs 1.0.1__py3-none-any.whl → 1.7.17__py3-none-any.whl

aiecs/tools/task_tools/chart_tool.py

@@ -6,17 +6,20 @@ import logging
 from typing import Dict, Any, List, Optional, Union, Tuple
 from enum import Enum
 
-from pydantic import BaseModel, Field, field_validator, ConfigDict
+from pydantic import BaseModel, Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
 import numpy as np
-import pandas as pd
+import pandas as pd  # type: ignore[import-untyped]
 import matplotlib.pyplot as plt
-import seaborn as sns
+import seaborn as sns  # type: ignore[import-untyped]
 
 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"
@@ -24,6 +27,7 @@ class ExportFormat(str, Enum):
     EXCEL = "excel"
     MARKDOWN = "markdown"
 
+
 class VisualizationType(str, Enum):
     HISTOGRAM = "histogram"
     BOXPLOT = "boxplot"
@@ -33,177 +37,147 @@ class VisualizationType(str, Enum):
     HEATMAP = "heatmap"
     PAIR = "pair"
 
-@register_tool('chart')
+
+@register_tool("chart")
 class ChartTool(BaseTool):
     """Chart and visualization tool: creates charts and exports data in various formats."""
 
     # Configuration schema
-    class Config(BaseModel):
-        model_config = ConfigDict()
-        """Configuration for the chart tool"""
+    class Config(BaseSettings):
+        """Configuration for the chart tool
+        
+        Automatically reads from environment variables with CHART_TOOL_ prefix.
+        Example: CHART_TOOL_EXPORT_DIR -> export_dir
+        """
+
+        model_config = SettingsConfigDict(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"
+            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"
+            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"
+            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"
-        )
+    class Read_dataSchema(BaseModel):
+        """Schema for read_data operation"""
+
+        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')
+        @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')
+        @field_validator("export_path")
         @classmethod
         def validate_export_path(cls, v, info):
-            if v and 'export_format' not in info.data:
+            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"
-        )
+    class VisualizeSchema(BaseModel):
+        """Schema for visualize operation"""
+
+        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"
+            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')
+        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')
+        @field_validator("export_path")
         @classmethod
         def validate_export_path(cls, v, info):
-            if v and 'export_format' not in info.data:
+            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"
-        )
+    class Export_dataSchema(BaseModel):
+        """Schema for export_data operation"""
+
+        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"
+            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')
+        @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')
+        @field_validator("export_path")
         @classmethod
         def validate_export_path(cls, v, info):
-            if v and 'export_format' not in info.data:
+            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):
+    def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
         """
         Initialize the chart tool
 
         Args:
             config: Optional configuration for the tool
+            **kwargs: Additional arguments passed to BaseTool (e.g., tool_name)
+
+        Configuration is automatically loaded by BaseTool from:
+        1. Explicit config dict (highest priority)
+        2. YAML config files (config/tools/chart.yaml)
+        3. Environment variables (via dotenv from .env files)
+        4. Tool defaults (lowest priority)
         """
-        super().__init__(config)
+        super().__init__(config, **kwargs)
 
-        # Parse configuration
-        self.config = self.Config(**(config or {}))
+        # Configuration is automatically loaded by BaseTool into self._config_obj
+        # Access config via self._config_obj (BaseSettings instance)
+        self.config = self._config_obj if self._config_obj else self.Config()
 
         # Create export directory if it doesn't exist
         os.makedirs(self.config.export_dir, exist_ok=True)
@@ -212,9 +186,14 @@ class ChartTool(BaseTool):
         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:
+        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
 
@@ -230,34 +209,37 @@ class ChartTool(BaseTool):
         ext = os.path.splitext(file_path)[1].lower()
 
         try:
-            if ext == '.sav':
-                import pyreadstat
+            if ext == ".sav":
+                import pyreadstat  # type: ignore[import-untyped]
+
                 df, meta = pyreadstat.read_sav(file_path)
                 return df
-            elif ext == '.sas7bdat':
-                import pyreadstat
+            elif ext == ".sas7bdat":
+                import pyreadstat  # type: ignore[import-untyped]
+
                 df, meta = pyreadstat.read_sas7bdat(file_path)
                 return df
-            elif ext == '.por':
-                import pyreadstat
+            elif ext == ".por":
+                import pyreadstat  # type: ignore[import-untyped]
+
                 df, meta = pyreadstat.read_por(file_path)
                 return df
-            elif ext == '.csv':
+            elif ext == ".csv":
                 return pd.read_csv(file_path, nrows=nrows)
-            elif ext in ['.xlsx', '.xls']:
+            elif ext in [".xlsx", ".xls"]:
                 return pd.read_excel(file_path, sheet_name=sheet_name, nrows=nrows)
-            elif ext == '.json':
+            elif ext == ".json":
                 return pd.read_json(file_path)
-            elif ext == '.parquet':
+            elif ext == ".parquet":
                 return pd.read_parquet(file_path)
-            elif ext == '.feather':
+            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:
+    def _export_result(self, result: Dict[str, Any], path: str, format: ExportFormat) -> str:
         """
         Export results to the specified format
 
@@ -279,10 +261,10 @@ class ChartTool(BaseTool):
                     elif isinstance(obj, np.ndarray):
                         return obj.tolist()
                     elif isinstance(obj, pd.DataFrame):
-                        return obj.to_dict(orient='records')
+                        return obj.to_dict(orient="records")
                     return str(obj)
 
-                with open(path, 'w') as f:
+                with open(path, "w") as f:
                     json.dump(result, f, default=json_serialize, indent=2)
 
             elif format == ExportFormat.CSV:
@@ -300,12 +282,12 @@ class ChartTool(BaseTool):
                     data_to_export.to_csv(path, index=False)
                 else:
                     # Fallback: convert the entire result to a flat structure
-                    flat_data = {}
+                    flat_data: Dict[str, Any] = {}
                     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:
+                    with open(path, "w", newline="") as f:
                         writer = csv.writer(f)
                         writer.writerow(flat_data.keys())
                         writer.writerow(flat_data.values())
@@ -326,21 +308,22 @@ class ChartTool(BaseTool):
                         html_content += f"<p>{value}</p>"
                 html_content += "</body></html>"
 
-                with open(path, 'w') as f:
+                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):
-                            value.to_excel(writer, sheet_name=key[:31])  # Excel sheet names limited to 31 chars
+                            # 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')
+                            pd.DataFrame({key: [value]}).to_excel(writer, sheet_name="Summary")
 
             elif format == ExportFormat.MARKDOWN:
-                with open(path, 'w') as f:
+                with open(path, "w") as f:
                     f.write("# Chart Results\n\n")
                     for key, value in result.items():
                         f.write(f"## {key}\n\n")
@@ -357,11 +340,19 @@ class ChartTool(BaseTool):
         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:
+    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
 
@@ -418,7 +409,7 @@ class ChartTool(BaseTool):
                     corr = df[variables].corr()
                 else:
                     corr = df.corr()
-                sns.heatmap(corr, annot=True, cmap='coolwarm', fmt=".2f")
+                sns.heatmap(corr, annot=True, cmap="coolwarm", fmt=".2f")
 
             elif plot_type == VisualizationType.PAIR:
                 if variables:
@@ -455,9 +446,7 @@ class ChartTool(BaseTool):
         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)}"
-            )
+            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]:
         """
@@ -471,12 +460,12 @@ class ChartTool(BaseTool):
         """
         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')
+            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')
+                result = result.dt.strftime("%Y-%m-%d %H:%M:%S")
             return result.to_dict()
         elif isinstance(result, dict):
             # Handle numpy types and datetime objects
@@ -497,10 +486,14 @@ class ChartTool(BaseTool):
         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]:
+    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
 
@@ -528,11 +521,12 @@ class ChartTool(BaseTool):
 
         # Create result
         result = {
-            'variables': df.columns.tolist(),
-            'observations': len(df),
-            'dtypes': {col: str(dtype) for col, dtype in df.dtypes.items()},
-            'memory_usage': df.memory_usage(deep=True).sum() / (1024 * 1024),  # MB
-            'preview': df.head(5).to_dict(orient='records')
+            "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
@@ -541,18 +535,26 @@ class ChartTool(BaseTool):
                 export_path = os.path.join(self.config.export_dir, export_path)
 
             self._export_result(result, export_path, export_format)
-            result['exported_to'] = export_path
+            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]:
+    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
 
@@ -600,15 +602,25 @@ class ChartTool(BaseTool):
 
         # Create visualization
         output_path = self._create_visualization(
-            df, plot_type, x, y, hue, variables, title, figsize, output_path, dpi
+            df,
+            plot_type,
+            x,
+            y,
+            hue,
+            variables,
+            title,
+            figsize,
+            output_path,
+            dpi,
         )
 
-        # Create result
+        # Create result - filter out None values from variables
+        var_list = variables or [v for v in [x, y, hue] if v is not None]
         result = {
-            'plot_type': plot_type,
-            'output_path': output_path,
-            'variables': variables or [x, y, hue],
-            'title': title or f"{plot_type.capitalize()} Plot"
+            "plot_type": plot_type,
+            "output_path": output_path,
+            "variables": var_list,
+            "title": title or f"{plot_type.capitalize()} Plot",
         }
 
         # Handle export if requested
@@ -617,15 +629,19 @@ class ChartTool(BaseTool):
                 export_path = os.path.join(self.config.export_dir, export_path)
 
             self._export_result(result, export_path, export_format)
-            result['exported_to'] = export_path
+            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]:
+    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
 
@@ -671,7 +687,7 @@ class ChartTool(BaseTool):
         # Export data
         try:
             if format == ExportFormat.JSON:
-                df.to_json(export_path, orient='records', indent=2)
+                df.to_json(export_path, orient="records", indent=2)
             elif format == ExportFormat.CSV:
                 df.to_csv(export_path, index=False)
             elif format == ExportFormat.HTML:
@@ -679,18 +695,18 @@ class ChartTool(BaseTool):
             elif format == ExportFormat.EXCEL:
                 df.to_excel(export_path, index=False)
             elif format == ExportFormat.MARKDOWN:
-                with open(export_path, 'w') as f:
+                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()
+            "format": format,
+            "path": export_path,
+            "rows": len(df),
+            "columns": len(df.columns),
+            "variables": df.columns.tolist(),
         }
 
         # Handle export if requested
@@ -699,6 +715,6 @@ class ChartTool(BaseTool):
                 export_path = os.path.join(self.config.export_dir, export_path)
 
             self._export_result(result, export_path, export_format)
-            result['exported_to'] = export_path
+            result["exported_to"] = export_path
 
         return result