Error recovery agent

When SQL execution fails, the Error Recovery Agent analyzes the error and attempts to fix the query. It retries up to 3 times before giving up.

Error recovery loop

How errors trigger the recovery and retry cycle.

text2sql_agent.py
python
def error_agent(state: AgentState) -> AgentState:
    """Handle errors and attempt to fix the SQL query"""
    error = state["error"]
    sql_query = state["sql_query"]
    question = state["question"]
    iteration = state.get("iteration", 0)

    # Give up after 3 tries
    if iteration > 3:
        state["final_answer"] = f"I couldn't generate a correct query. Error: {error}"
        return state

    prompt = f"""The following SQL query failed. Please fix it.

{SCHEMA_INFO}

Original Question: {question}
Failed SQL Query: {sql_query}
Error: {error}

Generate a corrected SQL query. Return ONLY the SQL, no explanation."""

    response = completion(
        model=DEFAULT_MODEL,
        messages=[
            {"role": "system", "content": AGENT_CONFIGS["error_agent"]["system_prompt"]},
            {"role": "user", "content": prompt}
        ],
        temperature=0
    )

    corrected = response.choices[0].message.content.strip()
    corrected = corrected.replace("```sql", "").replace("```", "").strip()

    state["sql_query"] = corrected
    state["error"] = ""  # Clear error for retry
    state["iteration"] = iteration + 1

    return state

The Error Agent uses the error message to fix queries.

After 3 failed attempts, the agent gives up gracefully and returns an error message to the user. Some queries are genuinely impossible given the schema. In production, you might log these failures to improve the system prompt or add new examples to reduce future errors.

Flashcards: Flashcards

Loading practice…

The Error Recovery Agent makes the system resilient to SQL failures. One more agent to go: the Analysis Agent that turns raw results into human-readable answers.