From 67b5fe14cbefbe592ffbd0080c81db58ee0de288 Mon Sep 17 00:00:00 2001 From: Ming-Jer Lee Date: Tue, 11 Aug 2026 08:10:37 -0700 Subject: [PATCH 1/2] docs: add design specs for graph utilization gaps in description generation and text2sql --- ...ESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md | 306 ++++++++++++++++ plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md | 329 ++++++++++++++++++ 2 files changed, 635 insertions(+) create mode 100644 plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md create mode 100644 plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md diff --git a/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md b/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md new file mode 100644 index 0000000..cbdc201 --- /dev/null +++ b/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md @@ -0,0 +1,306 @@ +# Design: Description Generation — Graph Utilization Gaps + +## Overview + +`generate_all_descriptions()` is architecturally sound: it walks computed columns in +topological order and builds each prompt from the column's direct lineage sources, so +descriptions cascade through the graph. This document specifies fixes for four gaps +that keep it from fully using the graph: + +| ID | Gap | Severity | +|----|-----|----------| +| D1 | Source-table columns are never described (cold-start at the graph boundary) | High | +| D2 | Incoming edges found by O(E) linear scan instead of the adjacency index | Medium (perf) | +| D3 | Rule-based fallback descriptions are indistinguishable from LLM output and cascade downstream | High | +| D4 | Queries without a destination table are skipped, unlike metadata propagation | Medium | + +Each gap is specified independently; D1 and D3 share a change to +`build_description_prompt()` and should land together or in sequence D3 → D1. + +--- + +## D1: Source-Column Cold Start + +### Problem Statement + +`MetadataManager.generate_all_descriptions()` only processes columns that satisfy +`col.is_computed()` on a query's `destination_table` +(`metadata_manager.py:90-100`). Columns of raw source tables (e.g. `raw.users.email`) +are never candidates. `build_description_prompt()` (`column.py:152-160`) only includes +source columns *that already have a description*, so for the first computed layer the +"Source columns" section is silently empty — the cascade starts from nothing exactly +where context matters most. + +### Current Behavior + +```python +pipeline = Pipeline(queries, dialect="bigquery") +pipeline.llm = llm +pipeline.generate_all_descriptions() +# raw.users.email -> no description, never processed +# staging.users.email_norm -> prompt contains only: column name, table, SQL expression +``` + +### Design + +Describe source columns using **forward usage context** — the graph knows how every +source column is consumed downstream even when nothing upstream exists. + +1. **New prompt builder** in `column.py`: + + ```python + def build_source_description_prompt(column: ColumnNode, pipeline: "Pipeline") -> str: + ``` + + Contents (all values through `sanitize_for_prompt` / `sanitize_sql_for_prompt`): + - `Column:` / `Table:` lines as today. + - `Sibling columns:` up to 15 other column names in the same table (domain signal). + - `Used downstream as:` up to 5 entries derived from outgoing edges + (`pipeline._get_outgoing_edges(column.full_name)`), each formatted as + ` = `. + - Same instruction block as `build_description_prompt` (one sentence, ≤15 words, + no SQL jargon). + +2. **New opt-in parameter** on `MetadataManager.generate_all_descriptions()` and the + `Pipeline` wrapper: + + ```python + def generate_all_descriptions(..., include_sources: bool = False) + ``` + + When `True`, source-table columns (columns of tables where + `table_graph.tables[t].is_source`) are processed **before** the topological walk of + computed columns, so their descriptions feed the first computed layer in the same + run. Existing skip rules apply unchanged: authored (`DescriptionSource.SOURCE`) + descriptions are never overwritten unless `overwrite=True`. + + Opt-in default rationale: the library is released (PyPI); flipping the default + changes LLM call volume and export diffs for existing users. Revisit the default in + a minor release. + +3. `generate_description()` dispatches on layer: source columns (no incoming edges, + `is_computed()` is False) use `build_source_description_prompt`; computed columns + keep `build_description_prompt`. + +### Test Cases + +- Two-layer fixture (`raw.users` → `staging.users` → `mart.user_stats`): + - `include_sources=True` produces descriptions on `raw.users.*` with + `description_source == GENERATED`. + - The prompt built for `staging.users.email_norm` afterwards contains the + `raw.users.email` description (assert via `build_description_prompt`). +- `include_sources=False` (default): behavior identical to today (regression test). +- Source column with an authored SQL-comment description is skipped unless + `overwrite=True`. +- Source column consumed by zero described targets still yields a valid prompt + (name + siblings only). + +### Risks and Mitigations + +- **Cost**: more LLM calls. Mitigated by opt-in flag; log the source-column count in + the existing progress logging. +- **Hallucination on thin context**: a bare column name invites guessing. The prompt + instruction already demands ≤15 words; usage context constrains the model further. + Output still passes `_validate_description_output`. + +--- + +## D2: O(V·E) Incoming-Edge Scans + +### Problem Statement + +Three functions in `column.py` find incoming edges with a full scan of every edge in +the pipeline: + +- `build_description_prompt` — `column.py:152` +- `propagate_metadata_backward` — `column.py:225` +- `propagate_metadata` — `column.py:268` + +```python +incoming_edges = [e for e in pipeline.edges if e.to_node == column] +``` + +`PipelineLineageGraph._incoming_index` (`column.py:315`, populated at `column.py:327` +keyed by `edge.to_node.full_name`) exists precisely for this lookup, and +`Pipeline._get_incoming_edges(full_name)` (`pipeline.py:193`) already wraps it. +Bulk operations (`generate_all_descriptions`, `propagate_all_metadata`) therefore run +in O(V·E) instead of O(V+E). + +### Design + +1. Replace all three scans with: + + ```python + incoming_edges = pipeline._get_incoming_edges(column.full_name) + ``` + +2. Promote the accessor to a public method `Pipeline.get_incoming_edges(full_name)` + (keep the underscore alias for backward compatibility), since `column.py` module + functions are public API and should not depend on a private method. + +3. **Precondition to verify in implementation**: `_incoming_index` is keyed by + `to_node.full_name`; confirm every caller passes `column.full_name` for a node that + is registered in `column_graph` (an unregistered node returns `[]` where the scan + would also return `[]` — equivalent, but assert this in tests, not assumptions). + +### Test Cases + +- Equivalence: for every column in an existing multi-query fixture pipeline, + `pipeline.get_incoming_edges(col.full_name)` equals the linear-scan result + (set equality on `(from_node.full_name, to_node.full_name)`). +- Existing `generate_all_descriptions` / `propagate_all_metadata` test suites pass + unchanged (behavior-preserving refactor). +- Optional benchmark note: synthetic pipeline with ~1k columns / ~5k edges shows the + bulk description pass no longer scales with E per column. + +### Risks and Mitigations + +- **Index staleness** if edges are appended without going through the graph's + `add_edge` path. Audit write paths for `column_graph.edges` during implementation; + the equivalence test above catches divergence. + +--- + +## D3: Fallback Descriptions Are Indistinguishable and Cascade + +### Problem Statement + +When the LLM call fails or validation rejects its output, +`_generate_fallback_description` (`column.py:187-193`) writes a humanized column name +("customer_ltv" → "Customer Ltv") and stamps it +`description_source = DescriptionSource.GENERATED` — the same value a real model +description gets. Consequences: + +1. Nothing persisted distinguishes placeholder text from model output; the + `return False` signal (added in PR #75) is transient. +2. `generate_all_descriptions` filters on `not col.description` + (`metadata_manager.py:97`), so a later re-run — e.g. after fixing LLM + configuration — **skips** every fallback column instead of retrying it. +3. Downstream prompts include the fallback text as source context + (`column.py:156-160`), feeding noise into the cascade. + +### Design + +1. **New enum member** in `models.py`: + + ```python + class DescriptionSource(Enum): + SOURCE = "source" + GENERATED = "generated" + PROPAGATED = "propagated" + FALLBACK = "fallback" # rule-based placeholder, no model involved + ``` + + `_generate_fallback_description` stamps `FALLBACK`. + +2. **Retry semantics**: the candidate filter in + `MetadataManager.generate_all_descriptions` treats `FALLBACK` as "no description": + + ```python + needs_description = ( + not col.description or col.description_source == DescriptionSource.FALLBACK + ) + ``` + + A re-run with a working LLM upgrades placeholders to real descriptions. This is an + intentional behavior change; call it out in the changelog. + +3. **Prompt hygiene** in `build_description_prompt`: list **all** direct source + columns by `full_name` (fixes the silent drop of undescribed sources), and append + `: ` only when a description exists and its source is not `FALLBACK`: + + ``` + Source columns: + - raw.users.email: User's primary email address + - raw.users.signup_ts + ``` + +4. **Serialization**: exporters emit `description_source` via `.value`; the new + `"fallback"` string flows through JSON/CSV export automatically. Verify the diff + tooling treats it as an ordinary value; if any deserializer whitelists enum + values, add `"fallback"`. + +### Test Cases + +- LLM raises → column gets fallback text with `description_source == FALLBACK`; + function returns `False`. +- Re-running `generate_all_descriptions` with a working LLM overwrites `FALLBACK` + columns (and still skips `SOURCE`/`GENERATED` unless `overwrite=True`). +- `build_description_prompt` for a column whose sources are (a) described, + (b) fallback-described, (c) undescribed lists all three by name and attaches text + only to (a). +- JSON export round-trip preserves `"fallback"`. + +### Risks and Mitigations + +- **Behavior change** (retry semantics): users relying on fallback text persisting + across runs will see regeneration attempts. Changelog + the `on_error="raise"` path + already exists for users who prefer hard failures. + +--- + +## D4: Queries Without a Destination Table Are Skipped + +### Problem Statement + +The two bulk operations disagree about which columns belong to a query: + +- `generate_all_descriptions` — only `query.destination_table` + (`metadata_manager.py:92-100`); terminal plain `SELECT` queries contribute nothing. +- `propagate_all_metadata` — `query.destination_table or f"{query_id}_result"` + (`metadata_manager.py:156`). + +Computed columns of a pipeline-final `SELECT` therefore receive propagated metadata +but never receive generated descriptions. + +### Design + +1. Extract one shared helper on `MetadataManager` (or module level): + + ```python + def _target_table(query: ParsedQuery) -> str: + return query.destination_table or f"{query.query_id}_result" + ``` + +2. Use it in both `generate_all_descriptions` and `propagate_all_metadata` so the + candidate column sets are computed identically (description generation keeps its + additional `needs_description` and `is_computed()` filters). + +### Test Cases + +- Pipeline ending in a plain `SELECT`: its computed output columns receive + descriptions. +- Parity test: the set of `(table, column)` pairs visited by description generation + equals the `is_computed()` subset of pairs visited by metadata propagation. +- Regression: pipelines where every query has a destination table are unaffected. + +### Risks and Mitigations + +- **Cost**: additional columns processed. Marginal — only terminal SELECTs; noted in + changelog. + +--- + +## Implementation Phases + +1. **Phase 1 — D2 (index lookups)**: behavior-preserving, unblocks perf; smallest + review surface. +2. **Phase 2 — D3 (FALLBACK source + prompt hygiene)**: changes `models.py`, + `column.py`, `metadata_manager.py`; includes the "list all sources" prompt change + D1 builds on. +3. **Phase 3 — D4 (target-table helper)**: small, isolated. +4. **Phase 4 — D1 (source-column descriptions)**: new prompt builder + `include_sources` + flag; depends on Phase 2's prompt shape. + +Each phase: `make pre-commit` (ruff check/format + tests) green; separate PRs +following `feat:`/`fix:` conventions. + +## Success Metrics + +- First-layer computed columns' prompts contain source context when + `include_sources=True` (D1). +- Bulk description generation does no per-column full-edge scans (D2). +- Zero fallback strings served as source context in any prompt; re-runs retry + fallback columns (D3). +- Description coverage includes terminal SELECT outputs; candidate sets of the two + bulk operations agree (D4). diff --git a/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md b/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md new file mode 100644 index 0000000..95ded06 --- /dev/null +++ b/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md @@ -0,0 +1,329 @@ +# Design: Text2SQL — Deeper Graph Context + +## Overview + +`GenerateSQLTool` (`tools/sql.py`) currently uses the graph mostly as a schema +catalog: tables, columns, descriptions, PII flags, and table-level "derived from" +lines. The deep graph assets — transitive column traversal, join-key knowledge, +final-table identification — are either gated behind the non-default `two_stage` +strategy or not surfaced at all. This document specifies five gaps: + +| ID | Gap | Severity | +|----|-----|----------| +| T1 | Column lineage absent from `direct` mode prompts | Medium | +| T2 | `expand_with_lineage` is 1-hop, not transitive | Medium | +| T3 | Join keys never surfaced (only "derived from" lines) | High | +| T4 | Final-table identification unused; intermediates presented as equally queryable | Medium | +| T5 | Non-LLM table-selection fallback ignores the graph | Low | + +T1–T3 all modify how the `relationship_section` prompt slot is assembled; T3 is the +highest-value change for generated-SQL correctness (join quality). + +### Shared change: graph-context assembly point + +Both `_generate_direct` (`sql.py:182`) and `_generate_two_stage` (`sql.py:238`) +currently fill `relationship_section` with a single builder call. Refactor both to a +shared assembly: + +```python +def _build_graph_context(self, builder: ContextBuilder, tables: List[str]) -> str: + parts = [ + builder.build_relationship_context(tables), # existing, T4 adjusts labels + builder.build_lineage_context(tables), # T1 wires into direct mode + builder.build_join_context(tables), # T3, new + ] + return "\n\n".join(p for p in parts if p) +``` + +The existing `GENERATE_SQL_PROMPT` / `GENERATE_SQL_WITH_EXPLANATION_PROMPT` templates +keep their `{relationship_section}` slot; no template signature change. Sanitization +stays where it is today (each section passed through `sanitize_for_prompt` at the +`prompt.format(...)` call sites). + +--- + +## T1: Column Lineage in Direct Mode + +### Problem Statement + +`_generate_direct` fills `relationship_section` with table-level +`build_relationship_context()` only (`sql.py:194`). `build_lineage_context()` — which +performs real transitive traversal via `trace_column_backward` (`context.py:451-483`) +— is only invoked in `_generate_two_stage` (`sql.py:255`). The default strategy never +shows the model how output columns relate to source columns, which is exactly the +information needed to pick the right table for a metric. + +### Design + +1. Use the shared `_build_graph_context` above in **both** strategies, so `direct` + mode gains the `## Column Lineage` section. +2. Make the existing hard-coded caps configurable on `ContextConfig`: + + ```python + max_lineage_columns_per_table: int = 10 # today: literal 10 (context.py:470) + max_lineage_lines: int = 20 # today: literal 20 (context.py:483) + ``` + +3. `include_lineage=False` in `ContextConfig` disables the section (already the + behavior of `build_lineage_context`; unchanged). + +### Test Cases + +- Direct-mode prompt for a two-layer fixture contains `## Column Lineage` with + `mart.x <- raw.y` lines. +- Caps respected: a table with 30 output columns contributes at most + `max_lineage_columns_per_table` lines; total lines ≤ `max_lineage_lines`. +- `include_lineage=False` removes the section from both strategies. + +### Risks and Mitigations + +- **Prompt growth** on wide pipelines — bounded by the two caps; both configurable. + +--- + +## T2: Transitive Lineage Expansion + +### Problem Statement + +`expand_with_lineage` (`context.py:274-303`) adds only the *direct* parent tables of +each selected table. In `two_stage` mode, a question whose answer requires a +grandparent table (e.g. join key only present two levels up) produces context that +omits it, and the generated SQL references a table the model never saw or invents +joins. + +### Current Behavior + +```python +# a -> b -> c (c selected) +builder.expand_with_lineage(["c"]) # ["c", "b"] — never "a" +``` + +### Design + +1. BFS to a configurable depth with a hard size cap: + + ```python + def expand_with_lineage(self, tables: List[str], depth: Optional[int] = None) -> List[str]: + ``` + + - `depth=None` reads `ContextConfig.lineage_expansion_depth` (new field, + default `2`). + - Traversal: frontier = selected tables; each round maps table → + `table_graph.tables[t].created_by` → `queries[qid].source_tables`; stop at + `depth` rounds or fixpoint. Maintain a visited set (the table graph is a DAG, + but self-referencing tables exist — `ParsedQuery.self_referenced_tables` — so + guard anyway). + - **Priority on truncation**: if the expanded set exceeds + `ContextConfig.max_tables`, keep shallower tables first (original selection, + then depth-1 parents, then depth-2, ...). + +2. Backward compatibility: `depth=1` reproduces today's behavior; the default of `2` + is a deliberate improvement documented in the changelog. + +### Test Cases + +- Chain `a -> b -> c`, select `["c"]`: depth 1 → `{b, c}`; depth 2 → `{a, b, c}`; + `None` → config default. +- Self-referencing query (`INSERT INTO t SELECT ... FROM t`) terminates. +- Truncation keeps the originally selected tables and nearest ancestors. + +### Risks and Mitigations + +- **Context bloat** in deep DAGs — depth default of 2 plus `max_tables` cap with + shallow-first priority. + +--- + +## T3: Join-Key Surfacing + +### Problem Statement + +The only relationship the model sees is `"- X is derived from Y"` +(`context.py:447`). Which columns actually join two tables — the single most +error-prone part of text2sql — is left for the model to guess from column names, +even though the pipeline holds this knowledge in two forms: observed join predicates +in the parsed SQL, and shared upstream ancestry in the column graph. + +### Design + +Two mechanisms, one output section. + +#### 1. Observed joins (from query ASTs) + +`ParsedQuery.ast` retains the sqlglot AST (`models.py:901`). New builder method: + +```python +def get_observed_joins(self, tables: Optional[List[str]] = None) -> List[Dict[str, Any]]: + # [{left_table, left_column, right_table, right_column, query_id}, ...] +``` + +Implementation sketch: +- For each `ParsedQuery`, walk `ast.find_all(exp.Join)`. +- From each join's `on` condition, collect top-level `exp.EQ` nodes whose both sides + are `exp.Column`; treat multiple EQs under one `AND` as a composite key (emit one + entry per column pair, same `query_id`). +- Resolve table aliases to real table names using the FROM/JOIN `exp.Table` nodes of + that query; consult `ParsedQuery.self_ref_aliases` for self-referencing tables. + **Phase 1 scope**: skip predicates whose alias resolves to a CTE-internal name + rather than a pipeline table; log at debug level. (CTE mapping via + `query_lineage` is a follow-up.) +- `USING (col)` joins (`exp.Join` args) emit `left.col = right.col`. +- Non-equi joins and `ON` conditions that are not column-to-column equality are + skipped. + +#### 2. Candidate joins (from shared lineage sources) + +For table pairs never joined in the pipeline (e.g. two marts), infer candidates: + +- For each output column of each context table, call + `pipeline.trace_column_backward(table, column)` once and record + `ultimate_source_full_name -> [(table, column), ...]`. +- Any source mapped to columns in ≥2 distinct context tables yields candidate pairs. +- Deduplicate against observed joins; cap total candidates + (`ContextConfig.max_join_hints: int = 15`, shared with observed joins, + observed first). + +#### 3. Prompt section + +```python +def build_join_context(self, tables: List[str]) -> str: +``` + +``` +## Join Hints + +- orders.customer_id = customers.id (observed in query_3) +- candidate: mart_ltv.user_id = mart_churn.user_id (both derive from raw.users.id) +``` + +Wired into `_build_graph_context` (see Overview) for both strategies. Candidates are +explicitly labeled `candidate:` so the model can weigh them below observed joins. + +### Test Cases + +- Fixture with `JOIN ... ON o.customer_id = c.id` and aliases → observed join with + real table names and `query_id`. +- Composite key (`ON a.x = b.x AND a.y = b.y`) → two entries, same query. +- `USING (id)` → resolved to both tables. +- Two marts sharing `raw.users.id` ancestry, never joined directly → one + `candidate:` line; no candidates between unrelated tables. +- Non-equi join (`ON a.ts > b.ts`) produces nothing. +- Cap: hints truncated to `max_join_hints`, observed joins prioritized. + +### Risks and Mitigations + +- **Alias/CTE resolution complexity** — phase 1 restricts to directly resolvable + aliases and logs skips; correctness over coverage. +- **False-positive candidates** (shared source ≠ valid join) — labeled `candidate:`, + deduplicated, capped; observed joins always listed first. +- **O(cols × trace) cost** — one backward trace per context column, bounded by + `max_lineage_columns_per_table`; traces are already the cost profile of + `build_lineage_context`. + +--- + +## T4: Final-Table Steering + +### Problem Statement + +`table_graph.get_final_tables()` (`table.py:304`, `read_by == []`) is never consulted +by `ContextBuilder`. `_format_table_context` annotates only `(Source table)` +(`context.py:398`); intermediates and marts look identical, so the model may answer +from a staging table when a mart exists — or from an intermediate that a later query +overwrites semantically. + +### Design + +1. **Role annotation** in `_format_table_context`: + - `(Source table)` — unchanged. + - `(Final table)` — `len(table_node.read_by) == 0` and not a source. + - `(Intermediate table)` — everything else. + Controlled by `ContextConfig.annotate_table_roles: bool = True`. + +2. **Prompt instruction**: add one line to the `## Instructions` block of both SQL + templates (via the existing `extra_instructions` slot assembly, applied in both + strategies): + + ``` + - Prefer final tables when they answer the question; use intermediate tables only when required + ``` + +3. **Truncation priority** in `build_schema_context` (`context.py:356-361`): when + trimming to `max_tables`, keep order final > intermediate > source (today it is + derived > source with no final/intermediate distinction). + +### Test Cases + +- Fixture `raw -> staging -> mart`: context labels `raw` source, `staging` + intermediate, `mart` final; instruction line present. +- `annotate_table_roles=False` restores current output (regression). +- Truncation with `max_tables=2` keeps `mart` and `staging`, drops `raw`. + +### Risks and Mitigations + +- **Fragment pipelines**: analyzing a subset of a real pipeline makes mid-DAG tables + look final. Acceptable — the annotation reflects the graph as loaded; document in + the tool docstring. + +--- + +## T5: Graph-Aware Selection Fallback + +### Problem Statement + +When the LLM table-selection call fails, `select_tables_by_keywords` +(`context.py:489-544`) scores tables purely lexically, and its `min_tables` padding +appends arbitrary tables in dict-insertion order. A question matching a mart by name +can miss the parent table required for its join, while padding adds noise tables. + +### Design + +1. **Score diffusion (1 hop)**: after lexical scoring, each table with score > 0 + contributes `0.5 × score` to its direct graph neighbors (parents via + `created_by`/`source_tables`, children via `read_by`). Single pass, applied on the + pre-diffusion scores (no iteration/feedback). +2. **Graph-aware padding**: to reach `min_tables`, prefer (in order) unselected + graph neighbors of already-selected tables, then final tables, then the current + arbitrary order. Deterministic tie-break: alphabetical. +3. Diffusion factor as a module constant (`_NEIGHBOR_SCORE_FACTOR = 0.5`); not + config-exposed until there is evidence tuning matters. + +### Test Cases + +- Mart matches keywords, its parent (needed for the join, zero lexical score) is + selected via diffusion before an unrelated lexically-weak table. +- Padding prefers neighbors of selected tables over unrelated tables; result order + deterministic across runs. +- Pure-lexical results unchanged when the graph has a single table. + +### Risks and Mitigations + +- Minimal — fallback path only. Diffusion is one pass and cannot cycle. + +--- + +## Implementation Phases + +1. **Phase 1 — T1 (lineage in direct mode) + shared `_build_graph_context`**: + creates the assembly point T3 plugs into; config caps added. +2. **Phase 2 — T4 (role annotation + instruction + truncation priority)**: isolated + in `context.py` + one instruction line in `sql.py`. +3. **Phase 3 — T2 (transitive expansion)**: `expand_with_lineage(depth)` + + `lineage_expansion_depth` config. +4. **Phase 4 — T3 (join hints)**: largest change; observed joins first, candidate + inference second (can ship as two PRs). +5. **Phase 5 — T5 (fallback diffusion)**: independent, lowest priority. + +Each phase: `make pre-commit` green; separate PRs (`feat:` prefix); docs-site update +for new `ContextConfig` fields after the API settles. + +## Success Metrics + +- Default (`direct`) prompts contain column lineage and join hints for pipelines + where they exist (T1, T3). +- `two_stage` context includes all ancestor tables within the configured depth (T2). +- Every observed equi-join in the pipeline's SQL appears in `## Join Hints` for + in-context tables, with zero fabricated observed joins (T3). +- Final/intermediate/source roles labeled in every schema context; truncation never + drops a final table while keeping a source table (T4). +- Fallback selection includes join-required parent tables for the fixture suite (T5). From 919d531c21e8d5d89eb374a358073e5a8fa87923 Mon Sep 17 00:00:00 2001 From: Ming-Jer Lee Date: Tue, 11 Aug 2026 11:31:30 -0700 Subject: [PATCH 2/2] docs: address review feedback on graph design specs - D1: dispatch source prompts on table_graph is_source, not is_computed() - T1/shared: resolve one capped table set for schema, graph context, notes - T3: USING joins only with single-table left input; chained-USING test - T3: candidates require identity-preserving paths via edge_type allowlist --- ...ESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md | 21 ++++- plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md | 86 +++++++++++++++---- 2 files changed, 87 insertions(+), 20 deletions(-) diff --git a/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md b/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md index cbdc201..e505fa1 100644 --- a/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md +++ b/plans/DESIGN_DESCRIPTION_GENERATION_GRAPH_GAPS.md @@ -78,9 +78,20 @@ source column is consumed downstream even when nothing upstream exists. changes LLM call volume and export diffs for existing users. Revisit the default in a minor release. -3. `generate_description()` dispatches on layer: source columns (no incoming edges, - `is_computed()` is False) use `build_source_description_prompt`; computed columns - keep `build_description_prompt`. +3. `generate_description()` dispatches on **table role, not `is_computed()`**: + + ```python + is_source_column = pipeline.table_graph.tables[column.table_name].is_source + ``` + + `is_computed()` cannot make this distinction: it returns `True` whenever + `query_id` is set (`models.py:570-578`), and `_add_query_columns()` assigns + `query_id` to every parsed node — so in `raw.users -> staging.users`, + `raw.users.email` has `layer="input"` and no incoming edges yet + `is_computed()` is still `True`. Source columns + (`table_node.is_source`) use `build_source_description_prompt`; all others keep + `build_description_prompt`. `layer == "input"` and the absence of incoming edges + are validation assertions in tests, not the dispatch criterion. ### Test Cases @@ -94,6 +105,10 @@ source column is consumed downstream even when nothing upstream exists. `overwrite=True`. - Source column consumed by zero described targets still yields a valid prompt (name + siblings only). +- **Dispatch test**: spy on both prompt builders in a `raw.users -> staging.users` + fixture; assert `raw.users.email` (which has `query_id` set and + `is_computed() == True`) is routed to `build_source_description_prompt` and + `staging.users.*` to `build_description_prompt`. ### Risks and Mitigations diff --git a/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md b/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md index 95ded06..0dce35c 100644 --- a/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md +++ b/plans/DESIGN_TEXT2SQL_GRAPH_CONTEXT.md @@ -19,11 +19,24 @@ strategy or not surfaced at all. This document specifies five gaps: T1–T3 all modify how the `relationship_section` prompt slot is assembled; T3 is the highest-value change for generated-SQL correctness (join quality). -### Shared change: graph-context assembly point +### Shared change: one capped table set, one graph-context assembly point -Both `_generate_direct` (`sql.py:182`) and `_generate_two_stage` (`sql.py:238`) -currently fill `relationship_section` with a single builder call. Refactor both to a -shared assembly: +**Current inconsistency**: `_generate_direct` builds `` via +`build_schema_context()`, which applies `max_tables` internally +(`context.py:356-361`), but then fills the relationship section from +`get_table_names()` — the *uncapped* list (`sql.py:192-194`). On pipelines larger +than `max_tables`, relationship/lineage/join sections would reference tables absent +from ``, and omitted tables can consume section caps before included tables +get their hints. + +Refactor both strategies to resolve the table set **once**: + +```python +def resolve_context_tables(self, tables: Optional[List[str]] = None) -> List[str]: + """Ordered, capped table list — the single source of truth for a prompt.""" + # selection (all tables / two-stage selection) -> role-priority truncation + # to config.max_tables (T4 defines the priority order) +``` ```python def _build_graph_context(self, builder: ContextBuilder, tables: List[str]) -> str: @@ -35,11 +48,20 @@ def _build_graph_context(self, builder: ContextBuilder, tables: List[str]) -> st return "\n\n".join(p for p in parts if p) ``` +The resolved list feeds **every** consumer: `build_context_for_tables()` (schema), +`_build_graph_context()` (relationships/lineage/joins), `_build_notes()` (PII), and +the `tables_used` result field. `build_schema_context()` becomes a thin wrapper that +resolves and delegates, so external callers keep their behavior. + The existing `GENERATE_SQL_PROMPT` / `GENERATE_SQL_WITH_EXPLANATION_PROMPT` templates keep their `{relationship_section}` slot; no template signature change. Sanitization stays where it is today (each section passed through `sanitize_for_prompt` at the `prompt.format(...)` call sites). +**Regression test**: fixture with more than `max_tables` tables — every table named +in the relationship, lineage, join, and notes sections must appear in ``, +and `tables_used` must equal the resolved list. + --- ## T1: Column Lineage in Direct Mode @@ -167,22 +189,46 @@ Implementation sketch: **Phase 1 scope**: skip predicates whose alias resolves to a CTE-internal name rather than a pipeline table; log at debug level. (CTE mapping via `query_lineage` is a follow-up.) -- `USING (col)` joins (`exp.Join` args) emit `left.col = right.col`. +- `USING (col)` joins emit `left.col = right.col` **only when the join's left input + resolves to exactly one physical table**. In a chain like + `a JOIN b USING (id) JOIN c USING (id)`, the second join's left input is the + composite relation `(a JOIN b)` — there is no unique physical `left_table`, and + picking `a` or `b` would fabricate an observed pair. Such joins are skipped and + debug-logged, consistent with the zero-fabrication goal. - Non-equi joins and `ON` conditions that are not column-to-column equality are skipped. -#### 2. Candidate joins (from shared lineage sources) - -For table pairs never joined in the pipeline (e.g. two marts), infer candidates: - -- For each output column of each context table, call - `pipeline.trace_column_backward(table, column)` once and record - `ultimate_source_full_name -> [(table, column), ...]`. -- Any source mapped to columns in ≥2 distinct context tables yields candidate pairs. -- Deduplicate against observed joins; cap total candidates +#### 2. Candidate joins (from shared lineage sources — identity-preserving paths only) + +For table pairs never joined in the pipeline (e.g. two marts), infer candidates. +**Shared ultimate ancestry alone does not prove an equality join**: +`trace_column_backward()` returns ultimate leaves and discards transformation and +grain, so `mart.user_counts.user_count = COUNT(raw.users.id)` and +`mart.user_ids.user_id = raw.users.id` both trace to `raw.users.id` — naive +inference would suggest the nonsensical `user_count = user_id`, and a `candidate:` +label does not make that safe for SQL generation. + +Restriction: a column qualifies as a candidate endpoint only if its **entire +backward path is identity-preserving**: + +- Use `pipeline.trace_column_backward_full(table, column)` (returns nodes **and** + edges) instead of the leaf-only variant. +- Every edge on the path from the column to the shared source must have + `edge_type` in the allowlist `{"direct", "star_passthrough", "cross_query"}` + (`ColumnEdge.edge_type`, `models.py:598-600`). Any `transform`, `aggregate`, + `join`, or unrecognized edge type disqualifies the path — unknown types fail + closed. +- Both endpoints of a candidate pair must qualify; the pair maps to the same shared + source column. +- Deduplicate against observed joins; cap total hints (`ContextConfig.max_join_hints: int = 15`, shared with observed joins, observed first). +If implementation finds `edge_type` granularity insufficient to prove identity +preservation (e.g. renames recorded as `transform`), ship observed joins **without** +inferred candidates rather than emit unproven ones — observed joins alone deliver +most of T3's value. + #### 3. Prompt section ```python @@ -204,9 +250,15 @@ explicitly labeled `candidate:` so the model can weigh them below observed joins - Fixture with `JOIN ... ON o.customer_id = c.id` and aliases → observed join with real table names and `query_id`. - Composite key (`ON a.x = b.x AND a.y = b.y`) → two entries, same query. -- `USING (id)` → resolved to both tables. -- Two marts sharing `raw.users.id` ancestry, never joined directly → one - `candidate:` line; no candidates between unrelated tables. +- `USING (id)` with a single-table left input → resolved to both tables. +- **Chained `USING`**: `a JOIN b USING (id) JOIN c USING (id)` → exactly one + observed hint (`a.id = b.id`); the second join emits nothing (composite left + input), and no `a.id = c.id` / `b.id = c.id` pair is fabricated. +- Two marts sharing `raw.users.id` ancestry via pass-through paths, never joined + directly → one `candidate:` line; no candidates between unrelated tables. +- **Aggregate counterexample**: `mart.user_counts.user_count = COUNT(raw.users.id)` + and `mart.user_ids.user_id = raw.users.id` both trace to `raw.users.id` → **no** + candidate emitted (the `aggregate` edge disqualifies the path). - Non-equi join (`ON a.ts > b.ts`) produces nothing. - Cap: hints truncated to `max_join_hints`, observed joins prioritized.