AskTable Agent Self-Correction Mechanism: AI Systems That Learn from Mistakes

AI-generated SQL may fail to execute due to incorrect table names, non-existent fields, syntax errors, etc. How can we make Agent learn from errors and automatically correct and retry?

AskTable's Agent self-correction mechanism achieves continuous optimization through a closed loop of error detection + Prompt adjustment + Case collection.

---

1. Why Do We Need Self-Correction?

1. Common AI-Generated Code Errors

Wrong Table Name:

-- AI Generated
SELECT * FROM order  -- Wrong: should be orders

-- Error Message
Table 'order' doesn't exist

Field Does Not Exist:

-- AI Generated
SELECT user_name FROM users  -- Wrong: should be username

-- Error Message
Unknown column 'user_name' in 'field list'

Syntax Error:

-- AI Generated
SELECT * FROM orders WHERE  -- Wrong: missing condition after WHERE

-- Error Message
You have an error in your SQL syntax

2. Limitations of Traditional Approaches

Direct Error Return:

• •❌ Poor user experience
• •❌ Cannot auto-fix
• •❌ Wastes user time

Infinite Retry:

• •❌ May fall into infinite loop
• •❌ Consumes large amount of Tokens
• •❌ High latency

---

2. Self-Correction Process

---

3. Core Implementation

1. Error Detection and Classification

class SQLError(Enum):
    TABLE_NOT_FOUND = "table_not_found"
    COLUMN_NOT_FOUND = "column_not_found"
    SYNTAX_ERROR = "syntax_error"
    PERMISSION_DENIED = "permission_denied"
    TIMEOUT = "timeout"
    UNKNOWN = "unknown"

def classify_error(error_message: str) -> SQLError:
    """Classify SQL errors"""
    if "doesn't exist" in error_message.lower():
        return SQLError.TABLE_NOT_FOUND
    elif "unknown column" in error_message.lower():
        return SQLError.COLUMN_NOT_FOUND
    elif "syntax" in error_message.lower():
        return SQLError.SYNTAX_ERROR
    elif "permission denied" in error_message.lower():
        return SQLError.PERMISSION_DENIED
    elif "timeout" in error_message.lower():
        return SQLError.TIMEOUT
    else:
        return SQLError.UNKNOWN

2. Retry Mechanism

async def execute_with_retry(
    question: str,
    datasource: DataSourceAdmin,
    max_retries: int = 3,
) -> QueryResult:
    """SQL execution with retry"""
    error_history: list[dict] = []

    for attempt in range(max_retries):
        try:
            # 1. Generate SQL
            sql = await generate_sql(
                question,
                datasource,
                error_history=error_history,
            )

            # 2. Execute SQL
            df = await datasource.execute_sql(sql)

            # 3. Success, save Good Case
            await save_good_case(question, sql, df)

            return QueryResult(sql=sql, dataframe=df)

        except Exception as e:
            # 4. Failure, record error
            error_type = classify_error(str(e))
            error_history.append({
                "sql": sql,
                "error": str(e),
                "error_type": error_type,
                "attempt": attempt + 1,
            })

            # 5. Last retry failed
            if attempt == max_retries - 1:
                await save_bad_case(question, sql, str(e))
                raise

            # 6. Continue retry
            log.warning(f"SQL execution failed (attempt {attempt + 1}): {e}")

3. Dynamic Prompt Adjustment

async def generate_sql(
    question: str,
    datasource: DataSourceAdmin,
    error_history: list[dict] | None = None,
) -> str:
    """Generate SQL, adjust Prompt based on error history"""
    # Base Prompt
    prompt = f"""
    Database: {datasource.to_markdown()}
    Question: {question}
    Please generate SQL query.
    """

    # If there is error history, add correction hints
    if error_history:
        prompt += "\n\n## Error History\n"
        for error in error_history:
            prompt += f"""
            Attempt {error['attempt']}:
            SQL: {error['sql']}
            Error: {error['error']}
            Error Type: {error['error_type']}

            Please correct the SQL based on the error message.
            """

    # Call LLM
    response = await llm.generate(prompt)
    return response["sql"]

4. Good/Bad Case Collection

async def save_good_case(
    question: str,
    sql: str,
    dataframe: pd.DataFrame,
):
    """Save successful case"""
    await db.execute(
        """
        INSERT INTO training_pairs (question, sql, status, created_at)
        VALUES (:question, :sql, 'good', NOW())
        """,
        {"question": question, "sql": sql},
    )

async def save_bad_case(
    question: str,
    sql: str,
    error: str,
):
    """Save failed case"""
    await db.execute(
        """
        INSERT INTO training_pairs (question, sql, error, status, created_at)
        VALUES (:question, :sql, :error, 'bad', NOW())
        """,
        {"question": question, "sql": sql, "error": error},
    )

---

4. Practical Cases

Case 1: Table Name Correction

# First attempt
sql_1 = "SELECT * FROM order"
error_1 = "Table 'order' doesn't exist"

# Second attempt (adjusted based on error)
sql_2 = "SELECT * FROM orders"  # Success!

Case 2: Field Name Correction

# First attempt
sql_1 = "SELECT user_name FROM users"
error_1 = "Unknown column 'user_name'"

# Second attempt
sql_2 = "SELECT username FROM users"  # Success!

---

5. Performance Optimization

1. Limit Retry Count

max_retries = 3  # Maximum 3 retries

Effect:

• •Avoid infinite retry
• •Control Token consumption

2. Error Caching

error_cache: dict[str, str] = {}

def get_cached_fix(sql: str, error: str) -> str | None:
    """Get cached fix for error"""
    cache_key = f"{sql}:{error}"
    return error_cache.get(cache_key)

Effect:

• •Same error uses cached fix directly
• •Reduce LLM calls

---

6. Summary

AskTable Agent's self-correction mechanism achieves:

✅ Auto-Fix: Common errors automatically corrected ✅ Continuous Optimization: Learning from mistakes ✅ User Experience: Reduce manual intervention ✅ Cost Control: Limit retry count

---

Related Reading:

• •AskTable Schema Linking

Technical Exchange:

• •Website: asktable.com