querychat multi-table standalone demo: Brilliant Retail (fine-jewelry) — diamonds, settings, customers, orders, order_items

app.py

# /// script
# dependencies = [
#     "querychat[polars,viz]",
#     "anthropic",
# ]
# ///
"""Multi-table querychat demo: a synthetic fine-jewelry retailer.

Five related tables -- diamonds, settings, customers, orders, and
order_items -- are registered with a single QueryChat instance so the LLM
can query any one of them or write joins across them (e.g. an "Engagement
Ring" order item joins to both its diamond and its setting).

Data is generated once and cached to ./data/*.parquet; later runs read the
cache instead of regenerating (see generate_data.py).

Usage:
    uv run app.py

Requires ANTHROPIC_API_KEY to be set (change client= below to use a
different provider).
"""

from pathlib import Path

from generate_data import load_tables
from querychat import QueryChat

GREETING = """\
## Brilliant Retail: Jewelry Explorer 💎

Welcome! You're chatting with a synthetic fine-jewelry retailer's data --
**diamonds**, **settings** (ring mountings), **customers**, **orders**, and
**order_items**, all linked together.

**Get started**

<ul>
<li><span class="suggestion">Show me a few example rows from each table</span></li>
<li><span class="suggestion">What's the average price of a round diamond over 1 carat?</span></li>
</ul>

**Explore across tables**

<ul>
<li><span class="suggestion">Which setting style is most often paired with princess-cut diamonds in engagement rings?</span></li>
<li><span class="suggestion">What's the average order value by acquisition channel?</span></li>
<li><span class="suggestion">Show total revenue by month over the last year</span></li>
</ul>
"""

DATA_DESCRIPTION = """\
Brilliant Retail is an online fine-jewelry retailer. This dataset covers its
diamond and setting catalog alongside customer order history.
"""

# Works around a ggsql bug (https://github.com/posit-dev/ggsql/issues/<TBD>):
# a date/timestamp column produced by an expression like DATE_TRUNC(...) fails
# ggsql's `SCALE ... VIA date` coercion unless it's explicitly cast to DATE.
EXTRA_INSTRUCTIONS = """\
When a `querychat_visualize` query buckets a date or timestamp column (e.g.
via DATE_TRUNC), always CAST the resulting expression to DATE in the SQL, for
example:

    SELECT CAST(DATE_TRUNC('month', order_date) AS DATE) AS month, ...

A column left as a raw TIMESTAMP will fail when referenced by
`SCALE ... VIA date` with a "Cannot coerce column ... to date" error.
"""

tables = load_tables(Path(__file__).parent / "data")

qc = QueryChat(
    tables["diamonds"],
    "diamonds",
    greeting=GREETING,
    data_description=DATA_DESCRIPTION,
    data_dict=Path(__file__).parent / "data-dict.yaml",
    extra_instructions=EXTRA_INSTRUCTIONS,
    client="anthropic",
)
qc.add_table(tables["settings"], "settings")
qc.add_table(tables["customers"], "customers")
qc.add_table(tables["orders"], "orders")
qc.add_table(tables["order_items"], "order_items")

app = qc.app()

if __name__ == "__main__":
    app.run()

data-dict.yaml

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.

generate_data.py

# /// script
# dependencies = ["polars"]
# ///
"""Generate (or load cached) synthetic data for a fine-jewelry retailer demo.

Five related tables are produced with a fixed random seed so the data is
identical across runs: diamonds, settings (ring mountings), customers,
orders, and order_items.

Run directly to force-regenerate the cache:
    uv run generate_data.py
"""

from __future__ import annotations

import random
import shutil
from datetime import date, timedelta
from pathlib import Path

import polars as pl

SEED = 42
TODAY = date(2026, 9, 1)
TABLE_NAMES = ["diamonds", "settings", "customers", "orders", "order_items"]


def load_tables(data_dir: Path) -> dict[str, pl.DataFrame]:
    """Return the 5 demo tables, generating and caching them on first call."""
    paths = {name: data_dir / f"{name}.parquet" for name in TABLE_NAMES}

    if all(path.exists() for path in paths.values()):
        return {name: pl.read_parquet(path) for name, path in paths.items()}

    tables = generate_tables()
    data_dir.mkdir(parents=True, exist_ok=True)
    for name, table in tables.items():
        table.write_parquet(paths[name])
    return tables


