feat(py): opt-in "maps" capability for visualizations

pkg-py/CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * **Faster iterative visualization, and support for layer-level data sources.** Visualization queries now execute as a whole through `ggsql.CachingReader`: the underlying SQL runs against your data source once and the result is memoized for the chat session, so re-rendering a chart with a different encoding no longer re-queries the source (a win for remote databases especially). Visualization layers can also draw from other relations in the query — e.g. raw points from the final `SELECT` with a summary line from a CTE — via `DRAW ... MAPPING ... FROM <source>`. Caching can be tuned or disabled with the `GGSQL_CACHE_TTL`, `GGSQL_CACHE_MAX_BYTES`, and `GGSQL_CACHE_DISABLED` environment variables. (#321)
 
+* New opt-in `"maps"` capability: pass `tools=("default", "maps")` to let the LLM draw choropleths (via the builtin `ggsql:world` dataset) and point maps through ggsql's spatial support. The `"default"` tools entry expands to the default tool set, so extra capabilities no longer require restating every tool. (#323)
+
 ### Bug fixes
 
 * JSON pins now fall back to `pin_read()` when DuckDB's JSON extension cannot be installed or loaded, including on builds that disable extension loading. Native DuckDB JSON reading is retained when the extension is available.
@@ -59,6 +61,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ```python
   qc = QueryChat(None, table_name="my_table")
 
+
   def server(input, output, session):
       qc.server(data_source=conn.table("my_table"), client=chat_client)
   ```
@@ -87,8 +90,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
   # Or, register all tables from a SQLAlchemy engine or Ibis backend at once:
   qc = QueryChat()
-  qc.add_tables(engine)       # SQLAlchemy engine
-  qc.add_tables(ibis_backend) # Ibis backend
+  qc.add_tables(engine)  # SQLAlchemy engine
+  qc.add_tables(ibis_backend)  # Ibis backend
 
   qc_vals = qc.server()
   qc_vals.table("orders").df()

pkg-py/docs/visualize.qmd

@@ -114,6 +114,20 @@ Tune or disable the cache with environment variables (set before the app starts)
 
 Learn more about the ggsql grammar at [ggsql.org](https://ggsql.org/).
 
+## Maps
+
+querychat can draw choropleths and point maps through ggsql's spatial support.
+Maps are opt-in because enabling them downloads the DuckDB `spatial` extension (a one-time, cached download) the first time they are set up:
+
+```{.python}
+qc = QueryChat(titanic(), "titanic", tools=("default", "maps"))
+```
+
+The `"default"` entry expands to the default tool set, so you don't have to restate every tool to add maps.
+
+With maps enabled, the LLM can answer geographic questions: "map revenue by country" joins your data with the builtin `ggsql:world` country boundaries (on ISO codes or country names), and "map our store locations" draws lon/lat points over a world outline.
+If a map query is attempted while maps are not enabled, the tool explains how to enable them.
+
 ## See also
 
 - [Tools](tools.qmd) — Understand what querychat can do under the hood

pkg-py/examples/13-maps-app.py

@@ -0,0 +1,87 @@
+# QueryChat maps capability demo — chat-forward, using Shiny Express.
+#
+# Two real-world tables are registered:
+#
+# - ``earthquakes``: every magnitude 2.5+ earthquake worldwide from the past
+#   month, from the USGS real-time feed (lon/lat point data).
+# - ``gdp``: GDP by country (World Bank, via the plotly datasets repo) with
+#   ISO 3166-1 alpha-3 codes, suitable for choropleths joined to the builtin
+#   ``ggsql:world`` country boundaries.
+#
+# The opt-in ``"maps"`` capability lets the LLM draw choropleths and point maps
+# through ggsql's spatial support. Try asking:
+#
+# - "Map the earthquakes from the past month, sized by magnitude"
+# - "Draw a choropleth of GDP by country"
+#
+# Usage:
+#     cd pkg-py
+#     uv run shiny run examples/13-maps-app.py
+
+import pandas as pd
+from querychat.express import QueryChat
+from shiny.express import render, ui
+from shinychat import chat_nav_panel
+
+# Real-time USGS feed: all magnitude 2.5+ earthquakes, past 30 days
+earthquakes = pd.read_csv(
+    "https://earthquake.usgs.gov/earthquakes/feed/v1.0/summary/2.5_month.csv"
+)
+
+# World Bank GDP (current US$) by country with ISO alpha-3 codes
+gdp = pd.read_csv(
+    "https://raw.githubusercontent.com/plotly/datasets/master/2014_world_gdp_with_codes.csv"
+)
+gdp.columns = ["country", "gdp_billions_usd", "iso_a3"]
+
+greeting = r"""
+Welcome! I'm chatting with two live datasets:
+
+* **earthquakes** — every magnitude 2.5+ earthquake worldwide from the past
+  month (USGS real-time feed)
+* **gdp** — GDP by country (World Bank), with ISO country codes
+
+Maps are enabled! Try asking:
+
+* <span class='suggestion'>Map the earthquakes from the past month, sized by magnitude</span>
+* <span class='suggestion'>Draw a choropleth of GDP by country</span>
+"""
+
+# 1. Provide data sources to QueryChat, with the opt-in "maps" capability
+qc = QueryChat(earthquakes, "earthquakes", greeting=greeting, tools=("default", "maps"))
+qc.add_table(gdp, "gdp")
+
+# 2. Hold the reactive data views for the secondary pages (qc.page() owns the
+#    entire page, so outputs can't live at the top level)
+with ui.hold() as earthquakes_view, ui.card(fill=True):
+
+    @render.data_frame
+    def earthquakes_table():
+        return qc.table("earthquakes").df()
+
+with ui.hold() as gdp_view, ui.card(fill=True):
+
+    @render.data_frame
+    def gdp_table():
+        return qc.table("gdp").df()
+
+# 3. Create a chat-first page (the chat owns the full window)
+qc.page(
+    "QueryChat Maps Demo",
+    pages_navbar=[
+        chat_nav_panel(
+            "Earthquakes",
+            earthquakes_view,
+            value="earthquakes",
+            sidebar=False,
+            content_width="100%",
+        ),
+        chat_nav_panel(
+            "GDP",
+            gdp_view,
+            value="gdp",
+            sidebar=False,
+            content_width="100%",
+        ),
+    ],
+)

pkg-py/src/querychat/_querychat_base.py

@@ -58,7 +58,7 @@
     from ._data_dict import DataDict
     from ._viz_tools import VisualizeData
 
-TOOL_GROUPS = Literal["filter", "update", "query", "visualize"]
+TOOL_GROUPS = Literal["filter", "update", "query", "visualize", "maps", "default"]
 DEFAULT_TOOLS: tuple[TOOL_GROUPS, ...] = ("filter", "query", "visualize")
 TABLE_NAME_PATTERN = r"^[a-zA-Z][a-zA-Z0-9_]*$"
 
@@ -312,7 +312,12 @@ def _create_session_client(
         if "visualize" in resolved_tools:
             viz_fn = visualize or (lambda _: None)
             chat.register_tool(
-                tool_visualize(executor, viz_fn, multi_table=multi_table)
+                tool_visualize(
+                    executor,
+                    viz_fn,
+                    multi_table=multi_table,
+                    maps="maps" in resolved_tools,
+                )
             )
 
         return chat
@@ -826,6 +831,17 @@ def normalize_tools(
         resolved = set(tools)
     if resolved is not None:
         resolved = {"update" if t == "filter" else t for t in resolved}
+        # Expand the "default" sentinel to the caller's default tool set.
+        if "default" in resolved:
+            resolved.discard("default")
+            resolved |= set(default) if default is not None else set()
+            resolved = {"update" if t == "filter" else t for t in resolved}
+        # "maps" is a capability modifier of the visualize tool.
+        if "maps" in resolved and "visualize" not in resolved:
+            raise ValueError(
+                'The "maps" capability requires the "visualize" tool. '
+                'Use tools=("default", "maps") or add "visualize".'
+            )
     if not check_deps:
         return resolved
     if resolved is not None and has_viz_tool(resolved) and not has_viz_deps():
@@ -836,6 +852,7 @@ def normalize_tools(
             stacklevel=2,
         )
         resolved.discard("visualize")
+        resolved.discard("maps")
     return resolved
 
 

pkg-py/src/querychat/_system_prompt.py

@@ -212,6 +212,9 @@ def render(
             "has_tool_update": "update" in tools if tools else False,
             "has_tool_query": "query" in tools if tools else False,
             "has_tool_visualize": has_viz_tool(tools),
+            "has_tool_maps": tools is not None and "maps" in tools,
+            "maps_disabled": has_viz_tool(tools)
+            and not (tools is not None and "maps" in tools),
             "include_query_guidelines": len(tools or ()) > 0,
             "multi_table": len(self._data_sources) > 1,
             "handoff_available": handoff_available,

pkg-py/src/querychat/_viz_ggsql.py

@@ -8,6 +8,7 @@
 from ._utils import as_narwhals
 
 if TYPE_CHECKING:
+    import ggsql
     import pyarrow as pa
 
     from ._datasource import DataSource
@@ -57,3 +58,90 @@ def unregister(self, name) -> None:
             "write to the primary reader, which is a no-op by design.",
             stacklevel=2,
         )
+
+
+SPATIAL_SMOKE_QUERY = "VISUALISE FROM ggsql:world DRAW spatial FILTER iso_a3 = 'FRA'"
+
+
+def install_extension(reader: ggsql.CachingReader, name: str) -> None:
+    """
+    Install and load a DuckDB extension on the reader's cache connection.
+
+    All ggsql-generated derived SQL runs on the cache connection, never on
+    the user's (typically locked-down) data source, so extensions must be
+    installed here. The first ``INSTALL`` downloads the binary (network
+    required); afterwards it is cached on disk.
+    """
+    reader.execute_sql_cached(f"INSTALL {name}")
+    reader.execute_sql_cached(f"LOAD {name}")
+
+
+def prewarm_icu(reader: ggsql.CachingReader) -> None:
+    """
+    Set up the DuckDB ICU extension (time zones, locale-aware collations).
+
+    DuckDB's CLI and Python client load ICU by default, so SQL written
+    against "DuckDB" — by users or LLMs — assumes it. Failure is
+    non-fatal: warns and continues.
+    """
+    try:
+        install_extension(reader, "icu")
+    except Exception as e:
+        warnings.warn(
+            "Could not set up the DuckDB ICU extension on the visualization "
+            f"engine ({e}). Continuing without it: time zone operations "
+            "(e.g. interval arithmetic on timestamptz columns) and "
+            "locale-aware collations will be unavailable in visualizations.",
+            stacklevel=2,
+        )
+
+
+def prewarm_spatial(reader: ggsql.CachingReader) -> bool:
+    """
+    Set up the DuckDB spatial extension for the ``"maps"`` capability.
+
+    The smoke query exercises the full spatial path (extension load,
+    builtin ``ggsql:world`` registration, spatial layer execution).
+    Failure warns and returns ``False`` so the caller can treat ``"maps"``
+    as unavailable.
+    """
+    try:
+        install_extension(reader, "spatial")
+        reader.execute(SPATIAL_SMOKE_QUERY)
+    except Exception as e:
+        warnings.warn(
+            "The 'maps' capability could not be enabled (spatial extension "
+            f"setup failed: {e}). Continuing without maps.",
+            stacklevel=2,
+        )
+        return False
+    return True
+
+
+def translate_spatial_error(error: Exception) -> str | None:
+    """
+    Translate ggsql's spatial-extension failure into actionable guidance.
+
+    Only call this when maps was requested but spatial setup failed earlier
+    in the session (see ``prewarm_spatial``). Returns the translated
+    message, or ``None`` when the error is not a spatial setup failure.
+
+    Matches DuckDB error wording as of DuckDB 1.5.5 — both the
+    extension-not-found form and the catalog error from calling an
+    ``st_*`` function after a failed LOAD. If a DuckDB upgrade rewords
+    these errors, update the patterns and tests together.
+    """
+    msg = str(error).lower()
+    is_extension_missing = "spatial" in msg and "extension" in msg
+    is_st_function_missing = (
+        "scalar function" in msg and "st_" in msg and "does not exist" in msg
+    )
+    if not (is_extension_missing or is_st_function_missing):
+        return None
+    return (
+        "This query requires map/spatial support, but spatial extension "
+        "setup failed earlier in this session (see the warning shown at "
+        "setup time, likely a failed extension download). Do not retry "
+        "with DRAW spatial, PROJECT TO <map>, or ggsql:world; suggest a "
+        "non-map visualization instead."
+    )

pkg-py/src/querychat/_viz_tools.py

@@ -22,7 +22,12 @@
 from ._tool_names import TOOL_VISUALIZE
 from ._utils import querychat_tool_starts_open, read_prompt_template, truncate_error
 from ._viz_altair_widget import AltairWidget, fit_chart_to_container
-from ._viz_ggsql import ExecutorReader
+from ._viz_ggsql import (
+    ExecutorReader,
+    prewarm_icu,
+    prewarm_spatial,
+    translate_spatial_error,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -62,6 +67,7 @@ def tool_visualize(
     update_fn: Callable[[VisualizeData], None],
     *,
     multi_table: bool = False,
+    maps: bool = False,
 ) -> Tool:
     """
     Create a tool that executes a ggsql query and renders the visualization.
@@ -74,14 +80,17 @@ def tool_visualize(
         Callback function to call with VisualizeData when visualization succeeds
     multi_table
         Whether multiple tables are registered.
+    maps
+        Whether the maps capability is enabled (pre-warms the DuckDB
+        spatial extension on the viz reader).
 
     Returns
     -------
     Tool
         A tool that can be registered with chatlas
 
     """
-    impl = visualize_impl(executor, update_fn)
+    impl = visualize_impl(executor, update_fn, maps=maps)
     impl.__doc__ = read_prompt_template(
         "tool-visualize.md",
         db_type=executor.get_db_type(),
@@ -186,6 +195,8 @@ def warn_viz_url_bookmarking() -> bool:
 def visualize_impl(
     executor: QueryExecutor,
     update_fn: Callable[[VisualizeData], None],
+    *,
+    maps: bool = False,
 ) -> Callable[[str, str], ContentToolResult]:
     """Create the visualize implementation function."""
     from ggsql import CachingReader, VegaLiteWriter, validate
@@ -195,13 +206,14 @@ def visualize_impl(
     # the reader's cache is freed when the client is dropped. Created lazily
     # so clients that never visualize don't pay for an idle DuckDB instance.
     reader = None
+    maps_available = maps
 
     def visualize(
         ggsql: str,
         title: str,
     ) -> ContentToolResult:
         """Execute a ggsql query and render the visualization."""
-        nonlocal reader
+        nonlocal reader, maps_available
         markdown = f"```sql\n{ggsql}\n```"
 
         try:
@@ -219,6 +231,9 @@ def visualize(
 
             if reader is None:
                 reader = CachingReader(ExecutorReader(executor))
+                prewarm_icu(reader)
+                if maps_available:
+                    maps_available = prewarm_spatial(reader)
             spec = reader.execute(ggsql)
 
             raw_chart = VegaLiteWriter().render_chart(spec)
@@ -244,7 +259,13 @@ def visualize(
             )
 
         except Exception as e:
-            error_msg = truncate_error(str(e))
+            # Only the setup-failed state needs translation: maps-disabled
+            # is covered by the prompt's maps-disabled note, and a spatial
+            # error with maps working is a real bug worth seeing raw.
+            translated = (
+                translate_spatial_error(e) if (maps and not maps_available) else None
+            )
+            error_msg = translated if translated is not None else truncate_error(str(e))
             markdown += f"\n\n> Error: {error_msg}"
             return ContentToolResult(value=markdown, error=Exception(error_msg))