diff --git a/pandasai/core/code_execution/code_executor.py b/pandasai/core/code_execution/code_executor.py
index b158ea8be..0bea78e6e 100644
--- a/pandasai/core/code_execution/code_executor.py
+++ b/pandasai/core/code_execution/code_executor.py
@@ -24,25 +24,16 @@ def add_to_env(self, key: str, value: Any) -> None:
         """
         self._environment[key] = value
 
-    def execute(self, code: str) -> dict:
-        try:
-            exec(code, self._environment)
-        except Exception as e:
-            raise CodeExecutionError("Code execution failed") from e
-        return self._environment
-
-    def execute_and_return_result(self, code: str) -> Any:
+    def execute(
+        self, code: str, context: Optional[Dict[str, Any]] = None
+    ) -> CodeExecutionResult:
         """
-        Executes the return updated environment
-        """
-        self.execute(code)
-
-        # Get the result
-        if "result" not in self._environment:
-            raise NoResultFoundError("No result returned")
+        Execute the code and return the result.
 
-        return self._environment.get("result", None)
+        Args:
+            code (str): The code to execute.
+            context (Dict[str, Any], optional): The context to execute the code in.
 
-    @property
-    def environment(self) -> dict:
-        return self._environment
+        Returns:
+            CodeExecutionResult: The result of the code execution.
+        """