Repository: gventuri/pandas-ai
File: pandasai/agent/base.py (Python)

Code before:
                for df in self._state.dfs:
                    con.register(df.schema.name, df)

                # Execute the query and fetch the result as a pandas DataFrame
                result = con.sql(query).df()

            return result
        except duckdb.Error as e:
            raise RuntimeError(f"SQL execution failed: {e}") from e

    def _execute_sql_query(self, query: str) -> pd.DataFrame:
        """
        Executes an SQL query on registered DataFrames.

        Args:
            query (str): The SQL query to execute.

        Returns:
            pd.DataFrame: The result of the SQL query as a pandas DataFrame.
        """
        if not self._state.dfs:
            raise ValueError("No DataFrames available to register for query execution.")

        if (
            self._state.dfs[0].schema.source
            and self._state.dfs[0].schema.source.type in LOCAL_SOURCE_TYPES
        ):
            return self._execute_local_sql_query(query)
        else:
            return self._state.dfs[0].execute_sql_query(query)

    def execute_with_retries(self, code: str) -> Any:
        """Execute the code with retry logic."""
        max_retries = self._state.config.max_retries
        attempts = 0

        while attempts <= max_retries:
            try:
                result = self.execute_code(code)
                return self._response_parser.parse(result, code)
            except CodeExecutionError as e:
                attempts += 1
                if attempts > max_retries:
                    self._state.logger.log(f"Max retries reached. Error: {e}")
                    raise
                self._state.logger.log(
                    f"Retrying execution ({attempts}/{max_retries})..."
                )
                code = self._regenerate_code_after_error(code, e)

    def train(

Review comment (suggestion):
Here given the repetition and how nested it is, we might declare a variable `source` or similar for readability

Code after:
                for df in self._state.dfs:
                    con.register(df.schema.name, df)

                # Execute the query and fetch the result as a pandas DataFrame
                result = con.sql(query).df()

            return result
        except duckdb.Error as e:
            raise RuntimeError(f"SQL execution failed: {e}") from e

    def _execute_sql_query(self, query: str) -> pd.DataFrame:
        """
        Executes an SQL query on registered DataFrames.

        Args:
            query (str): The SQL query to execute.

        Returns:
            pd.DataFrame: The result of the SQL query as a pandas DataFrame.
        """
        if not self._state.dfs:
            raise ValueError("No DataFrames available to register for query execution.")

        df0 = self._state.dfs[0]
        source = df0.schema.source or None

        if source and source.type in LOCAL_SOURCE_TYPES:
            return self._execute_local_sql_query(query)
        else:
            return df0.execute_sql_query(query)

    def execute_with_retries(self, code: str) -> Any:
        """Execute the code with retry logic."""
        max_retries = self._state.config.max_retries
        attempts = 0

        while attempts <= max_retries:
            try:
                result = self.execute_code(code)
                return self._response_parser.parse(result, code)
            except CodeExecutionError as e:
                attempts += 1
                if attempts > max_retries:
                    self._state.logger.log(f"Max retries reached. Error: {e}")
                    raise
                self._state.logger.log(
                    f"Retrying execution ({attempts}/{max_retries})..."
                )
                code = self._regenerate_code_after_error(code, e)

    def train(