sql_agent.py

import re

FORBIDDEN_SQL = re.compile(
    r"\b(DELETE|UPDATE|DROP|ALTER|INSERT|TRUNCATE|CREATE|REPLACE|ATTACH|DETACH|PRAGMA|VACUUM)\b",
    re.IGNORECASE,
)


def _chat_with_ollama(prompt):
    try:
        import ollama
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "The ollama Python package is not installed. Run: pip install ollama"
        ) from exc

    return ollama.chat(
        model="qwen2.5:7b",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ]
    )


def _clean_sql(sql):

    sql = sql.strip()

    sql = sql.replace("```sql", "")
    sql = sql.replace("```", "")
    sql = sql.strip()

    matches = re.findall(
        r"SELECT\b.*?(?:;|$)",
        sql,
        re.IGNORECASE | re.DOTALL
    )

    if not matches:
        raise ValueError("No valid SELECT query found.")

    sql = matches[0].strip()

    sql = sql.rstrip(";").strip() + ";"

    return sql


def validate_select_sql(sql):

    if FORBIDDEN_SQL.search(sql):
        raise ValueError("Unsafe SQL blocked.")

    cleaned = _clean_sql(sql)

    if FORBIDDEN_SQL.search(cleaned):
        raise ValueError("Unsafe SQL blocked.")

    if cleaned.count(";") > 1:
        raise ValueError("Only one SQL statement allowed.")

    if not cleaned.upper().startswith("SELECT"):
        raise ValueError("Only SELECT queries are allowed.")

    return cleaned


def select_relevant_tables(question, table_names, conversation_context=""):
    """Identify which database tables are needed to answer the user's question.
    If there are few tables (<= 15), returns all table names directly without filtering.
    """
    if len(table_names) <= 15:
        return table_names

    prompt = f"""
Identify which database tables are required to answer the user's question.
Respond ONLY with a comma-separated list of the relevant table names from the list below. Do not include any explanations, markdown, or extra words.

Table Names:
{", ".join(table_names)}

User Question:
{question}

Relevant Tables:
"""
    try:
        response = _chat_with_ollama(prompt)
        content = response["message"]["content"].strip()
        # Clean potential markdown or extra symbols
        content = content.replace("`", "").replace("[", "").replace("]", "").replace("'", "").replace('"', "")
        selected = [t.strip() for t in content.split(",") if t.strip() in table_names]
        if not selected:
            return table_names
        return selected
    except Exception:
        return table_names


def generate_sql(question, schema_text, relationship_text="", conversation_context=""):
    prompt = f"""
You are an expert MySQL query generator.

Database Relationships:
{relationship_text}

Database Schema:
{schema_text}

Conversation Context:
{conversation_context or "No previous context."}

Rules:

1. Return ONLY a SELECT statement.
2. Do not include any explanations, markdown, or comments.
3. Use only tables and columns present in the schema.
4. Use the provided conversation context for follow-up queries.
5. Generate valid MySQL syntax only. Do NOT use SQLite-specific syntax like COLLATE NOCASE.
6. Return exactly one SELECT query.
7. Do not generate UPDATE, DELETE, INSERT, DROP, ALTER, or other non-SELECT statements.

Examples:

Question:
Who works in sales?
SQL:
SELECT name FROM employees WHERE department='sales';

Question:
Who has the highest salary?
SQL:
SELECT name, salary FROM employees ORDER BY salary DESC LIMIT 1;

Question:
{question}
"""
    print("\n===================")
    print("QUESTION:", question)
    print("SCHEMA LENGTH:", len(schema_text))
    print("PROMPT LENGTH:", len(prompt))
    print("===================\n")

    response = _chat_with_ollama(prompt)

    return validate_select_sql(
        response["message"]["content"]
    )


def is_database_question(
    question,
    schema_text,
    conversation_context=""
):

    prompt = f"""
Database Schema:
{schema_text}

Conversation Context:
{conversation_context or "No previous context."}

Can the following question be answered using this database?

Question:
{question}

Respond ONLY:

YES

or

NO
"""

    response = _chat_with_ollama(prompt)

    answer = response["message"]["content"].strip().upper()

    return "YES" in answer


def generate_corrected_sql(
    question,
    schema_text,
    relationship_text="",
    failed_sql="",
    error_message="",
    conversation_context=""
):
    prompt = f"""
You are an expert MySQL query generator.
A previously generated SQL query failed with an execution error. Please correct the SQL query to resolve the error.

Database Relationships:
{relationship_text}

Database Schema:
{schema_text}

Conversation Context:
{conversation_context or "No previous context."}

User Question:
{question}

Failed SQL:
{failed_sql}

Error Message:
{error_message}

Rules:
1. Correct the SQL query to fix the error.
2. Return ONLY a SELECT statement.
3. Do not include any explanations, markdown, or comments.
4. Use only tables and columns present in the schema.
5. Generate valid MySQL syntax only. Do NOT use SQLite-specific syntax.
6. Return exactly one SELECT query.
7. Do not generate UPDATE, DELETE, INSERT, DROP, ALTER, or other non-SELECT statements.

Corrected SQL:
"""
    print("\n===================")
    print("CORRECTING SQL FOR QUESTION:", question)
    print("FAILED SQL:", failed_sql)
    print("ERROR:", error_message)
    print("===================\n")

    response = _chat_with_ollama(prompt)

    return validate_select_sql(
        response["message"]["content"]
    )