diff --git a/lumen/ai/agents/sql.py b/lumen/ai/agents/sql.py index 931fa8eea..dfd92ff7d 100644 --- a/lumen/ai/agents/sql.py +++ b/lumen/ai/agents/sql.py @@ -377,6 +377,18 @@ async def _validate_sql( feedback = f"{type(e).__name__}: {e!s}" if "KeyError" in feedback: feedback += " The data does not exist; select from available data sources." + elif "not found" in str(e).lower(): + # A referenced table does not exist (e.g. the model used the + # source name instead of one of its tables, common for + # multi-table sources like XArraySQLSource). Surface the + # actual table names so the retry stops guessing. + available = source.get_tables() + if available: + names = ", ".join(repr(t) for t in available[:50]) + feedback += ( + f" Use one of these exact table names, not the " + f"source name: {names}." + ) retry_result = await self.revise( feedback, messages, context, spec=sql_query, language=f"sql.{source.dialect}", diff --git a/lumen/tests/ai/test_sql_validation.py b/lumen/tests/ai/test_sql_validation.py new file mode 100644 index 000000000..1f88226b8 --- /dev/null +++ b/lumen/tests/ai/test_sql_validation.py @@ -0,0 +1,87 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +try: + import lumen.ai # noqa +except ModuleNotFoundError: + pytest.skip( + "lumen.ai could not be imported, skipping tests.", + allow_module_level=True, + ) + +from lumen.ai.agents.sql import SQLAgent + + +@pytest.mark.asyncio +async def test_validate_sql_surfaces_table_names_on_not_found(): + """On a 'table not found' error the retry feedback should list the + source's real table names so the model stops using the source name + (the failure mode for multi-table sources like XArraySQLSource).""" + agent = SQLAgent() + + source = MagicMock() + source.dialect = "duckdb" + source.get_tables.return_value = ["CMI_C01", "DQF_C01"] + + attempts = {"n": 0} + + def execute(sql, *a, **k): + attempts["n"] += 1 + if attempts["n"] == 1: + # First attempt: model referenced the source name as a table. + raise ValueError("Error during planning: table 'MySource' not found") + return None # Corrected query validates. + + source.execute.side_effect = execute + + captured = {} + + async def fake_revise(feedback, *a, **k): + captured["feedback"] = feedback + return "SELECT * FROM CMI_C01" + + agent.revise = fake_revise + + result = await agent._validate_sql( + context={}, + sql_query='SELECT * FROM "MySource"', + expr_slug="x", + source=source, + messages=[], + step=MagicMock(), + max_retries=2, + ) + + # Recovered using the corrected SQL. + assert "CMI_C01" in result + # The retry feedback surfaced the real table names + the hint. + assert "CMI_C01" in captured["feedback"] + assert "DQF_C01" in captured["feedback"] + assert "not the source name" in captured["feedback"] + + +@pytest.mark.asyncio +async def test_validate_sql_valid_query_no_retry(): + """A valid query validates on the first try with no retry/feedback.""" + agent = SQLAgent() + + source = MagicMock() + source.dialect = "duckdb" + source.execute.return_value = None # valid + + agent.revise = AsyncMock( + side_effect=AssertionError("revise should not be called") + ) + + result = await agent._validate_sql( + context={}, + sql_query="SELECT * FROM CMI_C01", + expr_slug="x", + source=source, + messages=[], + step=MagicMock(), + max_retries=2, + ) + assert "CMI_C01" in result + source.get_tables.assert_not_called()