diff --git a/pandasai/__init__.py b/pandasai/__init__.py index f3539f89d..3d9e2f9d0 100644 --- a/pandasai/__init__.py +++ b/pandasai/__init__.py @@ -98,22 +98,22 @@ class PandasAI: """ _dl: SmartDatalake = None - _config: Config + _config: [Config | dict] def __init__( - self, - llm=None, - conversational=False, - verbose=False, - enforce_privacy=False, - save_charts=False, - save_charts_path="", - enable_cache=True, - middlewares=None, - custom_whitelisted_dependencies=None, - enable_logging=True, - non_default_prompts: Optional[Dict[str, Type[Prompt]]] = None, - callback: Optional[BaseCallback] = None, + self, + llm=None, + conversational=False, + verbose=False, + enforce_privacy=False, + save_charts=False, + save_charts_path="", + enable_cache=True, + middlewares=None, + custom_whitelisted_dependencies=None, + enable_logging=True, + non_default_prompts: Optional[Dict[str, Type[Prompt]]] = None, + callback: Optional[BaseCallback] = None, ): """ __init__ method of the Class PandasAI @@ -142,8 +142,10 @@ def __init__( # noinspection PyArgumentList # https://stackoverflow.com/questions/61226587/pycharm-does-not-recognize-logging-basicconfig-handlers-argument - warnings.warn("`PandasAI` (class) is deprecated since v1.0 and will be removed " - "in a future release. Please use `SmartDataframe` instead.") + warnings.warn( + "`PandasAI` (class) is deprecated since v1.0 and will be removed " + "in a future release. Please use `SmartDataframe` instead." + ) self._config = Config( conversational=conversational, @@ -161,12 +163,12 @@ def __init__( ) def run( - self, - data_frame: Union[pd.DataFrame, List[pd.DataFrame]], - prompt: str, - show_code: bool = False, - anonymize_df: bool = True, - use_error_correction_framework: bool = True, + self, + data_frame: Union[pd.DataFrame, List[pd.DataFrame]], + prompt: str, + show_code: bool = False, + anonymize_df: bool = True, + use_error_correction_framework: bool = True, ) -> Union[str, pd.DataFrame]: """ Run the PandasAI to make Dataframes Conversational. @@ -198,12 +200,12 @@ def run( return self._dl.chat(prompt) def __call__( - self, - data_frame: Union[pd.DataFrame, List[pd.DataFrame]], - prompt: str, - show_code: bool = False, - anonymize_df: bool = True, - use_error_correction_framework: bool = True, + self, + data_frame: Union[pd.DataFrame, List[pd.DataFrame]], + prompt: str, + show_code: bool = False, + anonymize_df: bool = True, + use_error_correction_framework: bool = True, ) -> Union[str, pd.DataFrame]: """ __call__ method of PandasAI class. It calls the `run` method. diff --git a/pandasai/config.py b/pandasai/config.py index 5e13628b5..bd482b8e5 100644 --- a/pandasai/config.py +++ b/pandasai/config.py @@ -1,10 +1,15 @@ import json +import logging +from typing import Optional + from . import llm, middlewares, callbacks from .helpers.path import find_closest from .schemas.df_config import Config +logger = logging.getLogger(__name__) + -def load_config(override_config: Config = None): +def load_config(override_config: Optional[Config | dict] = None): config = {} if override_config is None: @@ -27,7 +32,7 @@ def load_config(override_config: Config = None): if config.get("callback") and not override_config.get("callback"): config["callback"] = getattr(callbacks, config["callback"])() except Exception: - pass + logger.error("Could not load configuration", exc_info=True) if override_config: config.update(override_config) diff --git a/pandasai/helpers/code_manager.py b/pandasai/helpers/code_manager.py index ac9e97f71..11b6b6b91 100644 --- a/pandasai/helpers/code_manager.py +++ b/pandasai/helpers/code_manager.py @@ -21,7 +21,7 @@ class CodeManager: _dfs: List _middlewares: List[Middleware] = [ChartsMiddleware()] - _config: Config + _config: [Config | dict] _logger: Logger = None _additional_dependencies: List[dict] = [] @@ -30,12 +30,12 @@ class CodeManager: def __init__( self, dfs: List, - config: Config, + config: [Config | dict], logger: Logger, ): """ Args: - config (Config, optional): Config to be used. Defaults to None. + config ([Config | dict], optional): Config to be used. Defaults to None. logger (Logger, optional): Logger to be used. Defaults to None. """ diff --git a/pandasai/smart_dataframe/__init__.py b/pandasai/smart_dataframe/__init__.py index 611a8a03e..ee5ad4f87 100644 --- a/pandasai/smart_dataframe/__init__.py +++ b/pandasai/smart_dataframe/__init__.py @@ -31,7 +31,7 @@ from ..helpers.logger import Logger from ..helpers.df_config_manager import DfConfigManager from ..helpers.from_google_sheets import from_google_sheets -from typing import List, Union +from typing import List, Union, Optional from ..middlewares.base import Middleware from ..helpers.df_info import DataFrameType, df_type from .abstract_df import DataframeAbstract @@ -53,7 +53,7 @@ def __init__( df: DataFrameType, name: str = None, description: str = None, - config: Config = None, + config: Optional[Config | dict] = None, sample_head: pd.DataFrame = None, logger: Logger = None, ): @@ -62,7 +62,7 @@ def __init__( df (Union[pd.DataFrame, pl.DataFrame]): Pandas or Polars dataframe name (str, optional): Name of the dataframe. Defaults to None. description (str, optional): Description of the dataframe. Defaults to "". - config (Config, optional): Config to be used. Defaults to None. + config ([Config | dict], optional): Config to be used. Defaults to None. logger (Logger, optional): Logger to be used. Defaults to None. """ self._original_import = df diff --git a/pandasai/smart_datalake/__init__.py b/pandasai/smart_datalake/__init__.py index 8a6d612b5..dd4a16353 100644 --- a/pandasai/smart_datalake/__init__.py +++ b/pandasai/smart_datalake/__init__.py @@ -44,7 +44,7 @@ class SmartDatalake: _dfs: List[DataFrameType] - _config: Config + _config: [Config | dict] _llm: LLM _cache: Cache = None _logger: Logger @@ -60,14 +60,14 @@ class SmartDatalake: def __init__( self, dfs: List[Union[DataFrameType, Any]], - config: Config = None, + config: Optional[Config | dict] = None, logger: Logger = None, memory: Memory = None, ): """ Args: dfs (List[Union[DataFrameType, Any]]): List of dataframes to be used - config (Config, optional): Config to be used. Defaults to None. + config ([Config | dict], optional): Config to be used. Defaults to None. logger (Logger, optional): Logger to be used. Defaults to None. """ @@ -135,12 +135,12 @@ def _load_dfs(self, dfs: List[Union[DataFrameType, Any]]): smart_dfs.append(df) self._dfs = smart_dfs - def _load_config(self, config: Config): + def _load_config(self, config: [Config | dict]): """ Load a config to be used to run the queries. Args: - config (Config): Config to be used + config ([Config | dict]): Config to be used """ config = load_config(config)