adding plotly

instructing the template with the opportunity to use plotly

Author: matteocacciola

pandasai/core/code_execution/environment.py

@@ -29,6 +29,7 @@ def get_environment() -> dict:
         "pd": import_dependency("pandas"),
         "plt": import_dependency("matplotlib.pyplot"),
         "np": import_dependency("numpy"),
+        "px": import_dependency("plotly.express"),
     }
 
     return env

pandasai/core/code_generation/code_cleaning.py

@@ -136,9 +136,10 @@ def clean_code(self, code: str) -> str:
             tuple: Cleaned code as a string and a list of additional dependencies.
         """
         code = self._replace_output_filenames_with_temp_chart(code)
+        code = self._replace_output_filenames_with_temp_json_chart(code)
 
-        # If plt.show is in the code, remove that line
-        code = re.sub(r"plt.show\(\)", "", code)
+        # If plt.show or fig.show is in the code, remove that line
+        code = re.sub(r"[a-z].show\(\)", "", code)
 
         tree = ast.parse(code)
         new_body = []
@@ -166,3 +167,16 @@ def _replace_output_filenames_with_temp_chart(self, code: str) -> str:
             lambda m: f"{m.group(1)}{chart_path}{m.group(1)}",
             code,
         )
+
+    def _replace_output_filenames_with_temp_json_chart(self, code: str) -> str:
+        """
+        Replace output file names with "temp_chart.json" (in case of usage of plotly).
+        """
+        _id = uuid.uuid4()
+        chart_path = os.path.join(DEFAULT_CHART_DIRECTORY, f"temp_chart_{_id}.json")
+        chart_path = chart_path.replace("\\", "\\\\")
+        return re.sub(
+            r"""(['"])([^'"]*\.json)\1""",
+            lambda m: f"{m.group(1)}{chart_path}{m.group(1)}",
+            code,
+        )

pandasai/core/prompts/templates/generate_python_code_with_sql.tmpl

@@ -10,6 +10,9 @@ def execute_sql_query(sql_query: str) -> pd.Dataframe
     """This method connects to the database, executes the sql query and returns the dataframe"""
 </function>
 
+For the charts, you can either use `matplotlib.pyplot` or `plotly.express` to generate the charts.
+If you use `plotly.express`, you have to save each chart as a dictionary into a JSON file.
+
 {% if last_code_generated != "" and context.memory.count() > 0 %}
 {{ last_code_generated }}
 {% else %}
@@ -31,4 +34,4 @@ At the end, declare "result" variable as a dictionary of type and value.
 
 Generate python code and return full updated code:
 
-### Note: Use only relevant table for query and do aggregation, sorting, joins and grouby through sql query
+### Note: Use only relevant table for query and do aggregation, sorting, joins and group by through sql query

pandasai/core/prompts/templates/shared/output_type_template.tmpl

@@ -1,5 +1,5 @@
 {% if not output_type %}
-type (possible values "string", "number", "dataframe", "plot"). No other type available. Examples: { "type": "string", "value": f"The highest salary is {highest_salary}." } or { "type": "number", "value": 125 } or { "type": "dataframe", "value": pd.DataFrame({...}) } or { "type": "plot", "value": "temp_chart.png" }
+type (possible values "string", "number", "dataframe", "plot", "iplot"). No other type available. "plot" is when "matplotlib" is used; "iplot" when "plotly" si used. Examples: { "type": "string", "value": f"The highest salary is {highest_salary}." } or { "type": "number", "value": 125 } or { "type": "dataframe", "value": pd.DataFrame({...}) } or { "type": "plot", "value": "temp_chart.png" } or { "type": "iplot", "value": "temp_chart.json" }
 {% elif output_type == "number" %}
 type (must be "number"), value must int. Example: { "type": "number", "value": 125 }
 {% elif output_type == "string" %}
@@ -8,4 +8,6 @@ type (must be "string"), value must be string. Example: { "type": "string", "val
 type (must be "dataframe"), value must be pd.DataFrame or pd.Series. Example: { "type": "dataframe", "value": pd.DataFrame({...}) }
 {% elif output_type == "plot" %}
 type (must be "plot"), value must be string. Example: { "type": "plot", "value": "temp_chart.png" }
+{% elif output_type == "iplot" %}
+type (must be "iplot"), value must be string. Example: { "type": "iplot", "value": "temp_chart.json" }
 {% endif %}

pandasai/core/response/__init__.py

@@ -2,6 +2,7 @@
 from .chart import ChartResponse
 from .dataframe import DataFrameResponse
 from .error import ErrorResponse
+from .interactive_chart import InteractiveChartResponse
 from .number import NumberResponse
 from .parser import ResponseParser
 from .string import StringResponse
@@ -10,6 +11,7 @@
     "ResponseParser",
     "BaseResponse",
     "ChartResponse",
+    "InteractiveChartResponse",
     "DataFrameResponse",
     "NumberResponse",
     "StringResponse",

pandasai/core/response/interactive_chart.py

@@ -0,0 +1,34 @@
+import json
+import os
+from typing import Any
+
+from .base import BaseResponse
+
+
+class InteractiveChartResponse(BaseResponse):
+    def __init__(self, value: Any, last_code_executed: str):
+        super().__init__(value, "ichart", last_code_executed)
+
+    def _get_chart(self) -> dict:
+        if isinstance(self.value, dict):
+            return self.value
+
+        if isinstance(self.value, str):
+            if os.path.exists(self.value):
+                with open(self.value, "rb") as f:
+                    return json.load(f)
+
+            return json.loads(self.value)
+
+        raise ValueError("Invalid value type for InteractiveChartResponse. Expected dict or str.")
+
+    def save(self, path: str):
+        img = self._get_chart()
+        with open(path, "w") as f:
+            json.dump(img, f)
+
+    def __str__(self) -> str:
+        return self.value if isinstance(self.value, str) else json.dumps(self.value)
+
+    def get_dict_image(self) -> dict:
+        return self._get_chart()

pandasai/core/response/parser.py

@@ -8,6 +8,7 @@
 from .base import BaseResponse
 from .chart import ChartResponse
 from .dataframe import DataFrameResponse
+from .interactive_chart import InteractiveChartResponse
 from .number import NumberResponse
 from .string import StringResponse
 
@@ -26,6 +27,8 @@ def _generate_response(self, result: dict, last_code_executed: str = None):
             return DataFrameResponse(result["value"], last_code_executed)
         elif result["type"] == "plot":
             return ChartResponse(result["value"], last_code_executed)
+        elif result["type"] == "iplot":
+            return InteractiveChartResponse(result["value"], last_code_executed)
         else:
             raise InvalidOutputValueMismatch(f"Invalid output type: {result['type']}")
 
@@ -72,4 +75,16 @@ def _validate_response(self, result: dict):
                     "Invalid output: Expected a plot save path str but received an incompatible type."
                 )
 
+        elif result["type"] == "iplot":
+            if not isinstance(result["value"], (str, dict)):
+                raise InvalidOutputValueMismatch(
+                    "Invalid output: Expected a plot save path str but received an incompatible type."
+                )
+
+            path_to_plot_pattern = r"^(\/[\w.-]+)+(/[\w.-]+)*$|^[^\s/]+(/[\w.-]+)*$"
+            if not bool(re.match(path_to_plot_pattern, result["value"])):
+                raise InvalidOutputValueMismatch(
+                    "Invalid output: Expected a plot save path str but received an incompatible type."
+                )
+
         return True

poetry.lock

@@ -28,6 +28,7 @@ seaborn = "^0.12.2"
 sqlglot = "^25.0.3" 
 pyarrow = "^14.0.1"
 pyyaml = "^6.0.2"
+plotly = "^6.1.1"
 
 [tool.poetry.group.dev]
 optional = true

pyproject.toml

@@ -18,6 +18,7 @@ def test_get_environment_with_secure_mode(self, mock_import_dependency):
         self.assertIn("pd", env)
         self.assertIn("plt", env)
         self.assertIn("np", env)
+        self.assertIn("px", env)
 
     @patch("pandasai.core.code_execution.environment.import_dependency")
     def test_get_environment_without_secure_mode(self, mock_import_dependency):
@@ -28,6 +29,7 @@ def test_get_environment_without_secure_mode(self, mock_import_dependency):
         self.assertIn("pd", env)
         self.assertIn("plt", env)
         self.assertIn("np", env)
+        self.assertIn("px", env)
         self.assertIsInstance(env["pd"], MagicMock)
 
     @patch("pandasai.core.code_execution.environment.importlib.import_module")