From a068ee8f3f49e7127ee7503bedd6e1b69cb687d8 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 22 Sep 2026 10:15:12 -0500 Subject: [PATCH 01/24] Add standalone multi-table querychat demo (fine-jewelry retailer) A single-folder, uv-runnable Python example for demoing querychat's multi-table support without needing the repo checked out or set up -- useful for customer demos. Data is synthetic (seeded, deterministic) and generated once to a local parquet cache on first run. --- .gitignore | 1 + .../examples/13-multi-table-standalone/app.py | 70 +++++ .../13-multi-table-standalone/data-dict.yaml | 178 +++++++++++++ .../generate_data.py | 252 ++++++++++++++++++ 4 files changed, 501 insertions(+) create mode 100644 pkg-py/examples/13-multi-table-standalone/app.py create mode 100644 pkg-py/examples/13-multi-table-standalone/data-dict.yaml create mode 100644 pkg-py/examples/13-multi-table-standalone/generate_data.py 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/examples/13-multi-table-standalone/app.py b/pkg-py/examples/13-multi-table-standalone/app.py new file mode 100644 index 000000000..e48cec3e5 --- /dev/null +++ b/pkg-py/examples/13-multi-table-standalone/app.py @@ -0,0 +1,70 @@ +# /// script +# dependencies = [ +# "querychat[pandas]", +# "anthropic", +# "pyarrow", +# ] +# /// +"""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** +Show me a few example rows from each table +What's the average price of a round diamond over 1 carat? + +**Explore across tables** +Which setting style is most often paired with princess-cut diamonds in engagement rings? +What's the average order value by acquisition channel? +Show total revenue by month over the last year +""" + +DATA_DESCRIPTION = """\ +Brilliant Retail is an online fine-jewelry retailer. This dataset covers its +diamond and setting catalog alongside customer order history. +""" + +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", + 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..7bbc50ef6 --- /dev/null +++ b/pkg-py/examples/13-multi-table-standalone/generate_data.py @@ -0,0 +1,252 @@ +# /// script +# dependencies = ["pandas", "pyarrow"] +# /// +"""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 pandas as pd + +SEED = 42 +TODAY = date(2026, 9, 1) +TABLE_NAMES = ["diamonds", "settings", "customers", "orders", "order_items"] + + +def load_tables(data_dir: Path) -> dict[str, pd.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: pd.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.to_parquet(paths[name], index=False) + return tables + + +def generate_tables() -> dict[str, pd.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) -> pd.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 pd.DataFrame(rows) + + +def generate_settings(rng: random.Random, n: int = 50) -> pd.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 pd.DataFrame(rows) + + +def generate_customers(rng: random.Random, n: int = 200) -> pd.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 pd.DataFrame(rows) + + +def generate_orders(rng: random.Random, customers: pd.DataFrame, n: int = 350) -> pd.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"].tolist() + + 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 pd.DataFrame(rows) + + +def generate_order_items( + rng: random.Random, + orders: pd.DataFrame, + diamonds: pd.DataFrame, + settings: pd.DataFrame, + n: int = 450, +) -> pd.DataFrame: + item_types = ["Engagement Ring", "Loose Diamond", "Wedding Band", "Earrings", "Necklace"] + item_type_weights = [35, 15, 20, 15, 15] + order_ids = orders["id"].tolist() + diamond_ids = diamonds["id"].tolist() + setting_ids = settings["id"].tolist() + 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, + } + ) + df = pd.DataFrame(rows) + df["diamond_id"] = df["diamond_id"].astype("Int64") + df["setting_id"] = df["setting_id"].astype("Int64") + return df + + +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}") From 64a2cf15e3d21165dabd9df0f0240bd441401183 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 22 Sep 2026 10:23:01 -0500 Subject: [PATCH 02/24] Enable visualization tool in multi-table standalone demo Pulls in querychat's viz extra so the chat can generate charts, not just tables -- previously the visualize tool was silently disabled with a warning. --- pkg-py/examples/13-multi-table-standalone/app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg-py/examples/13-multi-table-standalone/app.py b/pkg-py/examples/13-multi-table-standalone/app.py index e48cec3e5..c1a7c6826 100644 --- a/pkg-py/examples/13-multi-table-standalone/app.py +++ b/pkg-py/examples/13-multi-table-standalone/app.py @@ -1,6 +1,6 @@ # /// script # dependencies = [ -# "querychat[pandas]", +# "querychat[pandas,viz]", # "anthropic", # "pyarrow", # ] From ac4041a26cc129961c4ef5010d5e4efae9242cf8 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 22 Sep 2026 10:26:00 -0500 Subject: [PATCH 03/24] Switch multi-table standalone demo from pandas to polars Drops pandas as a dependency entirely -- polars' native parquet I/O and nullable Int64 inference handle everything generate_data.py needs without it. --- .../examples/13-multi-table-standalone/app.py | 3 +- .../generate_data.py | 49 +++++++++---------- 2 files changed, 24 insertions(+), 28 deletions(-) diff --git a/pkg-py/examples/13-multi-table-standalone/app.py b/pkg-py/examples/13-multi-table-standalone/app.py index c1a7c6826..ad1d667bb 100644 --- a/pkg-py/examples/13-multi-table-standalone/app.py +++ b/pkg-py/examples/13-multi-table-standalone/app.py @@ -1,8 +1,7 @@ # /// script # dependencies = [ -# "querychat[pandas,viz]", +# "querychat[polars,viz]", # "anthropic", -# "pyarrow", # ] # /// """Multi-table querychat demo: a synthetic fine-jewelry retailer. diff --git a/pkg-py/examples/13-multi-table-standalone/generate_data.py b/pkg-py/examples/13-multi-table-standalone/generate_data.py index 7bbc50ef6..989bf322d 100644 --- a/pkg-py/examples/13-multi-table-standalone/generate_data.py +++ b/pkg-py/examples/13-multi-table-standalone/generate_data.py @@ -1,5 +1,5 @@ # /// script -# dependencies = ["pandas", "pyarrow"] +# dependencies = ["polars"] # /// """Generate (or load cached) synthetic data for a fine-jewelry retailer demo. @@ -18,28 +18,28 @@ from datetime import date, timedelta from pathlib import Path -import pandas as pd +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, pd.DataFrame]: +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: pd.read_parquet(path) for name, path in paths.items()} + 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.to_parquet(paths[name], index=False) + table.write_parquet(paths[name]) return tables -def generate_tables() -> dict[str, pd.DataFrame]: +def generate_tables() -> dict[str, pl.DataFrame]: rng = random.Random(SEED) diamonds = generate_diamonds(rng) settings = generate_settings(rng) @@ -55,7 +55,7 @@ def generate_tables() -> dict[str, pd.DataFrame]: } -def generate_diamonds(rng: random.Random, n: int = 400) -> pd.DataFrame: +def generate_diamonds(rng: random.Random, n: int = 400) -> pl.DataFrame: shapes = [ "Round", "Princess", "Cushion", "Oval", "Emerald", "Radiant", "Asscher", "Marquise", "Pear", "Heart", @@ -100,10 +100,10 @@ def generate_diamonds(rng: random.Random, n: int = 400) -> pd.DataFrame: "price": price, } ) - return pd.DataFrame(rows) + return pl.DataFrame(rows) -def generate_settings(rng: random.Random, n: int = 50) -> pd.DataFrame: +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, @@ -129,10 +129,10 @@ def generate_settings(rng: random.Random, n: int = 50) -> pd.DataFrame: "base_price": base_price, } ) - return pd.DataFrame(rows) + return pl.DataFrame(rows) -def generate_customers(rng: random.Random, n: int = 200) -> pd.DataFrame: +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", @@ -164,15 +164,15 @@ def generate_customers(rng: random.Random, n: int = 200) -> pd.DataFrame: "acquisition_channel": rng.choice(channels), } ) - return pd.DataFrame(rows) + return pl.DataFrame(rows) -def generate_orders(rng: random.Random, customers: pd.DataFrame, n: int = 350) -> pd.DataFrame: +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"].tolist() + customer_ids = customers["id"].to_list() rows = [] for i in range(1, n + 1): @@ -189,21 +189,21 @@ def generate_orders(rng: random.Random, customers: pd.DataFrame, n: int = 350) - "channel": rng.choice(channels), } ) - return pd.DataFrame(rows) + return pl.DataFrame(rows) def generate_order_items( rng: random.Random, - orders: pd.DataFrame, - diamonds: pd.DataFrame, - settings: pd.DataFrame, + orders: pl.DataFrame, + diamonds: pl.DataFrame, + settings: pl.DataFrame, n: int = 450, -) -> pd.DataFrame: +) -> pl.DataFrame: item_types = ["Engagement Ring", "Loose Diamond", "Wedding Band", "Earrings", "Necklace"] item_type_weights = [35, 15, 20, 15, 15] - order_ids = orders["id"].tolist() - diamond_ids = diamonds["id"].tolist() - setting_ids = settings["id"].tolist() + 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"])) @@ -238,10 +238,7 @@ def generate_order_items( "price": price, } ) - df = pd.DataFrame(rows) - df["diamond_id"] = df["diamond_id"].astype("Int64") - df["setting_id"] = df["setting_id"].astype("Int64") - return df + return pl.DataFrame(rows) if __name__ == "__main__": From 40f328e4c9992f350cbf9549d3cfa129deb639c3 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 22 Sep 2026 10:41:49 -0500 Subject: [PATCH 04/24] Work around a ggsql date-coercion bug in the multi-table standalone demo DATE_TRUNC(...)-derived columns hit a ggsql bug where SCALE ... VIA date can't coerce a Timestamp column (posit-dev/ggsql#559). Add extra_instructions telling the LLM to CAST such expressions to DATE, so monthly/weekly revenue charts keep working until the upstream fix lands. --- .../examples/13-multi-table-standalone/app.py | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/pkg-py/examples/13-multi-table-standalone/app.py b/pkg-py/examples/13-multi-table-standalone/app.py index ad1d667bb..7450ff279 100644 --- a/pkg-py/examples/13-multi-table-standalone/app.py +++ b/pkg-py/examples/13-multi-table-standalone/app.py @@ -34,13 +34,19 @@ **order_items**, all linked together. **Get started** -Show me a few example rows from each table -What's the average price of a round diamond over 1 carat? + +
    +
  • Show me a few example rows from each table
    • +
  • What's the average price of a round diamond over 1 carat?
    • +
