diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 914a27b9..6f680289 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### New features +* Data dictionaries now support more of the data-dict spec: table-level assertion `constraints`, column `label`, `time_zone`, and `examples`, and relationship `aliases`/`conflicts`. Labels from value-to-label maps are shown to the LLM alongside the raw values, so it can interpret codes without writing them into SQL incorrectly. Columns marked `display: restricted` are hidden from the schema card shown to users (the LLM still sees them, marked as restricted). (#316 follow-up) + * **Faster iterative visualization, and support for layer-level data sources.** Visualization queries now execute as a whole through `ggsql.CachingReader`: the underlying SQL runs against your data source once and the result is memoized for the chat session, so re-rendering a chart with a different encoding no longer re-queries the source (a win for remote databases especially). Visualization layers can also draw from other relations in the query — e.g. raw points from the final `SELECT` with a summary line from a CTE — via `DRAW ... MAPPING ... FROM `. Caching can be tuned or disabled with the `GGSQL_CACHE_TTL`, `GGSQL_CACHE_MAX_BYTES`, and `GGSQL_CACHE_DISABLED` environment variables. (#321) ### Bug fixes diff --git a/pkg-py/docs/context.qmd b/pkg-py/docs/context.qmd index bd0603d3..ede9af52 100644 --- a/pkg-py/docs/context.qmd +++ b/pkg-py/docs/context.qmd @@ -84,6 +84,8 @@ glossary: Each entry under `tables` describes one table. Its `name` must match the table name you pass to `QueryChat` or `.add_table()`. - **`description`**: What this table represents (one sentence is usually enough). +- **`label`**: Optional display title for the table, included in the LLM's context. +- **`constraints`**: Optional table-level assertions (e.g. `- assert: end_date >= start_date`), conveyed to the LLM when it fetches the table's schema. - **`columns`**: A list of column annotations. Each column can have: - `name`: Column name (must match the actual column) - `type`: Semantic type hint — `string`, `number`, `number(id)`, `number(quantity)`, `date`, `enum` @@ -91,6 +93,9 @@ Each entry under `tables` describes one table. Its `name` must match the table n - `description`: What this column means in plain English - `values`: For `enum` columns, the possible values — a list, or a value-to-label map (e.g. `{M: Male, F: Female}`) - `range`: Optional `[min, max]` bounds. Either end may be left open with the `-.inf`/`.inf` sentinels (or `null`) — if you leave a bound open, describe the range in prose in the column's `description`. Replaces live min/max queries for the column + - `label` / `time_zone`: Optional display title, and time zone for `datetime` columns — both shown in the schema view + - `examples`: Optional sample values (up to 5 are shown) + - `display`: Set to `restricted` for columns that shouldn't appear in user-facing output. The LLM still sees the column (marked as restricted) and can query it, but it's hidden from the schema card shown to users Columns listed in the data dictionary are excluded from the auto-generated schema (since your description supersedes the auto-detected metadata). Columns not listed are still auto-detected as usual. @@ -101,6 +106,8 @@ The `relationships` section tells the LLM how to join tables. Each entry has: - `description`: A plain-English description of the relationship - `cardinality`: `one-to-one`, `one-to-many`, or `many-to-one` - `join`: The join condition (e.g., `orders.customer_id = customers.id`) +- `aliases`: For self-joins, a map of the alias each side joins under to the table it refers to (e.g. `{mother: otters, pup: otters}`) +- `conflicts`: Column names that appear on both sides of the join with different meanings #### Glossary diff --git a/pkg-py/src/querychat/_data_dict.py b/pkg-py/src/querychat/_data_dict.py index 90c0bcc4..7714fa4f 100644 --- a/pkg-py/src/querychat/_data_dict.py +++ b/pkg-py/src/querychat/_data_dict.py @@ -83,6 +83,13 @@ class ColumnSpec(BaseModel): open range. Replaces live min/max statistics queries when set. examples Representative sample values shown to the LLM as context. + label + Display title for the column. + time_zone + Time zone for ``datetime`` columns, shown in the schema view. + display + When ``"restricted"``, the column is flagged as not for user-facing + output. """ @@ -95,6 +102,9 @@ class ColumnSpec(BaseModel): values: list[Any] | dict[str, str] | None = None range: list[Any] | None = None examples: list[Any] | None = None + label: str | None = None + time_zone: str | None = None + display: Literal["restricted"] | None = None @field_validator("values", mode="before") @classmethod @@ -141,6 +151,11 @@ class TableSpec(BaseModel): columns Per-column specifications. Columns not listed here are documented using live statistics inferred from the data. + label + Display title for the table. + constraints + Table-level assertions conveyed to the LLM in the ``get_schema`` + response. """ @@ -148,6 +163,8 @@ class TableSpec(BaseModel): description: str | None = None details: str | None = None columns: list[ColumnSpec] = [] + label: str | None = None + constraints: list[str | ConstraintAssertion] = [] class RelationshipSpec(BaseModel): @@ -162,12 +179,20 @@ class RelationshipSpec(BaseModel): Cardinality string (e.g. ``"one-to-many"``). join SQL JOIN clause or expression that links the tables. + aliases + Alias-to-table map naming the role each join side plays (required + for self-joins). + conflicts + Column names that appear on both sides of the join with different + meanings. """ description: str | None = None cardinality: str | None = None join: str + aliases: dict[str, str] | None = None + conflicts: list[str] | None = None _LEGACY_TABLES_MESSAGE = """\ @@ -325,7 +350,17 @@ def to_prompt_dict(self) -> dict[str, Any]: result["description"] = self.description if self.tables: result["tables"] = { - t.name: ({"description": t.description} if t.description else None) + t.name: ( + { + k: v + for k, v in { + "description": t.description, + "label": t.label, + }.items() + if v is not None + } + or None + ) for t in self.tables } if self.relationships: @@ -356,24 +391,7 @@ def get_table_schema( for meta in metas: spec = documented.get(meta.name) if spec is not None: - if spec.type is not None: - meta.sql_type = spec.type - if spec.range is not None: - meta.min_val = spec.range[0] - meta.max_val = spec.range[1] - if spec.values is not None: - vals = ( - list(spec.values) - if isinstance(spec.values, dict) - else spec.values - ) - meta.categories = [str(v) for v in vals] - if spec.description is not None: - meta.description = spec.description - if spec.units is not None: - meta.units = spec.units - if spec.constraints: - meta.constraints = [str(c) for c in spec.constraints] + apply_column_spec(meta, spec) else: undocumented.append(meta) @@ -404,3 +422,32 @@ def from_yaml(cls, path: Path | str) -> DataDict: if dd.name is None: dd = dd.model_copy(update={"name": path.stem}) return dd + + +def apply_column_spec(meta: ColumnMeta, spec: ColumnSpec) -> None: + """Overlay a :class:`ColumnSpec` onto a live :class:`ColumnMeta`.""" + if spec.type is not None: + meta.sql_type = spec.type + if spec.range is not None: + meta.min_val = spec.range[0] + meta.max_val = spec.range[1] + if spec.values is not None: + if isinstance(spec.values, dict): + meta.categories = [str(k) for k in spec.values] + meta.value_labels = dict(spec.values) + else: + meta.categories = [str(v) for v in spec.values] + if spec.description is not None: + meta.description = spec.description + if spec.units is not None: + meta.units = spec.units + if spec.constraints: + meta.constraints = [str(c) for c in spec.constraints] + if spec.label is not None: + meta.label = spec.label + if spec.time_zone is not None: + meta.time_zone = spec.time_zone + if spec.display is not None: + meta.display = spec.display + if spec.examples is not None and spec.display != "restricted": + meta.examples = [str(e) for e in spec.examples] diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index 05fb716d..5e45ae94 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/pkg-py/src/querychat/_datasource.py @@ -60,6 +60,25 @@ class ColumnMeta: constraints: list[str] = field(default_factory=list) """Free-text constraints (e.g. 'non-negative').""" + examples: list[str] = field(default_factory=list) + """Sample values from the data dictionary (withheld when restricted).""" + + value_labels: dict[str, str] | None = None + """Optional value -> label map from map-form `values` in the data dictionary.""" + + display: str | None = None + """Set to 'restricted' for columns not intended for user-facing output.""" + + label: str | None = None + """Display title from the data dictionary.""" + + time_zone: str | None = None + """Time zone for datetime columns, from the data dictionary.""" + + +MAX_SCHEMA_EXAMPLES = 5 +"""Maximum number of data-dict example values shown per column in schema text.""" + def format_schema(table_name: str, columns: list[ColumnMeta]) -> str: """Format column metadata into schema string.""" @@ -71,6 +90,12 @@ def format_schema(table_name: str, columns: list[ColumnMeta]) -> str: header += f" [{col.units}]" lines.append(header) + if col.display == "restricted": + lines.append(" [restricted: do not display in user-facing output]") + if col.label: + lines.append(f" Label: {col.label}") + if col.time_zone: + lines.append(f" Time zone: {col.time_zone}") if col.description: lines.append(f" Description: {col.description}") if col.constraints: @@ -85,6 +110,17 @@ def format_schema(table_name: str, columns: list[ColumnMeta]) -> str: elif col.categories: cats = ", ".join(f"'{v}'" for v in col.categories) lines.append(f" Categorical values: {cats}") + if col.value_labels: + labels = ", ".join(f"{k}={v}" for k, v in col.value_labels.items()) + lines.append(f" Labels: {labels}") + if col.examples: + shown = col.examples[:MAX_SCHEMA_EXAMPLES] + more = ( + f" … ({len(col.examples) - MAX_SCHEMA_EXAMPLES} more)" + if len(col.examples) > MAX_SCHEMA_EXAMPLES + else "" + ) + lines.append(f" Examples: {', '.join(shown)}{more}") return "\n".join(lines) diff --git a/pkg-py/src/querychat/tools.py b/pkg-py/src/querychat/tools.py index 21de714a..7209ce3e 100644 --- a/pkg-py/src/querychat/tools.py +++ b/pkg-py/src/querychat/tools.py @@ -68,6 +68,11 @@ def get_schema(table_name: str) -> ContentToolResult: columns = executor.get_column_details(table_name, categorical_threshold) schema_text = format_schema(table_name, columns) + if dd is not None: + tspec = dd.get_table(table_name) + if tspec is not None and tspec.constraints: + rendered = "\n".join(f"- {c}" for c in tspec.constraints) + schema_text += f"\n\nTable constraints:\n{rendered}" return GetSchemaResult( value=schema_text, table_name=table_name, @@ -127,6 +132,9 @@ def tool_get_schema( def _schema_table(columns: list[ColumnMeta]) -> Tag: + # Restricted columns are LLM-visible (annotated in the schema text) but + # hidden from user-facing output per the data-dict spec. + columns = [c for c in columns if c.display != "restricted"] headers = ("Column", "Type", "Range / Values", "Description", "Constraints") rows: list[Tag] = [] diff --git a/pkg-py/tests/test_data_dict.py b/pkg-py/tests/test_data_dict.py index e8a827be..1ee3acb2 100644 --- a/pkg-py/tests/test_data_dict.py +++ b/pkg-py/tests/test_data_dict.py @@ -363,6 +363,121 @@ def test_to_prompt_dict_includes_glossary() -> None: assert d["glossary"]["ARR"] == "Annual Recurring Revenue" +def test_new_field_parsing(tmp_path: Path) -> None: + yaml_content = """\ +tables: + - name: orders + label: Orders + constraints: + - assert: end_date >= start_date + description: Valid date range + - primary_key + columns: + - name: status + type: enum + label: Order status + values: {W: Web, R: Retail} + - name: ssn + type: string + display: restricted + - name: created_at + type: datetime + time_zone: America/Chicago +relationships: + - join: mother.otter_no = pup.pup_number + cardinality: one-to-many + aliases: {mother: otters, pup: otters} + conflicts: [description] +""" + f = tmp_path / "spec.yaml" + f.write_text(yaml_content) + dd = DataDict.from_yaml(f) + + t = dd.tables[0] + assert t.label == "Orders" + assert len(t.constraints) == 2 + table_constraint = t.constraints[0] + assert not isinstance(table_constraint, str) + assert table_constraint.assert_ == "end_date >= start_date" + assert table_constraint.description == "Valid date range" + assert t.constraints[1] == "primary_key" + + status, ssn, created_at = t.columns + assert status.label == "Order status" + assert ssn.display == "restricted" + assert created_at.time_zone == "America/Chicago" + + rel = dd.relationships[0] + assert rel.aliases == {"mother": "otters", "pup": "otters"} + assert rel.conflicts == ["description"] + + +def test_get_table_schema_new_column_fields() -> None: + dd = DataDict( + tables=[ + TableSpec( + name="orders", + columns=[ + ColumnSpec( + name="status", + label="Order status", + values={"W": "Web", "R": "Retail"}, + examples=["W", "R"], + ), + ColumnSpec( + name="ssn", + display="restricted", + examples=["000-00-0000"], + ), + ColumnSpec(name="created_at", time_zone="America/Chicago"), + ], + ) + ], + ) + df = pl.DataFrame( + { + "status": ["W", "R"], + "ssn": ["x", "y"], + "created_at": [1, 2], + } + ) + executor = _make_executor(df, "orders") + cols = dd.get_table_schema("orders", executor, categorical_threshold=10) + col_map = {c.name: c for c in cols} + + status = col_map["status"] + assert status.label == "Order status" + assert status.categories == ["W", "R"] + assert status.value_labels == {"W": "Web", "R": "Retail"} + assert status.examples == ["W", "R"] + + ssn = col_map["ssn"] + assert ssn.display == "restricted" + assert ssn.examples == [] # withheld on restricted columns + + assert col_map["created_at"].time_zone == "America/Chicago" + + +def test_to_prompt_dict_new_pass_throughs() -> None: + from querychat._data_dict import RelationshipSpec + + dd = DataDict( + tables=[TableSpec(name="orders", description="Orders", label="Orders title")], + relationships=[ + RelationshipSpec( + join="mother.otter_no = pup.pup_number", + cardinality="one-to-many", + aliases={"mother": "otters", "pup": "otters"}, + conflicts=["description"], + ) + ], + ) + d = dd.to_prompt_dict() + assert d["tables"]["orders"]["label"] == "Orders title" + assert d["relationships"][0]["aliases"] == {"mother": "otters", "pup": "otters"} + assert d["relationships"][0]["conflicts"] == ["description"] + + def test_constraint_assertion_constructible_by_name_and_round_trips() -> None: from querychat._data_dict import ConstraintAssertion diff --git a/pkg-py/tests/test_datasource.py b/pkg-py/tests/test_datasource.py index 12b555a7..32584850 100644 --- a/pkg-py/tests/test_datasource.py +++ b/pkg-py/tests/test_datasource.py @@ -586,6 +586,39 @@ def test_sqlalchemy_cleanup_does_not_dispose_caller_owned_engine( assert len(source.get_data()) > 0 +def test_format_schema_new_annotations() -> None: + cols = [ + ColumnMeta( + name="status", + sql_type="STRING", + kind="text", + categories=["W", "R"], + value_labels={"W": "Web", "R": "Retail"}, + label="Order status", + examples=["W", "R", "W", "R", "W", "R", "W"], + ), + ColumnMeta( + name="ssn", + sql_type="STRING", + kind="text", + display="restricted", + ), + ColumnMeta( + name="created_at", + sql_type="TIMESTAMP", + kind="date", + time_zone="America/Chicago", + ), + ] + text = format_schema("orders", cols) + + assert " Label: Order status" in text + assert " Categorical values: 'W', 'R'" in text + assert " Labels: W=Web, R=Retail" in text + # examples capped at 5 + assert " Examples: W, R, W, R, W … (2 more)" in text + assert "[restricted: do not display in user-facing output]" in text + assert " Time zone: America/Chicago" in text def test_format_schema_open_ended_range() -> None: cols = [ ColumnMeta(name="a", sql_type="INTEGER", kind="numeric", min_val=0), diff --git a/pkg-py/tests/test_tools.py b/pkg-py/tests/test_tools.py index b1305dc2..6420e4e2 100644 --- a/pkg-py/tests/test_tools.py +++ b/pkg-py/tests/test_tools.py @@ -331,6 +331,35 @@ def test_get_schema_tool_uses_singular_column_preview() -> None: assert result.extra["display"].value_preview == "1 column" +def test_schema_card_hides_restricted_columns() -> None: + from querychat.tools import _schema_table + + cols = [ + ColumnMeta(name="id", sql_type="INTEGER", kind="numeric"), + ColumnMeta( + name="ssn", sql_type="STRING", kind="text", display="restricted" + ), + ] + html = str(_schema_table(cols)) + assert "id" in html + assert "ssn" not in html + + +def test_get_schema_includes_table_constraints() -> None: + df = pl.DataFrame({"amount": [1, 2]}) + executor, table_names = _make_executor_and_table(df, "orders") + dd = DataDict( + tables=[ + TableSpec( + name="orders", + constraints=[{"assert": "amount >= 0", "description": "Non-negative"}], + ) + ], + ) + tool = tool_get_schema([dd], executor, table_names, categorical_threshold=10) + result = tool.func(table_name="orders") + assert "Table constraints:" in result.value + assert "- amount >= 0 (Non-negative)" in result.value def test_schema_card_shows_open_ended_range() -> None: from querychat.tools import _schema_table diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 2b3dc65d..5ac2fbc0 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -4,6 +4,10 @@ * Data dictionaries now follow the current [data-dict spec](https://data-dict.tidyverse.org/spec.html): `tables` must be a list of tables each with a `name` field (previously a mapping of name to definition), column `range` is a two-element `[min, max]` list, and `values` may be a value-to-label map. Update older dictionary files accordingly; `data-dict validate-spec` (from the data-dict CLI) checks a file against the spec. Files in the older mapping format, or with a `$version` that is invalid or predates the first spec version, fail with a descriptive error; files declaring a newer `$version` warn at load time but are still read on a best-effort basis. (#316) +## New features + +* Data dictionaries now support more of the [data-dict spec](https://data-dict.tidyverse.org/spec.html): table-level assertion `constraints`, column `label`, `time_zone`, and `examples`, and relationship `aliases`/`conflicts`. Labels from value-to-label maps are shown to the LLM alongside the raw values, so it can interpret codes without writing them into SQL incorrectly. Columns marked `display: restricted` are hidden from the schema card shown to users (the LLM still sees them, marked as restricted). (#316 follow-up) + ## Bug fixes * JSON pins now fall back to `pins::pin_read()` when DuckDB's JSON extension cannot be installed or loaded, including on builds that disable extension loading. Native DuckDB JSON reading is retained when the extension is available. diff --git a/pkg-r/R/DBISource.R b/pkg-r/R/DBISource.R index f78fa538..bc778ef9 100644 --- a/pkg-r/R/DBISource.R +++ b/pkg-r/R/DBISource.R @@ -131,7 +131,8 @@ DBISource <- R6::R6Class( list( text = format_schema_from_details( as.character(DBI::dbQuoteIdentifier(private$.conn, self$table_name)), - details + details, + table_spec = table_spec ), columns = details ) @@ -232,6 +233,26 @@ DBISource <- R6::R6Class( ) ) +# Render spec `constraints` entries (barewords and assertion objects) to text. +format_constraint_entries <- function(constraints) { + vapply( + constraints, + function(c) { + if (is.list(c)) { + # Assertion object: {assert, language?, description?} + txt <- c[["assert"]] %||% "" + if (!is.null(c[["description"]])) { + txt <- paste0(txt, " (", c[["description"]], ")") + } + txt + } else { + as.character(c) + } + }, + character(1) + ) +} + # Returns a list of named lists (one per column) with all schema metadata. # Used by both format_schema_from_details() (LLM text) and JSON serialization (UI). build_column_details_impl <- function( @@ -273,7 +294,9 @@ build_column_details_impl <- function( c("integer", "numeric", "double", "Date", "POSIXct", "POSIXt") ) { numeric_columns <- c(numeric_columns, col) - if (is.null(spec[["range"]])) { + # Spec `range` and `values` both replace live min/max stats, so skip + # the aggregates (a wasted scan on remote tables) when either is set. + if (is.null(spec[["range"]]) && is.null(spec[["values"]])) { select_parts <- c( select_parts, paste0( @@ -371,19 +394,9 @@ build_column_details_impl <- function( min_val <- NA_character_ max_val <- NA_character_ categories <- list() + value_labels <- NULL constraints <- if (length(spec[["constraints"]]) > 0) { - lapply(spec[["constraints"]], function(c) { - if (is.list(c)) { - # Assertion object: {assert, language?, description?} - txt <- c[["assert"]] %||% "" - if (!is.null(c[["description"]])) { - txt <- paste0(txt, " (", c[["description"]], ")") - } - txt - } else { - as.character(c) - } - }) + as.list(format_constraint_entries(spec[["constraints"]])) } else { list() } @@ -404,7 +417,9 @@ build_column_details_impl <- function( } min_val <- as.character(lo %||% "?") max_val <- as.character(hi %||% "?") - } else { + # Spec `values` replace live stats for this column (handled below), + # so only query stats when neither `range` nor `values` is given. + } else if (is.null(spec[["values"]])) { min_key <- paste0(col, "__min") max_key <- paste0(col, "__max") if ( @@ -418,18 +433,21 @@ build_column_details_impl <- function( max_val <- as.character(column_stats[[max_key]]) } } - } else if (col %in% text_columns) { - if (!is.null(spec[["values"]])) { - vals <- spec[["values"]] - # Map form (value -> label): the values are the names - if (!is.null(names(vals))) { - vals <- as.list(names(vals)) - } - categories <- as.list(as.character(vals)) - } else if (col %in% names(categorical_values)) { - vals <- categorical_values[[col]] - if (length(vals) > 0) categories <- as.list(as.character(vals)) + } + # Spec `values` replace categorical inference for any column kind (a + # numeric code can have labeled values just like a text column can). + if (!is.null(spec[["values"]])) { + vals <- spec[["values"]] + # Map form (value -> label): values are the names, labels the values + if (!is.null(names(vals))) { + value_labels <- as.list(as.character(vals)) + names(value_labels) <- names(vals) + vals <- as.list(names(vals)) } + categories <- as.list(as.character(vals)) + } else if (col %in% text_columns && col %in% names(categorical_values)) { + vals <- categorical_values[[col]] + if (length(vals) > 0) categories <- as.list(as.character(vals)) } list( @@ -437,6 +455,18 @@ build_column_details_impl <- function( sql_type = sql_type, units = spec[["units"]] %||% NA_character_, description = spec[["description"]] %||% NA_character_, + label = spec[["label"]] %||% NA_character_, + time_zone = spec[["time_zone"]] %||% NA_character_, + display = spec[["display"]] %||% NA_character_, + examples = if ( + !is.null(spec[["examples"]]) && + !identical(spec[["display"]], "restricted") + ) { + as.list(as.character(spec[["examples"]])) + } else { + list() + }, + value_labels = value_labels, min_val = min_val, max_val = max_val, categories = categories, @@ -445,14 +475,32 @@ build_column_details_impl <- function( }) } -# Format a list of column details (from build_column_details_impl) into LLM schema text. -format_schema_from_details <- function(table_name_display, details) { +# Format a list of column details (from build_column_details_impl) into LLM +# schema text. When `table_spec` is given, table-level constraints are appended. +format_schema_from_details <- function( + table_name_display, + details, + table_spec = NULL +) { schema_lines <- c(paste("Table:", table_name_display), "Columns:") for (col in details) { line <- paste0("- ", col$name, " (", col$sql_type, ")") if (!is.na(col$units)) { line <- paste0(line, " [", col$units, "]") } + if (!is.na(col$display) && identical(col$display, "restricted")) { + line <- paste( + line, + " [restricted: do not display in user-facing output]", + sep = "\n" + ) + } + if (!is.na(col$label)) { + line <- paste(line, paste0(" Label: ", col$label), sep = "\n") + } + if (!is.na(col$time_zone)) { + line <- paste(line, paste0(" Time zone: ", col$time_zone), sep = "\n") + } if (!is.na(col$description)) { line <- paste( line, @@ -481,9 +529,41 @@ format_schema_from_details <- function(table_name_display, details) { sep = "\n" ) } + if (length(col$value_labels) > 0) { + labels_str <- paste0( + names(col$value_labels), + "=", + as.character(col$value_labels), + collapse = ", " + ) + line <- paste(line, paste0(" Labels: ", labels_str), sep = "\n") + } + if (length(col$examples) > 0) { + shown <- col$examples[seq_len(min(5L, length(col$examples)))] + more <- if (length(col$examples) > 5) { + # (\u2026 is an ellipsis; R sources must stay ASCII for R CMD check) + paste0(" \u2026 (", length(col$examples) - 5, " more)") + } else { + "" + } + line <- paste( + line, + paste0(" Examples: ", paste(shown, collapse = ", "), more), + sep = "\n" + ) + } schema_lines <- c(schema_lines, line) } - paste(schema_lines, collapse = "\n") + text <- paste(schema_lines, collapse = "\n") + if (length(table_spec[["constraints"]]) > 0) { + rendered <- format_constraint_entries(table_spec[["constraints"]]) + text <- paste0( + text, + "\n\nTable constraints:\n", + paste0("- ", rendered, collapse = "\n") + ) + } + text } get_schema_impl <- function( @@ -505,7 +585,8 @@ get_schema_impl <- function( ) format_schema_from_details( as.character(DBI::dbQuoteIdentifier(conn, table_name)), - details + details, + table_spec = table_spec ) } diff --git a/pkg-r/R/DataDict.R b/pkg-r/R/DataDict.R index ea047b5e..f4497132 100644 --- a/pkg-r/R/DataDict.R +++ b/pkg-r/R/DataDict.R @@ -147,11 +147,10 @@ data_dict_to_prompt_list <- function(dd) { tables <- dd_tables_by_name(dd) if (length(tables) > 0) { result[["tables"]] <- lapply(tables, function(ts) { - if (is.null(ts[["description"]])) { - NULL - } else { - list(description = ts[["description"]]) - } + entry <- compact( + list(description = ts[["description"]], label = ts[["label"]]) + ) + if (length(entry) == 0) NULL else entry }) } if (length(dd[["relationships"]]) > 0) { @@ -160,7 +159,9 @@ data_dict_to_prompt_list <- function(dd) { list( join = rs[["join"]], description = rs[["description"]], - cardinality = rs[["cardinality"]] + cardinality = rs[["cardinality"]], + aliases = rs[["aliases"]], + conflicts = rs[["conflicts"]] ) ) }) diff --git a/pkg-r/R/QueryExecutor.R b/pkg-r/R/QueryExecutor.R index 375257ee..cac616a5 100644 --- a/pkg-r/R/QueryExecutor.R +++ b/pkg-r/R/QueryExecutor.R @@ -155,7 +155,8 @@ DuckDBExecutor <- R6::R6Class( list( text = format_schema_from_details( as.character(DBI::dbQuoteIdentifier(private$conn, table_name)), - details + details, + table_spec = table_spec ), columns = details ) diff --git a/pkg-r/R/TblSqlSource.R b/pkg-r/R/TblSqlSource.R index 63846bf5..94e741be 100644 --- a/pkg-r/R/TblSqlSource.R +++ b/pkg-r/R/TblSqlSource.R @@ -130,7 +130,8 @@ TblSqlSource <- R6::R6Class( list( text = format_schema_from_details( as.character(DBI::dbQuoteIdentifier(private$.conn, self$table_name)), - details + details, + table_spec = table_spec ), columns = details ) diff --git a/pkg-r/R/querychat_tools.R b/pkg-r/R/querychat_tools.R index b7e0b075..497b7c3e 100644 --- a/pkg-r/R/querychat_tools.R +++ b/pkg-r/R/querychat_tools.R @@ -358,6 +358,12 @@ querychat_tool_result <- function( } schema_table <- function(columns) { + # Restricted columns are LLM-visible (annotated in the schema text) but + # hidden from user-facing output per the data-dict spec. + columns <- Filter( + function(column) !identical(column$display, "restricted"), + columns + ) headers <- c( "Column", "Type", diff --git a/pkg-r/tests/testthat/test-DataDict.R b/pkg-r/tests/testthat/test-DataDict.R index 5d62162b..0a2e3b95 100644 --- a/pkg-r/tests/testthat/test-DataDict.R +++ b/pkg-r/tests/testthat/test-DataDict.R @@ -226,6 +226,27 @@ describe("data_dict_to_prompt_list()", { expect_false("glossary" %in% names(result)) }) + it("passes through table label and relationship aliases/conflicts", { + dd <- list( + tables = list( + list(name = "orders", description = "Orders", label = "Orders title") + ), + relationships = list( + list( + join = "mother.otter_no = pup.pup_number", + cardinality = "one-to-many", + aliases = list(mother = "otters", pup = "otters"), + conflicts = list("description") + ) + ) + ) + result <- data_dict_to_prompt_list(dd) + expect_equal(result[["tables"]][["orders"]][["label"]], "Orders title") + rel <- result[["relationships"]][[1]] + expect_equal(rel[["aliases"]], list(mother = "otters", pup = "otters")) + expect_equal(rel[["conflicts"]], list("description")) + }) + it("includes table entry as NULL when table has no description", { dd <- list(tables = list(list(name = "no_desc"))) result <- data_dict_to_prompt_list(dd) diff --git a/pkg-r/tests/testthat/test-querychat_tools.R b/pkg-r/tests/testthat/test-querychat_tools.R index a27ab510..51961efe 100644 --- a/pkg-r/tests/testthat/test-querychat_tools.R +++ b/pkg-r/tests/testthat/test-querychat_tools.R @@ -517,6 +517,75 @@ describe("tool_update_dashboard_impl()", { describe("tool_get_schema()", { skip_if_no_dataframe_engine() + it("renders new dict fields and hides restricted columns from the card", { + df <- data.frame( + id = 1:5, + category = c("A", "B", "A", "B", "A"), + code = c(1L, 2L, 1L, 2L, 1L), + ssn = c("x", "y", "z", "w", "v"), + stringsAsFactors = FALSE + ) + df_source <- local_data_frame_source(df) + executor <- local_executor(df_source) + table_spec <- list( + name = "test_table", + constraints = list( + list(assert = "id >= 0", description = "Non-negative id") + ), + columns = list( + list( + name = "id", + description = "Primary key", + label = "Identifier", + time_zone = "UTC", + constraints = c("positive", "required"), + range = list(1, 5), + examples = as.list(1:7) + ), + list( + name = "category", + values = list(A = "Alpha", B = "Beta") + ), + list( + name = "code", + values = list("1" = "Low", "2" = "High") + ), + list( + name = "ssn", + display = "restricted", + examples = list("000-00-0000") + ) + ) + ) + data_dicts <- list(list(tables = list(table_spec))) + schema_result <- executor$get_schema_result( + "test_table", + 20, + table_spec = table_spec + ) + text <- schema_result$text + + expect_match(text, "Label: Identifier", fixed = TRUE) + expect_match(text, "Time zone: UTC", fixed = TRUE) + expect_match(text, "Categorical values: 'A', 'B'", fixed = TRUE) + expect_match(text, "Labels: A=Alpha, B=Beta", fixed = TRUE) + # Spec values apply to numeric-coded columns too (not just text) + expect_match(text, "Categorical values: '1', '2'", fixed = TRUE) + expect_match(text, "Labels: 1=Low, 2=High", fixed = TRUE) + expect_match(text, "… (2 more)", fixed = TRUE) + expect_match( + text, + "[restricted: do not display in user-facing output]", + fixed = TRUE + ) + expect_match(text, "Table constraints:", fixed = TRUE) + expect_match(text, "- id >= 0 (Non-negative id)", fixed = TRUE) + + # Restricted column is hidden from the user-facing card + card_html <- as.character(schema_table(schema_result$columns)) + expect_false(grepl("ssn", card_html)) + }) + it("renders spec range, value keys, and assertion constraints", { df <- data.frame( id = 1:5, @@ -552,9 +621,10 @@ describe("tool_get_schema()", { # Open-ended ranges keep their known bound expect_match(text, "Range: 1 to ?", fixed = TRUE) expect_match(text, "Range: ? to 5", fixed = TRUE) - # Map-form values render the values (keys), not the labels + # Map-form values keep the values (keys) as categories; labels are + # rendered separately as a `Labels:` annotation expect_match(text, "Categorical values: 'A', 'B'", fixed = TRUE) - expect_false(grepl("Alpha", text, fixed = TRUE)) + expect_match(text, "Labels: A=Alpha, B=Beta", fixed = TRUE) # Assertion objects render as 'assert (description)' alongside barewords expect_match(text, "grp <> '' (Non-empty)", fixed = TRUE) expect_match(text, "required", fixed = TRUE) diff --git a/pkg-r/vignettes/context.Rmd b/pkg-r/vignettes/context.Rmd index 1f4237e8..bac08208 100644 --- a/pkg-r/vignettes/context.Rmd +++ b/pkg-r/vignettes/context.Rmd @@ -100,6 +100,8 @@ glossary: Each entry under `tables` describes one table. Its `name` must match the table name you pass to `QueryChat$new()` or `$add_table()`. - **`description`**: What this table represents (one sentence is usually enough). +- **`label`**: Optional display title for the table, included in the LLM's context. +- **`constraints`**: Optional table-level assertions (e.g. `- assert: end_date >= start_date`), conveyed to the LLM when it fetches the table's schema. - **`columns`**: A list of column annotations. Each column can have: - `name`: Column name (must match the actual column) - `type`: Semantic type hint — `string`, `number`, `number(id)`, `number(quantity)`, `date`, `enum` @@ -107,6 +109,9 @@ Each entry under `tables` describes one table. Its `name` must match the table n - `description`: What this column means in plain English - `values`: For `enum` columns, the possible values — a list, or a value-to-label map (e.g. `{M: Male, F: Female}`) - `range`: Optional `[min, max]` bounds. Either end may be left open with the `-.inf`/`.inf` sentinels (or `null`) — if you leave a bound open, describe the range in prose in the column's `description`. Replaces live min/max queries for the column + - `label` / `time_zone`: Optional display title, and time zone for `datetime` columns — both shown in the schema view + - `examples`: Optional sample values (up to 5 are shown) + - `display`: Set to `restricted` for columns that shouldn't appear in user-facing output. The LLM still sees the column (marked as restricted) and can query it, but it's hidden from the schema card shown to users Columns listed in the data dictionary are excluded from the auto-generated schema (since your description supersedes the auto-detected metadata). Columns not listed are still auto-detected as usual. @@ -117,6 +122,8 @@ The `relationships` section tells the LLM how to join tables. Each entry has: - `description`: A plain-English description of the relationship - `cardinality`: `one-to-one`, `one-to-many`, or `many-to-one` - `join`: The join condition (e.g., `orders.customer_id = customers.id`) +- `aliases`: For self-joins, a map of the alias each side joins under to the table it refers to (e.g. `{mother: otters, pup: otters}`) +- `conflicts`: Column names that appear on both sides of the join with different meanings #### Glossary