add boxplot in plot list as well

Author: ArslanSaleem

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"]
         )
 

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"