Repository: gventuri/pandas-ai
File: pandasai/smart_datalake/__init__.py (Python)

Code before:

        self._query_exec_tracker.publish()

        return result

    def _add_result_to_memory(self, result: dict):
        """
        Add the result to the memory.

        Args:
            result (dict): The result to add to the memory
        """
        if not isinstance(result, dict):
            self.logger.log(
                f"Inappropriate type of 'result' produced by generated code. "
                f"Expected 'dict', got '{result.__class__.__name__}'",
                level=logging.WARNING,
            )
            return

        if "type" not in result or "value" not in result:
            self.logger.log(
                f"Both 'type' and 'value' items should be present in 'result' "
                f"produced by generated code. Instead it contains the next "
                f"content:\n{result}"
            )
            return

        if result["type"] in ["string", "number"]:
            self._memory.add(result["value"], False)
        elif result["type"] in ["dataframe", "plot"]:
            self._memory.add("Ok here it is", False)

    def _retry_run_code(self, code: str, e: Exception) -> List:
        """
        A method to retry the code execution with error correction framework.

        Args:
            code (str): A python code
            e (Exception): An exception
            dataframes

        Returns (str): A python code
        """

        self.logger.log(f"Failed with error: {e}. Retrying", logging.ERROR)

        default_values = {
            "engine": self._dfs[0].engine,
            "code": code,
            "error_returned": e,

Review comment (suggestion):
Great catch. I'm not entirely sure this is the place to do this check.

What if before we run `self._add_result_to_memory(result)` on line 508, we execute another method `self._validate_result` and in there we:
- check if both type and value are present
- check if the type is one of the ones accepted

If any of the conditions is not true, we raise an error.

What do you think? @nautics889 

Code after:

        result = self._query_exec_tracker.execute_func(
            self._response_parser.parse, result
        )

        self._query_exec_tracker.success = True

        self._query_exec_tracker.publish()

        return result

    def _add_result_to_memory(self, result: dict):
        """
        Add the result to the memory.

        Args:
            result (dict): The result to add to the memory
        """
        if result["type"] in ["string", "number"]:
            self._memory.add(result["value"], False)
        elif result["type"] in ["dataframe", "plot"]:
            self._memory.add("Ok here it is", False)

    def _retry_run_code(self, code: str, e: Exception) -> List:
        """
        A method to retry the code execution with error correction framework.

        Args:
            code (str): A python code
            e (Exception): An exception
            dataframes

        Returns (str): A python code
        """

        self.logger.log(f"Failed with error: {e}. Retrying", logging.ERROR)

        default_values = {
            "engine": self._dfs[0].engine,
            "code": code,
            "error_returned": e,
        }
        error_correcting_instruction = self._get_prompt(
            "correct_error",
            default_prompt=CorrectErrorPrompt,
            default_values=default_values,
        )

        result = self._llm.generate_code(error_correcting_instruction)
        if self._config.callback is not None:
            self._config.callback.on_code(result[0])