diff --git a/.gitignore b/.gitignore index f03a3e8bb..25fc3c564 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ README.html test-results/ shiny_bookmarks/ python-package/examples/titanic.db +pkg-py/examples/13-multi-table-standalone/data/ .quarto *.db diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 757d609c0..c5e66f6f4 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- 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. + ## [0.9.0] - 2026-09-13 ### New features diff --git a/pkg-py/docs/visualize.qmd b/pkg-py/docs/visualize.qmd index e8d720b68..4f3c863aa 100644 --- a/pkg-py/docs/visualize.qmd +++ b/pkg-py/docs/visualize.qmd @@ -96,8 +96,38 @@ Each chart has controls in its footer: Note that visualization queries are independent of any active dashboard filter set by the `update` tool. They always run against the full dataset. +### Caching + +To keep iterative chart refinement fast, querychat memoizes the data portion of visualization queries: when the LLM issues identical SQL within 5 minutes (for example, when only the visual encoding changes), the cached result is reused instead of re-reading your data source. +Each chat session has its own cache, bounded at 512 MiB, which is freed when the session ends. + +This assumes the data doesn't change mid-session. +If it can — for example, update queries are enabled (`QUERYCHAT_ENABLE_UPDATE_QUERIES=true`), or another process writes to the database — charts may lag those changes by up to the TTL below. + +Tune or disable the cache with environment variables (set before the app starts): + +| Variable | Default | Purpose | +|---|---|---| +| `GGSQL_CACHE_TTL` | `300` | Seconds a cached result is treated as fresh | +| `GGSQL_CACHE_MAX_BYTES` | `512mb` | Per-session memory budget before least-recently-used eviction | +| `GGSQL_CACHE_DISABLED` | unset | Set to any non-empty value except `0` to always re-read the source | + 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-multi-table-standalone/app.py b/pkg-py/examples/13-multi-table-standalone/app.py new file mode 100644 index 000000000..7450ff279 --- /dev/null +++ b/pkg-py/examples/13-multi-table-standalone/app.py @@ -0,0 +1,90 @@ +# /// script +# dependencies = [ +# "querychat[polars,viz]", +# "anthropic", +# ] +# /// +"""Multi-table querychat demo: a synthetic fine-jewelry retailer. + +Five related tables -- diamonds, settings, customers, orders, and +order_items -- are registered with a single QueryChat instance so the LLM +can query any one of them or write joins across them (e.g. an "Engagement +Ring" order item joins to both its diamond and its setting). + +Data is generated once and cached to ./data/*.parquet; later runs read the +cache instead of regenerating (see generate_data.py). + +Usage: + uv run app.py + +Requires ANTHROPIC_API_KEY to be set (change client= below to use a +different provider). +""" + +from pathlib import Path + +from generate_data import load_tables +from querychat import QueryChat + +GREETING = """\ +## Brilliant Retail: Jewelry Explorer 💎 + +Welcome! You're chatting with a synthetic fine-jewelry retailer's data -- +**diamonds**, **settings** (ring mountings), **customers**, **orders**, and +**order_items**, all linked together. + +**Get started** + + + +**Explore across tables** + + +""" + +DATA_DESCRIPTION = """\ +Brilliant Retail is an online fine-jewelry retailer. This dataset covers its +diamond and setting catalog alongside customer order history. +""" + +# Works around a ggsql bug (https://github.com/posit-dev/ggsql/issues/): +# a date/timestamp column produced by an expression like DATE_TRUNC(...) fails +# ggsql's `SCALE ... VIA date` coercion unless it's explicitly cast to DATE. +EXTRA_INSTRUCTIONS = """\ +When a `querychat_visualize` query buckets a date or timestamp column (e.g. +via DATE_TRUNC), always CAST the resulting expression to DATE in the SQL, for +example: + + SELECT CAST(DATE_TRUNC('month', order_date) AS DATE) AS month, ... + +A column left as a raw TIMESTAMP will fail when referenced by +`SCALE ... VIA date` with a "Cannot coerce column ... to date" error. +""" + +tables = load_tables(Path(__file__).parent / "data") + +qc = QueryChat( + tables["diamonds"], + "diamonds", + greeting=GREETING, + data_description=DATA_DESCRIPTION, + data_dict=Path(__file__).parent / "data-dict.yaml", + extra_instructions=EXTRA_INSTRUCTIONS, + client="anthropic", +) +qc.add_table(tables["settings"], "settings") +qc.add_table(tables["customers"], "customers") +qc.add_table(tables["orders"], "orders") +qc.add_table(tables["order_items"], "order_items") + +app = qc.app() + +if __name__ == "__main__": + app.run() diff --git a/pkg-py/examples/13-multi-table-standalone/data-dict.yaml b/pkg-py/examples/13-multi-table-standalone/data-dict.yaml new file mode 100644 index 000000000..22cfb8583 --- /dev/null +++ b/pkg-py/examples/13-multi-table-standalone/data-dict.yaml @@ -0,0 +1,178 @@ +version: "0.1.0" + +tables: + diamonds: + description: > + Loose diamond inventory. One row per diamond, graded on the 4 Cs + (carat, cut, color, clarity) and certified by an independent lab. + columns: + - name: id + type: number(id) + constraints: [primary_key] + description: Unique diamond identifier. + - name: shape + type: enum + values: [Round, Princess, Cushion, Oval, Emerald, Radiant, Asscher, Marquise, Pear, Heart] + description: Cut shape of the diamond. + - name: carat + type: number(quantity) + units: carat + description: Weight of the diamond. + - name: cut + type: enum + values: [Super Ideal, Ideal, Very Good, Good] + description: Cut quality grade -- how well the diamond's facets interact with light. + - name: color + type: enum + values: [D, E, F, G, H, I, J] + description: Color grade on the standard D (colorless) to J (faint color) scale. + - name: clarity + type: enum + values: [FL, IF, VVS1, VVS2, VS1, VS2, SI1, SI2] + description: > + Clarity grade, from FL (flawless) to SI2 (slight inclusions visible + under magnification). + - name: certified_by + type: enum + values: [GIA, IGI, AGS] + description: Independent grading lab that certified the diamond. + - name: price + type: number(currency) + units: USD + description: Price of the loose diamond. + + settings: + description: Ring mountings that a diamond can be set into. One row per setting. + columns: + - name: id + type: number(id) + constraints: [primary_key] + description: Unique setting identifier. + - name: style + type: enum + values: [Solitaire, Halo, Pave, Three-Stone, Vintage, Bezel] + description: Setting style. + - name: metal + type: enum + values: [Platinum, 18K White Gold, 18K Yellow Gold, 18K Rose Gold, Palladium] + description: Metal the setting is made of. + - name: band_width_mm + type: number(quantity) + units: mm + description: Width of the ring band. + - name: base_price + type: number(currency) + units: USD + description: Price of the setting alone, without a diamond. + + customers: + description: One row per customer. + columns: + - name: id + type: number(id) + constraints: [primary_key] + description: Unique customer identifier. + - name: name + type: string + description: Customer's full name. + - name: state + type: string + description: US state abbreviation of the customer's shipping address. + - name: join_date + type: date + description: Date the customer created their account. + - name: acquisition_channel + type: enum + values: [Organic Search, Paid Social, Referral, Email, Direct, Showroom Walk-in] + description: How the customer first found the retailer. + + orders: + description: One row per order. + columns: + - name: id + type: number(id) + constraints: [primary_key] + description: Unique order identifier. + - name: customer_id + type: number(id) + constraints: [foreign_key] + description: Customer who placed the order. Joins to customers.id. + - name: order_date + type: date + description: Date the order was placed. + - name: status + type: enum + values: [Delivered, Shipped, Processing, Returned] + description: Current fulfillment status of the order. + - name: channel + type: enum + values: [Online, Showroom] + description: Where the order was placed. + + order_items: + description: > + Line items within an order. One row per item purchased. An item's + diamond_id and setting_id are populated according to its item_type: + "Engagement Ring" has both, "Loose Diamond" has only diamond_id, + "Wedding Band" has only setting_id, and "Earrings"/"Necklace" have + neither -- their price is not decomposed into components. + columns: + - name: id + type: number(id) + constraints: [primary_key] + description: Unique order item identifier. + - name: order_id + type: number(id) + constraints: [foreign_key] + description: Order this item belongs to. Joins to orders.id. + - name: item_type + type: enum + values: [Engagement Ring, Loose Diamond, Wedding Band, Earrings, Necklace] + description: Kind of item purchased. + - name: diamond_id + type: number(id) + constraints: [foreign_key] + description: > + Diamond included in this item, if any. Joins to diamonds.id. Null + unless item_type is "Engagement Ring" or "Loose Diamond". + - name: setting_id + type: number(id) + constraints: [foreign_key] + description: > + Setting included in this item, if any. Joins to settings.id. Null + unless item_type is "Engagement Ring" or "Wedding Band". + - name: price + type: number(currency) + units: USD + description: Price of this line item. + +relationships: + - description: Each order belongs to one customer. + cardinality: many-to-one + join: orders.customer_id = customers.id + - description: Each order has one or more line items. + cardinality: one-to-many + join: orders.id = order_items.order_id + - description: An order item may include one diamond (Engagement Ring or Loose Diamond). + cardinality: many-to-one + join: order_items.diamond_id = diamonds.id + - description: An order item may include one setting (Engagement Ring or Wedding Band). + cardinality: many-to-one + join: order_items.setting_id = settings.id + +glossary: + 4 Cs: The standard diamond grading criteria -- carat (weight), cut, color, and clarity. + carat: Unit of diamond weight; 1 carat = 200 milligrams. Larger carat generally means higher price. + cut: > + How well a diamond's facets interact with light, independent of its + shape. Graded Super Ideal > Ideal > Very Good > Good in this dataset. + color: > + Diamond color grade from D (colorless, most valuable) to J (faint color, + least valuable) in this dataset. + clarity: > + Presence of internal/external flaws, from FL (flawless) down to SI2 + (inclusions visible under 10x magnification). + setting: > + The metal mounting that holds a diamond, e.g. a ring band. Sold + separately as-is, or combined with a diamond into an Engagement Ring + order item. diff --git a/pkg-py/examples/13-multi-table-standalone/generate_data.py b/pkg-py/examples/13-multi-table-standalone/generate_data.py new file mode 100644 index 000000000..989bf322d --- /dev/null +++ b/pkg-py/examples/13-multi-table-standalone/generate_data.py @@ -0,0 +1,249 @@ +# /// script +# dependencies = ["polars"] +# /// +"""Generate (or load cached) synthetic data for a fine-jewelry retailer demo. + +Five related tables are produced with a fixed random seed so the data is +identical across runs: diamonds, settings (ring mountings), customers, +orders, and order_items. + +Run directly to force-regenerate the cache: + uv run generate_data.py +""" + +from __future__ import annotations + +import random +import shutil +from datetime import date, timedelta +from pathlib import Path + +import polars as pl + +SEED = 42 +TODAY = date(2026, 9, 1) +TABLE_NAMES = ["diamonds", "settings", "customers", "orders", "order_items"] + + +def load_tables(data_dir: Path) -> dict[str, pl.DataFrame]: + """Return the 5 demo tables, generating and caching them on first call.""" + paths = {name: data_dir / f"{name}.parquet" for name in TABLE_NAMES} + + if all(path.exists() for path in paths.values()): + return {name: pl.read_parquet(path) for name, path in paths.items()} + + tables = generate_tables() + data_dir.mkdir(parents=True, exist_ok=True) + for name, table in tables.items(): + table.write_parquet(paths[name]) + return tables + + +def generate_tables() -> dict[str, pl.DataFrame]: + rng = random.Random(SEED) + diamonds = generate_diamonds(rng) + settings = generate_settings(rng) + customers = generate_customers(rng) + orders = generate_orders(rng, customers) + order_items = generate_order_items(rng, orders, diamonds, settings) + return { + "diamonds": diamonds, + "settings": settings, + "customers": customers, + "orders": orders, + "order_items": order_items, + } + + +def generate_diamonds(rng: random.Random, n: int = 400) -> pl.DataFrame: + shapes = [ + "Round", "Princess", "Cushion", "Oval", "Emerald", + "Radiant", "Asscher", "Marquise", "Pear", "Heart", + ] + shape_weights = [30, 12, 10, 12, 8, 8, 4, 4, 6, 6] + cuts = ["Super Ideal", "Ideal", "Very Good", "Good"] + cut_weights = [10, 40, 35, 15] + cut_multiplier = {"Super Ideal": 1.15, "Ideal": 1.05, "Very Good": 0.95, "Good": 0.85} + colors = ["D", "E", "F", "G", "H", "I", "J"] + color_multiplier = {"D": 1.2, "E": 1.15, "F": 1.1, "G": 1.0, "H": 0.9, "I": 0.8, "J": 0.7} + clarities = ["FL", "IF", "VVS1", "VVS2", "VS1", "VS2", "SI1", "SI2"] + clarity_multiplier = { + "FL": 1.3, "IF": 1.25, "VVS1": 1.15, "VVS2": 1.1, + "VS1": 1.0, "VS2": 0.95, "SI1": 0.85, "SI2": 0.75, + } + certifiers = ["GIA", "IGI", "AGS"] + certifier_weights = [60, 30, 10] + + rows = [] + for i in range(1, n + 1): + carat = round(rng.uniform(0.3, 3.5), 2) + cut = rng.choices(cuts, weights=cut_weights)[0] + color = rng.choice(colors) + clarity = rng.choice(clarities) + price = round( + 3000 + * (carat**1.8) + * cut_multiplier[cut] + * color_multiplier[color] + * clarity_multiplier[clarity], + -1, + ) + rows.append( + { + "id": i, + "shape": rng.choices(shapes, weights=shape_weights)[0], + "carat": carat, + "cut": cut, + "color": color, + "clarity": clarity, + "certified_by": rng.choices(certifiers, weights=certifier_weights)[0], + "price": price, + } + ) + return pl.DataFrame(rows) + + +def generate_settings(rng: random.Random, n: int = 50) -> pl.DataFrame: + styles = ["Solitaire", "Halo", "Pave", "Three-Stone", "Vintage", "Bezel"] + style_multiplier = { + "Solitaire": 1.0, "Halo": 1.4, "Pave": 1.3, + "Three-Stone": 1.5, "Vintage": 1.6, "Bezel": 1.1, + } + metal_price = { + "Platinum": 1400, "18K White Gold": 900, "18K Yellow Gold": 850, + "18K Rose Gold": 850, "Palladium": 1000, + } + metals = list(metal_price) + + rows = [] + for i in range(1, n + 1): + style = rng.choice(styles) + metal = rng.choice(metals) + base_price = round(metal_price[metal] * style_multiplier[style] * rng.uniform(0.9, 1.1), -1) + rows.append( + { + "id": i, + "style": style, + "metal": metal, + "band_width_mm": round(rng.uniform(1.5, 3.5), 1), + "base_price": base_price, + } + ) + return pl.DataFrame(rows) + + +def generate_customers(rng: random.Random, n: int = 200) -> pl.DataFrame: + first_names = [ + "Olivia", "Liam", "Emma", "Noah", "Ava", "Ethan", "Sophia", "Mason", + "Isabella", "James", "Mia", "Benjamin", "Charlotte", "Lucas", "Amelia", + "Henry", "Harper", "Alexander", "Evelyn", "Michael", + ] + last_names = [ + "Smith", "Johnson", "Williams", "Brown", "Jones", "Garcia", "Miller", + "Davis", "Rodriguez", "Martinez", "Chen", "Kim", "Patel", "Nguyen", + "Anderson", "Taylor", "Thomas", "Moore", "Jackson", "Martin", + ] + states = [ + "CA", "NY", "TX", "FL", "WA", "IL", "MA", "CO", "GA", "NC", + "VA", "PA", "OH", "AZ", "OR", "NJ", "MD", "MN", "CT", "NV", + ] + channels = [ + "Organic Search", "Paid Social", "Referral", "Email", "Direct", "Showroom Walk-in", + ] + start = date(2022, 1, 1) + span_days = (TODAY - start).days + + rows = [] + for i in range(1, n + 1): + rows.append( + { + "id": i, + "name": f"{rng.choice(first_names)} {rng.choice(last_names)}", + "state": rng.choice(states), + "join_date": start + timedelta(days=rng.randint(0, span_days)), + "acquisition_channel": rng.choice(channels), + } + ) + return pl.DataFrame(rows) + + +def generate_orders(rng: random.Random, customers: pl.DataFrame, n: int = 350) -> pl.DataFrame: + statuses = ["Delivered", "Shipped", "Processing", "Returned"] + status_weights = [70, 15, 10, 5] + channels = ["Online", "Showroom"] + join_date_by_customer = dict(zip(customers["id"], customers["join_date"])) + customer_ids = customers["id"].to_list() + + rows = [] + for i in range(1, n + 1): + customer_id = rng.choice(customer_ids) + join_date = join_date_by_customer[customer_id] + max_offset = max((TODAY - join_date).days, 1) + order_date = join_date + timedelta(days=rng.randint(0, max_offset)) + rows.append( + { + "id": i, + "customer_id": customer_id, + "order_date": order_date, + "status": rng.choices(statuses, weights=status_weights)[0], + "channel": rng.choice(channels), + } + ) + return pl.DataFrame(rows) + + +def generate_order_items( + rng: random.Random, + orders: pl.DataFrame, + diamonds: pl.DataFrame, + settings: pl.DataFrame, + n: int = 450, +) -> pl.DataFrame: + item_types = ["Engagement Ring", "Loose Diamond", "Wedding Band", "Earrings", "Necklace"] + item_type_weights = [35, 15, 20, 15, 15] + order_ids = orders["id"].to_list() + diamond_ids = diamonds["id"].to_list() + setting_ids = settings["id"].to_list() + price_by_diamond = dict(zip(diamonds["id"], diamonds["price"])) + price_by_setting = dict(zip(settings["id"], settings["base_price"])) + + rows = [] + for i in range(1, n + 1): + item_type = rng.choices(item_types, weights=item_type_weights)[0] + diamond_id: int | None = None + setting_id: int | None = None + + if item_type == "Engagement Ring": + diamond_id = rng.choice(diamond_ids) + setting_id = rng.choice(setting_ids) + price = price_by_diamond[diamond_id] + price_by_setting[setting_id] + elif item_type == "Loose Diamond": + diamond_id = rng.choice(diamond_ids) + price = price_by_diamond[diamond_id] + elif item_type == "Wedding Band": + setting_id = rng.choice(setting_ids) + price = round(price_by_setting[setting_id] * rng.uniform(0.7, 1.0), -1) + elif item_type == "Earrings": + price = round(rng.uniform(400, 4000), -1) + else: # Necklace + price = round(rng.uniform(500, 6000), -1) + + rows.append( + { + "id": i, + "order_id": rng.choice(order_ids), + "item_type": item_type, + "diamond_id": diamond_id, + "setting_id": setting_id, + "price": price, + } + ) + return pl.DataFrame(rows) + + +if __name__ == "__main__": + target_dir = Path(__file__).parent / "data" + if target_dir.exists(): + shutil.rmtree(target_dir) + load_tables(target_dir) + print(f"Wrote {', '.join(TABLE_NAMES)} to {target_dir}") diff --git a/pkg-py/src/querychat/_handoff_prompt.py b/pkg-py/src/querychat/_handoff_prompt.py index 10ed4bc74..03f0c4215 100644 --- a/pkg-py/src/querychat/_handoff_prompt.py +++ b/pkg-py/src/querychat/_handoff_prompt.py @@ -2,7 +2,7 @@ import json from pathlib import Path -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal import chevron from pydantic import BaseModel, Field, create_model, field_validator @@ -116,12 +116,12 @@ def handoff_result_model( ) -> type[HandoffResult]: # Static typing cannot express Literal values created from runtime names. table_name_type = Literal[tuple(table_names)] # type: ignore[valid-type] - referenced_tables = ( + referenced_tables: tuple[Any, Any] = ( list[table_name_type], # type: ignore[valid-type] Field(description="Registered table names used by the handoff source."), ) language_type = Literal[tuple(languages)] # type: ignore[valid-type] - language = ( + language: tuple[Any, Any] = ( language_type, # type: ignore[valid-type] Field(description="Programming language used by the handoff."), ) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 0d2535b78..b541ba8de 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/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 3d7656d0a..5906d86cd 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -17,7 +17,7 @@ from ._querychat_core import warn_multi_table_flat_accessor from ._table_accessor import TableAccessor from ._viz_altair_widget import AltairWidget -from ._viz_ggsql import execute_ggsql +from ._viz_ggsql import caching_reader, execute_ggsql from ._viz_utils import has_viz_tool, preload_viz_deps_server, preload_viz_deps_ui if TYPE_CHECKING: @@ -509,12 +509,16 @@ def restore_viz_widgets( restored: list[VizWidgetEntry] = [] + # One reader for the whole restore pass: restored widgets often share a + # base read, so the memo avoids re-executing it for each entry. + reader = caching_reader(executor) + for entry in saved_widgets: widget_id = entry["widget_id"] ggsql_str = entry["ggsql"] try: - validated = validate(ggsql_str) - spec = execute_ggsql(executor, validated) + validate(ggsql_str) # early syntax check; errors skip the entry + spec = execute_ggsql(reader, ggsql_str) altair_widget = AltairWidget.from_ggsql(spec, widget_id=widget_id) register_widget(widget_id, altair_widget.widget) restored.append(entry) diff --git a/pkg-py/src/querychat/_system_prompt.py b/pkg-py/src/querychat/_system_prompt.py index b53626827..5552d50b4 100644 --- a/pkg-py/src/querychat/_system_prompt.py +++ b/pkg-py/src/querychat/_system_prompt.py @@ -212,6 +212,7 @@ 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, "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 9e166c151..c7cbf1c58 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -2,8 +2,8 @@ from __future__ import annotations -import re -from typing import TYPE_CHECKING +import warnings +from typing import TYPE_CHECKING, Literal from ._utils import to_polars @@ -13,90 +13,163 @@ from ._query_executor import QueryExecutor -def execute_ggsql(executor: QueryExecutor, validated: ggsql.Validated) -> ggsql.Spec: +class ExecutorReader: """ - Execute a pre-validated ggsql query against a QueryExecutor, returning a Spec. + ggsql reader protocol over a querychat executor (source reads only). - Executes the SQL portion through the executor (preserving database pushdown), - then feeds the result into a ggsql DuckDBReader to produce a Spec. + Designed exclusively for composition with ``ggsql.CachingReader``: the + caching reader stages every derived table into its own in-memory cache, + so this object only ever serves base reads of the user's data, executed + through the querychat executor (preserving database pushdown). + + ``register``/``unregister`` are part of ggsql's reader protocol but are + never called in this composition — ``CachingReader`` routes all writes + to its private cache — so they are no-ops here. + """ + + def __init__(self, executor: QueryExecutor) -> None: + self._executor = executor + + def execute_sql(self, sql: str): + """Run a base read through the executor, as a pyarrow Table.""" + return to_polars(self._executor.execute_query(sql)).to_arrow() + + def register(self, name, table, replace) -> None: + # Writes belong to CachingReader's cache, never the primary. + pass + + def unregister(self, name) -> None: + pass + + +def caching_reader(executor: QueryExecutor) -> ggsql.CachingReader: + """ + Create a ``ggsql.CachingReader`` serving base reads through ``executor``. + + The reader owns a private in-memory cache that memoizes source reads, so + its lifetime should match the span over which repeated reads pay off: + create one per chat client and reuse it across viz calls. When the client + is dropped, the reader (and its cache) is freed with it. Callers own the + reader — there is deliberately no shared registry here, since a + module-level reader keeps its executor (and cache) alive for the life of + the process. + + .. note:: + + Memoized reads assume the underlying data doesn't change: an + identical query within the cache TTL (5 minutes by default) is + served from the cache rather than re-read from the source. If the + data can change mid-session — e.g. update queries are enabled + (``QUERYCHAT_ENABLE_UPDATE_QUERIES=true``), or a live database is + written to by another process — charts may lag those changes by up + to the TTL. Mutations made through querychat's query tool do not + invalidate the memo. Tune or disable the cache via the + ``GGSQL_CACHE_TTL``, ``GGSQL_CACHE_MAX_BYTES``, and + ``GGSQL_CACHE_DISABLED`` environment variables (read when each + reader is created). Parameters ---------- executor - The querychat QueryExecutor to execute the SQL portion against. - validated - A pre-validated ggsql query (from ``ggsql.validate()``). + The querychat QueryExecutor (or bare DataSource) to read through. Returns ------- - ggsql.Spec - The writer-independent plot specification. + ggsql.CachingReader + A caching reader wrapping the executor. """ - from ggsql import DuckDBReader - - visual = validated.visual() - if has_layer_level_source(visual): - # Short term, querychat only supports visual layers that can be replayed - # from one final SQL result. Long term, the cleaner fix is likely to use - # ggsql's native remote-reader execution path (for example via ODBC-backed - # Readers) instead of reconstructing multi-relation scope here. - raise ValueError( - "Layer-specific sources are not currently supported in querychat visual " - "queries. Rewrite the query so that all layers come from the final SQL " - "result." - ) + from ggsql import CachingReader - pl_df = to_polars(executor.execute_query(validated.sql())) + return CachingReader(ExecutorReader(executor)) - reader = DuckDBReader("duckdb://memory") - table = extract_visualise_table(visual) - if table is not None: - # VISUALISE [mappings] FROM — register data under the - # referenced table name and execute the visual part directly. - name = table[1:-1] if table.startswith('"') and table.endswith('"') else table - reader.register(name, pl_df) - return reader.execute(visual) - else: - # SELECT ... VISUALISE — no FROM in VISUALISE clause, so register - # under a synthetic name and prepend a SELECT. - reader.register("_data", pl_df) - return reader.execute(f"SELECT * FROM _data {visual}") +SPATIAL_SMOKE_QUERY = ( + "INSTALL spatial; LOAD spatial; " + "VISUALISE FROM ggsql:world DRAW spatial FILTER iso_a3 = 'FRA'" +) -def extract_visualise_table(visual: str) -> str | None: +def prewarm_spatial(reader: ggsql.CachingReader) -> bool: """ - Extract the table name from ``VISUALISE ... FROM
`` if present. + Install and load the DuckDB spatial extension for the reader's cache. - This reimplements a small part of ggsql's parsing because the current - Python bindings do not expose the top-level VISUALISE source directly. + The first ``INSTALL`` downloads the extension binary (network required); + afterwards it is cached on disk and this is cheap. Returns ``True`` when + spatial support is ready. On failure, warns and returns ``False`` so the + caller can treat the ``"maps"`` capability as unavailable. """ - draw_pos = re.search(r"\bDRAW\b", visual, re.IGNORECASE) - vis_clause = visual[: draw_pos.start()] if draw_pos else visual - m = re.search(r'\bFROM\s+("[^"]+?"|\S+)', vis_clause, re.IGNORECASE) - return m.group(1) if m else None + try: + 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 has_layer_level_source(visual: str) -> bool: +def translate_spatial_error( + error: Exception, *, maps_state: Literal["enabled", "disabled", "setup_failed"] +) -> str | None: """ - Return ``True`` when a DRAW clause defines its own ``FROM ``. + Translate ggsql's spatial-extension failure into actionable guidance. - Querychat currently replays the VISUALISE portion against a single local - relation, so layer-specific sources cannot be preserved reliably. + Returns the translated message, or ``None`` when the error is not a + spatial setup failure (or maps are fully enabled, making the message + moot). ``maps_state`` is ``"enabled"`` when the pre-warm succeeded, + ``"setup_failed"`` when it ran but failed, and ``"disabled"`` when the + session never opted into maps. """ - clauses = re.split( - r"(?=\b(?:DRAW|SCALE|PROJECT|FACET|PLACE|LABEL|THEME)\b)", - visual, - flags=re.IGNORECASE, + msg = str(error).lower() + if maps_state == "enabled" or "spatial" not in msg or "extension" not in msg: + return None + if maps_state == "setup_failed": + 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." + ) + return ( + "This query requires map/spatial support, which is not enabled in " + "this session. Do not retry with DRAW spatial, PROJECT TO , or " + "ggsql:world; suggest a non-map visualization instead. The app " + 'developer can enable maps with tools=("default", "maps").' ) - for clause in clauses: - if not re.match(r"^\s*DRAW\b", clause, re.IGNORECASE): - continue - if re.search( - r'\bMAPPING\b[\s\S]*?\bFROM\s+("[^"]+?"|\S+)', - clause, - re.IGNORECASE, - ): - return True - return False + + +def execute_ggsql(reader: ggsql.CachingReader, query: str) -> ggsql.Spec: + """ + Execute a ggsql query through a caching reader, returning a Spec. + + The whole query (SQL + VISUALISE) executes through the reader: base reads + go to the underlying executor (preserving database pushdown) and are + memoized for the cache TTL, while ggsql-generated derived tables live in + the reader's private in-memory cache. Layer-level sources (``DRAW ... + FROM
``) work because every layer source is staged through that + cache. + + Callers that want validation feedback should run ``ggsql.validate()`` + first themselves; the raw query string is what gets executed (a + ``Validated`` splits ``VISUALISE ... FROM `` into overlapping SQL + and visual parts that cannot be concatenated back into a valid query). + + Parameters + ---------- + reader + A ``ggsql.CachingReader`` from ``caching_reader()``. Reuse one reader + across related viz calls so memoized base reads are shared. + query + The full ggsql query string (SQL + VISUALISE). + + Returns + ------- + ggsql.Spec + The writer-independent plot specification. + + """ + return reader.execute(query) diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index fa6306d8c..4191d5c27 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 execute_ggsql +from ._viz_ggsql import ( + caching_reader, + execute_ggsql, + 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 when the viz reader is first created). 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,15 +195,25 @@ 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 VegaLiteWriter, validate + # One CachingReader per chat client (this closure is created once per + # client): the memo only pays off across viz calls within a session, and + # 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, maps_available markdown = f"```sql\n{ggsql}\n```" try: @@ -210,7 +229,11 @@ def visualize( "\n".join(error["message"] for error in validated.errors()) ) - spec = execute_ggsql(executor, validated) + if reader is None: + reader = caching_reader(executor) + if maps_available: + maps_available = prewarm_spatial(reader) + spec = execute_ggsql(reader, ggsql) raw_chart = VegaLiteWriter().render_chart(spec) altair_widget = AltairWidget(copy.deepcopy(raw_chart)) @@ -235,7 +258,15 @@ def visualize( ) except Exception as e: - error_msg = truncate_error(str(e)) + maps_state = ( + "enabled" + if maps_available + else ("setup_failed" if maps else "disabled") + ) + translated = translate_spatial_error(e, maps_state=maps_state) + 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/_viz_utils.py b/pkg-py/src/querychat/_viz_utils.py index 204d34486..9da578ece 100644 --- a/pkg-py/src/querychat/_viz_utils.py +++ b/pkg-py/src/querychat/_viz_utils.py @@ -14,6 +14,11 @@ def has_viz_tool(tools: set[str] | None) -> bool: return tools is not None and "visualize" in tools +def has_maps_tool(tools: set[str] | None) -> bool: + """Check if the maps capability is among the configured tools.""" + return tools is not None and "maps" in tools + + def has_viz_deps() -> bool: """Check whether visualization dependencies (ggsql, altair, shinywidgets, vl-convert-python) are installed.""" return all( 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..17d35a94f --- /dev/null +++ b/pkg-py/src/querychat/prompts/ggsql-spatial.md @@ -0,0 +1,36 @@ +## Maps with ggsql (spatial) + +The `maps` capability is enabled: the DuckDB `spatial` extension is already installed and loaded. Do NOT write `INSTALL spatial`, `LOAD spatial`, or `ST_*` function calls — querychat has handled spatial setup, and spatial SQL functions may not exist on the user's database. Keep spatial logic in ggsql clauses. + +### 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. + +### 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 +``` + +### 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. + +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..b0fc77b2d 100644 --- a/pkg-py/src/querychat/prompts/prompt.md +++ b/pkg-py/src/querychat/prompts/prompt.md @@ -196,6 +196,12 @@ Match the chart type to what the user is trying to understand: {{> ggsql-syntax}} +{{#has_tool_maps}} + + +{{> ggsql-spatial}} + +{{/has_tool_maps}} {{#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 e7072808b..d11c3b02f 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -6,48 +6,7 @@ import pytest from querychat._datasource import DataFrameSource from querychat._viz_altair_widget import AltairWidget -from querychat._viz_ggsql import ( - execute_ggsql, - extract_visualise_table, - has_layer_level_source, -) - - -class TestExtractVisualiseTable: - """Tests for extract_visualise_table() parsing.""" - - def test_bare_identifier(self): - assert ( - extract_visualise_table("VISUALISE x, y FROM mytable DRAW point") - == "mytable" - ) - - def test_quoted_identifier(self): - assert ( - extract_visualise_table('VISUALISE x FROM "my table" DRAW point') - == '"my table"' - ) - - def test_no_from_returns_none(self): - assert extract_visualise_table("VISUALISE x, y DRAW point") is None - - def test_ignores_draw_level_from(self): - visual = "VISUALISE x, y DRAW bar MAPPING z AS fill FROM summary" - assert extract_visualise_table(visual) is None - - -class TestHasLayerLevelSource: - def test_detects_draw_level_from(self): - visual = "VISUALISE x, y DRAW bar MAPPING z AS fill FROM summary" - assert has_layer_level_source(visual) - - def test_ignores_visualise_from(self): - visual = "VISUALISE x, y FROM sales DRAW point MAPPING z AS color" - assert not has_layer_level_source(visual) - - def test_ignores_scale_from(self): - visual = "VISUALISE x, y DRAW point MAPPING z AS color SCALE x FROM [0, 10]" - assert not has_layer_level_source(visual) +from querychat._viz_ggsql import caching_reader, execute_ggsql class TestGgsqlValidate: @@ -146,7 +105,7 @@ def test_full_pipeline(self): nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = execute_ggsql(caching_reader(ds), query) altair_widget = AltairWidget.from_ggsql(spec) result = altair_widget.widget.chart.to_dict() assert "$schema" in result @@ -158,7 +117,7 @@ def test_with_filtered_query(self): ) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data WHERE x > 2 VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = execute_ggsql(caching_reader(ds), query) assert spec.metadata()["rows"] == 3 @pytest.mark.ggsql @@ -166,7 +125,7 @@ def test_spec_has_visual(self): nw_df = nw.from_native(pl.DataFrame({"x": [1, 2], "y": [3, 4]})) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = execute_ggsql(caching_reader(ds), query) assert "VISUALISE" in spec.visual() @pytest.mark.ggsql @@ -174,7 +133,7 @@ def test_visualise_from_path(self): nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) ds = DataFrameSource(nw_df, "test_data") query = "VISUALISE x, y FROM test_data DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = execute_ggsql(caching_reader(ds), query) assert spec.metadata()["rows"] == 3 assert "VISUALISE" in spec.visual() @@ -185,7 +144,7 @@ def test_with_pandas_dataframe(self): nw_df = nw.from_native(pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = execute_ggsql(caching_reader(ds), query) altair_widget = AltairWidget.from_ggsql(spec) result = altair_widget.widget.chart.to_dict() assert "$schema" in result @@ -206,49 +165,11 @@ def test_supports_uppercase_column_references_without_renaming(self): "VISUALISE ROOM_TYPE AS x, listings AS y " "DRAW bar" ) - spec = execute_ggsql(ds, ggsql.validate(query)) - assert "VISUALISE" in spec.visual() - - def test_passes_original_column_names_through_to_ggsql(self, monkeypatch): - nw_df = nw.from_native( - pl.DataFrame( - { - "ROOM_TYPE": ["Entire home", "Private room"], - "listings": [10, 20], - } - ) - ) - ds = DataFrameSource(nw_df, "upper_table") - validated = type( - "Validated", - (), - { - "sql": lambda self: "SELECT ROOM_TYPE, listings FROM upper_table", - "visual": lambda self: ( - "VISUALISE ROOM_TYPE AS x, listings AS y DRAW bar" - ), - }, - )() - captured_columns = None - - class FakeReader: - def register(self, name, df): - nonlocal captured_columns - captured_columns = list(df.columns) - - def execute(self, query): - return type("Spec", (), {"visual": lambda self: query})() - - import ggsql - - monkeypatch.setattr(ggsql, "DuckDBReader", lambda *_args: FakeReader()) - - spec = execute_ggsql(ds, validated) + spec = execute_ggsql(caching_reader(ds), query) assert "VISUALISE" in spec.visual() - assert captured_columns == ["ROOM_TYPE", "listings"] @pytest.mark.ggsql - def test_rejects_layer_level_from_sources_with_clear_error(self): + def test_supports_layer_level_from_sources(self): nw_df = nw.from_native( pl.DataFrame( { @@ -272,11 +193,12 @@ def test_rejects_layer_level_from_sources_with_clear_error(self): DRAW bar MAPPING region AS x, total AS y FROM summary """ - with pytest.raises( - ValueError, - match="Layer-specific sources are not currently supported", - ): - execute_ggsql(ds, ggsql.validate(query)) + # Previously raised "Layer-specific sources are not currently + # supported"; CachingReader stages each layer source through its + # cache, so the query now runs end to end. + spec = execute_ggsql(caching_reader(ds), query) + assert spec is not None + assert "VISUALISE" in spec.visual() @pytest.mark.ggsql def test_supports_single_relation_raw_plus_summary_overlay(self): @@ -320,6 +242,182 @@ def test_supports_single_relation_raw_plus_summary_overlay(self): DRAW line MAPPING category AS color FILTER layer_type = 'summary' """ - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = execute_ggsql(caching_reader(ds), query) assert spec.metadata()["rows"] == 4 assert "VISUALISE" in spec.visual() + + +class TestCachingBehavior: + """CachingReader shares memoized source reads across viz executions.""" + + @pytest.mark.ggsql + def test_identical_queries_share_source_reads(self, monkeypatch): + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + + calls = 0 + original = ds.execute_query + + def counting(query): + nonlocal calls + calls += 1 + return original(query) + + monkeypatch.setattr(ds, "execute_query", counting) + + query = "SELECT * FROM test_data VISUALISE x, y DRAW point" + reader = caching_reader(ds) + execute_ggsql(reader, query) + execute_ggsql(reader, query) + assert calls == 1 + + @pytest.mark.ggsql + def test_independent_readers_do_not_share_memo(self, monkeypatch): + """Readers own their caches; nothing is shared implicitly.""" + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + + calls = 0 + original = ds.execute_query + + def counting(query): + nonlocal calls + calls += 1 + return original(query) + + monkeypatch.setattr(ds, "execute_query", counting) + + query = "SELECT * FROM test_data VISUALISE x, y DRAW point" + execute_ggsql(caching_reader(ds), query) + execute_ggsql(caching_reader(ds), query) + assert calls == 2 + + +class TestRestoreVizWidgets: + """ + Bookmark restore executes through a CachingReader (PR #321 follow-up). + + Regression guard: an earlier draft passed the executor itself to + execute_ggsql(), so every widget silently failed to restore (the + AttributeError was swallowed by restore_viz_widgets' except Exception). + """ + + @pytest.fixture + def register_spy(self, monkeypatch): + """Stub out the two places restore_viz_widgets needs a Shiny session.""" + import ipywidgets + import shinywidgets + + registered = [] + monkeypatch.setattr( + shinywidgets, + "register_widget", + lambda widget_id, widget: registered.append(widget_id), + ) + # shinywidgets hooks ipywidgets construction (Widget.on_widget_constructed) + # and requires an active session; unhook so widgets build session-free. + monkeypatch.setattr(ipywidgets.Widget, "_widget_construction_callback", None) + return registered + + @pytest.mark.ggsql + def test_restores_and_registers_widgets(self, register_spy): + from querychat._shiny_module import restore_viz_widgets + + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + entries = [ + { + "widget_id": "w1", + "ggsql": "SELECT * FROM test_data VISUALISE x, y DRAW point", + }, + { + "widget_id": "w2", + "ggsql": "SELECT * FROM test_data WHERE x > 1 VISUALISE x, y DRAW point", + }, + ] + + restored = restore_viz_widgets(ds, entries) + + assert [e["widget_id"] for e in restored] == ["w1", "w2"] + assert register_spy == ["w1", "w2"] + + @pytest.mark.ggsql + def test_shared_base_reads_are_memoized(self, monkeypatch, register_spy): + """One reader per restore pass: identical queries share one source read.""" + from querychat._shiny_module import restore_viz_widgets + + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + + calls = 0 + original = ds.execute_query + + def counting(query): + nonlocal calls + calls += 1 + return original(query) + + monkeypatch.setattr(ds, "execute_query", counting) + + query = "SELECT * FROM test_data VISUALISE x, y DRAW point" + entries = [ + {"widget_id": "w1", "ggsql": query}, + {"widget_id": "w2", "ggsql": query}, + ] + + restored = restore_viz_widgets(ds, entries) + + assert len(restored) == 2 + assert calls == 1 + + @pytest.mark.ggsql + def test_invalid_entries_are_skipped_with_warning(self, register_spy): + from querychat._shiny_module import restore_viz_widgets + + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + entries = [ + { + "widget_id": "bad", + "ggsql": "SELECT * FROM nope VISUALISE x, y DRAW point", + }, + { + "widget_id": "good", + "ggsql": "SELECT * FROM test_data VISUALISE x, y DRAW point", + }, + ] + + with pytest.warns(UserWarning, match="Failed to restore"): + restored = restore_viz_widgets(ds, entries) + + assert [e["widget_id"] for e in restored] == ["good"] + assert register_spy == ["good"] + + +class TestPrewarmSpatial: + def test_success_runs_setup_query(self): + from querychat._viz_ggsql import prewarm_spatial + + calls = [] + + class FakeReader: + def execute(self, query): + calls.append(query) + return object() + + assert prewarm_spatial(FakeReader()) is True + assert len(calls) == 1 + assert "INSTALL spatial" in calls[0] + assert "LOAD spatial" in calls[0] + + def test_failure_warns_and_returns_false(self): + from querychat._viz_ggsql import prewarm_spatial + + class FailReader: + 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..f84504267 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,30 @@ 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 + assert "ggsql:world" not in prompt + + @pytest.fixture def sample_df(): return pl.DataFrame( @@ -376,3 +447,83 @@ 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_spatial_once_when_maps_enabled(self, data_source, monkeypatch): + import querychat._viz_tools as viz_tools + + calls = [] + monkeypatch.setattr( + viz_tools, "prewarm_spatial", lambda reader: 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(calls) == 1 + + @pytest.mark.ggsql + def test_no_prewarm_when_maps_disabled(self, data_source, monkeypatch): + import querychat._viz_tools as viz_tools + + calls = [] + monkeypatch.setattr( + viz_tools, "prewarm_spatial", lambda reader: 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 calls == [] + + +class TestTranslateSpatialError: + SPATIAL_ERR = Exception( + 'IO Error: Extension "/home/u/.duckdb/extensions/v1.5.2/linux_amd64/' + 'spatial.duckdb_extension" not found' + ) + + def test_disabled_message_points_to_tools(self): + from querychat._viz_ggsql import translate_spatial_error + + msg = translate_spatial_error(self.SPATIAL_ERR, maps_state="disabled") + assert msg is not None + assert "maps" in msg + assert 'tools=("default", "maps")' in msg + + def test_setup_failed_message(self): + from querychat._viz_ggsql import translate_spatial_error + + msg = translate_spatial_error(self.SPATIAL_ERR, maps_state="setup_failed") + assert msg is not None + assert "setup failed" in msg + + def test_enabled_or_nonspatial_returns_none(self): + from querychat._viz_ggsql import translate_spatial_error + + assert translate_spatial_error(self.SPATIAL_ERR, maps_state="enabled") is None + assert ( + translate_spatial_error(Exception("boom"), maps_state="disabled") is None + ) + + @pytest.mark.ggsql + def test_tool_translates_spatial_error_when_maps_disabled( + self, data_source, monkeypatch + ): + import querychat._viz_tools as viz_tools + + def boom(reader, query): + raise RuntimeError(str(self.SPATIAL_ERR)) + + monkeypatch.setattr(viz_tools, "execute_ggsql", boom) + 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 'tools=("default", "maps")' in str(result.error) diff --git a/pkg-r/DESCRIPTION b/pkg-r/DESCRIPTION index 4dba23c39..e5fe614f0 100644 --- a/pkg-r/DESCRIPTION +++ b/pkg-r/DESCRIPTION @@ -1,6 +1,6 @@ Package: querychat Title: Filter and Query Data Frames in 'shiny' Using an LLM Chat Interface -Version: 0.4.0 +Version: 0.4.1.9000 Authors@R: c( person("Garrick", "Aden-Buie", , "garrick@posit.co", role = c("aut", "cre"), comment = c(ORCID = "0000-0002-7111-0077")), @@ -60,7 +60,7 @@ Suggests: withr VignetteBuilder: knitr -Config/roxygen2/version: 8.1.0 +Config/roxygen2/version: 8.1.1 Config/testthat/edition: 3 Config/testthat/parallel: true Encoding: UTF-8 diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 2759c6f42..0b2ceb326 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -1,3 +1,11 @@ +# querychat (development version) + +# querychat 0.4.1 + +## Bug fixes + +* Tests no longer fail on systems where Suggests packages like duckdb or RSQLite aren't installed (as on some CRAN check flavors): test fixtures now skip when a required database engine is missing instead of erroring or silently falling back to a different engine. (#317) + # querychat 0.4.0 ## New features diff --git a/pkg-r/R/querychat-package.R b/pkg-r/R/querychat-package.R index c09a99379..737b09342 100644 --- a/pkg-r/R/querychat-package.R +++ b/pkg-r/R/querychat-package.R @@ -77,10 +77,7 @@ NULL NULL release_bullets <- function() { - c( - "Run `staticimports::import()` to update static imports", - "Enable `development.mode: auto` in `_pkgdown.yml` and remove this release bullet." - ) + c("Run `staticimports::import()` to update static imports") } suppress_rcmdcheck <- function() { diff --git a/pkg-r/cran-comments.md b/pkg-r/cran-comments.md index 1e2ec3c52..3211fb10b 100644 --- a/pkg-r/cran-comments.md +++ b/pkg-r/cran-comments.md @@ -1,7 +1,21 @@ +## Resubmission + +This is a resubmission of an archived package. querychat was archived on +2026-09-27 because R CMD check failed on `r-devel-linux-x86_64-fedora-clang`: +the test suite assumed that duckdb (a package listed in Suggests) was +installed on the checking machine. + +The failing tests have been fixed: tests that require a package listed in +`Suggests` now skip when that package isn't installed, rather than erroring +or silently falling back to another engine. No changes were made to runtime +package code. + ## R CMD check results 0 errors | 0 warnings | 1 note +* This is a new submission of a previously archived package. + * The CRAN incoming feasibility check flags https://platform.openai.com/ (linked from README.md) with status 403. The site blocks automated requests; the link works in a browser. diff --git a/pkg-r/tests/testthat/helper-fixtures.R b/pkg-r/tests/testthat/helper-fixtures.R index 9b254fa07..7f49235f3 100644 --- a/pkg-r/tests/testthat/helper-fixtures.R +++ b/pkg-r/tests/testthat/helper-fixtures.R @@ -76,11 +76,33 @@ local_sqlite_connection <- function( list(conn = conn, path = temp_db) } -# Skip test if no DataFrameSource engine is available -skip_if_no_dataframe_engine <- function() { - if (!rlang::is_installed("duckdb") && !rlang::is_installed("RSQLite")) { - skip("Neither duckdb nor RSQLite is installed") +# Skip test if no DataFrameSource engine is available. +# +# When `engine` is NULL, skip only if neither duckdb nor RSQLite is installed +# (mirroring how DataFrameSource resolves a default engine). When a specific +# engine is requested, skip unless that engine's package is installed, so +# tests never silently fall back to a different engine (e.g. SQLite) and then +# fail their DuckDB-specific expectations. +skip_if_no_dataframe_engine <- function(engine = NULL) { + if (is.null(engine)) { + engine <- getOption("querychat.DataFrameSource.engine", NULL) } + + if (is.null(engine)) { + if (!rlang::is_installed("duckdb") && !rlang::is_installed("RSQLite")) { + skip("Neither duckdb nor RSQLite is installed") + } + return(invisible()) + } + + engine <- tolower(engine) + if (engine == "duckdb") { + skip_if_not_installed("duckdb") + } else if (engine == "sqlite") { + skip_if_not_installed("RSQLite") + } + + invisible() } # Create a DataFrameSource with automatic cleanup @@ -90,6 +112,7 @@ local_data_frame_source <- function( engine = "duckdb", env = parent.frame() ) { + skip_if_no_dataframe_engine(engine) df_source <- DataFrameSource$new(data, table_name, engine = engine) withr::defer(df_source$cleanup(), envir = env) df_source @@ -119,6 +142,7 @@ local_recording_data_frame_source <- function( engine = "duckdb", env = parent.frame() ) { + skip_if_no_dataframe_engine(engine) state <- new.env(parent = emptyenv()) state$get_data_calls <- 0L state$get_data_error <- NULL @@ -339,6 +363,11 @@ local_querychat <- function( ..., env = parent.frame() ) { + # Plain data frames are wrapped in a DataFrameSource using the default + # engine, which requires duckdb or RSQLite to be installed. + if (is.data.frame(data_source)) { + skip_if_no_dataframe_engine() + } qc <- QueryChat$new(data_source, table_name, ...) withr::defer(qc$cleanup(), envir = env) qc diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index f3a4720ae..2499c17af 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -222,6 +222,8 @@ describe("QueryChat integration with DBISource", { }) describe("QueryChat$cleanup()", { + skip_if_no_dataframe_engine() + it("cleans up data source resources", { test_df <- new_test_df() qc <- QueryChat$new(test_df, greeting = "Test") @@ -235,6 +237,8 @@ describe("QueryChat$cleanup()", { }) describe("QueryChat$system_prompt", { + skip_if_no_dataframe_engine() + it("returns the system prompt from the client", { test_df <- new_test_df() qc <- QueryChat$new(test_df, greeting = "Test") @@ -328,6 +332,8 @@ describe("QueryChat$data_source", { }) describe("QueryChat$client()", { + skip_if_no_dataframe_engine() + it("uses default tools when tools = NA", { qc <- QueryChat$new( new_test_df(), @@ -678,6 +684,8 @@ describe("QueryChat$client()", { }) test_that("QueryChat$generate_greeting() generates a greeting using the LLM client", { + skip_if_no_dataframe_engine() + client <- mock_ellmer_chat_client( public = list( chat = function(message, ...) { @@ -698,6 +706,7 @@ test_that("QueryChat$generate_greeting() generates a greeting using the LLM clie }) test_that("QueryChat$server() errors when called outside Shiny context", { + skip_if_no_dataframe_engine() withr::local_envvar(OPENAI_API_KEY = "boop") test_df <- new_test_df() @@ -710,6 +719,8 @@ test_that("QueryChat$server() errors when called outside Shiny context", { }) test_that("QueryChat$new() validates history and stores it verbatim", { + skip_if_no_dataframe_engine() + test_df <- new_test_df() qc_default <- QueryChat$new(test_df, greeting = "Test") @@ -835,6 +846,8 @@ test_that("QueryChat$app_obj() infers Shiny bookmarking from history's restore_m }) describe("QueryChat internal client handoff availability", { + skip_if_no_dataframe_engine() + local_mocked_r6_class( QueryChat, public = list( @@ -914,6 +927,8 @@ describe("querychat()", { }) describe("QueryChat$console()", { + skip_if_no_dataframe_engine() + local_mocked_r6_class( QueryChat, public = list( @@ -1086,6 +1101,8 @@ test_that("querychat_app() only cleans up data frame sources on exit", { }) describe("QueryChat$server() client override", { + skip_if_no_dataframe_engine() + it("accepts a client parameter", { withr::local_envvar(OPENAI_API_KEY = "boop") test_df <- new_test_df() diff --git a/pkg-r/tests/testthat/test-QueryChatSystemPrompt.R b/pkg-r/tests/testthat/test-QueryChatSystemPrompt.R index 1138311e7..1954aea09 100644 --- a/pkg-r/tests/testthat/test-QueryChatSystemPrompt.R +++ b/pkg-r/tests/testthat/test-QueryChatSystemPrompt.R @@ -125,6 +125,8 @@ describe("QueryChatSystemPrompt$new()", { }) describe("QueryChatSystemPrompt$render()", { + skip_if_no_dataframe_engine() + it("renders handoff guidance only when available", { df <- new_test_df() ds <- DataFrameSource$new(df, "test_table") @@ -268,6 +270,8 @@ describe("QueryChatSystemPrompt$render()", { }) it("includes db_type in rendered output", { + skip_if_not_installed("duckdb") + df <- new_test_df() ds <- DataFrameSource$new(df, "test_table") withr::defer(ds$cleanup()) @@ -370,6 +374,8 @@ describe("QueryChatSystemPrompt$render()", { }) it("detects DuckDB correctly", { + skip_if_not_installed("duckdb") + df <- new_test_df() ds <- DataFrameSource$new(df, "test_table") withr::defer(ds$cleanup()) @@ -405,6 +411,8 @@ describe("QueryChatSystemPrompt$render()", { }) describe("QueryChatSystemPrompt with full prompt.md template", { + skip_if_no_dataframe_engine() + it("renders full template with data_description", { df <- new_test_df(3) ds <- DataFrameSource$new(df, "test_table") @@ -428,6 +436,8 @@ describe("QueryChatSystemPrompt with full prompt.md template", { }) it("includes DuckDB-specific content for DuckDB sources", { + skip_if_not_installed("duckdb") + df <- new_test_df() ds <- DataFrameSource$new(df, "test_table") withr::defer(ds$cleanup()) diff --git a/pyproject.toml b/pyproject.toml index 56ca89e43..acc98c42f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,7 +62,8 @@ gradio = ["gradio>=6.0"] dash = ["dash-ag-grid>=31.0", "dash[async]>=3.1", "dash-bootstrap-components>=2.0", "pandas"] # Visualization with ggsql viz = [ - "ggsql>=0.3.2", + # CachingReader (required by _viz_ggsql.py) first appears in 0.5.2 + "ggsql>=0.5.2", "altair>=6.0", "shinywidgets>=0.8.0", "vl-convert-python>=1.9.0", @@ -117,7 +118,8 @@ dev = [ "ibis-framework[duckdb]>=9.0.0", # see note on the `ibis` extra above "sqlglot<30.18.0", - "ggsql>=0.3.2", + # CachingReader (required by _viz_ggsql.py) first appears in 0.5.2 + "ggsql>=0.5.2", "altair>=6.0", "shinywidgets>=0.8.0", "vl-convert-python>=1.9.0",