BIRD: linking tables, inferred joins, lenient plan intent and contact columns

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.

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)
 

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]
 
 

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

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 `<name>_id` columns without a foreign key: orders.customer_id -> customers.id."""
+    """Guess joins the database does not declare, from names:
+
+    - a `<name>_id` column joins the primary key of the table `<name>`: 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
 
 

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

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))

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
     ],
 )

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"

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))