Розробка AI-генерації SQL-запитів за текстовим описом (Text-to-SQL)

AI-генерація SQL з текстових запитів (Text-to-SQL)

Text-to-SQL дозволяє бізнес-аналітикам та менеджерам напрямку задавати питання базі даних без знання SQL. Ключовий технічний виклик — не генерація SQL синтаксису (LLM справляється легко), а передача точного контексту схеми: які таблиці існують, як вони пов'язані, які значення допустимі в enum-полях.

Архітектура Text-to-SQL системи

from anthropic import Anthropic
import psycopg2
import json
from typing import Optional
from dataclasses import dataclass

client = Anthropic()

@dataclass
class QueryResult:
    sql: str
    explanation: str
    rows: list[dict]
    error: Optional[str] = None

class TextToSQLEngine:

    def __init__(self, connection_string: str):
        self.conn = psycopg2.connect(connection_string)
        self.schema_cache: dict = {}

    def get_schema(self, tables: list[str] = None) -> str:
        """Отримує DDL схему з PostgreSQL"""

        query = """
        SELECT
            t.table_name,
            c.column_name,
            c.data_type,
            c.is_nullable,
            c.column_default,
            tc.constraint_type,
            kcu.column_name as fk_column,
            ccu.table_name as fk_table
        FROM information_schema.tables t
        JOIN information_schema.columns c ON t.table_name = c.table_name
        LEFT JOIN information_schema.key_column_usage kcu
            ON c.table_name = kcu.table_name AND c.column_name = kcu.column_name
        LEFT JOIN information_schema.table_constraints tc
            ON kcu.constraint_name = tc.constraint_name
        LEFT JOIN information_schema.constraint_column_usage ccu
            ON tc.constraint_name = ccu.constraint_name
        WHERE t.table_schema = 'public'
        """

        if tables:
            placeholders = ",".join(["%s"] * len(tables))
            query += f" AND t.table_name IN ({placeholders})"

        with self.conn.cursor() as cur:
            cur.execute(query, tables or [])
            rows = cur.fetchall()

        # Форматуємо як DDL
        tables_dict = {}
        for row in rows:
            table_name = row[0]
            if table_name not in tables_dict:
                tables_dict[table_name] = {"columns": [], "foreign_keys": []}

            col_def = f"  {row[1]} {row[2].upper()}"
            if row[3] == "NO":
                col_def += " NOT NULL"
            if row[4]:
                col_def += f" DEFAULT {row[4]}"
            if row[5] == "PRIMARY KEY":
                col_def += " PRIMARY KEY"

            tables_dict[table_name]["columns"].append(col_def)

            if row[5] == "FOREIGN KEY" and row[7]:
                tables_dict[table_name]["foreign_keys"].append(
                    f"  FOREIGN KEY ({row[6]}) REFERENCES {row[7]}"
                )

        ddl_parts = []
        for table, info in tables_dict.items():
            ddl = f"CREATE TABLE {table} (\n"
            ddl += ",\n".join(info["columns"])
            if info["foreign_keys"]:
                ddl += ",\n" + ",\n".join(info["foreign_keys"])
            ddl += "\n);"
            ddl_parts.append(ddl)

        return "\n\n".join(ddl_parts)

    def get_sample_values(self, important_columns: dict[str, list[str]]) -> str:
        """Отримує приклади значень для enum/category полів"""
        samples = []

        with self.conn.cursor() as cur:
            for table_col, _ in important_columns.items():
                table, col = table_col.split(".")
                try:
                    cur.execute(
                        f"SELECT DISTINCT {col} FROM {table} LIMIT 10"
                    )
                    values = [str(row[0]) for row in cur.fetchall()]
                    samples.append(f"-- {table}.{col}: {', '.join(values)}")
                except Exception:
                    pass

        return "\n".join(samples)

    def generate_sql(self, question: str, context_tables: list[str] = None) -> QueryResult:
        """Генерує SQL з текстового запитання"""

        schema = self.get_schema(context_tables)

        # Додатковий контекст: приклади значень для рядкових полів
        sample_values = self._get_relevant_samples(question)

        response = client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=2048,
            system="""Ти — експерт по SQL та PostgreSQL.
Генеруй точні, оптимізовані SQL запити на основі схеми БД.

Правила:
- Використовуй тільки існуючі таблиці та колонки з схеми
- Переважай JOIN замість підзапитів де можливо
- Додавай LIMIT 1000 для запитів без агрегації
- Для дат використовуй PostgreSQL функції: DATE_TRUNC, NOW(), EXTRACT
- Завжди додавай ORDER BY для передбачуваності результатів
- Якщо запитання неоднозначне — вибирай найймовірнішу інтерпретацію

Поверни JSON:
{
  "sql": "<SQL запит>",
  "explanation": "<пояснення що робить запит, 1-2 речення>",
  "assumptions": ["<допущення 1 якщо були>"]
}""",
            messages=[{
                "role": "user",
                "content": f"""Запитання: {question}

Схема бази даних:
```sql
{schema}

{f"Приклади значень:{chr(10)}{sample_values}" if sample_values else ""}""" }] )

    text = response.content[0].text
    try:
        # Парсимо JSON ответ
        start = text.find("{")
        end = text.rfind("}") + 1
        data = json.loads(text[start:end])

        sql = data["sql"]
        explanation = data.get("explanation", "")

        # Виконуємо запит
        rows = self._execute_safe(sql)

        return QueryResult(sql=sql, explanation=explanation, rows=rows)

    except Exception as e:
        return QueryResult(sql="", explanation="", rows=[], error=str(e))

