Text to SQL Agent with LangChain

@tool
def get_database_schema(table_name: str = None) -> str:
    """Get database schema information for SQL query generation.
    Use this first to understand table structure before creating queries."""
    print(f"🔍 Getting schema for: {table_name if table_name else 'all tables'}")

    if table_name:
        try:
            tables = db.get_usable_table_names()
            if table_name.lower() in [t.lower() for t in tables]:
                result = db.get_table_info([table_name])
                print(f"✔ Retrieved schema for table: {table_name}")
                return result
            else:
                return f"Error: Table '{table_name}' not found. Available tables: {', '.join(tables)}"
        except Exception as e:
            return f"Error getting table info: {e}"
    else:
        print("✔ Retrieved full database schema")
        return SCHEMA

result = get_database_schema.invoke({'table_name': 'departments'})
print(result)

🔍 Getting schema for: departments
✔ Retrieved schema for table: departments

CREATE TABLE departments (
	dept_no CHAR(4) NOT NULL,
	dept_name VARCHAR(40) NOT NULL,
	PRIMARY KEY (dept_no),
	UNIQUE (dept_name)
)

/*
3 rows from departments table:
dept_no	dept_name
d009	Customer Service
d005	Development
d002	Finance
*/

@tool
def generate_sql_query(question: str, schema_info: str = None) -> str:
    """Generate a SQL SELECT query from a natural language question using database schema.
    Always use this after getting schema information."""
    print(f"🔺 Generating SQL for: {question[:100]}...")

    schema_to_use = schema_info if schema_info else SCHEMA

    prompt = f"""
                Based on this database schema:
                {schema_to_use}

                Generate a SQL query to answer this question: {question}

                Rules:
                - Use only SELECT statements
                - Include only existing columns and tables
                - Add appropriate WHERE, GROUP BY, ORDER BY clauses as needed
                - Limit results to 10 rows unless specified otherwise
                - Use proper SQL syntax for SQLite

                Return only the SQL query, nothing else.
                """

    try:
        response = llm.invoke(prompt)
        query = response.content.strip()
        print(f"✔ Generated SQL query")
        return query
    except Exception as e:
        return f"Error generating query: {e}"

@tool
def validate_sql_query(query: str) -> str:
    """Validate SQL query for safety and syntax before execution.
    Returns 'Valid: <query>' if safe or 'Error: <message>' if unsafe."""
    print(f"🔍 Validating SQL: {query[:100]}...")

    clean_query = query.strip()

    # Remove markdown SQL code block markers (```sql ... ```)
    clean_query = re.sub(r'```sql\s*', '', clean_query, flags=re.IGNORECASE)
    clean_query = re.sub(r'```\s*', '', clean_query)
    clean_query = clean_query.strip().rstrip(";")

    # Check 1: Must be a SELECT statement
    if not clean_query.lower().startswith("select"):
        return "Error: Only SELECT statements are allowed"

    # Check 2: Block dangerous SQL keywords
    dangerous_keywords = ['INSERT', 'UPDATE', 'DELETE', 'ALTER', 'DROP', 'CREATE', 'TRUNCATE']
    query_upper = clean_query.upper()

    for keyword in dangerous_keywords:
        if keyword in query_upper:
            return f"Error: {keyword} operations are not allowed"

    print("✔ Query validation passed")
    return f"Valid: {clean_query}"

@tool
def execute_sql_query(query: str) -> str:
    """Execute a validated SQL query and return results.
    Only use this after validating the query for safety."""
    print(f"🚀 Executing SQL: {query[:100]}...")

    try:
        clean_query = query.strip()
        if clean_query.startswith("Valid: "):
            clean_query = clean_query[7:]  # Remove "Valid: " prefix

        clean_query = re.sub(r'```sql\s*', '', clean_query, flags=re.IGNORECASE)
        clean_query = re.sub(r'```\s*', '', clean_query)
        clean_query = clean_query.strip().rstrip(";")

        result = db.run(clean_query)
        print("✔ Query executed successfully")

        if result:
            return f"Query Results:\n{result}"
        else:
            return "Query executed successfully but returned no results."

    except Exception as e:
        error_msg = f"Execution Error: {str(e)}"
        print(f"✘ {error_msg}")
        return error_msg

@tool
def fix_sql_error(original_query: str, error_message: str, question: str) -> str:
    """Fix a failed SQL query by analyzing the error and generating a corrected version.
    Use this when validation or execution fails."""
    print(f"🔺 Fixing SQL error: {error_message[:100]}...")

    fix_prompt = f"""
                    The following SQL query failed:
                    Query: {original_query}
                    Error: {error_message}
                    Original Question: {question}

                    Database Schema:
                    {SCHEMA}

                    Analyze the error and provide a corrected SQL query that:
                    1. Fixes the specific error mentioned
                    2. Still answers the original question
                    3. Uses only valid table and column names from the schema
                    4. Follows SQLite syntax rules

                    Return only the corrected SQL query, nothing else.
                    """

    try:
        response = llm.invoke(fix_prompt)
        fixed_query = response.content.strip()
        print("✔ Generated fixed SQL query")
        return fixed_query
    except Exception as e:
        return f"Error generating fix: {e}"

