aws-samples / aws-samples/sample-data-agent-on-duckdb

Add a semantic layer as a pipeline stage (metric registry + compiler)

Open
#1 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1
Forks
1
PR merge metrics
No merged PRs in 30d

Description

## Context

`docs/design.md` names a semantic layer as the next layer three times (:13, :434,
:469), and `tests/unit/test_pipeline.py:59` already pins the extension point: a
stage inserted at index 1 yields `gate → semantic → rewrite → execute` through
the real tool. This issue describes the layer that fills that slot.

The problem it addresses: the model currently writes every aggregate itself.
`docs/design.md` lists the figures the eval asserts (401,346.77 ETH inbound), but
those definitions live in the prompt and in the eval fixtures, not in code, so
nothing stops two sessions from computing "inbound ETH" two ways — one dividing
by 1e18, one not; one counting reverted traces, one not. A semantic layer moves
the definition into code and out of the model's hands, which is the posture the
governance layer already takes towards authorization.

The design below is implemented and green against the current `main`: 353 unit
tests, up from 303, and `ruff check` clean. I am opening this before the PR to
get feedback on the surface and the scope.

## The approach

The layer is not a new entry point. It is one more step in the statement
pipeline.

1. **The definition goes into code.** A registry holds one frozen entry per
number: which table it reads, how it aggregates, which rows count, which
dimensions it supports, the unit, the owner, the version.
2. **The model gets a call, not a second tool.** A metric is called inside
ordinary SQL — `metric('inbound_eth', …)` — so it can be joined, wrapped and
aggregated further.
3. **A stage does the translation.** It replaces the call with the SQL the
definition fixes, rewriting the statement in place.
4. **The stage sits after the read-only gate and before governance.** The
semantic layer expands and governance narrows, so expansion goes first and
governance applies its policy to the real tables and columns.
5. **The layers behind it need no change.** They receive ordinary SQL and never
learn that the semantic layer exists.
6. **Anything the registry does not cover is refused there.** An undefined
metric, dimension or window is rejected with the available names listed,
rather than approximated.
7. **The definition travels back with the number.** The result carries the scope
and version for the model to cite, and the audit line records which
definition produced the figure.

The rest of this issue is the detail behind those seven steps.

## The surface

A table-function-shaped call inside ordinary SQL, expanded by the stage before
governance sees the statement:

SELECT * FROM metric('inbound_eth', address => '0x...', time_window => 'last_7d')
SELECT * FROM metric('unblended_spend', start_date => '2025-01',
end_date => '2025-02', group_by => ['service'])

The stage replaces the call with an aliased subquery, so it composes with the
rest of the statement (JOIN, outer GROUP BY, CTAS into a session working set) and
with result shaping unchanged. Two calls in one statement each keep their alias.
SQL containing no call is returned byte-identical, not re-serialized.

The form is verified against this repo's dependencies.
`duckdb.extract_statements` parses it, `compute.gate.check` allows it, and
sqlglot parses it to `exp.Table → exp.Anonymous(this='metric')` with the named
arguments intact, so the stage reads its arguments off the AST rather than
parsing text. Both `=>` (`exp.Kwarg`) and `:=` (`exp.PropertyEQ`) are accepted,
since DuckDB takes either.

Argument names avoid DuckDB reserved words on purpose. `window => 'x'` and
`end => 'x'` both make DuckDB's parser raise a syntax error, and
`gate.single_statement` returns `True` when the parse fails, so such a statement
passes the gate and the error surfaces from the engine instead. The argument
names are `time_window`, `start_date`, `end_date`, `group_by`, `filters`.

## Where it plugs in

`STAGES` becomes `[Gate(), Semantic(), Govern(), Cost(), Execute()]`. The order
is forced, not chosen: the semantic layer **expands** (metric name → SQL over
base tables) and governance **narrows** (RLS predicate injection, CLS column
exclusion). Expanding after governance would let SQL no policy had seen reach the
engine. A test covers the consequence: a restricted principal's row filter lands
inside the metric's own subquery, so the metric sums only rows the caller may
read.

Governance therefore never sees the `metric(...)` node, and the table-function
allowlist in `governance/rewrite.py:226` needs no exception for it. When
expansion fails the stage rejects rather than passing the statement on.
Governance would refuse an unexpanded call anyway, which makes that the second
line of defence rather than the first.

The stage emits no trace event for a statement that calls no metric, following
`Cost`'s `if st.cost["sources"]:` pattern. Plain SQL still traces
`gate → rewrite → execute`, and the existing UI needs no change.

## Constraints the existing layers impose

These came out of reading `governance/rewrite.py` and `context/cost.py`, and they
shape the design more than the surface does.

1. **Partition predicates must be literal.** `context/cost.py:86` classifies
`date = (subquery)` as non-literal: the statement satisfies the required-filter
gate but gets no path expansion, which is a whole-prefix LIST by another name.
Relative windows are therefore resolved at compile time into
`BETWEEN 'a' AND 'b'`, a form `cost._literal_partitions` expands. Every metric
call then carries a literal, bounded partition predicate by construction
rather than by the model remembering to write one. Resolution is string
arithmetic over a configured anchor date, so it costs no engine round-trip.
2. **Table-level deny-by-default** (`rewrite.py:268-278`): every table a compiled
metric names must be in the caller's policy. The metric catalogue in the
system prompt is therefore filtered by `governance.visible_tables()` and
`governance.denied_columns()`, and by `storage.registry.enabled()`, the way
`context/datamap.py` already filters L0 and `search_schema`. A CLS-denied
column is dropped from the metric's dimension list rather than offered and
then refused.
3. **No local name may shadow a governed table** (`rewrite.py:251-257`): the
compiler emits no CTE or alias named after a governed table, and never the
reserved `_r_` prefix.
4. **Denied column names are matched by name, without scope**
(`rewrite.py:304-323`).
5. **Statement roots are an allowlist** (`rewrite.py:47`): compiled output is a
single `SELECT`.

