def _process_result()

in databao/visualizers/vega_chat.py [0:0]

• 45 lines of code
• 11 McCabe index (conditional complexity)

    def _process_result(self, state: VegaChatState, spec_df: pd.DataFrame) -> VegaChatResult:
        # Use the possibly transformed dataframe tied to the generated spec
        model_out = state["messages"][-1]
        text = model_out.message.text()
        meta = {"messages": state["messages"]}  # Full history. Also used for edit follow ups.
        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! Output: {text}",
                meta=meta,
                plot=None,
                code=spec_json,
                spec=spec,
                spec_df=spec_df,
                visualizer=self,
            )

        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, spec_df)
            elif (png_bytes := vl_to_png_bytes(spec, spec_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=spec_df,
                    visualizer=self,
                )
        elif self._return_interactive_chart:
            plot = VegaVisTool(spec, spec_df)
        else:
            plot = to_altair_chart(spec, spec_df)

        return VegaChatResult(
            text=text,
            meta=meta,
            plot=plot,
            code=spec_json,
            spec=spec,
            spec_df=spec_df,
            visualizer=self,
        )