SQL_SYSTEM_PROMPT = f"""You are an expert SQL analyst working with an employees database.

Database Schema:
{SCHEMA}

Your workflow for answering questions:
1. Use `get_database_schema` first to understand available tables and columns (if needed)
2. Use `generate_sql_query` to create SQL based on the question
3. Use `validate_sql_query` to check the query for safety and syntax
4. Use `execute_sql_query` to run the validated query
5. If there's an error, use `fix_sql_error` to correct it and try again (up to 3 times)
6. Provide a clear answer based on the query results

Rules:
- Always follow the workflow step by step
- If a query fails, use the fix tool and try again
- Provide clear, informative answers
- Be precise with table and column names
- Handle errors gracefully and try to fix them
- If you fail after 3 attempts, explain what went wrong

Available tools for each step:
- get_database_schema: Get table structure info
- generate_sql_query: Create SQL from question
- validate_sql_query: Check query safety/syntax
- execute_sql_query: Run the query
- fix_sql_error: Fix failed queries

Remember: Always validate queries before executing them for safety.
"""

def ask_sql(question: str):
    """Ask the SQL agent a question using the full workflow."""
    print(f"\n{'='*60}")
    print(f"SQL AGENT - Question: {question}")
    print('='*60)

    for event in sql_agent.stream({"messages": question}, stream_mode="values"):
        msg = event["messages"][-1]

        if hasattr(msg, 'tool_calls') and msg.tool_calls:
            for tc in msg.tool_calls:
                print(f"\n🔺 Using Tool: {tc['name']}")
                print(f"Args: {str(tc['args'])[:200]}")

        elif hasattr(msg, 'content') and msg.content:
            print(f"\n💼 Answer:\n{msg.content}")

ask_sql("What is the average salary of employees in the Sales department?")

============================================================
SQL AGENT - Question: What is the average salary of employees in the Sales department?
============================================================

💼 Answer:
What is the average salary of employees in the Sales department?

🔺 Using Tool: get_database_schema
Args: {'table_name': 'departments'}
🔍 Getting schema for: departments
✔ Retrieved schema for table: departments

💼 Answer:
CREATE TABLE departments (
	dept_no CHAR(4) NOT NULL,
	dept_name VARCHAR(40) NOT NULL,
	...

🔺 Using Tool: generate_sql_query
Args: {'question': 'What is the average salary of employees in the Sales department?', 'schema_info': '...'}
🔺 Generating SQL for: What is the average salary of employees in the Sales department?...
✔ Generated SQL query

💼 Answer:
SELECT AVG(salary) AS average_salary
FROM employees
JOIN departments ON employees.dept_no = departments.dept_no
WHERE departments.dept_name = 'Sales';

🔺 Using Tool: validate_sql_query
Args: {'query': 'SELECT AVG(salary) AS average_salary\nFROM employees\nJOIN departments ON...'}
🔍 Validating SQL: SELECT AVG(salary) AS average_salary...
✔ Query validation passed

💼 Answer:
Valid: SELECT AVG(salary) AS average_salary
FROM employees
JOIN departments ON employees.dept_no = departments.dept_no
WHERE departments.dept_name = 'Sales'

🔺 Using Tool: execute_sql_query
Args: {'query': 'SELECT AVG(salary) AS average_salary\nFROM employees\nJOIN dept_emp...'}
🚀 Executing SQL: SELECT AVG(salary) AS average_salary
FROM employees
JOIN dept_emp ON employees.emp_no = dept_emp.emp_no...
✘ Execution Error: (sqlite3.OperationalError) near "JOIN": syntax error
[SQL: SELECT AVG(salary)...JOIN salaries ON employees.emp_no = salaries.emp_no]

💼 Answer:
Execution Error: near "JOIN": syntax error...

🔺 Using Tool: fix_sql_error
Args: {'error_message': 'near "JOIN": syntax error', 'original_query': '...'}

User Question: "What is the average salary in Sales?"
     │
     ▼
get_database_schema("departments")    ← Understand table structure
     │
     ▼
generate_sql_query(question, schema)  ← LLM writes the SELECT query
     │
     ▼
validate_sql_query(query)             ← Safety check — must be SELECT only
     │         │
     │  [FAIL] └──── Error: dangerous keyword → STOP
     │
     ▼
execute_sql_query(query)              ← Run against SQLite
     │         │
     │  [FAIL] └──── Execution error → fix_sql_error() → retry (max 3x)
     │
     ▼
LLM synthesizes final answer from query results