Spaces:
Sleeping
Sleeping
| """Tests for schema linking — TableIndex, find_relevant_tables, and the | |
| text_to_sql safety cap. Provider-agnostic: assertions hold under both the | |
| sentence-transformer and lexical-fallback embedding providers. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from lexsi_ds.agent.context import AgentContext, ColumnInfo, DatasetHandle, TableInfo | |
| from lexsi_ds.agent.table_index import TableIndex | |
| from lexsi_ds.agent.tools.find_relevant_tables import FindRelevantTablesArgs, _run | |
| from lexsi_ds.agent.tools.text_to_sql import _render_schema | |
| def _tbl(name, cols, comment=""): | |
| return TableInfo(name=name, columns=[ColumnInfo(c, "VARCHAR") for c in cols], | |
| n_rows=10, comment=comment) | |
| def _handle(tables, ds_id="syn"): | |
| return DatasetHandle(id=ds_id, kind="attached", duckdb_path=Path("/tmp/none.duckdb"), | |
| tables=tables) | |
| def _ctx(tables, ds_id="syn"): | |
| return AgentContext(dataset=_handle(tables, ds_id), run_id="t") | |
| _THREE = [ | |
| _tbl("loan_book", ["loan_id", "principal", "interest", "default_flag"], "loans issued to customers"), | |
| _tbl("weather_daily", ["date", "temperature", "rainfall"], "daily weather readings"), | |
| _tbl("customer_profile", ["customer_id", "age", "city"], "customer demographics"), | |
| ] | |
| def test_table_index_ranks_relevant_first(): | |
| idx = TableIndex.build(_THREE) | |
| hits = idx.search("average principal of loans that defaulted", top_k=3) | |
| assert hits, "expected at least one hit" | |
| assert hits[0][0].name == "loan_book" | |
| def test_table_index_empty_query(): | |
| assert TableIndex.build(_THREE).search(" ") == [] | |
| def test_find_relevant_tables_tool(): | |
| ctx = _ctx(_THREE) | |
| res = _run(FindRelevantTablesArgs(question="loan principal default", top_k=2), ctx) | |
| assert res.ok | |
| assert res.payload["total_tables"] == 3 | |
| assert len(res.payload["tables"]) <= 2 | |
| assert "loan_book" in res.payload["tables"] | |
| def test_find_relevant_tables_no_tables(): | |
| res = _run(FindRelevantTablesArgs(question="x"), _ctx([])) | |
| assert not res.ok and res.error == "no_tables" | |
| def test_index_cached_on_ctx(): | |
| ctx = _ctx(_THREE) | |
| _run(FindRelevantTablesArgs(question="loans"), ctx) | |
| assert "table_index:syn" in ctx.cache | |
| # ---- text_to_sql safety cap ---- | |
| def _many(n): | |
| tables = [_tbl(f"t{i:02d}", ["col_a", "col_b"]) for i in range(n)] | |
| tables.append(_tbl("sales_revenue", ["revenue", "region", "quarter"], "sales revenue by region")) | |
| return tables | |
| def _count_rendered(schema: str) -> int: | |
| return schema.count("### TABLE") | |
| def test_small_schema_renders_all_no_cap(): | |
| schema = _render_schema(_ctx(_THREE), question="loans") | |
| assert _count_rendered(schema) == 3 | |
| assert "auto-selected" not in schema | |
| def test_large_schema_auto_capped(): | |
| ctx = _ctx(_many(50), ds_id="big") | |
| schema = _render_schema(ctx, question="total revenue by region") | |
| n = _count_rendered(schema) | |
| assert n <= 21 and n < 51 # capped well below the 51 total | |
| assert "sales_revenue" in schema # the relevant table survived the cap | |
| assert "auto-selected" in schema # cap note present | |
| def test_explicit_tables_subset_scopes_exactly(): | |
| ctx = _ctx(_many(50), ds_id="big2") | |
| schema = _render_schema(ctx, question="anything", tables=["t00", "sales_revenue"]) | |
| assert _count_rendered(schema) == 2 | |
| assert "sales_revenue" in schema and "### TABLE `t00`" in schema | |
| def test_bad_explicit_subset_falls_back(): | |
| # nonexistent names → don't render an empty schema; fall back gracefully | |
| ctx = _ctx(_THREE, ds_id="fb") | |
| schema = _render_schema(ctx, question="loans", tables=["does_not_exist"]) | |
| assert _count_rendered(schema) == 3 | |