feat: Surface newer data-dict spec fields in prompts and schema views (#316 follow-up)

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 <source>`. Caching can be tuned or disabled with the `GGSQL_CACHE_TTL`, `GGSQL_CACHE_MAX_BYTES`, and `GGSQL_CACHE_DISABLED` environment variables. (#321)
 
 ### Bug fixes

pkg-py/docs/context.qmd

@@ -84,13 +84,18 @@ 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`
   - `constraints`: Optional list — barewords like `primary_key`/`foreign_key`, or assertion objects (`assert` with an optional `language` and `description`)
   - `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
 

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,13 +151,20 @@ 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.
 
     """
 
     name: str
     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]

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)
 

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] = []
 

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