def generate_tables() -> dict[str, pl.DataFrame]:
    rng = random.Random(SEED)
    diamonds = generate_diamonds(rng)
    settings = generate_settings(rng)
    customers = generate_customers(rng)
    orders = generate_orders(rng, customers)
    order_items = generate_order_items(rng, orders, diamonds, settings)
    return {
        "diamonds": diamonds,
        "settings": settings,
        "customers": customers,
        "orders": orders,
        "order_items": order_items,
    }


def generate_diamonds(rng: random.Random, n: int = 400) -> pl.DataFrame:
    shapes = [
        "Round", "Princess", "Cushion", "Oval", "Emerald",
        "Radiant", "Asscher", "Marquise", "Pear", "Heart",
    ]
    shape_weights = [30, 12, 10, 12, 8, 8, 4, 4, 6, 6]
    cuts = ["Super Ideal", "Ideal", "Very Good", "Good"]
    cut_weights = [10, 40, 35, 15]
    cut_multiplier = {"Super Ideal": 1.15, "Ideal": 1.05, "Very Good": 0.95, "Good": 0.85}
    colors = ["D", "E", "F", "G", "H", "I", "J"]
    color_multiplier = {"D": 1.2, "E": 1.15, "F": 1.1, "G": 1.0, "H": 0.9, "I": 0.8, "J": 0.7}
    clarities = ["FL", "IF", "VVS1", "VVS2", "VS1", "VS2", "SI1", "SI2"]
    clarity_multiplier = {
        "FL": 1.3, "IF": 1.25, "VVS1": 1.15, "VVS2": 1.1,
        "VS1": 1.0, "VS2": 0.95, "SI1": 0.85, "SI2": 0.75,
    }
    certifiers = ["GIA", "IGI", "AGS"]
    certifier_weights = [60, 30, 10]

    rows = []
    for i in range(1, n + 1):
        carat = round(rng.uniform(0.3, 3.5), 2)
        cut = rng.choices(cuts, weights=cut_weights)[0]
        color = rng.choice(colors)
        clarity = rng.choice(clarities)
        price = round(
            3000
            * (carat**1.8)
            * cut_multiplier[cut]
            * color_multiplier[color]
            * clarity_multiplier[clarity],
            -1,
        )
        rows.append(
            {
                "id": i,
                "shape": rng.choices(shapes, weights=shape_weights)[0],
                "carat": carat,
                "cut": cut,
                "color": color,
                "clarity": clarity,
                "certified_by": rng.choices(certifiers, weights=certifier_weights)[0],
                "price": price,
            }
        )
    return pl.DataFrame(rows)


def generate_settings(rng: random.Random, n: int = 50) -> pl.DataFrame:
    styles = ["Solitaire", "Halo", "Pave", "Three-Stone", "Vintage", "Bezel"]
    style_multiplier = {
        "Solitaire": 1.0, "Halo": 1.4, "Pave": 1.3,
        "Three-Stone": 1.5, "Vintage": 1.6, "Bezel": 1.1,
    }
    metal_price = {
        "Platinum": 1400, "18K White Gold": 900, "18K Yellow Gold": 850,
        "18K Rose Gold": 850, "Palladium": 1000,
    }
    metals = list(metal_price)

    rows = []
    for i in range(1, n + 1):
        style = rng.choice(styles)
        metal = rng.choice(metals)
        base_price = round(metal_price[metal] * style_multiplier[style] * rng.uniform(0.9, 1.1), -1)
        rows.append(
            {
                "id": i,
                "style": style,
                "metal": metal,
                "band_width_mm": round(rng.uniform(1.5, 3.5), 1),
                "base_price": base_price,
            }
        )
    return pl.DataFrame(rows)