**Explore across tables** -Which setting style is most often paired with princess-cut diamonds in engagement rings? -What's the average order value by acquisition channel? -Show total revenue by month over the last year + +
    +
  • Which setting style is most often paired with princess-cut diamonds in engagement rings?
    • +
  • What's the average order value by acquisition channel?
    • +
  • Show total revenue by month over the last year
    • +
""" DATA_DESCRIPTION = """\ @@ -48,6 +54,20 @@ 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( @@ -56,6 +76,7 @@ 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") From 74878bd3eeee97224861ee833652fcb44178e183 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 13:17:53 -0500 Subject: [PATCH 05/24] test: layer-level sources succeed; shared source reads (failing) --- pkg-py/tests/test_ggsql.py | 118 ++++++++++--------------------------- 1 file changed, 32 insertions(+), 86 deletions(-) diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index e7072808b..54ff51ce7 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 execute_ggsql class TestGgsqlValidate: @@ -209,46 +168,8 @@ def test_supports_uppercase_column_references_without_renaming(self): 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) - 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(ds, ggsql.validate(query)) + assert spec is not None + assert "VISUALISE" in spec.visual() @pytest.mark.ggsql def test_supports_single_relation_raw_plus_summary_overlay(self): @@ -323,3 +245,27 @@ def test_supports_single_relation_raw_plus_summary_overlay(self): spec = execute_ggsql(ds, ggsql.validate(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" + execute_ggsql(ds, ggsql.validate(query)) + execute_ggsql(ds, ggsql.validate(query)) + assert calls == 1 From bc0b8047a9718e3f83a0183b9616c2d8005b26e1 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 13:20:37 -0500 Subject: [PATCH 06/24] feat: route ggsql viz execution through ggsql.CachingReader execute_ggsql now executes the whole query (SQL + VISUALISE) through a CachingReader wrapping the querychat executor, replacing the hand-rolled split path (executor SQL -> polars -> fresh DuckDBReader per call). Base reads keep database pushdown and are memoized across viz calls within a session (one reader per executor, weakly held); derived tables stage in the reader's private cache, so layer-level DRAW...FROM sources now work instead of raising. Signature change: execute_ggsql(executor, query: str) takes the raw query string. A Validated's sql()/visual() parts overlap for VISUALISE...FROM queries (validate normalizes the FROM into both) and cannot be concatenated back into a valid query. --- pkg-py/src/querychat/_shiny_module.py | 4 +- pkg-py/src/querychat/_viz_ggsql.py | 139 ++++++++++++-------------- pkg-py/src/querychat/_viz_tools.py | 2 +- 3 files changed, 67 insertions(+), 78 deletions(-) diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 3d7656d0a..212424efc 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -513,8 +513,8 @@ def restore_viz_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(executor, 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/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index 9e166c151..14bf8da0a 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -2,7 +2,7 @@ from __future__ import annotations -import re +import weakref from typing import TYPE_CHECKING from ._utils import to_polars @@ -13,19 +13,74 @@ 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: # noqa: ARG002 + # Writes belong to CachingReader's cache, never the primary. + pass + + def unregister(self, name) -> None: # noqa: ARG002 + pass + + +# One CachingReader per executor: the memo only pays off across viz calls +# within a chat session, so the reader (and its cache) must outlive any +# single execute_ggsql() call. Weak keys let finished sessions' executors — +# and their caches — be garbage collected. +_readers: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + +def _caching_reader(executor: QueryExecutor) -> ggsql.CachingReader: + from ggsql import CachingReader + + reader = _readers.get(executor) + if reader is None: + reader = CachingReader(ExecutorReader(executor)) + _readers[executor] = reader + return reader + + +def execute_ggsql(executor: QueryExecutor, query: str) -> ggsql.Spec: + """ + Execute a ggsql query against an executor, returning a Spec. + + The whole query (SQL + VISUALISE) executes through a + ``ggsql.CachingReader`` wrapping the executor: base reads go to the + 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 ---------- 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. + query + The full ggsql query string (SQL + VISUALISE). Returns ------- @@ -33,70 +88,4 @@ def execute_ggsql(executor: QueryExecutor, validated: ggsql.Validated) -> ggsql. The writer-independent plot specification. """ - 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." - ) - - pl_df = to_polars(executor.execute_query(validated.sql())) - - 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}") - - -def extract_visualise_table(visual: str) -> str | None: - """ - Extract the table name from ``VISUALISE ... FROM
`` if present. - - This reimplements a small part of ggsql's parsing because the current - Python bindings do not expose the top-level VISUALISE source directly. - """ - 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 - - -def has_layer_level_source(visual: str) -> bool: - """ - Return ``True`` when a DRAW clause defines its own ``FROM ``. - - Querychat currently replays the VISUALISE portion against a single local - relation, so layer-specific sources cannot be preserved reliably. - """ - clauses = re.split( - r"(?=\b(?:DRAW|SCALE|PROJECT|FACET|PLACE|LABEL|THEME)\b)", - visual, - flags=re.IGNORECASE, - ) - 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 + return _caching_reader(executor).execute(query) diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index fa6306d8c..9ec265813 100644 --- a/pkg-py/src/querychat/_viz_tools.py +++ b/pkg-py/src/querychat/_viz_tools.py @@ -210,7 +210,7 @@ def visualize( "\n".join(error["message"] for error in validated.errors()) ) - spec = execute_ggsql(executor, validated) + spec = execute_ggsql(executor, ggsql) raw_chart = VegaLiteWriter().render_chart(spec) altair_widget = AltairWidget(copy.deepcopy(raw_chart)) From 6cc5ae41e6fe17aa8b9af8b69834e44c5055c538 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 16:03:27 -0500 Subject: [PATCH 07/24] test: pass raw query strings to execute_ggsql Completes bc0b8047: execute_ggsql now takes the whole query string (a Validated splits layer-level sources into overlapping parts that cannot be concatenated back), but these call sites still passed ggsql.validate(query). --- pkg-py/tests/test_ggsql.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index 54ff51ce7..eb4519695 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -105,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(ds, query) altair_widget = AltairWidget.from_ggsql(spec) result = altair_widget.widget.chart.to_dict() assert "$schema" in result @@ -117,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(ds, query) assert spec.metadata()["rows"] == 3 @pytest.mark.ggsql @@ -125,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(ds, query) assert "VISUALISE" in spec.visual() @pytest.mark.ggsql @@ -133,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(ds, query) assert spec.metadata()["rows"] == 3 assert "VISUALISE" in spec.visual() @@ -144,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(ds, query) altair_widget = AltairWidget.from_ggsql(spec) result = altair_widget.widget.chart.to_dict() assert "$schema" in result @@ -165,7 +165,7 @@ 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)) + spec = execute_ggsql(ds, query) assert "VISUALISE" in spec.visual() @pytest.mark.ggsql @@ -196,7 +196,7 @@ def test_supports_layer_level_from_sources(self): # 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(ds, ggsql.validate(query)) + spec = execute_ggsql(ds, query) assert spec is not None assert "VISUALISE" in spec.visual() @@ -242,7 +242,7 @@ 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(ds, query) assert spec.metadata()["rows"] == 4 assert "VISUALISE" in spec.visual() @@ -266,6 +266,6 @@ def counting(query): monkeypatch.setattr(ds, "execute_query", counting) query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - execute_ggsql(ds, ggsql.validate(query)) - execute_ggsql(ds, ggsql.validate(query)) + execute_ggsql(ds, query) + execute_ggsql(ds, query) assert calls == 1 From 314a80cffed0e9146e5321abc896758b0e51942b Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 16:08:54 -0500 Subject: [PATCH 08/24] deps: require ggsql>=0.5.2 for viz _viz_ggsql.py lazily imports ggsql.CachingReader at first viz call; older releases would fail with ImportError at runtime. CachingReader first ships in ggsql 0.5.2 (posit-dev/ggsql-python#10). --- pyproject.toml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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", From bb758ace513fdcc86f972daf89d68179bbee1fe3 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 16:12:57 -0500 Subject: [PATCH 09/24] refactor: own the CachingReader in the viz tool closure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The module-level WeakKeyDictionary never actually released readers: each value (CachingReader -> ExecutorReader) holds a strong reference to its key (the executor), so the weak key could never die and entries accumulated for the life of the process — the leak the weak keys were meant to prevent. The viz tool closure is created once per chat client and already has exactly the memo's desired lifetime, so the reader now lives there, created lazily on first viz call. execute_ggsql() takes the reader explicitly, making cache sharing visible at the call site; a new test guards against reintroducing an implicit shared registry. --- pkg-py/src/querychat/_viz_ggsql.py | 57 ++++++++++++++++++------------ pkg-py/src/querychat/_viz_tools.py | 13 +++++-- pkg-py/tests/test_ggsql.py | 44 +++++++++++++++++------ 3 files changed, 78 insertions(+), 36 deletions(-) diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index 14bf8da0a..45c05ac10 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -2,7 +2,6 @@ from __future__ import annotations -import weakref from typing import TYPE_CHECKING from ._utils import to_polars @@ -42,33 +41,44 @@ def unregister(self, name) -> None: # noqa: ARG002 pass -# One CachingReader per executor: the memo only pays off across viz calls -# within a chat session, so the reader (and its cache) must outlive any -# single execute_ggsql() call. Weak keys let finished sessions' executors — -# and their caches — be garbage collected. -_readers: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() +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. + + Parameters + ---------- + executor + The querychat QueryExecutor (or bare DataSource) to read through. + Returns + ------- + ggsql.CachingReader + A caching reader wrapping the executor. -def _caching_reader(executor: QueryExecutor) -> ggsql.CachingReader: + """ from ggsql import CachingReader - reader = _readers.get(executor) - if reader is None: - reader = CachingReader(ExecutorReader(executor)) - _readers[executor] = reader - return reader + return CachingReader(ExecutorReader(executor)) -def execute_ggsql(executor: QueryExecutor, query: str) -> ggsql.Spec: +def execute_ggsql(reader: ggsql.CachingReader, query: str) -> ggsql.Spec: """ - Execute a ggsql query against an executor, returning a Spec. + Execute a ggsql query through a caching reader, returning a Spec. - The whole query (SQL + VISUALISE) executes through a - ``ggsql.CachingReader`` wrapping the executor: base reads go to the - 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. + 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 @@ -77,8 +87,9 @@ def execute_ggsql(executor: QueryExecutor, query: str) -> ggsql.Spec: Parameters ---------- - executor - The querychat QueryExecutor (or bare DataSource) to read through. + 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). @@ -88,4 +99,4 @@ def execute_ggsql(executor: QueryExecutor, query: str) -> ggsql.Spec: The writer-independent plot specification. """ - return _caching_reader(executor).execute(query) + return reader.execute(query) diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index 9ec265813..c7c5a2906 100644 --- a/pkg-py/src/querychat/_viz_tools.py +++ b/pkg-py/src/querychat/_viz_tools.py @@ -22,7 +22,7 @@ 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 if TYPE_CHECKING: from collections.abc import Callable @@ -190,11 +190,18 @@ def visualize_impl( """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 + def visualize( ggsql: str, title: str, ) -> ContentToolResult: """Execute a ggsql query and render the visualization.""" + nonlocal reader markdown = f"```sql\n{ggsql}\n```" try: @@ -210,7 +217,9 @@ def visualize( "\n".join(error["message"] for error in validated.errors()) ) - spec = execute_ggsql(executor, ggsql) + if reader is None: + reader = caching_reader(executor) + spec = execute_ggsql(reader, ggsql) raw_chart = VegaLiteWriter().render_chart(spec) altair_widget = AltairWidget(copy.deepcopy(raw_chart)) diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index eb4519695..67b976199 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -6,7 +6,7 @@ import pytest from querychat._datasource import DataFrameSource from querychat._viz_altair_widget import AltairWidget -from querychat._viz_ggsql import execute_ggsql +from querychat._viz_ggsql import caching_reader, execute_ggsql class TestGgsqlValidate: @@ -105,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, 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 @@ -117,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, query) + spec = execute_ggsql(caching_reader(ds), query) assert spec.metadata()["rows"] == 3 @pytest.mark.ggsql @@ -125,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, query) + spec = execute_ggsql(caching_reader(ds), query) assert "VISUALISE" in spec.visual() @pytest.mark.ggsql @@ -133,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, query) + spec = execute_ggsql(caching_reader(ds), query) assert spec.metadata()["rows"] == 3 assert "VISUALISE" in spec.visual() @@ -144,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, 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 @@ -165,7 +165,7 @@ def test_supports_uppercase_column_references_without_renaming(self): "VISUALISE ROOM_TYPE AS x, listings AS y " "DRAW bar" ) - spec = execute_ggsql(ds, query) + spec = execute_ggsql(caching_reader(ds), query) assert "VISUALISE" in spec.visual() @pytest.mark.ggsql @@ -196,7 +196,7 @@ def test_supports_layer_level_from_sources(self): # 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(ds, query) + spec = execute_ggsql(caching_reader(ds), query) assert spec is not None assert "VISUALISE" in spec.visual() @@ -242,7 +242,7 @@ def test_supports_single_relation_raw_plus_summary_overlay(self): DRAW line MAPPING category AS color FILTER layer_type = 'summary' """ - spec = execute_ggsql(ds, query) + spec = execute_ggsql(caching_reader(ds), query) assert spec.metadata()["rows"] == 4 assert "VISUALISE" in spec.visual() @@ -266,6 +266,28 @@ def counting(query): monkeypatch.setattr(ds, "execute_query", counting) query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - execute_ggsql(ds, query) - execute_ggsql(ds, query) + 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 From 544530c45bee853d6b8aadb45bf73de0d5c0cd87 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 16:55:06 -0500 Subject: [PATCH 10/24] fix: restore viz widgets through a CachingReader restore_viz_widgets() passed the executor itself to execute_ggsql(), whose signature now takes a ggsql.CachingReader. The AttributeError was swallowed by the per-entry except, so every viz widget silently failed to restore from bookmarks. Create one reader per restore pass (restored widgets sharing a base read also reuse the memo) and add regression tests covering restore, memo sharing, and invalid-entry skipping. --- pkg-py/src/querychat/_shiny_module.py | 8 ++- pkg-py/tests/test_ggsql.py | 100 ++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 2 deletions(-) diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 212424efc..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: validate(ggsql_str) # early syntax check; errors skip the entry - spec = execute_ggsql(executor, ggsql_str) + 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/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index 67b976199..c10b5131a 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -291,3 +291,103 @@ def counting(query): 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"] From 4eac7eb5ca41a714a0d665b0126eae6aec6bf0b7 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 16:55:06 -0500 Subject: [PATCH 11/24] docs: viz cache behavior, staleness caveat, and GGSQL_CACHE_* env vars Memoized base reads assume the data doesn't change mid-session; with update queries enabled or an externally-written live database, charts can lag changes by up to the TTL. Document the caveat and the GGSQL_CACHE_TTL / GGSQL_CACHE_MAX_BYTES / GGSQL_CACHE_DISABLED knobs in the visualize guide and the caching_reader() docstring. --- pkg-py/docs/visualize.qmd | 16 ++++++++++++++++ pkg-py/src/querychat/_viz_ggsql.py | 14 ++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/pkg-py/docs/visualize.qmd b/pkg-py/docs/visualize.qmd index e8d720b68..1a708fb70 100644 --- a/pkg-py/docs/visualize.qmd +++ b/pkg-py/docs/visualize.qmd @@ -96,6 +96,22 @@ 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/). ## See also diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index 45c05ac10..fde86ee64 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -53,6 +53,20 @@ def caching_reader(executor: QueryExecutor) -> ggsql.CachingReader: 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 From 5d788896b4b219f7f851c031ec2372b256946119 Mon Sep 17 00:00:00 2001 From: Garrick Aden-Buie Date: Tue, 29 Sep 2026 17:35:44 -0400 Subject: [PATCH 12/24] r: Increment dev version 0.4.0.9000 --- pkg-r/DESCRIPTION | 2 +- pkg-r/NEWS.md | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pkg-r/DESCRIPTION b/pkg-r/DESCRIPTION index 4dba23c39..1d27069f0 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.0.9000 Authors@R: c( person("Garrick", "Aden-Buie", , "garrick@posit.co", role = c("aut", "cre"), comment = c(ORCID = "0000-0002-7111-0077")), diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 2759c6f42..4f27669ba 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -1,3 +1,5 @@ +# querychat (development version) + # querychat 0.4.0 ## New features From 36b6d0c293c37967b46b1cc319edd8a00862cf0c Mon Sep 17 00:00:00 2001 From: Garrick Aden-Buie Date: Tue, 29 Sep 2026 17:44:13 -0400 Subject: [PATCH 13/24] fix(pkg-r): skip tests that assume Suggests packages are installed (#317) * fix(pkg-r): skip tests that assume Suggests packages are installed CRAN's r-devel-linux-x86_64-fedora-clang flavor has RSQLite but not duckdb installed, which caused the test suite to ERROR: - 92 errors: tests requesting engine = "duckdb" (explicitly or via the local_data_frame_source() fixture default) hard-failed in check_installed("duckdb"), including one variant where the non-interactive check surfaced as "menu() cannot be used non-interactively" - 4 failures: tests using the default engine silently fell back to SQLite and then failed DuckDB-specific expectations skip_if_no_dataframe_engine() only skipped when *neither* engine was installed, so it never fired on that machine, and it had no way to express "this test requires a specific engine". It is now engine-aware: - with engine = NULL, keep the old any-engine semantics (mirroring how DataFrameSource resolves a default engine) - with a specific engine, skip unless that engine's package is installed, so tests no longer silently degrade to another engine local_data_frame_source() and local_recording_data_frame_source() now call it with their requested engine, and local_querychat() skips for plain data.frame inputs, so fixture-based tests skip automatically. Tests that construct DataFrameSource or QueryChat directly gained skip_if_no_dataframe_engine() guards (6 describe blocks and 3 test_that() blocks in test-QueryChat.R, 2 describe blocks in test-QueryChatSystemPrompt.R, plus 3 explicit skip_if_not_installed ("duckdb") calls in DuckDB-asserting tests). Audited all Suggests packages by running the full suite with every Suggests dependency (except the testthat/withr harness and later, which is a hard dependency of shiny) shadowed by broken stubs. That surfaced 46 additional tests failing on "no compatible database engine" with both engines absent, all fixed by the guards above. Verified by simulating the CRAN condition (FAIL 0, SKIP 98, PASS 716 for the duckdb-only case), a full Suggests blackout (FAIL 0, SKIP 199, PASS 1198), and with everything installed (FAIL 0, SKIP 0, PASS 2195, twice). * docs(pkg-r): add NEWS bullet for CRAN test skip fix --- pkg-r/NEWS.md | 4 ++ pkg-r/tests/testthat/helper-fixtures.R | 37 +++++++++++++++++-- pkg-r/tests/testthat/test-QueryChat.R | 17 +++++++++ .../testthat/test-QueryChatSystemPrompt.R | 10 +++++ 4 files changed, 64 insertions(+), 4 deletions(-) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 4f27669ba..520a87325 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -1,5 +1,9 @@ # querychat (development version) +## 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/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()) From 6ccec6bc42a995c1fd7061027906a63b89c593b6 Mon Sep 17 00:00:00 2001 From: Garrick Aden-Buie Date: Wed, 30 Sep 2026 12:31:03 -0400 Subject: [PATCH 14/24] r/v0.4.1 (#319) * v0.4.1 * docs(pkg-r): drop dev-mode bullet from release_bullets() * docs(pkg-r): update cran-comments for 0.4.1 resubmission --- pkg-r/DESCRIPTION | 2 +- pkg-r/NEWS.md | 2 +- pkg-r/R/querychat-package.R | 5 +---- pkg-r/cran-comments.md | 14 ++++++++++++++ 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/pkg-r/DESCRIPTION b/pkg-r/DESCRIPTION index 1d27069f0..407171252 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.9000 +Version: 0.4.1 Authors@R: c( person("Garrick", "Aden-Buie", , "garrick@posit.co", role = c("aut", "cre"), comment = c(ORCID = "0000-0002-7111-0077")), diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 520a87325..5f95c5b92 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -1,4 +1,4 @@ -# querychat (development version) +# querychat 0.4.1 ## Bug fixes 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. From b6bd3dbdee40994c00781bdf8cdb978cb16a5ae4 Mon Sep 17 00:00:00 2001 From: Garrick Aden-Buie Date: Wed, 30 Sep 2026 12:48:37 -0400 Subject: [PATCH 15/24] r: bump dev version to v0.4.1.9000 --- pkg-r/DESCRIPTION | 2 +- pkg-r/NEWS.md | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pkg-r/DESCRIPTION b/pkg-r/DESCRIPTION index 407171252..a147b9506 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.1 +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")), diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 5f95c5b92..0b2ceb326 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -1,3 +1,5 @@ +# querychat (development version) + # querychat 0.4.1 ## Bug fixes From c3f24b92497bf4071a21bcedd12e470259f5d60d Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 17:59:43 -0500 Subject: [PATCH 16/24] feat(prompts): add ggsql-spatial partial for maps capability --- pkg-py/src/querychat/prompts/ggsql-spatial.md | 36 +++++++++++++++++++ pkg-py/src/querychat/prompts/prompt.md | 6 ++++ 2 files changed, 42 insertions(+) create mode 100644 pkg-py/src/querychat/prompts/ggsql-spatial.md 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. From f4f87793b54dea171b52f8bcf2c1c8a6184973e3 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:00:45 -0500 Subject: [PATCH 17/24] feat(py): add default sentinel and maps modifier to tools --- pkg-py/src/querychat/_querychat_base.py | 14 +++++++- pkg-py/src/querychat/_viz_utils.py | 5 +++ pkg-py/tests/test_viz_tools.py | 46 +++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 0d2535b78..106527d45 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_]*$" @@ -826,6 +826,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 +847,7 @@ def normalize_tools( stacklevel=2, ) resolved.discard("visualize") + resolved.discard("maps") return resolved 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/tests/test_viz_tools.py b/pkg-py/tests/test_viz_tools.py index f5364be96..97ec4b5d4 100644 --- a/pkg-py/tests/test_viz_tools.py +++ b/pkg-py/tests/test_viz_tools.py @@ -58,6 +58,52 @@ 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 and "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="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 From a4bb7f73a19a4e56f7d6ac89fd1fa9355d7af401 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:01:19 -0500 Subject: [PATCH 18/24] feat(py): render ggsql-spatial partial when maps enabled --- pkg-py/src/querychat/_system_prompt.py | 1 + pkg-py/tests/test_viz_tools.py | 24 ++++++++++++++++++++++++ 2 files changed, 25 insertions(+) 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/tests/test_viz_tools.py b/pkg-py/tests/test_viz_tools.py index 97ec4b5d4..30f528a2c 100644 --- a/pkg-py/tests/test_viz_tools.py +++ b/pkg-py/tests/test_viz_tools.py @@ -146,6 +146,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( From d06507c5e7e5cdfc151387a22707b459ed5d9688 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:03:54 -0500 Subject: [PATCH 19/24] docs: maps capability --- pkg-py/CHANGELOG.md | 6 ++++++ pkg-py/docs/visualize.qmd | 14 ++++++++++++++ 2 files changed, 20 insertions(+) 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 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 From 143544163c371dec364b5e4a15efaf0218e346f3 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:18:45 -0500 Subject: [PATCH 20/24] feat(py): pre-warm spatial extension when maps enabled --- pkg-py/src/querychat/_querychat_base.py | 7 +++++- pkg-py/src/querychat/_viz_ggsql.py | 28 +++++++++++++++++++++ pkg-py/src/querychat/_viz_tools.py | 12 +++++++-- pkg-py/tests/test_ggsql.py | 29 ++++++++++++++++++++++ pkg-py/tests/test_viz_tools.py | 33 +++++++++++++++++++++++++ 5 files changed, 106 insertions(+), 3 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 106527d45..b541ba8de 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -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 diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index fde86ee64..03b42d8cd 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -2,6 +2,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING from ._utils import to_polars @@ -83,6 +84,33 @@ def caching_reader(executor: QueryExecutor) -> ggsql.CachingReader: return CachingReader(ExecutorReader(executor)) +SPATIAL_SMOKE_QUERY = ( + "INSTALL spatial; LOAD spatial; " + "VISUALISE FROM ggsql:world DRAW spatial FILTER iso_a3 = 'FRA'" +) + + +def prewarm_spatial(reader: ggsql.CachingReader) -> bool: + """ + Install and load the DuckDB spatial extension for the reader's cache. + + 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. + """ + 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 execute_ggsql(reader: ggsql.CachingReader, query: str) -> ggsql.Spec: """ Execute a ggsql query through a caching reader, returning a Spec. diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index c7c5a2906..d04dc1cee 100644 --- a/pkg-py/src/querychat/_viz_tools.py +++ b/pkg-py/src/querychat/_viz_tools.py @@ -22,7 +22,7 @@ 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 caching_reader, execute_ggsql +from ._viz_ggsql import caching_reader, execute_ggsql, prewarm_spatial if TYPE_CHECKING: from collections.abc import Callable @@ -62,6 +62,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 +75,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 +85,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 +190,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 VegaLiteWriter, validate @@ -219,6 +225,8 @@ def visualize( if reader is None: reader = caching_reader(executor) + if maps: + prewarm_spatial(reader) spec = execute_ggsql(reader, ggsql) raw_chart = VegaLiteWriter().render_chart(spec) diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index c10b5131a..3288422d5 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -391,3 +391,32 @@ 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 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 Exception( + "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 30f528a2c..8dd5e121b 100644 --- a/pkg-py/tests/test_viz_tools.py +++ b/pkg-py/tests/test_viz_tools.py @@ -446,3 +446,36 @@ 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 == [] From 84d21d4e9f704dbbc7112218f30f1f2564c35f70 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:19:15 -0500 Subject: [PATCH 21/24] feat(py): translate spatial errors when maps disabled --- pkg-py/src/querychat/_viz_ggsql.py | 33 ++++++++++++++++++++- pkg-py/src/querychat/_viz_tools.py | 24 +++++++++++---- pkg-py/tests/test_viz_tools.py | 47 ++++++++++++++++++++++++++++++ 3 files changed, 98 insertions(+), 6 deletions(-) diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index 03b42d8cd..dd5f3fd2c 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -3,7 +3,7 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal from ._utils import to_polars @@ -111,6 +111,37 @@ def prewarm_spatial(reader: ggsql.CachingReader) -> bool: return True +def translate_spatial_error( + error: Exception, *, maps_state: Literal["enabled", "disabled", "setup_failed"] +) -> str | None: + """ + Translate ggsql's spatial-extension failure into actionable guidance. + + 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. + """ + 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").' + ) + + def execute_ggsql(reader: ggsql.CachingReader, query: str) -> ggsql.Spec: """ Execute a ggsql query through a caching reader, returning a Spec. diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index d04dc1cee..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 caching_reader, execute_ggsql, prewarm_spatial +from ._viz_ggsql import ( + caching_reader, + execute_ggsql, + prewarm_spatial, + translate_spatial_error, +) if TYPE_CHECKING: from collections.abc import Callable @@ -201,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: @@ -225,8 +231,8 @@ def visualize( if reader is None: reader = caching_reader(executor) - if maps: - prewarm_spatial(reader) + if maps_available: + maps_available = prewarm_spatial(reader) spec = execute_ggsql(reader, ggsql) raw_chart = VegaLiteWriter().render_chart(spec) @@ -252,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/tests/test_viz_tools.py b/pkg-py/tests/test_viz_tools.py index 8dd5e121b..d44ab9084 100644 --- a/pkg-py/tests/test_viz_tools.py +++ b/pkg-py/tests/test_viz_tools.py @@ -479,3 +479,50 @@ def test_no_prewarm_when_maps_disabled(self, data_source, monkeypatch): 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 Exception(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) From b89789d67e40ca1962f92f5c2a439ed6db89bf89 Mon Sep 17 00:00:00 2001 From: cpsievert Date: Mon, 12 Oct 2026 23:21:49 +0000 Subject: [PATCH 22/24] `devtools::document()` (GitHub Actions) --- pkg-r/DESCRIPTION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg-r/DESCRIPTION b/pkg-r/DESCRIPTION index a147b9506..e5fe614f0 100644 --- a/pkg-r/DESCRIPTION +++ b/pkg-r/DESCRIPTION @@ -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 From 8adcb0696300df29e4e94583f879f08dde0ec338 Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:22:12 -0500 Subject: [PATCH 23/24] fix(py): widen create_model field spec annotations for pyright --- pkg-py/src/querychat/_handoff_prompt.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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."), ) From de7d660de45f2adce0d3aacddb44dd162cd24a1d Mon Sep 17 00:00:00 2001 From: Carson Date: Mon, 12 Oct 2026 18:29:58 -0500 Subject: [PATCH 24/24] fix(py): ruff lint cleanups (test style, drop stale noqa from #321) --- pkg-py/src/querychat/_viz_ggsql.py | 4 ++-- pkg-py/tests/test_ggsql.py | 5 +++-- pkg-py/tests/test_viz_tools.py | 7 ++++--- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index dd5f3fd2c..c7cbf1c58 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/pkg-py/src/querychat/_viz_ggsql.py @@ -34,11 +34,11 @@ 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: # noqa: ARG002 + def register(self, name, table, replace) -> None: # Writes belong to CachingReader's cache, never the primary. pass - def unregister(self, name) -> None: # noqa: ARG002 + def unregister(self, name) -> None: pass diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index 3288422d5..d11c3b02f 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -294,7 +294,8 @@ def counting(query): class TestRestoreVizWidgets: - """Bookmark restore executes through a CachingReader (PR #321 follow-up). + """ + 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 @@ -414,7 +415,7 @@ def test_failure_warns_and_returns_false(self): class FailReader: def execute(self, query): - raise Exception( + raise RuntimeError( "IO Error: Extension spatial.duckdb_extension not found" ) diff --git a/pkg-py/tests/test_viz_tools.py b/pkg-py/tests/test_viz_tools.py index d44ab9084..f84504267 100644 --- a/pkg-py/tests/test_viz_tools.py +++ b/pkg-py/tests/test_viz_tools.py @@ -76,13 +76,14 @@ def test_default_sentinel_alone(self): from querychat._querychat_base import DEFAULT_TOOLS, normalize_tools resolved = normalize_tools("default", default=DEFAULT_TOOLS, check_deps=False) - assert "visualize" in resolved and "query" in resolved + 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="maps.*visualize"): + 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): @@ -517,7 +518,7 @@ def test_tool_translates_spatial_error_when_maps_disabled( import querychat._viz_tools as viz_tools def boom(reader, query): - raise Exception(self.SPATIAL_ERR) + raise RuntimeError(str(self.SPATIAL_ERR)) monkeypatch.setattr(viz_tools, "execute_ggsql", boom) tool = tool_visualize(data_source, lambda _: None)