diff --git a/README.md b/README.md index 3670e41..23a49bc 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ question ─▶ schema retrieval ─▶ LLM: plan + SQL ─▶ AST validation | Validation | deterministic | `sqlglot` AST checks, row limit enforced (see below). The SQL that runs is regenerated from the validated tree. | | Execution | deterministic | Read-only connection, statement timeout, client-side row cap. | | Repair | LLM, bounded | A repairable failure (validator rejection, SQL error, timeout, bad result, malformed model reply) is sent back to the model at most `MAX_RETRIES` times; every repair is validated again. Unsafe intent is never retried. | -| Result checks | deterministic | Empty result, aggregate over nothing, ranking led by NULL, columns concatenated into one, NULL-only column, row limit reached, duplicates. Missing-value probes catch filter values and `LIKE` patterns that match nothing (e.g. `'Ivory Coast'`, a time written `'0:01:54%'` where `'1:54.123'` is stored) and show how stored values look. | +| Result checks | deterministic | Empty result, aggregate over nothing, ranking led by NULL, top N that returned fewer than N rows, columns concatenated into one, NULL-only column, row limit reached, duplicates. Missing-value probes catch filter values and `LIKE` patterns that match nothing (e.g. `'Ivory Coast'`, a time written `'0:01:54%'` where `'1:54.123'` is stored) and show how stored values look. | | Chart | deterministic | Chosen from column kinds and values: stat tile, line, bar, scatter or none. The model's suggestion is only a tie-breaker. | | Answer | LLM + check | 1–3 sentences. Every number must come from the rows; otherwise a template answer built from the rows is used. | diff --git a/backend/app/agent/result_checks.py b/backend/app/agent/result_checks.py index eb88f05..38469d4 100644 --- a/backend/app/agent/result_checks.py +++ b/backend/app/agent/result_checks.py @@ -116,6 +116,17 @@ def check_result( ) ) + wanted = _top_n(tree, limit_action) if tree is not None else None + if wanted is not None and result.row_count < wanted: + checks.append( + ResultCheck( + code="fewer_than_requested", + severity="warning", + message=f"The query asks for the top {wanted} but returned {result.row_count}; " + "some may have no data for the filters used.", + ) + ) + if result.row_count > 1: for index, column in enumerate(result.columns): if all(row[index] is None for row in result.rows): @@ -149,6 +160,48 @@ def check_result( return checks +def _top_n(tree: exp.Expr, limit_action: str) -> int | None: + """N of a top-N query: an ORDER BY ... LIMIT N the model wrote, in the final SELECT or in a + CTE or derived table that the final SELECT returns row for row (no grouping, aggregates or + DISTINCT). A LIMIT the validator added or clamped is the row cap, not a request.""" + if not isinstance(tree, exp.Select): + return None + if limit_action == "kept" and tree.args.get("order"): + return _limit_count(tree) + if ( + tree.args.get("group") + or tree.args.get("distinct") + or any(projection.find(exp.AggFunc) for projection in tree.expressions) + ): + return None + ctes = {cte.alias_or_name: cte.this for cte in tree.ctes} + sources = [tree.args.get("from_"), *(tree.args.get("joins") or [])] + for source in sources: + node = source.this if source is not None else None + if isinstance(node, exp.Subquery): + inner = node.this + elif isinstance(node, exp.Table) and not node.db and node.name in ctes: + inner = ctes[node.name] + else: + continue + if isinstance(inner, exp.Select) and inner.args.get("order"): + count = _limit_count(inner) + if count is not None: + return count + return None + + +def _limit_count(query: exp.Select) -> int | None: + fetch = query.args.get("limit") + if isinstance(fetch, exp.Fetch): # FETCH FIRST [n] ROWS ONLY; no count means one row + value = fetch.args.get("count") or exp.Literal.number(1) + elif isinstance(fetch, exp.Limit): + value = fetch.expression + else: + return None + return int(value.this) if isinstance(value, exp.Literal) and value.is_int else None + + def _hashable(value: Any) -> Any: return repr(value) if isinstance(value, (list, dict)) else value diff --git a/backend/tests/test_result_checks.py b/backend/tests/test_result_checks.py index 25ada5d..b362976 100644 --- a/backend/tests/test_result_checks.py +++ b/backend/tests/test_result_checks.py @@ -73,6 +73,49 @@ def test_limit_reached_only_when_the_limit_was_not_asked_for() -> None: assert codes(result(["a"], rows), limit=10, action="kept") == [] # "top 10" asked for 10 rows +TOP5_IN_CTE = ( + "WITH top5 AS (SELECT id, name FROM c ORDER BY population DESC LIMIT 5) " + "SELECT t.name, g.total FROM top5 AS t JOIN g ON g.id = t.id WHERE g.total IS NOT NULL " + "ORDER BY t.name LIMIT 1000" +) +TOP5 = "WITH top5 AS (SELECT co2 FROM t ORDER BY co2 DESC LIMIT 5)" +FOUR_ROWS = [["A", 1], ["B", 2], ["C", 3], ["D", 4]] + + +@pytest.mark.parametrize( + ("sql", "action"), + [ + ("SELECT name, co2 FROM t ORDER BY co2 DESC LIMIT 5", "kept"), + ("SELECT name, co2 FROM t ORDER BY co2 DESC FETCH FIRST 5 ROWS ONLY", "kept"), + (TOP5_IN_CTE, "added"), # Q052: top 5 picked in a CTE, then a filter drops one + ("SELECT s.name, s.co2 FROM (SELECT name, co2 FROM t ORDER BY co2 DESC LIMIT 5) AS s", "added"), + ], +) +def test_top_n_with_fewer_rows_warns(sql: str, action: str) -> None: + checks = check_result(result(["name", "v"], FOUR_ROWS), sql, "postgres", 1000, action) + found = [(c.code, c.severity, c.repairable) for c in checks] + assert found == [("fewer_than_requested", "warning", False)] + assert checks[0].message.startswith("The query asks for the top 5 but returned 4;") + + +@pytest.mark.parametrize( + ("sql", "action", "limit"), + [ + ("SELECT name, co2 FROM t ORDER BY co2 DESC LIMIT 1000", "added", 1000), # only the row cap + ("SELECT name, co2 FROM t ORDER BY co2 DESC LIMIT 1000", "clamped", 1000), # asked for too many + ("SELECT name, co2 FROM t LIMIT 5", "kept", 5), # no ranking + ("SELECT name, co2 FROM t ORDER BY co2 DESC LIMIT 4", "kept", 4), # got all it asked for + # The outer query aggregates the top 5: one row is the expected shape. + (f"{TOP5} SELECT SUM(co2) FROM top5", "added", 1000), + (f"{TOP5} SELECT DISTINCT co2 FROM top5", "added", 1000), + # A top-N used only as a filter does not set the number of output rows. + ("SELECT name, co2 FROM t WHERE id IN (SELECT id FROM t ORDER BY co2 DESC LIMIT 5)", "added", 1000), + ], +) +def test_no_top_n_warning_otherwise(sql: str, action: str, limit: int) -> None: + assert "fewer_than_requested" not in codes(result(["name", "v"], FOUR_ROWS), sql, limit, action) + + def test_duplicate_rows() -> None: assert codes(result(["a"], [[1], [1], [2]])) == ["duplicate_rows"] diff --git a/frontend/src/components/Verification.test.tsx b/frontend/src/components/Verification.test.tsx new file mode 100644 index 0000000..46f654a --- /dev/null +++ b/frontend/src/components/Verification.test.tsx @@ -0,0 +1,66 @@ +import { render, screen, within } from '@testing-library/react' +import { describe, expect, it } from 'vitest' +import type { AgentResult } from '../api/types' +import { Verification } from './Verification' + +const ASSUMPTIONS = [ + "Use the 'co2' column representing annual fossil and industry CO2 emissions (million tonnes).", + "Only entities with entity_type='country' are considered.", +] + +const RESULT: AgentResult = { + status: 'success', + question: 'Which 5 countries emitted the most CO2 in 2023?', + resolved_question: null, + answer: 'China emitted the most.', + answer_source: 'llm', + clarification_question: null, + explanation: null, + plan: { + intent: 'ranking', + tables: ['co2_emissions', 'countries'], + metrics: ['co2'], + filters: ['year = 2023', "entity_type = 'country'"], + group_by: [], + order_by: 'co2 desc', + limit: 5, + assumptions: ASSUMPTIONS, + }, + sql: 'SELECT name, co2 FROM co2_emissions ORDER BY co2 DESC LIMIT 5', + columns: ['name', 'co2'], + rows: [['China', 11472.4]], + column_units: {}, + chart_suggestion: null, + chart: null, + checks: [], + error: null, + trace: [], + metadata: { + database_id: 'owid', + dialect: 'postgresql', + model: 'gemini-3.6-flash', + prompt_version: 'sql-generator/3', + tables_used: ['co2_emissions'], + execution_time_ms: 41, + row_count: 1, + truncated: false, + retry_count: 0, + request_id: null, + }, +} + +describe('Verification plan', () => { + it('lists sentence-style assumptions one per line instead of joining them', () => { + render() + + const items = within(screen.getByRole('list')).getAllByRole('listitem') + expect(items.map((item) => item.textContent)).toEqual(ASSUMPTIONS) + expect(screen.queryByText(/\.;/)).toBeNull() + }) + + it('still joins short filter expressions on one line', () => { + render() + + expect(screen.getByText("year = 2023; entity_type = 'country'")).toBeTruthy() + }) +}) diff --git a/frontend/src/components/Verification.tsx b/frontend/src/components/Verification.tsx index 7997fd2..bcc5092 100644 --- a/frontend/src/components/Verification.tsx +++ b/frontend/src/components/Verification.tsx @@ -95,14 +95,24 @@ function Section({ title, action, children }: { title: string; action?: ReactNod } function PlanList({ plan }: { plan: QueryPlan }) { - const rows: [string, string][] = [['Intent', plan.intent]] + const rows: [string, ReactNode][] = [['Intent', plan.intent]] if (plan.tables.length) rows.push(['Tables', plan.tables.join(', ')]) if (plan.metrics.length) rows.push(['Metrics', plan.metrics.join(', ')]) if (plan.filters.length) rows.push(['Filters', plan.filters.join('; ')]) if (plan.group_by.length) rows.push(['Grouped by', plan.group_by.join(', ')]) if (plan.order_by) rows.push(['Ordered by', plan.order_by]) if (plan.limit) rows.push(['Limit', String(plan.limit)]) - if (plan.assumptions.length) rows.push(['Assumptions', plan.assumptions.join('; ')]) + // Assumptions are usually whole sentences: one per line, since joining them would print ".;". + if (plan.assumptions.length) { + rows.push([ + 'Assumptions', + , + ]) + } return (