Handle invalid Vega-Lite specs that are still drawable (#75)

Author: mare5x

databao/core/visualizer.py

@@ -38,9 +38,19 @@ def _repr_mimebundle_(self, include: Any = None, exclude: Any = None) -> Any:
         if hasattr(self.plot, "_repr_mimebundle_"):
             return self.plot._repr_mimebundle_(include, exclude)
 
-        plot_html = self._get_plot_html()
-        if plot_html is not None:
-            return {"text/html": plot_html}
+        mimebundle = {}
+        if (plot_html := self._get_plot_html()) is not None:
+            mimebundle["text/html"] = plot_html
+
+        # TODO Handle all _repr_*_ methods
+        # These are mostly for fallback representations
+        if hasattr(self.plot, "_repr_png_"):
+            mimebundle["image/png"] = self.plot._repr_png_()
+        if hasattr(self.plot, "_repr_jpeg_"):
+            mimebundle["image/jpeg"] = self.plot._repr_jpeg_()
+
+        if len(mimebundle) > 0:
+            return mimebundle
         return None
 
     def _get_plot_html(self) -> str | None:

databao/visualizers/vega_chat.py

@@ -1,19 +1,26 @@
 import dataclasses
+import io
 import json
+import logging
 from typing import Any
 
 import altair
 import pandas as pd
+from edaplot.image_utils import vl_to_png_bytes
 from edaplot.llms import LLMConfig as VegaLLMConfig
 from edaplot.vega import to_altair_chart
 from edaplot.vega_chat.vega_chat import VegaChat, VegaChatConfig
+from PIL import Image
 
 from databao.configs.llm import LLMConfig
 from databao.core import ExecutionResult, VisualisationResult, Visualizer
 from databao.visualizers.vega_vis_tool import VegaVisTool
 
+logger = logging.getLogger(__name__)
+
 
 class VegaChatResult(VisualisationResult):
+    plot: VegaVisTool | altair.Chart | Image.Image | None = None
     spec: dict[str, Any] | None = None
     spec_df: pd.DataFrame | None = None
 
@@ -33,6 +40,14 @@ def altair(self) -> altair.Chart | None:
             return None
         return to_altair_chart(self.spec, self.spec_df)
 
+    def image(self) -> Image.Image | None:
+        """Return a static PIL.Image.Image."""
+        if self.spec is None or self.spec_df is None:
+            return None
+        if (png_bytes := vl_to_png_bytes(self.spec, self.spec_df)) is not None:
+            return Image.open(io.BytesIO(png_bytes))
+        return None
+
 
 def _convert_llm_config(llm_config: LLMConfig) -> VegaLLMConfig:
     # N.B. The two config classes are nearly identical.
@@ -74,28 +89,49 @@ def visualize(self, request: str | None, data: ExecutionResult) -> VegaChatResul
         model = VegaChat.from_config(config=self._vega_config, df=data.df)
         model_out = model.query_sync(request)
 
+        # Use the possibly transformed dataframe tied to the generated spec
+        preprocessed_df = model.dataframe
+        text = model_out.message.text()
+        meta = dataclasses.asdict(model_out)
         spec = model_out.spec
+        spec_json = json.dumps(spec, indent=2) if spec is not None else None
         if spec is None or not model_out.is_drawable or model_out.is_empty_chart:
             return VegaChatResult(
-                text=f"Failed to visualize request {request}",
-                meta=dataclasses.asdict(model_out),
+                text=f"Failed to visualize request! Output: {text}",
+                meta=meta,
                 plot=None,
-                code=None,
+                code=spec_json,
+                spec=spec,
+                spec_df=preprocessed_df,
             )
 
-        text = model_out.message.text()
-        spec_json = json.dumps(spec, indent=2)
-
-        # Use the possibly transformed dataframe tied to the generated spec
-        preprocessed_df = model.dataframe
-        if self._return_interactive_chart:
+        if not model_out.is_valid_schema and model_out.is_drawable:
+            # Vega-Lite specs can be invalid (so cannot be used with altair), but they might still be drawable with
+            # another backend.
+            logger.warning("Generated Vega-Lite spec is not valid, but it is still drawable: %s", spec_json)
+            if self._return_interactive_chart:
+                # The VegaVisTool backend uses vega-embed so it can handle corrupt specs
+                plot = VegaVisTool(spec, preprocessed_df)
+            elif (png_bytes := vl_to_png_bytes(spec, preprocessed_df)) is not None:
+                # Try to convert to an Image that can still be displayed in Jupyter notebooks
+                plot = Image.open(io.BytesIO(png_bytes))
+            else:
+                return VegaChatResult(
+                    text=f"Failed to visualize request! Output: {text}",
+                    meta=meta,
+                    plot=None,
+                    code=spec_json,
+                    spec=spec,
+                    spec_df=preprocessed_df,
+                )
+        elif self._return_interactive_chart:
             plot = VegaVisTool(spec, preprocessed_df)
         else:
             plot = to_altair_chart(spec, preprocessed_df)
 
         return VegaChatResult(
             text=text,
-            meta=dataclasses.asdict(model_out),
+            meta=meta,
             plot=plot,
             code=spec_json,
             spec=spec,