databao/visualizers/vega_chat.py

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 MessageInfo, VegaChatConfig, VegaChatGraph, VegaChatState
from langchain_core.runnables import RunnableConfig
from PIL import Image

from databao.configs.llm import LLMConfig
from databao.core import ExecutionResult, VisualisationResult, Visualizer
from databao.executors.base import GraphExecutor
from databao.visualizers.vega_vis_tool import VegaVisTool

logger = logging.getLogger(__name__)


class VegaChatResult(VisualisationResult):
    spec: dict[str, Any] | None = None
    spec_df: pd.DataFrame | None = None

    # TODO expose as part of the VisualisationResult API
    def interactive(self) -> VegaVisTool | None:
        """Return an interactive UI wizard for the Vega-Lite chart.

        The returned chart object can be rendered in interactive notebooks."""
        if self.spec is None or self.spec_df is None:
            return None
        return VegaVisTool(self.spec, self.spec_df)

    def altair(self) -> altair.Chart | None:
        """Return an interactive Altair chart.

        The returned chart object can be rendered in interactive notebooks."""
        if self.spec is None or self.spec_df is 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.
    return VegaLLMConfig(
        name=llm_config.name,
        temperature=llm_config.temperature,
        max_tokens=llm_config.max_tokens,
        reasoning_effort=llm_config.reasoning_effort,
        cache_system_prompt=llm_config.cache_system_prompt,
        timeout=llm_config.timeout,
        api_base_url=llm_config.api_base_url,
        use_responses_api=llm_config.use_responses_api,
        ollama_pull_model=llm_config.ollama_pull_model,
        model_kwargs=llm_config.model_kwargs,
    )


class VegaChatVisualizer(Visualizer):
    def __init__(self, llm_config: LLMConfig, *, return_interactive_chart: bool = False):
        vega_llm = _convert_llm_config(llm_config)
        self._vega_config = VegaChatConfig(
            llm_config=vega_llm,
            data_normalize_column_names=True,  # To deal with column names that have special characters
        )
        self._return_interactive_chart = return_interactive_chart

    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,
        )

    def _run_vega_chat(
        self, request: str, df: pd.DataFrame, *, messages: list[MessageInfo] | None = None, stream: bool = False
    ) -> VegaChatResult:
        vega_chat = VegaChatGraph(self._vega_config, df=df)
        start_state = vega_chat.get_start_state(request, messages=messages)
        compiled_graph = vega_chat.compile_graph(is_async=False)
        # Use an empty `config` instead of `None` due to a bug in the "AI Agents Debugger" PyCharm plugin.
        final_state: VegaChatState = GraphExecutor._invoke_graph_sync(
            compiled_graph, start_state, config=RunnableConfig(), stream=stream
        )
        processed_df = vega_chat.dataframe
        return self._process_result(final_state, processed_df)

    def visualize(self, request: str | None, data: ExecutionResult, *, stream: bool = False) -> VegaChatResult:
        if data.df is None:
            return VegaChatResult(text="Nothing to visualize", meta={}, plot=None, code=None, visualizer=self)
        if request is None:
            # We could also call the ChartRecommender module, but since we want a
            # single output plot, we'll just use a simple prompt.
            request = "I don't know what the data is about. Show me an interesting plot."
        return self._run_vega_chat(request, data.df, stream=stream)

    def edit(self, request: str, visualization: VisualisationResult, *, stream: bool = False) -> VegaChatResult:
        if not isinstance(visualization, VegaChatResult):
            raise ValueError(f"{self.__class__.__name__} can only edit {VegaChatResult.__name__} objects")
        if visualization.spec_df is None:
            raise ValueError("No dataframe found in the provided visualization")
        messages = visualization.meta.get("messages", None)
        if messages is None:
            raise ValueError("No message history found in the provided visualization")
        return self._run_vega_chat(request, visualization.spec_df, messages=messages, stream=stream)