diff --git a/EVALUATION_PLAN.md b/EVALUATION_PLAN.md index 2215f8b..90b1234 100644 --- a/EVALUATION_PLAN.md +++ b/EVALUATION_PLAN.md @@ -164,12 +164,13 @@ all), with cryptic column names, dirty values and no column comments. Difficulty this is the setting its owner would choose. The app's default stays `safe`, so results on BIRD describe a database profiled with `full`. * Sensitive columns: the app hides every column whose name looks like a birth date, e-mail, phone - and the like. BIRD connections allow the four birth-date columns (formula_1.drivers.dob, - financial.client.birth_date, thrombosis_prediction.patient.birthday, - european_football_2.player.birthday) through `allow_columns`, as an owner of this long-published - data would (evaluation/bird.py ALLOW_COLUMNS). E-mail and phone columns stay hidden; questions that - need them count as misses. Runs before this change hid `dob` and `birth_date` and showed - `birthday` and california_schools' `admemail1-3`. + and the like. BIRD connections allow every such column of BIRD (the four birth dates, + california_schools' `phone` and `admemail1-3`, student_club's member `phone` and `email`) through + `allow_columns`, as an owner of this long-published data would (evaluation/bird.py + ALLOW_COLUMNS): other systems on the leaderboard see all columns. The app's default is unchanged. + Runs up to commit 9ba4627 kept e-mail and phone hidden (6 Mini-Dev questions need them and + counted as misses); runs before that hid `dob` and `birth_date` and showed `birthday` and + california_schools' `admemail1-3`. * Evidence (--evidence): BIRD gives each question a hint ("evidence"), e.g. "eligible free rate = Free Meal Count / Enrollment". `on` passes it as the user's definitions (the API's optional `definitions` field, applied literally); `off` measures the agent without it. Published BIRD scores are usually with evidence. diff --git a/backend/app/agent/controller.py b/backend/app/agent/controller.py index f9e3f73..fa27ce7 100644 --- a/backend/app/agent/controller.py +++ b/backend/app/agent/controller.py @@ -205,8 +205,10 @@ def run( with trace.timed() as ms: profile = connection.profile() - # Earlier questions too: "what about China?" alone names no table or metric. - context = build_context(profile, " ".join([*(t.question for t in history), question])) + # Earlier questions too: "what about China?" alone names no table or metric. And the + # definitions: "finished the race refers to time IS NOT NULL" names the column to use. + terms = [*(t.question for t in history), question, definitions or ""] + context = build_context(profile, " ".join(terms)) meta.tables_used = context.table_names trace.record("schema_retrieval", duration_ms=ms[0], tables=context.table_names) diff --git a/backend/app/agent/schema_retriever.py b/backend/app/agent/schema_retriever.py index 9638a42..87299b3 100644 --- a/backend/app/agent/schema_retriever.py +++ b/backend/app/agent/schema_retriever.py @@ -1,7 +1,8 @@ """Select the parts of a database profile relevant to a question and render them for the LLM. Retrieval is deterministic keyword matching over table/column names, comments and -categorical values, plus the tables those matches join to. Small schemas are sent whole. +categorical values, plus the tables those matches join to and the tables that link two of them. +Small schemas are sent whole. The `retrieve` function is the seam where semantic (embedding) retrieval can be added later. Sensitive columns are never rendered. Sampled values are rendered as quoted data and the @@ -69,9 +70,14 @@ def retrieve(profile: DatabaseProfile, question: str) -> list[TableProfile]: matched = {name for name, score in scores.items() if score > 0} if not matched: return list(profile.tables) + targets: dict[str, set[str]] = {} for rel in profile.relationships: # bring in the tables needed to join the matches + targets.setdefault(rel.from_table, set()).add(rel.to_table) if rel.from_table in matched: matched.add(rel.to_table) + # A linking table (hero_power between superhero and superpower, race results between drivers + # and races) rarely shares a word with the question, but the answer runs through it. + matched |= {name for name, to in targets.items() if len(to & matched) >= 2} return [t for t in profile.tables if t.qualified_name in matched] diff --git a/backend/app/agent/sql_generator.py b/backend/app/agent/sql_generator.py index f575073..c982afd 100644 --- a/backend/app/agent/sql_generator.py +++ b/backend/app/agent/sql_generator.py @@ -10,9 +10,9 @@ from collections.abc import Sequence from dataclasses import dataclass -from typing import Annotated, Literal +from typing import Annotated, Any, Literal, get_args -from pydantic import BaseModel, Field, StringConstraints, model_validator +from pydantic import BaseModel, Field, StringConstraints, field_validator, model_validator from app.agent.schema_retriever import SchemaContext from app.llm.client import LLMCall, LLMClient @@ -85,8 +85,11 @@ _Text = Annotated[str, StringConstraints(max_length=300)] +Intent = Literal["lookup", "aggregate", "ranking", "trend", "comparison", "other"] + + class QueryPlan(BaseModel): - intent: Literal["lookup", "aggregate", "ranking", "trend", "comparison", "other"] = "other" + intent: Intent = "other" tables: list[_Text] = Field(default=[], max_length=20) metrics: list[_Text] = Field(default=[], max_length=20) filters: list[_Text] = Field(default=[], max_length=20) @@ -95,6 +98,13 @@ class QueryPlan(BaseModel): limit: int | None = None assumptions: list[_Text] = Field(default=[], max_length=10) + @field_validator("intent", mode="before") + @classmethod + def _unknown_intent_is_other(cls, value: Any) -> Any: + """The intent is only a label shown with the plan: an unknown one ("lookup|aggregate", copied + from the prompt's list) must not throw away a reply whose SQL is fine.""" + return value if value in get_args(Intent) else "other" + class Turn(BaseModel): """An earlier question in the conversation, sent back by the client (the server keeps no diff --git a/backend/app/database/profiler.py b/backend/app/database/profiler.py index 05155f8..81c20b4 100644 --- a/backend/app/database/profiler.py +++ b/backend/app/database/profiler.py @@ -74,7 +74,13 @@ def fingerprint(tables: Iterable[TableProfile], relationships: Iterable[Relation def infer_relationships(tables: list[TableProfile], declared: list[Relationship]) -> list[Relationship]: - """Guess joins for `_id` columns without a foreign key: orders.customer_id -> customers.id.""" + """Guess joins the database does not declare, from names: + + - a `_id` column joins the primary key of the table ``: orders.customer_id -> + customers.id; + - a column named like a declared single-column foreign key joins what that key references: + rulings.uuid -> cards.uuid is declared, so foreign_data.uuid -> cards.uuid too (dumps often + keep some foreign keys and drop others).""" by_name: dict[str, TableProfile] = {t.name.lower(): t for t in tables} declared_keys = {(r.from_table, tuple(r.from_columns)) for r in declared} inferred = [] @@ -100,6 +106,27 @@ def infer_relationships(tables: list[TableProfile], declared: list[Relationship] ) ) break + references: dict[str, set[tuple[str, str]]] = {} + for r in declared: + if len(r.from_columns) == 1: + references.setdefault(r.from_columns[0].lower(), set()).add((r.to_table, r.to_columns[0])) + known = declared_keys | {(r.from_table, tuple(r.from_columns)) for r in inferred} + for table in tables: + for column in table.columns: + targets = references.get(column.name.lower(), set()) + if len(targets) != 1 or column.primary_key or (table.qualified_name, (column.name,)) in known: + continue # no such key, or ambiguous (same name, different targets) + to_table, to_column = next(iter(targets)) + if to_table != table.qualified_name: + inferred.append( + Relationship( + from_table=table.qualified_name, + from_columns=[column.name], + to_table=to_table, + to_columns=[to_column], + inferred=True, + ) + ) return inferred diff --git a/backend/tests/test_agent.py b/backend/tests/test_agent.py index acc9569..89bbef8 100644 --- a/backend/tests/test_agent.py +++ b/backend/tests/test_agent.py @@ -81,6 +81,30 @@ def test_large_schemas_are_narrowed_to_relevant_tables() -> None: assert [t.name for t in retrieve(_profile(20), "total invoice amount")] == ["invoices"] +def test_a_table_linking_two_selected_tables_is_added() -> None: + def table(name: str, *columns: str) -> TableProfile: + cols = [ColumnProfile(name="id", type="INTEGER", nullable=False, primary_key=True)] + cols += [ColumnProfile(name=c, type="TEXT", nullable=True) for c in columns] + return TableProfile(schema_name="s", name=name, kind="table", columns=cols) + + def join(child: str, column: str, parent: str) -> Relationship: + return Relationship( + from_table=f"s.{child}", from_columns=[column], to_table=f"s.{parent}", to_columns=["id"] + ) + + tables = [table(f"filler_{i}") for i in range(10)] + tables += [table("hero", "alias"), table("power", "ability"), table("hero_power", "hero_id", "power_id")] + tables.append(table("sidekick", "hero_id")) # joins only one selected table: not a link + profile = DatabaseProfile( + database_id="x", dialect="postgresql", sampling=SamplingMode.SAFE, fingerprint="f", tables=tables, + relationships=[join("hero_power", "hero_id", "hero"), join("hero_power", "power_id", "power"), + join("sidekick", "hero_id", "hero")], + ) # fmt: skip + # The question names no word of hero_power: it joins the two matched tables, so it is added. + names = [t.name for t in retrieve(profile, "Which ability does the alias Abomination have?")] + assert names == ["hero", "power", "hero_power"] + + def test_sensitive_columns_are_never_rendered() -> None: context = build_context(_profile(1), "invoices") assert "amount" in context.text and "card_number" not in context.text diff --git a/backend/tests/test_profiler.py b/backend/tests/test_profiler.py index 2a94a69..56b9b03 100644 --- a/backend/tests/test_profiler.py +++ b/backend/tests/test_profiler.py @@ -106,6 +106,30 @@ def test_infer_relationships_from_names() -> None: assert all(r.inferred for r in inferred) +def test_a_declared_foreign_key_is_inferred_for_columns_of_the_same_name() -> None: + cards = _table("cards", ("id", True), ("uuid", False)) + rulings = _table("rulings", ("id", True), ("uuid", False)) + foreign_data = _table("foreign_data", ("id", True), ("uuid", False)) + declared = [ + Relationship(from_table="s.rulings", from_columns=["uuid"], to_table="s.cards", to_columns=["uuid"]) + ] + inferred = infer_relationships([cards, rulings, foreign_data], declared) + # foreign_data.uuid joins cards.uuid like rulings.uuid; cards.uuid is the target itself. + assert [(r.from_table, r.from_columns, r.to_table, r.to_columns, r.inferred) for r in inferred] == [ + ("s.foreign_data", ["uuid"], "s.cards", ["uuid"], True) + ] + + +def test_a_column_name_declared_with_two_targets_is_not_guessed() -> None: + a, b = _table("a", ("id", True)), _table("b", ("id", True)) + x, y, z = (_table(n, ("id", True), ("ref", False)) for n in ("x", "y", "z")) + declared = [ + Relationship(from_table="s.x", from_columns=["ref"], to_table="s.a", to_columns=["id"]), + Relationship(from_table="s.y", from_columns=["ref"], to_table="s.b", to_columns=["id"]), + ] + assert infer_relationships([a, b, x, y, z], declared) == [] + + def test_declared_foreign_keys_are_not_duplicated() -> None: customers = _table("customers", ("id", True)) orders = _table("orders", ("id", True), ("customer_id", False)) diff --git a/backend/tests/test_sql_generator.py b/backend/tests/test_sql_generator.py index 14345a9..369f3ed 100644 --- a/backend/tests/test_sql_generator.py +++ b/backend/tests/test_sql_generator.py @@ -37,11 +37,18 @@ def test_clarification_reply_is_parsed() -> None: assert parse_output(json.dumps(reply), GeneratedSQL).clarification_question == "Total or per capita?" +@pytest.mark.parametrize("intent", ["lookup|aggregate", "delete", None, 3]) +def test_an_unknown_intent_is_other_and_keeps_the_reply(intent: object) -> None: + # The intent is only a label shown with the plan; rejecting the reply would discard valid SQL. + reply = {"plan": {**PLAN, "intent": intent}, "sql": "SELECT 1", "explanation": "x"} + generated = parse_output(json.dumps(reply), GeneratedSQL) + assert generated.plan is not None and generated.plan.intent == "other" and generated.sql == "SELECT 1" + + @pytest.mark.parametrize( "reply", [ {"sql": "SELECT 1", "clarification_question": "Which year?", "explanation": "x"}, # both - {"plan": {**PLAN, "intent": "delete"}, "sql": "SELECT 1", "explanation": "x"}, # unknown intent {"plan": {**PLAN, "filters": ["x" * 301]}, "sql": "SELECT 1", "explanation": "x"}, # oversized ], ) diff --git a/evaluation/bird.py b/evaluation/bird.py index 77b1120..f396438 100644 --- a/evaluation/bird.py +++ b/evaluation/bird.py @@ -79,13 +79,21 @@ " -c max_parallel_workers_per_gather=0 -c synchronize_seqscans=off" ) -# Birth dates in BIRD's public, long-published data, allowed as a database owner would allow them -# in databases.toml (allow_columns). E-mail and phone columns stay hidden: questions that need them -# count as misses. Configuration only: the agent knows nothing about these names. +# Every column the sensitive-column policy would hide in BIRD's public, long-published data (birth +# dates, phone numbers, e-mail addresses), allowed as a database owner would allow them in +# databases.toml (allow_columns). Other systems on the leaderboard see every column, so hiding them +# would score questions that need them as misses for a policy reason. The app's default is unchanged. +# Configuration only: the agent knows nothing about these names. ALLOW_COLUMNS = [ + "california_schools.schools.admemail1", + "california_schools.schools.admemail2", + "california_schools.schools.admemail3", + "california_schools.schools.phone", "european_football_2.player.birthday", "financial.client.birth_date", "formula_1.drivers.dob", + "student_club.member.email", + "student_club.member.phone", "thrombosis_prediction.patient.birthday", ] @@ -254,7 +262,11 @@ def connect( def fingerprint(url: SecretStr) -> dict[str, int]: """Row count of every BIRD table: ground truth is only valid on the data it was built from.""" counts: dict[str, int] = {} - with psycopg.connect(_libpq(url), options="-c default_transaction_read_only=on") as conn: + # Its own statement timeout: under the account's (the agent's 5 s) a count on a busy server once + # timed out and ended a run (its answers were already saved). + timeout_ms = GROUND_TRUTH_TIMEOUT_SECONDS * 1000 + options = f"-c default_transaction_read_only=on -c statement_timeout={timeout_ms}" + with psycopg.connect(_libpq(url), options=options) as conn: tables = conn.execute( "SELECT schemaname, tablename FROM pg_tables" " WHERE schemaname NOT IN ('pg_catalog', 'information_schema') ORDER BY 1, 2" diff --git a/evaluation/tests/test_bird.py b/evaluation/tests/test_bird.py index 1aec419..163177b 100644 --- a/evaluation/tests/test_bird.py +++ b/evaluation/tests/test_bird.py @@ -10,7 +10,9 @@ import pytest from pydantic import SecretStr +from app.database.profiler import is_sensitive from evaluation.bird import ( + ALLOW_COLUMNS, bird_url, definitions_for, execution_accuracy, @@ -32,6 +34,14 @@ ] # fmt: skip +def test_allowed_columns_are_ones_the_policy_would_hide() -> None: + # Allowing a column the policy does not hide would mean nothing; a typo would silently allow none. + assert sorted(set(ALLOW_COLUMNS)) == ALLOW_COLUMNS + for name in ALLOW_COLUMNS: + parts = name.split(".") # schema.table.column + assert len(parts) == 3 and is_sensitive(parts[2]), name + + def test_questions_are_converted(tmp_path: Path) -> None: path = tmp_path / "q.json" path.write_text(json.dumps(RAW))