def generate_customers(rng: random.Random, n: int = 200) -> pl.DataFrame:
    first_names = [
        "Olivia", "Liam", "Emma", "Noah", "Ava", "Ethan", "Sophia", "Mason",
        "Isabella", "James", "Mia", "Benjamin", "Charlotte", "Lucas", "Amelia",
        "Henry", "Harper", "Alexander", "Evelyn", "Michael",
    ]
    last_names = [
        "Smith", "Johnson", "Williams", "Brown", "Jones", "Garcia", "Miller",
        "Davis", "Rodriguez", "Martinez", "Chen", "Kim", "Patel", "Nguyen",
        "Anderson", "Taylor", "Thomas", "Moore", "Jackson", "Martin",
    ]
    states = [
        "CA", "NY", "TX", "FL", "WA", "IL", "MA", "CO", "GA", "NC",
        "VA", "PA", "OH", "AZ", "OR", "NJ", "MD", "MN", "CT", "NV",
    ]
    channels = [
        "Organic Search", "Paid Social", "Referral", "Email", "Direct", "Showroom Walk-in",
    ]
    start = date(2022, 1, 1)
    span_days = (TODAY - start).days

    rows = []
    for i in range(1, n + 1):
        rows.append(
            {
                "id": i,
                "name": f"{rng.choice(first_names)} {rng.choice(last_names)}",
                "state": rng.choice(states),
                "join_date": start + timedelta(days=rng.randint(0, span_days)),
                "acquisition_channel": rng.choice(channels),
            }
        )
    return pl.DataFrame(rows)


def generate_orders(rng: random.Random, customers: pl.DataFrame, n: int = 350) -> pl.DataFrame:
    statuses = ["Delivered", "Shipped", "Processing", "Returned"]
    status_weights = [70, 15, 10, 5]
    channels = ["Online", "Showroom"]
    join_date_by_customer = dict(zip(customers["id"], customers["join_date"]))
    customer_ids = customers["id"].to_list()

    rows = []
    for i in range(1, n + 1):
        customer_id = rng.choice(customer_ids)
        join_date = join_date_by_customer[customer_id]
        max_offset = max((TODAY - join_date).days, 1)
        order_date = join_date + timedelta(days=rng.randint(0, max_offset))
        rows.append(
            {
                "id": i,
                "customer_id": customer_id,
                "order_date": order_date,
                "status": rng.choices(statuses, weights=status_weights)[0],
                "channel": rng.choice(channels),
            }
        )
    return pl.DataFrame(rows)


def generate_order_items(
    rng: random.Random,
    orders: pl.DataFrame,
    diamonds: pl.DataFrame,
    settings: pl.DataFrame,
    n: int = 450,
) -> pl.DataFrame:
    item_types = ["Engagement Ring", "Loose Diamond", "Wedding Band", "Earrings", "Necklace"]
    item_type_weights = [35, 15, 20, 15, 15]
    order_ids = orders["id"].to_list()
    diamond_ids = diamonds["id"].to_list()
    setting_ids = settings["id"].to_list()
    price_by_diamond = dict(zip(diamonds["id"], diamonds["price"]))
    price_by_setting = dict(zip(settings["id"], settings["base_price"]))

    rows = []
    for i in range(1, n + 1):
        item_type = rng.choices(item_types, weights=item_type_weights)[0]
        diamond_id: int | None = None
        setting_id: int | None = None

        if item_type == "Engagement Ring":
            diamond_id = rng.choice(diamond_ids)
            setting_id = rng.choice(setting_ids)
            price = price_by_diamond[diamond_id] + price_by_setting[setting_id]
        elif item_type == "Loose Diamond":
            diamond_id = rng.choice(diamond_ids)
            price = price_by_diamond[diamond_id]
        elif item_type == "Wedding Band":
            setting_id = rng.choice(setting_ids)
            price = round(price_by_setting[setting_id] * rng.uniform(0.7, 1.0), -1)
        elif item_type == "Earrings":
            price = round(rng.uniform(400, 4000), -1)
        else:  # Necklace
            price = round(rng.uniform(500, 6000), -1)

        rows.append(
            {
                "id": i,
                "order_id": rng.choice(order_ids),
                "item_type": item_type,
                "diamond_id": diamond_id,
                "setting_id": setting_id,
                "price": price,
            }
        )
    return pl.DataFrame(rows)


if __name__ == "__main__":
    target_dir = Path(__file__).parent / "data"
    if target_dir.exists():
        shutil.rmtree(target_dir)
    load_tables(target_dir)
    print(f"Wrote {', '.join(TABLE_NAMES)} to {target_dir}")