Updated the executor function to provide support for plotly that fixes #36

lida/components/executor.py

@@ -9,6 +9,7 @@
 
 import matplotlib.pyplot as plt
 import pandas as pd
+import plotly.io as pio
 
 from lida.datamodel import ChartExecutorResponse, Summary
 
@@ -235,7 +236,47 @@ def execute(
                             )
                         )
             return charts
+
+        elif library == "plotly":
+            for code in code_specs:
+                try:
+                    ex_locals = get_globals_dict(code, data)
+                    exec(code, ex_locals)
+                    chart = ex_locals["chart"]
+
+                    if pio:
+                        chart_bytes = pio.to_image(chart, 'png')
+                        plot_data = base64.b64encode(chart_bytes).decode('utf-8')
+
+                        charts.append(
+                            ChartExecutorResponse(
+                                spec=None,
+                                status=True,
+                                raster=plot_data,
+                                code=code,
+                                library=library,
+                            )
+                        )
+                except Exception as exception_error:
+                    print(code)
+                    print(traceback.format_exc())
+                    if return_error:
+                        charts.append(
+                            ChartExecutorResponse(
+                                spec=None,
+                                status=False,
+                                raster=None,
+                                code=code,
+                                library=library,
+                                error={
+                                    "message": str(exception_error),
+                                    "traceback": traceback.format_exc(),
+                                },
+                            )
+                        )
+            return charts
+
         else:
             raise Exception(
-                f"Unsupported library. Supported libraries are altair, matplotlib, seaborn, ggplot. You provided {library}"
+                f"Unsupported library. Supported libraries are altair, matplotlib, seaborn, ggplot, plotly. You provided {library}"
             )

lida/components/scaffold.py

@@ -83,6 +83,23 @@ def plot(data: pd.DataFrame):
     return chart
 chart = plot(data) # data already contains the data to be plotted.  Always include this line. No additional code beyond this line..
 """
+
+        elif library == "plotly":
+            instructions = {
+                "role": "system",
+                "content": f"If calculating metrics such as mean, median, mode, etc. ALWAYS use the option 'numeric_only=True' when applicable and available, AVOID visualizations that require nbformat library, ALWAYS return the chart object {mpl_pre}",
+            }
+            template = \
+                """
+import plotly.express as px
+<imports>
+def plot(data: pd.DataFrame):
+    <stub> # only modify this section
+    fig.show()
+    return chart
+chart = plot(data) # data already contains the data to be plotted.  Always include this line. No additional code beyond this line..
+"""
+
         else:
             raise ValueError(
                 "Unsupported library. Choose from 'matplotlib', 'seaborn', 'plotly', 'bokeh', 'ggplot', 'altair'."

pyproject.toml

@@ -36,7 +36,8 @@ dependencies = [
     "networkx",
     "geopandas",
     "matplotlib-venn",
-    "wordcloud", 
+    "wordcloud",
+    "kaleido"
 ]
 optional-dependencies = {web = ["fastapi", "uvicorn"], transformers = ["llmx[transformers]"], tools=["geopy", "basemap", "basemap-data-hires"], infographics=["peacasso"]}