Разработка AI-агента с доступом к базе данных

Разработка AI-агента с доступом к базе данных

AI-агент с доступом к БД может самостоятельно запрашивать данные, формировать SQL-запросы на основе вопросов на естественном языке и выполнять действия записи с необходимыми проверками. Text-to-SQL — ключевая технология, позволяющая нетехническим пользователям взаимодействовать с данными через диалог.

Text-to-SQL: преобразование вопросов в SQL

from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain.agents import create_sql_agent

# Подключение к PostgreSQL
db = SQLDatabase.from_uri(
    "postgresql://user:password@localhost:5432/company_db",
    include_tables=["orders", "customers", "products", "inventory"],
    sample_rows_in_table_info=3,  # Включаем примеры данных в схему
)

llm = ChatOpenAI(model="gpt-4o", temperature=0)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)

# SQL агент с автоматическим исправлением ошибок
agent = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=10,
)

# Примеры запросов
result = agent.invoke({"input": "Каковы топ-5 клиентов по выручке за последние 3 месяца?"})
result = agent.invoke({"input": "Покажи товары с остатком на складе менее 10 единиц"})

Безопасность: только READ-доступ

from langchain_community.utilities import SQLDatabase
from sqlalchemy import create_engine, text

# Read-only пользователь PostgreSQL
READ_ONLY_USER_URI = "postgresql://readonly_user:pass@localhost:5432/db"

# Дополнительная валидация: запрет DML-операций
def validate_sql_query(query: str) -> bool:
    """Проверяет, что запрос является только SELECT"""
    forbidden_keywords = ["INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER", "TRUNCATE"]
    query_upper = query.upper()
    for keyword in forbidden_keywords:
        if keyword in query_upper:
            return False
    return True

class SafeSQLTool:
    def __init__(self, db_uri: str):
        self.engine = create_engine(db_uri)

    def execute_query(self, query: str) -> str:
        if not validate_sql_query(query):
            return "ERROR: Only SELECT queries are allowed"

        # Ограничение количества строк
        if "LIMIT" not in query.upper():
            query = f"{query.rstrip(';')} LIMIT 100"

        with self.engine.connect() as conn:
            result = conn.execute(text(query))
            rows = result.fetchall()
            columns = result.keys()
            return str([dict(zip(columns, row)) for row in rows])

Схема базы с контекстом для LLM

Качество Text-to-SQL критически зависит от качества описания схемы:

SCHEMA_CONTEXT = """
Таблицы базы данных:

1. orders (заказы)
   - id: PK, INTEGER
   - customer_id: FK → customers.id
   - status: VARCHAR (pending, confirmed, shipped, delivered, cancelled)
   - total_amount: DECIMAL(12,2) — сумма заказа в рублях
   - created_at: TIMESTAMP
   - shipped_at: TIMESTAMP (NULL если не отгружен)

2. customers (клиенты)
   - id: PK
   - name: VARCHAR — наименование компании или ФИО
   - inn: VARCHAR(12) — ИНН юрлица/ИП
   - segment: VARCHAR (enterprise, mid, small) — сегмент клиента
   - manager_id: FK → employees.id — ответственный менеджер

3. products (товары)
   - sku: VARCHAR — артикул
   - name: VARCHAR
   - category: VARCHAR
   - price_rub: DECIMAL
   - cost_rub: DECIMAL — себестоимость

ВАЖНО: Статусы заказа: 'delivered' = успешно выполнен. 'cancelled' = отменён.
       Выручка = сумма total_amount заказов со статусом 'delivered'.
"""

system_prompt = f"""Ты — аналитик данных. Переводи вопросы в SQL-запросы.
Используй следующую схему базы данных:

{SCHEMA_CONTEXT}

Правила:
- Только SELECT запросы
- Всегда добавляй LIMIT (не более 1000)
- Используй русские алиасы для читаемости
- При агрегации — добавляй ORDER BY"""

Практический кейс: BI-агент для e-commerce

Задача: аналитический ассистент для коммерческого директора — анализ продаж, ABC-анализ ассортимента, воронка заказов, cohort retention.

БД: PostgreSQL, 15 таблиц, 3M заказов.

Примеры диалогов:

Пользователь: «Какова конверсия из pending в delivered за последние 30 дней по сегментам клиентов?»

Агент генерирует:

SELECT
    c.segment AS сегмент,
    COUNT(*) FILTER (WHERE o.status = 'pending') AS ожидающих,
    COUNT(*) FILTER (WHERE o.status = 'delivered') AS доставлено,
    ROUND(
        COUNT(*) FILTER (WHERE o.status = 'delivered')::decimal /
        NULLIF(COUNT(*), 0) * 100, 1
    ) AS конверсия_pct
FROM orders o
JOIN customers c ON o.customer_id = c.id
WHERE o.created_at >= NOW() - INTERVAL '30 days'
GROUP BY c.segment
ORDER BY конверсия_pct DESC
LIMIT 100;

Результаты:

• Время получения аналитики: 2 дня → 3 минуты
• Accuracy SQL (вопросы → корректный SQL): 87%
• Типичные ошибки: неверные JOIN при сложных запросах (решается через few-shot примеры в промпте)

Few-shot примеры для повышения точности

FEW_SHOT_EXAMPLES = """
Примеры корректных запросов:

Вопрос: Топ-10 товаров по марже за последний квартал
SQL:
SELECT p.name, p.sku,
       SUM(oi.quantity * (p.price_rub - p.cost_rub)) AS маржа_руб
FROM order_items oi
JOIN products p ON oi.sku = p.sku
JOIN orders o ON oi.order_id = o.id
WHERE o.status = 'delivered'
  AND o.created_at >= DATE_TRUNC('quarter', CURRENT_DATE) - INTERVAL '3 months'
GROUP BY p.name, p.sku
ORDER BY маржа_руб DESC
LIMIT 10;

Вопрос: Средний чек по месяцам за 2025 год
SQL:
SELECT DATE_TRUNC('month', created_at) AS месяц,
       ROUND(AVG(total_amount), 0) AS средний_чек,
       COUNT(*) AS кол_заказов
FROM orders
WHERE status = 'delivered'
  AND created_at BETWEEN '2025-01-01' AND '2025-12-31'
GROUP BY 1
ORDER BY 1;
"""

Сроки

• Разработка Text-to-SQL агента: 2–3 недели
• Настройка схемы и few-shot примеров: 1 неделя
• Тестирование на реальных запросах: 1–2 недели
• Итого: 4–6 недель