Repository: gventuri/pandas-ai
File: pandasai/helpers/shortcuts.py (Python)

Code before:
    ) -> None:
        """Plot a confusion matrix."""

        self.run(
            df,
            f"""
Plot a confusion matrix with the following y_true and y_pred:
y_true = {y_true}
y_pred = {y_pred}
""",
        )

    def plot_roc_curve(self, df: pd.DataFrame, y_true: list, y_pred: list) -> None:
        """Plot a ROC curve."""

        self.run(
            df,
            f"""
Plot a ROC curve with the following y_true and y_pred:
y_true = {y_true}
y_pred = {y_pred}
""",
        )

    def boxplot(self, df: pd.DataFrame, col: str | list[str] | None = None, by: str | list[str] | None = None, style: str | None = None):
        """Draw a box plot to show distributions with respect to categories. The user can provide style info """
        if not isinstance(col, (str, list, type(None))):
            raise TypeError("The 'col' argument must be a string, a list of strings, or None.")
        if not isinstance(by, (str, list, type(None))):
            raise TypeError("The 'by' argument must be a string, a list of strings, or None.")

        prompt = "Plot a box-and-whisker plot"

        if isinstance(col, str):
            prompt += f" for the variable '{col}'"
        elif isinstance(col, list):
            var_list = [f"'{v}'" for v in col]
            if len(var_list) > 1:
                prompt += f" for the variables {', '.join(var_list[:-1])} and {var_list[-1]}"
            else:
                prompt += f" for the variable {var_list[0]}"

        if by is not None:
            prompt += f" grouped by '{by}'"

        if style is not None:
            prompt += f"\nStyle: '''{style}'''"

        self.run(df, prompt)

    def rolling_mean(self, df: pd.DataFrame, column: str, window: int) -> pd.DataFrame:

Review comment (suggestion):
Let's add Args and Return here.

I.e:
```
def chat_completion(self, value: str) -> str:
        """
        Query the chat completion API

        Args:
            value (str): Prompt

        Returns:
            str: LLM response
        """
```

Code after:
            df (pd.DataFrame): The DataFrame containing the data.
            x (list): The x values for the bar chart.
            y (list): The y values for the bar chart.

        Returns:
            None
        """

        self.run(
            df,
            f"""
Plot a bar chart with the following x and y:
x = {x}
y = {y}
""",
        )

    def plot_histogram(self, df: pd.DataFrame, column: str) -> None:
        """
        Plot a histogram.

        Args:
            df (pd.DataFrame): The DataFrame containing the data.
            column (str): The column to plot the histogram for.

        Returns:
            None
        """

        self.run(df, f"Plot a histogram of the column {column}.")

    def plot_line_chart(self, df: pd.DataFrame, x: list, y: list) -> None:
        """
        Plot a line chart.

        Args:
            df (pd.DataFrame): The DataFrame containing the data.
            x (list): The x values for the line chart.
            y (list): The y values for the line chart.

        Returns:
            None
        """

        self.run(
            df,
            f"""
Plot a line chart with the following x and y:
x = {x}
y = {y}
""",