Return common.ai SQLToolset errors to the agent so it self-corrects#68117
Open
kaxil wants to merge 1 commit into
Open
Return common.ai SQLToolset errors to the agent so it self-corrects#68117kaxil wants to merge 1 commit into
SQLToolset errors to the agent so it self-corrects#68117kaxil wants to merge 1 commit into
Conversation
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.
5f903be to
6daef20
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
SQLToolset'squery,get_schema, andlist_tablestools raised on any database error, which failed the whole@task.agenttask. 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 asCOUNT(*) rows, orPARSE_JSONon a column that is already aVARIANT) failed the task instead of correcting itself. The same agent over Snowflake's managed MCPsql_exec_toolran clean, because that tool returns errors to the model.What changed
SQLToolset.call_toolnow surfaces any tool failure to the agent aspydantic_ai.ModelRetry, carrying the database's own error message. It does not inspect the error type or text.pydantic-aibounds the loop with the tool'smax_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_toolsbridge feedsModelRetryback to the model as tool output, so it now applies the samemax_retriesbound: 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).