def _execute_safe(self, sql: str) -> list[dict]:
    """Виконує тільки SELECT запити"""
    sql_upper = sql.strip().upper()
    if not sql_upper.startswith("SELECT") and not sql_upper.startswith("WITH"):
        raise ValueError("Only SELECT queries are allowed")

    with self.conn.cursor() as cur:
        cur.execute(sql)
        columns = [desc[0] for desc in cur.description]
        rows = cur.fetchall()
        return [dict(zip(columns, row)) for row in rows]

def _get_relevant_samples(self, question: str) -> str:
    """Проста евристика для визначення релевантних enum полів"""
    # У реальній системі — LLM визначає потрібні поля
    return ""

### Самокорегуючийся генератор

```python
class SelfCorrectingTextToSQL:
    """Ітеративно виправляє SQL при помилках виконання"""

    def __init__(self, engine: TextToSQLEngine):
        self.engine = engine

    def query(self, question: str, max_attempts: int = 3) -> QueryResult:
        """Генерує SQL з автоматичним виправленням помилок"""

        result = self.engine.generate_sql(question)
        if not result.error:
            return result

        # Ітеративно виправляємо
        messages = [{
            "role": "user",
            "content": f"Запитання: {question}\n\nСгенерував запит:\n```sql\n{result.sql}\n```\n\nПомилка: {result.error}\n\nВиправ запит."
        }]

        for attempt in range(max_attempts - 1):
            response = client.messages.create(
                model="claude-sonnet-4-5",
                max_tokens=1024,
                system="Ти — SQL експерт. Виправляй SQL запити по помилках виконання. Поверни тільки виправлений SQL.",
                messages=messages,
            )

            fixed_sql = response.content[0].text.strip()
            if "```sql" in fixed_sql:
                fixed_sql = fixed_sql.split("```sql")[1].split("```")[0].strip()

            try:
                rows = self.engine._execute_safe(fixed_sql)
                return QueryResult(sql=fixed_sql, explanation="Auto-corrected", rows=rows)
            except Exception as e:
                messages.append({"role": "assistant", "content": response.content[0].text})
                messages.append({"role": "user", "content": f"Все ще помилка: {e}"})

        return QueryResult(sql=result.sql, rows=[], error="Max attempts reached", explanation="")

NL інтерфейс з історією

class ConversationalDataAnalyst:
    """Діаловий інтерфейс для роботи з даними"""

    def __init__(self, connection_string: str):
        self.engine = TextToSQLEngine(connection_string)
        self.history: list[dict] = []
        self.last_sql: str = ""

    def ask(self, question: str) -> str:
        """Відповідає на запитання з урахуванням історії діалогу"""

        # Додаємо контекст попереднього запиту
        context = ""
        if self.last_sql:
            context = f"\nПопередній запит:\n```sql\n{self.last_sql}\n```"

        # Підтримка уточнювальних запитань
        if any(word in question.lower() for word in ["і також", "тепер", "додай", "також"]):
            enhanced = f"На основі попереднього запиту, {question}"
        else:
            enhanced = question

        result = self.engine.generate_sql(enhanced + context)

        if result.error:
            return f"Помилка виконання запиту: {result.error}"

        self.last_sql = result.sql
        self.history.append({"question": question, "sql": result.sql})

        # Форматуємо результат
        if not result.rows:
            return "Запит виконаний успішно, але даних не знайдено."

        response_text = f"{result.explanation}\n\n"
        response_text += f"SQL: `{result.sql}`\n\n"
        response_text += f"Результати ({len(result.rows)} рядків):\n"

        # Таблиця результатів
        if result.rows:
            headers = list(result.rows[0].keys())
            response_text += " | ".join(headers) + "\n"
            response_text += " | ".join(["---"] * len(headers)) + "\n"
            for row in result.rows[:10]:
                response_text += " | ".join(str(v) for v in row.values()) + "\n"
            if len(result.rows) > 10:
                response_text += f"... і ще {len(result.rows) - 10} рядків"

        return response_text

Практичний кейс: аналітика e-commerce

Задача: продакт-менеджери формували задачи аналітикам (2–5 днів очікування), тому що не знали SQL. База даних: PostgreSQL, 23 таблиці, ~50M записів.

Впровадження:

• Text-to-SQL інтерфейс в Slack: /data <запитання>
• White-list дозволених таблиць для продактів (без особистих даних)
• Кешування часто задаваних запитань

Метрики:

• ad-hoc запити від продактів без участі аналітиків: 0 → 23 на тиждень
• Час отримання ответу на простий запит: 2 дні → 30 секунд
• Точність генерованого SQL: 89% (не потребують правки)
• 11% запитів потребували ітеративного уточнення через діалог

Терміни

• Базовий двигун (екстракція схеми + генерація SQL): 1 тиждень
• Самокорегуючийся генератор + історія: 1-2 тижні
• Діаловий інтерфейс з Slack інтеграцією: 1 тиждень
• Кешування запитів + оптимізація продуктивності: 1 тиждень