diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 403df5432..dae1b7fdc 100644 --- a/pkg-py/CHANGELOG.md +++ b/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 `. 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() diff --git a/pkg-py/docs/visualize.qmd b/pkg-py/docs/visualize.qmd index 1a708fb70..4f3c863aa 100644 --- a/pkg-py/docs/visualize.qmd +++ b/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 diff --git a/pkg-py/examples/13-maps-app.py b/pkg-py/examples/13-maps-app.py new file mode 100644 index 000000000..9832990df --- /dev/null +++ b/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: + +* Map the earthquakes from the past month, sized by magnitude +* Draw a choropleth of GDP by country +""" + +# 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%", + ), + ], +) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 5abd1ad52..1d1568845 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/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 diff --git a/pkg-py/src/querychat/_system_prompt.py b/pkg-py/src/querychat/_system_prompt.py index b53626827..fdc67a21b 100644 --- a/pkg-py/src/querychat/_system_prompt.py +++ b/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, diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index ff7e9bc82..c96961bf2 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/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 , or ggsql:world; suggest a " + "non-map visualization instead." + ) diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index e38ba688f..684d33ade 100644 --- a/pkg-py/src/querychat/_viz_tools.py +++ b/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,6 +80,9 @@ 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 ------- @@ -81,7 +90,7 @@ def tool_visualize( 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)) diff --git a/pkg-py/src/querychat/prompts/ggsql-spatial.md b/pkg-py/src/querychat/prompts/ggsql-spatial.md new file mode 100644 index 000000000..0f9c2d5cd --- /dev/null +++ b/pkg-py/src/querychat/prompts/ggsql-spatial.md @@ -0,0 +1,65 @@ +## Maps with ggsql (spatial) + +The `maps` capability is enabled: the DuckDB `spatial` extension is already installed and loaded on the ggsql compute connection. Do NOT write `INSTALL spatial` or `LOAD spatial` — querychat has handled spatial setup. + +The two query paths differ in what SQL is available. `querychat_query` runs on the user's database, where `ST_*` spatial functions may not exist — keep spatial logic out of data queries. Visualization queries run on ggsql's DuckDB connection, where the spatial extension IS loaded, so `ST_*` functions (e.g. `ST_Within`, `ST_Point`) are available in the SELECT clause of a `VISUALISE` query when ggsql clauses alone can't express the logic (e.g. point-in-polygon joins). + +### Choropleths from `ggsql:world` + +`ggsql:world` is a builtin dataset of country boundaries (Natural Earth 1:110m) with columns: `name`, `iso_a3`, `continent`, `subregion`, `income_group`, `population`, `gdp`, `geom` (geometry). It is resolved by ggsql, not the user's database — but you can join it with the user's tables in the SELECT clause and ggsql will stage the join itself. + +To map the user's data, join it to `ggsql:world` on `iso_a3` (preferred) or `name`, then `DRAW spatial` with a fill mapping: + +```sql +SELECT w.name, w.geom, s.revenue +FROM ggsql:world w JOIN sales_by_country s ON w.iso_a3 = s.country_code +VISUALISE +DRAW spatial MAPPING revenue AS fill +``` + +The geometry column is auto-detected; declare `MAPPING geom AS geometry` explicitly only if detection fails. Before joining, inspect the user's identifier values (e.g. with the schema or query tool): country names vary ("United States" vs "USA"), and ISO codes only match if the user's column is also ISO 3166-1 alpha-3. Filter to a region with `FILTER continent = 'Europe'` on the DRAW clause when the data only covers one. + +`ggsql:world` is reachable only from visualization queries — the data-query tool cannot read it. To answer questions about map coverage (e.g. "which countries are missing from this map?"), don't refuse: use a visualization query with an anti-join (`LEFT JOIN ... WHERE ... IS NULL`) against `ggsql:world`. + +### Point maps + +For lon/lat columns, map them to `lon`/`lat` (or `x`/`y` with `PROJECT x, y TO `) and draw `point` over a `ggsql:world` background: + +```sql +SELECT city, lon, lat, passengers FROM airports +VISUALISE +DRAW spatial MAPPING null AS fill FROM ggsql:world +DRAW point MAPPING lon, lat, passengers AS size +PROJECT TO natural +``` + +**Always put `lon`/`lat`/`size` mappings on the `DRAW point` layer, never in the global `VISUALISE` mappings.** Global mappings are inherited by every layer, and the `spatial` background layer only accepts `geometry`/`fill`/`stroke`/`opacity` — a global `lon`/`lat`/`size` mapping fails with "Layer 'spatial' does not support the `lat`, `size`, `lon` mappings". If you see that error, move the point aesthetics out of `VISUALISE` and onto `DRAW point MAPPING ...`. + +### Spatial joins (point-in-polygon) + +To aggregate point data by country (e.g. earthquake counts per country), spatially join against `ggsql:world` in the SELECT clause. One gotcha: when `ggsql:world` is staged for a join, its `geom` column arrives as WKB `BLOB`, not `GEOMETRY` — wrap it with `ST_GeomFromWKB()` or the join fails with `No function matches ... 'ST_Within(GEOMETRY, BLOB)'`. Aggregate first, then join the geometry on last: + +```sql +WITH counts AS ( + SELECT w.iso_a3, COUNT(*) AS quakes + FROM ggsql:world w + JOIN earthquakes e + ON ST_Within(ST_Point(e.longitude, e.latitude), ST_GeomFromWKB(w.geom)) + GROUP BY w.iso_a3 +) +SELECT w.name, w.geom, c.quakes +FROM ggsql:world w +JOIN counts c ON w.iso_a3 = c.iso_a3 +VISUALISE +DRAW spatial MAPPING quakes AS fill +``` + +### Projections + +`PROJECT TO ` sets the map projection. Useful named projections: `geographic` (unprojected lon/lat), `mercator`, `natural` (Natural Earth), `robinson`, `equal_earth`, `orthographic`, `albers` (US-centric). Set the view center with `SETTING origin => (lon, lat)`, e.g. `PROJECT TO orthographic SETTING origin => (-98, 39)` for a US globe. For anything else, use `PROJECT TO crs SETTING target => 'EPSG:XXXX'` or a proj string. `SCALE lon`/`SCALE lat` limits transfer to the projection. + +`SETTING origin` re-centers the map but does NOT zoom — a regional map with only `origin` still shows the whole world. To frame a region, also bound the view with `SCALE lon FROM (min, max)` and `SCALE lat FROM (min, max)`. + +The ocean (sphere) background and graticule are drawn by the renderer and cannot be restyled — no setting is exposed for them, so explain the limitation rather than guessing one. + +Keep maps simple: one `spatial` layer for boundaries plus at most one or two data layers. Do not put geometry columns in `GROUP BY` — join them on last, in the final SELECT. diff --git a/pkg-py/src/querychat/prompts/prompt.md b/pkg-py/src/querychat/prompts/prompt.md index bf4734919..8f4f02b9d 100644 --- a/pkg-py/src/querychat/prompts/prompt.md +++ b/pkg-py/src/querychat/prompts/prompt.md @@ -196,6 +196,16 @@ Match the chart type to what the user is trying to understand: {{> ggsql-syntax}} +{{#has_tool_maps}} + + +{{> ggsql-spatial}} + +{{/has_tool_maps}} +{{#maps_disabled}} + +**Map visualizations are available but not enabled.** querychat can draw choropleths and point maps, but the `maps` capability is opt-in and this session has not enabled it. Do NOT attempt `DRAW spatial`, `ggsql:world`, `ST_*` functions, or `PROJECT TO ` — they will fail. If the user asks for a map, explain that maps are an opt-in capability the app developer can enable with `tools=("default", "maps")`, and offer a fallback such as a scatter plot of longitude/latitude with `DRAW point`. +{{/maps_disabled}} {{#has_tool_query}} **Avoid redundant expanded results.** If you run a preparatory query before visualizing, or if both a table and chart would show the same data, always pass `collapsed=true` on the query so the user sees the chart prominently, not a duplicate table above it. The user can still expand the table if they want the exact values. diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index 1f59718ca..e6f1c11f3 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -505,3 +505,66 @@ def test_invalid_entries_are_skipped_with_warning(self, register_spy): assert [e["widget_id"] for e in restored] == ["good"] assert register_spy == ["good"] + + +class TestPrewarmIcu: + def test_installs_and_loads_icu(self): + from querychat._viz_ggsql import prewarm_icu + + cached_calls = [] + + class FakeReader: + def execute_sql_cached(self, sql): + cached_calls.append(sql) + return object() + + assert prewarm_icu(FakeReader()) is None + assert cached_calls == ["INSTALL icu", "LOAD icu"] + + def test_icu_failure_warns_and_continues(self): + from querychat._viz_ggsql import prewarm_icu + + class FailReader: + def execute_sql_cached(self, sql): + raise RuntimeError("IO Error: could not download extension") + + with pytest.warns(UserWarning, match="ICU"): + assert prewarm_icu(FailReader()) is None + + +class TestPrewarmSpatial: + def test_success_runs_smoke_query(self): + from querychat._viz_ggsql import SPATIAL_SMOKE_QUERY, prewarm_spatial + + cached_calls = [] + execute_calls = [] + + class FakeReader: + def execute_sql_cached(self, sql): + cached_calls.append(sql) + return object() + + def execute(self, query): + execute_calls.append(query) + return object() + + assert prewarm_spatial(FakeReader()) is True + assert cached_calls == ["INSTALL spatial", "LOAD spatial"] + assert execute_calls == [SPATIAL_SMOKE_QUERY] + + def test_spatial_failure_warns_and_returns_false(self): + from querychat._viz_ggsql import prewarm_spatial + + class FailReader: + def execute_sql_cached(self, sql): + raise RuntimeError( + "IO Error: Extension spatial.duckdb_extension not found" + ) + + def execute(self, query): + raise RuntimeError( + "IO Error: Extension spatial.duckdb_extension not found" + ) + + with pytest.warns(UserWarning, match="maps"): + assert prewarm_spatial(FailReader()) is False diff --git a/pkg-py/tests/test_viz_tools.py b/pkg-py/tests/test_viz_tools.py index f5364be96..bb20e3897 100644 --- a/pkg-py/tests/test_viz_tools.py +++ b/pkg-py/tests/test_viz_tools.py @@ -58,6 +58,53 @@ def test_check_deps_false_skips_check(self, monkeypatch): assert result == {"visualize"} +class TestMapsToolModifier: + def test_default_sentinel_expands(self): + """The "default" entry expands to the default tool set.""" + from querychat._querychat_base import DEFAULT_TOOLS, normalize_tools + + resolved = normalize_tools( + ("default", "maps"), default=DEFAULT_TOOLS, check_deps=False + ) + assert "visualize" in resolved + assert "query" in resolved + assert "maps" in resolved + assert "default" not in resolved + + def test_default_sentinel_alone(self): + """tools="default" alone resolves to the default tool set.""" + from querychat._querychat_base import DEFAULT_TOOLS, normalize_tools + + resolved = normalize_tools("default", default=DEFAULT_TOOLS, check_deps=False) + assert "visualize" in resolved + assert "query" in resolved + + def test_maps_requires_visualize(self): + """"maps" without "visualize" is a hard error.""" + from querychat._querychat_base import DEFAULT_TOOLS, normalize_tools + + with pytest.raises(ValueError, match=r"maps.*visualize"): + normalize_tools(("query", "maps"), default=DEFAULT_TOOLS, check_deps=False) + + def test_maps_dropped_with_visualize_when_deps_missing(self, monkeypatch): + """When viz deps are missing, maps is dropped alongside visualize.""" + real_find_spec = importlib.util.find_spec + + def mock_find_spec(name, *args, **kwargs): + if name == "ggsql": + return None + return real_find_spec(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", mock_find_spec) + + from querychat._querychat_base import DEFAULT_TOOLS, normalize_tools + + with pytest.warns(UserWarning, match="pip install querychat\\[viz\\]"): + resolved = normalize_tools(("default", "maps"), default=DEFAULT_TOOLS) + assert "visualize" not in resolved + assert "maps" not in resolved + + def test_ggsql_syntax_reference_uses_range_not_errorbar(): syntax = (PROMPTS_DIR / "ggsql-syntax.md").read_text() assert "`range`" in syntax @@ -100,6 +147,66 @@ def test_main_prompt_render_includes_ggsql_reference(data_source): assert "{{> ggsql-syntax}}" not in prompt +def test_main_prompt_render_includes_spatial_reference_with_maps(data_source): + system_prompt = QueryChatSystemPrompt( + prompt_template=PROMPTS_DIR / "prompt.md", + data_source=data_source, + ) + + prompt = system_prompt.render({"update", "query", "visualize", "maps"}) + + assert "" in prompt + assert "ggsql:world" in prompt + + +def test_main_prompt_render_omits_spatial_reference_without_maps(data_source): + system_prompt = QueryChatSystemPrompt( + prompt_template=PROMPTS_DIR / "prompt.md", + data_source=data_source, + ) + + prompt = system_prompt.render({"update", "query", "visualize"}) + + assert "" not in prompt + + +def test_main_prompt_render_shows_maps_disabled_note_without_maps(data_source): + system_prompt = QueryChatSystemPrompt( + prompt_template=PROMPTS_DIR / "prompt.md", + data_source=data_source, + ) + + prompt = system_prompt.render({"update", "query", "visualize"}) + + assert "not enabled" in prompt + assert 'tools=("default", "maps")' in prompt + # The note must not leak the spatial syntax reference + assert "" not in prompt + + +def test_main_prompt_render_hides_maps_disabled_note_with_maps(data_source): + system_prompt = QueryChatSystemPrompt( + prompt_template=PROMPTS_DIR / "prompt.md", + data_source=data_source, + ) + + prompt = system_prompt.render({"update", "query", "visualize", "maps"}) + + assert "" in prompt + assert "Map visualizations" not in prompt + + +def test_main_prompt_render_hides_maps_disabled_note_without_visualize(data_source): + system_prompt = QueryChatSystemPrompt( + prompt_template=PROMPTS_DIR / "prompt.md", + data_source=data_source, + ) + + prompt = system_prompt.render({"update", "query"}) + + assert "Map visualizations" not in prompt + + @pytest.fixture def sample_df(): return pl.DataFrame( @@ -376,3 +483,125 @@ def test_compound_chart_returns_bytes(self): result = render_chart_to_png(chart) assert isinstance(result, bytes) assert result[:8] == b"\x89PNG\r\n\x1a\n" + + +class TestMapsPrewarmWiring: + @pytest.mark.ggsql + def test_prewarms_icu_and_spatial_once_when_maps_enabled( + self, data_source, monkeypatch + ): + import querychat._viz_tools as viz_tools + + icu_calls = [] + spatial_calls = [] + monkeypatch.setattr( + viz_tools, "prewarm_icu", lambda reader: icu_calls.append(reader) + ) + monkeypatch.setattr( + viz_tools, + "prewarm_spatial", + lambda reader: spatial_calls.append(reader) or True, + ) + tool = tool_visualize(data_source, lambda _: None, maps=True) + tool.func( + ggsql="SELECT x, y FROM test_data VISUALISE x, y DRAW point", title="t1" + ) + tool.func( + ggsql="SELECT x, y FROM test_data VISUALISE x, y DRAW point", title="t2" + ) + assert len(icu_calls) == 1 + assert len(spatial_calls) == 1 + + @pytest.mark.ggsql + def test_prewarms_icu_but_not_spatial_when_maps_disabled( + self, data_source, monkeypatch + ): + import querychat._viz_tools as viz_tools + + icu_calls = [] + spatial_calls = [] + monkeypatch.setattr( + viz_tools, "prewarm_icu", lambda reader: icu_calls.append(reader) + ) + monkeypatch.setattr( + viz_tools, + "prewarm_spatial", + lambda reader: spatial_calls.append(reader) or True, + ) + tool = tool_visualize(data_source, lambda _: None) + tool.func( + ggsql="SELECT x, y FROM test_data VISUALISE x, y DRAW point", title="t1" + ) + assert len(icu_calls) == 1 + assert spatial_calls == [] + + +class TestTranslateSpatialError: + SPATIAL_ERR = Exception( + 'IO Error: Extension "/home/u/.duckdb/extensions/v1.5.5/linux_amd64/' + 'spatial.duckdb_extension" not found' + ) + ST_FUNC_ERR = Exception( + "Catalog Error: Scalar Function with name st_within does not exist!" + ) + + def test_extension_error_message(self): + from querychat._viz_ggsql import translate_spatial_error + + msg = translate_spatial_error(self.SPATIAL_ERR) + assert msg is not None + assert "setup failed" in msg + + def test_st_function_error_message(self): + # A failed LOAD followed by an st_* call surfaces as a catalog + # error, not an extension error — same recovery guidance applies. + from querychat._viz_ggsql import translate_spatial_error + + msg = translate_spatial_error(self.ST_FUNC_ERR) + assert msg is not None + assert "setup failed" in msg + + def test_nonspatial_returns_none(self): + from querychat._viz_ggsql import translate_spatial_error + + assert translate_spatial_error(Exception("boom")) is None + + @pytest.mark.ggsql + def test_tool_translates_spatial_error_after_setup_failure( + self, data_source, monkeypatch + ): + import querychat._viz_tools as viz_tools + + monkeypatch.setattr(viz_tools, "prewarm_icu", lambda reader: None) + monkeypatch.setattr(viz_tools, "prewarm_spatial", lambda reader: False) + monkeypatch.setattr( + "ggsql.CachingReader.execute", + lambda reader, query: (_ for _ in ()).throw( + RuntimeError(str(self.SPATIAL_ERR)) + ), + ) + tool = tool_visualize(data_source, lambda _: None, maps=True) + result = tool.func( + ggsql="SELECT x, y FROM test_data VISUALISE x, y DRAW point", title="t" + ) + assert result.error is not None + assert "setup failed" in str(result.error) + + @pytest.mark.ggsql + def test_tool_does_not_translate_when_maps_disabled( + self, data_source, monkeypatch + ): + # With maps disabled, the system prompt tells the LLM not to + # attempt spatial features; a raw error is the acceptable backstop. + monkeypatch.setattr( + "ggsql.CachingReader.execute", + lambda reader, query: (_ for _ in ()).throw( + RuntimeError(str(self.SPATIAL_ERR)) + ), + ) + tool = tool_visualize(data_source, lambda _: None) + result = tool.func( + ggsql="SELECT x, y FROM test_data VISUALISE x, y DRAW point", title="t" + ) + assert result.error is not None + assert "setup failed" not in str(result.error)