Skip to content

Return common.ai SQLToolset errors to the agent so it self-corrects#68117

Open
kaxil wants to merge 1 commit into
apache:mainfrom
astronomer:sqltoolset-return-errors-to-agent
Open

Return common.ai SQLToolset errors to the agent so it self-corrects#68117
kaxil wants to merge 1 commit into
apache:mainfrom
astronomer:sqltoolset-return-errors-to-agent

Conversation

@kaxil
Copy link
Copy Markdown
Member

@kaxil kaxil commented Jun 6, 2026

SQLToolset's query, get_schema, and list_tables tools raised on any database error, which failed the whole @task.agent task. The agent never saw the error, so it could not fix its SQL and retry within the run. A capable model that composes a slightly imperfect query (a Snowflake reserved word used as an unquoted alias such as COUNT(*) rows, or PARSE_JSON on a column that is already a VARIANT) failed the task instead of correcting itself. The same agent over Snowflake's managed MCP sql_exec_tool ran clean, because that tool returns errors to the model.

What changed

SQLToolset.call_tool now surfaces any tool failure to the agent as pydantic_ai.ModelRetry, carrying the database's own error message. It does not inspect the error type or text. pydantic-ai bounds the loop with the tool's max_retries, so a fixable SQL error is corrected within the run while an unrecoverable one (a bad connection, an auth failure) exhausts the budget and fails the task for Airflow to retry.

The airflow_toolset_to_langchain_tools bridge feeds ModelRetry back to the model as tool output, so it now applies the same max_retries bound: a tool that keeps failing propagates once the budget is exhausted instead of looping forever (the counter resets after a success or after it propagates).

SQLToolset's query, get_schema, and list_tables tools raised on any
database error, failing the whole @task.agent task. The agent never saw
the error, so it could not correct its SQL and retry within the run.

call_tool now surfaces any tool failure to the agent as ModelRetry,
carrying the database's own message, without inspecting the error type or
text. pydantic-ai bounds the loop with the tool's max_retries, so a
fixable SQL error is corrected within the run while an unrecoverable one
(a bad connection, an auth failure) exhausts the budget and fails the
task for Airflow to retry.

The airflow_toolset_to_langchain_tools bridge feeds ModelRetry back to
the model as tool output; it now bounds that by the tool's max_retries
too, so a tool that keeps failing propagates instead of looping forever.

Removes the previous _is_retryable_query_error classifier and its
per-driver exception-tuple imports: matching exception class names and
error-message substrings is brittle and cannot cover every backend.
@kaxil kaxil force-pushed the sqltoolset-return-errors-to-agent branch from 5f903be to 6daef20 Compare June 6, 2026 02:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant