pandas-ai icon indicating copy to clipboard operation
pandas-ai copied to clipboard

Support for Figure object output type

Open Blubbaa opened this issue 6 months ago • 2 comments

🚀 The feature

I would love to have the ability of setting the ouput type to a plotly or matplotlib figure, instead of saving plots to PNG and returning filepaths, where the usefulness is quite limited.

Motivation, pitch

I recently started using pandasai for building custom data analysis apps and I like it quite a lot so far. I was wondering why it is limited to the four output datatypes and even more why the (cumbersome) way of saving images to disk and returning a filepath has been chosen. Maybe it has a security related reason or it is due to the client-server architecture of pandasai? Instead returning objects (like already implemented for the dataframe) would open much more potential, especially regarding plots and figures.

I have already tinkered with modifying the output_type_template.tmpl and output_validator.py in order to make pandasai return figure objects. However I do not really know about potential problems/implications of this and thus am proposing this here as a feature request, since my "hacky" implementation is probably not how it should be implemented.

Here you can see the resulting prompt used and the generated code, which works fine for now in a standalone app. Prompt used:

<dataframe>
dfs[0]:150x5
Sepal_Length,Sepal_Width,Petal_Length,Petal_Width,Class
7.2,3.4,6.4,0.3,Iris-setosa
4.5,4.1,3.5,1.7,Iris-virginica
6.0,2.6,5.9,2.5,Iris-versicolor
</dataframe>


Update this initial code:
"""python
# TODO: import the required dependencies
import pandas as pd

# Write code here

# Declare result var:
type (must be "figure"), value must be a matplotlib.figure or plotly.graph_objects.Figure. Example: { "type": "figure", "value": go.Figure(...) }   

"""



### QUERY
 Plot the sepal length and width of the data and color points by class

Variable `dfs: list[pd.DataFrame]` is already declared.

At the end, declare "result" variable as a dictionary of type and value.

If you are asked to plot a chart, use "plotly" for charts, save as png.


Generate python code and return full updated code:

Resulting Code:

df = dfs[0]
fig = px.scatter(df, x='Sepal_Length', y='Sepal_Width', color='Class', title='Sepal Length vs Sepal Width', labels={'Sepal_Length': 'Sepal Length', 'Sepal_Width': 'Sepal Width'})
result = {'type': 'figure', 'value': fig}

Alternatives

I know its also possible to convert plotly figures from/to json. So maybe this could be another option to return (or potentially also save) the figure as json instead.

Additional context

Final Result in Chatbot App: image

Blubbaa avatar Aug 21 '24 12:08 Blubbaa