## Metric definitions over the existing raw tables

The metrics compile against the tables the two scenarios already mount. There is
no pre-aggregated layer, no loader and no new data — `docs/design.md:581` is
explicit that the repository ships loaders and views and never data, and the
on-chain source is a public dataset. A metric entry therefore carries its own
`from` and any joins. The frozen definition is the point; whether the aggregation
was frozen at build time or at compile time is not.

Seven metrics, overlapping the eval's existing ground truth:

- on-chain (`eth_traces`, day granularity): `inbound_eth`, `outbound_eth`,
`counterparties`, `sanctioned_inflow` (joins `ofac_addresses`)
- cloud-ops: `unblended_spend` (`cur_line_items`, month), `failed_api_calls` and
`api_error_rate` (`cloudtrail_events`, day)

Each entry declares owner, definition text, unit, version, the dimensions it
supports and the rows in scope. An unknown metric, an unsupported dimension or an
undefined window is refused at compile time with the available names listed.
Windows are `last_7d`/`last_30d`/`last_90d` for day-partitioned metrics and
`last_1m`/`last_3m`/`last_6m` for month-partitioned ones; a window of the wrong
granularity is refused rather than converted.

Relative windows need an anchor. A metric may declare its own; otherwise
`SEMANTIC_ANCHOR` supplies it, and a metric with neither refuses a relative
window rather than resolving to today. On a static sample dataset "today" reads
empty, and an empty result looks like an answer.

One definition detail is worth review. The on-chain metrics restrict rows to
`status = 1 AND call_type = 'call'`, which I chose as the settled, value-moving
subset of traces. I have not validated that against the AWS Public Blockchain
Data schema, and it is the pair of constants the ETH figures depend on.

## What changes

- new `data_agent/semantic/` — `metrics.py` (registry) and `compile.py`
(compiler, SQL surface, prompt catalogue); depends on `storage` only, and reads
`governance` for catalogue filtering the way `context` does
- `pipeline.py`: one `Semantic` stage; the statement keeps both the submitted SQL
and the expanded SQL, and the compiled definitions are appended to the tool
result so the model can cite the scope of a figure it reports
- `context/prompt.py`: a generated metric catalogue section, filtered by role
- `audit.py`: a `semantic` field naming the metric and version behind each
figure, so an audit line records which version of a number the answer quoted;
`semantic` joins the documented `stage` values
- `tests/unit/test_pipeline.py`: the stage-list and phase-sequence assertions
gain the new stage, and the extension-point test's throwaway stage is renamed
so it does not collide with the real one
- `tests/unit/test_semantic.py`: 50 tests — compilation, refusals, the call
surface, the pipeline path, governance interaction, the audit line, catalogue
filtering
- `docs/design.md`, `README.md`, `.env.example`

## What does not change

`tools.py` and `ui/app.py` need no edit. `_phase_line` falls through to
`· {phase}` for an unrecognised phase (`ui/app.py:315`), so the trace renders as
is. A three-line case there would render the metric name and version instead, and
I can include that in the same PR.

## One finding from reading

`Govern.phase` is the string `"rewrite"`, but `docs/design.md` wrote `govern` in
the pipeline diagram (:36), in the stage-extension note (:470) and in the
streaming-trace section (:487); the tests use `rewrite`. The same three places
needed editing for the new stage, so the change corrects them. Taking that as a
separate PR instead is fine with me.

## Open decisions

1. **Call or views.** `metric(...)` keeps the dimension allowlist, the unit and
the rows in scope enforceable in code, and lets a refusal name what exists.
Governed views compose more naturally with ordinary SQL and need no new
surface, but the slicing then happens outside the definition, so a view cannot
refuse a dimension it does not support. I went with the call. Switching to
views changes the compiler's output, not the layer.
2. **Scope of the first PR.** The branch currently ships both scenarios' metrics.
Splitting it into on-chain first and cloud-ops second is straightforward.
3. **The anchor default.** `SEMANTIC_ANCHOR` is unset by default, so relative
windows are unavailable until an operator sets it, and explicit
`start_date`/`end_date` always work. The alternative is defaulting it to the
sample dataset's end date so `last_7d` works on a fresh deployment. I chose
the former because the dataset's end date is deployment-specific.

Contributor guide

Open the contributing guide

Research direction

Start with docs/design.md and the existing stage flow in pipeline.py, then inspect governance/rewrite.py and context/cost.py for the constraints described. Run tests/unit/test_pipeline.py and review tests/unit/test_semantic.py alongside the listed semantic, prompt, audit, and documentation changes. Done means the stage order, metric refusals, governance interaction, catalogue filtering, audit metadata, and documented decisions are agreed and covered by the tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, sql
Domain
ai, backend, data-engineering, databases
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.