feat(pkg-py): Add ggsql visualization integration

pkg-py/examples/10-viz-app.py

@@ -0,0 +1,21 @@
+from querychat import QueryChat
+from querychat.data import titanic
+
+from shiny import App, ui
+
+qc = QueryChat(
+    titanic(),
+    "titanic",
+    tools=("query", "visualize_query"),
+)
+
+app_ui = ui.page_fillable(
+    qc.ui(),
+)
+
+
+def server(input, output, session):
+    qc.server()
+
+
+app = App(app_ui, server)

pkg-py/src/querychat/__init__.py

@@ -2,9 +2,11 @@
 from ._deprecated import mod_server as server
 from ._deprecated import mod_ui as ui
 from ._shiny import QueryChat
+from .tools import VisualizeQueryData
 
 __all__ = (
     "QueryChat",
+    "VisualizeQueryData",
     # TODO(lifecycle): Remove these deprecated functions when we reach v1.0
     "greeting",
     "init",

pkg-py/src/querychat/_datasource.py

@@ -214,6 +214,7 @@ def __init__(self, df: nw.DataFrame, table_name: str):
         self._df_lib = native_namespace.__name__
 
         self._conn = duckdb.connect(database=":memory:")
+        # NOTE: if native representation is polars, pyarrow is required for registration
         self._conn.register(table_name, self._df.to_native())
         self._conn.execute("""
 -- extensions: lock down supply chain + auto behaviors

pkg-py/src/querychat/_ggsql.py

@@ -0,0 +1,80 @@
+"""Helpers for ggsql integration."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+import narwhals.stable.v1 as nw
+
+if TYPE_CHECKING:
+    import ggsql
+    import polars as pl
+    from narwhals.stable.v1.typing import IntoFrame
+
+    from ._datasource import DataSource
+
+
+def to_polars(data: IntoFrame) -> pl.DataFrame:
+    """Convert any narwhals-compatible frame to a polars DataFrame."""
+    nw_df = nw.from_native(data)
+    if isinstance(nw_df, nw.LazyFrame):
+        nw_df = nw_df.collect()
+    return nw_df.to_polars()
+
+
+def execute_ggsql(data_source: DataSource, query: str) -> ggsql.Spec:
+    """
+    Execute a full ggsql query against a DataSource, returning a Spec.
+
+    Uses ggsql.validate() to split SQL from VISUALISE, executes the SQL
+    through DataSource (preserving database pushdown), then feeds the result
+    into a ggsql DuckDBReader to produce a Spec.
+
+    Parameters
+    ----------
+    data_source
+        The querychat DataSource to execute the SQL portion against.
+    query
+        A full ggsql query (SQL + VISUALISE).
+
+    Returns
+    -------
+    ggsql.Spec
+        The writer-independent plot specification.
+
+    """
+    import ggsql as _ggsql
+
+    validated = _ggsql.validate(query)
+    pl_df = to_polars(data_source.execute_query(validated.sql()))
+
+    reader = _ggsql.DuckDBReader("duckdb://memory")
+    reader.register("_data", pl_df)
+    return reader.execute(f"SELECT * FROM _data {validated.visual()}")
+
+
+def spec_to_altair(spec: ggsql.Spec) -> ggsql.AltairChart:
+    """Render a ggsql Spec to an Altair chart via VegaLiteWriter."""
+    import ggsql as _ggsql
+
+    writer = _ggsql.VegaLiteWriter()
+    return writer.render_chart(spec, validate=False)
+
+
+def extract_title(spec: ggsql.Spec) -> str | None:
+    """
+    Extract the title from a ggsql Spec's rendered Vega-Lite JSON.
+
+    TODO: Replace with ``spec.title()`` once ggsql exposes this natively.
+    """
+    import ggsql as _ggsql
+
+    writer = _ggsql.VegaLiteWriter()
+    vl: dict[str, object] = json.loads(writer.render(spec))
+    title = vl.get("title")
+    if isinstance(title, str):
+        return title
+    if isinstance(title, dict):
+        return title.get("text")
+    return None

pkg-py/src/querychat/_icons.py

@@ -2,7 +2,14 @@
 
 from shiny import ui
 
-ICON_NAMES = Literal["arrow-counterclockwise", "funnel-fill", "terminal-fill", "table"]
+ICON_NAMES = Literal[
+    "arrow-counterclockwise",
+    "bar-chart-fill",
+    "funnel-fill",
+    "graph-up",
+    "terminal-fill",
+    "table",
+]
 
 
 def bs_icon(name: ICON_NAMES) -> ui.HTML:
@@ -14,7 +21,9 @@ def bs_icon(name: ICON_NAMES) -> ui.HTML:
 
 BS_ICONS = {
     "arrow-counterclockwise": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" class="bi bi-arrow-counterclockwise" style="height:1em;width:1em;fill:currentColor;vertical-align:-0.125em;" aria-hidden="true" role="img"><path fill-rule="evenodd" d="M8 3a5 5 0 1 1-4.546 2.914.5.5 0 0 0-.908-.417A6 6 0 1 0 8 2v1z"></path><path d="M8 4.466V.534a.25.25 0 0 0-.41-.192L5.23 2.308a.25.25 0 0 0 0 .384l2.36 1.966A.25.25 0 0 0 8 4.466z"></path></svg>',
+    "bar-chart-fill": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" class="bi bi-bar-chart-fill" style="height:1em;width:1em;fill:currentColor;vertical-align:-0.125em;" aria-hidden="true" role="img"><path d="M1 11a1 1 0 0 1 1-1h2a1 1 0 0 1 1 1v3a1 1 0 0 1-1 1H2a1 1 0 0 1-1-1zm5-4a1 1 0 0 1 1-1h2a1 1 0 0 1 1 1v7a1 1 0 0 1-1 1H7a1 1 0 0 1-1-1zm5-5a1 1 0 0 1 1-1h2a1 1 0 0 1 1 1v12a1 1 0 0 1-1 1h-2a1 1 0 0 1-1-1z"/></svg>',
     "funnel-fill": '<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" class="bi bi-funnel-fill" viewBox="0 0 16 16"><path d="M1.5 1.5A.5.5 0 0 1 2 1h12a.5.5 0 0 1 .5.5v2a.5.5 0 0 1-.128.334L10 8.692V13.5a.5.5 0 0 1-.342.474l-3 1A.5.5 0 0 1 6 14.5V8.692L1.628 3.834A.5.5 0 0 1 1.5 3.5z"/></svg>',
+    "graph-up": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" class="bi bi-graph-up" style="height:1em;width:1em;fill:currentColor;vertical-align:-0.125em;" aria-hidden="true" role="img"><path fill-rule="evenodd" d="M0 0h1v15h15v1H0zm14.817 3.113a.5.5 0 0 1 .07.704l-4.5 5.5a.5.5 0 0 1-.74.037L7.06 6.767l-3.656 5.027a.5.5 0 0 1-.808-.588l4-5.5a.5.5 0 0 1 .758-.06l2.609 2.61 4.15-5.073a.5.5 0 0 1 .704-.07"/></svg>',
     "terminal-fill": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" class="bi bi-terminal-fill " style="height:1em;width:1em;fill:currentColor;vertical-align:-0.125em;" aria-hidden="true" role="img" ><path d="M0 3a2 2 0 0 1 2-2h12a2 2 0 0 1 2 2v10a2 2 0 0 1-2 2H2a2 2 0 0 1-2-2V3zm9.5 5.5h-3a.5.5 0 0 0 0 1h3a.5.5 0 0 0 0-1zm-6.354-.354a.5.5 0 1 0 .708.708l2-2a.5.5 0 0 0 0-.708l-2-2a.5.5 0 1 0-.708.708L4.793 6.5 3.146 8.146z"></path></svg>',
     "table": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" class="bi bi-table " style="height:1em;width:1em;fill:currentColor;vertical-align:-0.125em;" aria-hidden="true" role="img" ><path d="M0 2a2 2 0 0 1 2-2h12a2 2 0 0 1 2 2v12a2 2 0 0 1-2 2H2a2 2 0 0 1-2-2V2zm15 2h-4v3h4V4zm0 4h-4v3h4V8zm0 4h-4v3h3a1 1 0 0 0 1-1v-2zm-5 3v-3H6v3h4zm-5 0v-3H1v2a1 1 0 0 0 1 1h3zm-4-4h4V8H1v3zm0-4h4V4H1v3zm5-3v3h4V4H6zm4 4H6v3h4V8z"></path></svg>',
 }

pkg-py/src/querychat/_preload_viz_deps.py

@@ -0,0 +1,88 @@
+"""
+Preload visualization JavaScript dependencies for faster first-chart render.
+
+When an Altair chart is first rendered through shinywidgets, the browser must:
+
+1. Parse **libembed-amd.js** (~3.9 MB), the Jupyter widget runtime bundled
+   with shinywidgets.
+2. Fetch **vega-embed**, **vega**, and **vega-lite** from a CDN (esm.sh).
+   These are imported by Altair's ``JupyterChart`` anywidget ESM and total
+   roughly 1-2 MB of additional JavaScript.
+
+Together these steps can add several seconds of latency to the first chart.
+Subsequent charts are fast because the browser caches the assets.
+
+This module eliminates that first-chart penalty by rendering a tiny, off-screen
+Altair chart at page-load time.  The browser eagerly walks the full widget
+pipeline, and by the time the user's first real chart arrives the JavaScript
+is already cached.
+
+Both :func:`preload_viz_deps_ui` and :func:`preload_viz_deps_server` must be
+called **outside** a Shiny module so the widget ID is not namespaced.
+Module-level flags ensure each function produces output at most once per
+process, so multiple ``QueryChat`` instances on the same page share a single
+preload widget.
+"""
+
+from __future__ import annotations
+
+PRELOAD_WIDGET_ID = "__querychat_preload_viz_deps__"
+
+_preload_ui_emitted = False
+"""Flag so the preload widget output is only emitted once per process."""
+
+
+def preload_viz_deps_ui():
+    """
+    Return an off-screen widget output that preloads visualization JS dependencies.
+
+    Must be called **outside** a Shiny module (e.g. from ``QueryChat.ui()``) so
+    the widget ID is not namespaced.  Returns ``None`` after the first call.
+
+    When paired with :func:`preload_viz_deps_server`, the browser walks the
+    full rendering pipeline (libembed-amd.js → anywidget ESM → vega-embed CDN),
+    caching everything before the user's first real chart request.
+    """
+    global _preload_ui_emitted  # noqa: PLW0603
+    if _preload_ui_emitted:
+        return None
+    _preload_ui_emitted = True
+
+    from htmltools import tags
+
+    try:
+        from shinywidgets import output_widget
+
+        return tags.div(
+            output_widget(PRELOAD_WIDGET_ID),
+            style="position:absolute; left:-9999px; width:1px; height:1px;",
+        )
+    except ImportError:
+        return None
+
+
+_preload_server_registered = False
+"""Flag so the preload chart is only registered once per process."""
+
+
+def preload_viz_deps_server() -> None:
+    """
+    Register a minimal Altair chart to trigger full JS dependency loading.
+
+    Must be called **outside** a Shiny module (e.g. from ``QueryChat.server()``)
+    so the widget ID matches the un-namespaced output created by
+    :func:`preload_viz_deps_ui`.  Subsequent calls are no-ops.
+    """
+    global _preload_server_registered  # noqa: PLW0603
+    if _preload_server_registered:
+        return
+    _preload_server_registered = True
+
+    try:
+        import altair as alt
+        from shinywidgets import register_widget
+
+        chart = alt.Chart({"values": [{"x": 0}]}).mark_point().encode(x="x:Q")
+        register_widget(PRELOAD_WIDGET_ID, chart)
+    except Exception:  # noqa: S110
+        pass  # Best-effort: don't break the app if deps are unavailable

pkg-py/src/querychat/_querychat_base.py

@@ -25,17 +25,44 @@
 from ._utils import MISSING, MISSING_TYPE, is_ibis_table
 from .tools import (
     UpdateDashboardData,
+    VisualizeQueryData,
     tool_query,
     tool_reset_dashboard,
     tool_update_dashboard,
+    tool_visualize_query,
 )
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
     from narwhals.stable.v1.typing import IntoFrame
 
-TOOL_GROUPS = Literal["update", "query"]
+TOOL_GROUPS = Literal["update", "query", "visualize_query"]
+DEFAULT_TOOLS: tuple[TOOL_GROUPS, ...] = ("update", "query")
+ALL_TOOLS: tuple[TOOL_GROUPS, ...] = (
+    "update",
+    "query",
+    "visualize_query",
+)
+
+VIZ_TOOLS: tuple[TOOL_GROUPS, ...] = ("visualize_query",)
+
+
+def check_viz_dependencies(tools: tuple[TOOL_GROUPS, ...] | None) -> None:
+    """Raise ImportError early if viz tools are requested but ggsql is not installed."""
+    if tools is None:
+        return
+    has_viz = any(t in VIZ_TOOLS for t in tools)
+    if not has_viz:
+        return
+    try:
+        import altair as alt  # noqa: F401
+        import ggsql  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            f"Visualization tools require ggsql and altair: {e}. "
+            "Install them with: pip install querychat[viz]"
+        ) from e
 
 
 class QueryChatBase(Generic[IntoFrameT]):
@@ -58,7 +85,7 @@ def __init__(
         *,
         greeting: Optional[str | Path] = None,
         client: Optional[str | chatlas.Chat] = None,
-        tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None = ("update", "query"),
+        tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None = DEFAULT_TOOLS,
         data_description: Optional[str | Path] = None,
         categorical_threshold: int = 20,
         extra_instructions: Optional[str | Path] = None,
@@ -72,7 +99,8 @@ def __init__(
                 "Table name must begin with a letter and contain only letters, numbers, and underscores",
             )
 
-        self.tools = normalize_tools(tools, default=("update", "query"))
+        self.tools = normalize_tools(tools, default=DEFAULT_TOOLS)
+        check_viz_dependencies(self.tools)
         self.greeting = greeting.read_text() if isinstance(greeting, Path) else greeting
 
         # Store init parameters for deferred system prompt building
@@ -132,18 +160,22 @@ def client(
         tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None | MISSING_TYPE = MISSING,
         update_dashboard: Callable[[UpdateDashboardData], None] | None = None,
         reset_dashboard: Callable[[], None] | None = None,
+        visualize_query: Callable[[VisualizeQueryData], None] | None = None,
     ) -> chatlas.Chat:
         """
         Create a chat client with registered tools.
 
         Parameters
         ----------
         tools
-            Which tools to include: `"update"`, `"query"`, or both.
+            Which tools to include: `"update"`, `"query"`, `"visualize_query"`,
+            or a combination.
         update_dashboard
             Callback when update_dashboard tool succeeds.
         reset_dashboard
             Callback when reset_dashboard tool is invoked.
+        visualize_query
+            Callback when visualize_query tool succeeds.
 
         Returns
         -------
@@ -172,6 +204,10 @@ def client(
         if "query" in tools:
             chat.register_tool(tool_query(data_source))
 
+        if "visualize_query" in tools:
+            query_viz_fn = visualize_query or (lambda _: None)
+            chat.register_tool(tool_visualize_query(self._data_source, query_viz_fn))
+
         return chat
 
     def generate_greeting(self, *, echo: Literal["none", "output"] = "none") -> str:
@@ -216,6 +252,10 @@ def data_source(self, value: IntoFrame | sqlalchemy.Engine) -> None:
         self._data_source = normalize_data_source(value, self._table_name)
         self._build_system_prompt()
 
+    def _has_viz_tool(self) -> bool:
+        """Check if visualize_query is among the configured tools."""
+        return self.tools is not None and "visualize_query" in self.tools
+
     def cleanup(self) -> None:
         """Clean up resources associated with the data source."""
         if self._data_source is not None: