add boxplot in plot list as well

Sinaptik-AI · Jun 14, 2024 · 4f1e527 · 4f1e527
1 parent c035cbd
commit 4f1e527
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/pandasai/ee/agents/semantic_agent/pipeline/code_generator.py b/pandasai/ee/agents/semantic_agent/pipeline/code_generator.py
@@ -86,7 +86,8 @@ def execute(self, input_data: Any, **kwargs) -> Any:
     def _get_type(self, input: dict) -> bool:
         return (
             "plot"
-            if input["type"] in ["bar", "line", "histogram", "pie", "scatter"]
+            if input["type"]
+            in ["bar", "line", "histogram", "pie", "scatter", "boxplot"]
             else input["type"]
         )
 

diff --git a/tests/unit_tests/ee/semantic_agent/test__semantic_code_generator.py b/tests/unit_tests/ee/semantic_agent/test__semantic_code_generator.py
@@ -307,11 +307,10 @@ def test_generate_matplolib_boxplot_chart_code(
 
         logic_unit = code_gen.execute(json_str, context=context, logger=logger)
         assert isinstance(logic_unit, LogicUnitOutput)
-        print(logic_unit.output)
         assert (
             logic_unit.output
             == """
-
+import matplotlib.pyplot as plt
 import pandas as pd
 
 sql_query="SELECT `orders`.`ship_country` AS ship_country, SUM(`orders`.`freight`) AS total_freight FROM `orders` GROUP BY ship_country"