Text2SQL

Сегодня хочу рассказать, как я добавил в нашу корпоративную ИИ-вики возможность общаться с продакшн-БД на русском языке.

Прежде чем добавлять Text2SQL как полноценный инструмент, хотелось протестировать, будет ли это вообще нормально работать на нашей продакшн-БД и типичных запросах пользователей. Я пошёл по этому гайду от LlamaIndex, он достаточно простой и, как оказалось, вполне рабочий. При этом, конечно, по теме Text2SQL есть куча статей со сложными идеями, хитрым парсингом схемы БД и так далее. Но мне хотелось быстро сделать что-то рабочее и полезное. Вкратце перескажу, в чём суть пайплайна.

Итоговый промпт для генерации SQL состоит из нескольких частей:

• Запрос пользователя - например, "Сколько исследований с патологией "Фиброз" было обработано в Челябинской области по ФЛГ за декабрь?"
• Описание таблиц БД - краткое саммери содержания, названия столбцов и их типы. Без этой информации будет непонятно, к каким таблицам обращаться, какие столбцы запрашивать, и какие операции на них делать. Если БД большая и сложная, то нужно сначала понять, какие именно таблицы релевантны запросу.

• Примеры строчек из таблиц - причём, желательно релевантные запросу. К примеру, название нужного кастомера по Челябинской области - "Челябинск flg". Как LLM это понять для генерации нужного запроса, если мы не покажем ей строчки с этим кастомером?

LlamaIndex из коробки предоставляет все эти возможности. В итоге для генерации промпта и SQL-запроса достаточно лишь:

• Выбрать эмбеддинг-модель для генерации эмбеддингов строчек и таблиц. Я использовал ту, что используется для индексации документов в нашей ИИ-вики
• Сгенерировать словесное саммери таблиц. Можно, с помощью LLM, но я предпочёл запариться и написать руками, чтоб добавить больше контекста
• С помощью эмбеддингов получить наиболее релевантные запросу таблицы. Количество - гмперпараметр
• С помощью эмбеддингов получить наиболее релевантные строки для выбранных таблиц
• Выбрать LLM-модель. Я использовал ту же, что уже используется в нашей ИИ-вики
• Распарсить ответ модели, чтоб получить итоговый SQL-запрос. LLMки иногда добавляют ненужные пояснения или префиксы, поэтому нужно написать функцию-парсер

Всё это заработало с ходу на простых запросах - уже прикольно. Но в более сложных случаях - когда нужно определённым образом заджойнить несколько таблиц или выдрать ответ из JSONB-колонки, иногда начинались проблемы - сгенерированный SQL падал с ошибкой.

Что бы я сделал в таком случае, если бы сам писал нужный запрос? Посмотрел бы на ошибку и задебажил бы свой запрос. Я решил поступить так же - если запрос неуспешен, я подаю на вход LLM обновлённый промпт, содержащий текст ошибки. Работает практически идеально, все мои тестовые запросы успешно отработали. По ходу ещё немного подредактировал изначальный промпт - например, попросил модель все текстовые сравнения проводить без учёта регистра.

Добавляем Text2SQLTool

Пришло время добавлять эту функциональность в ИИ-вики как полноценный инструмент. Что вообще такое "инструмент" в этом контексте? Это некий объект, у которого есть:

• Название - например, Text2SQLTool
• Описание на естественном языке - для чего нужен этот инструмент, в каких случаях его использовать. Это нужно, чтоб LLM на этапе выбора инструмента по запросу пользователя могла понять, какой инструмент сейчас наиболее релевантен. Я указал два критерия - либо если запрос о какой-то статистической информации об обработанных исследовниях, либо если пользователь прямо пишет, что это запрос к БД
• Сам код инструмента - на вход он принимает запрос пользователя и какую-то дополнительную мета-информацию (например, какие доступы есть у пользователя), а на выходе получается какая-то информация для LLM - в данном случае результаты работы SQL-запроса

Итоговый пайплайн работы выглядит так:

• LLM получает на вход запрос пользователя и набор доступных инструментов (например, поиск по внутренней документации, поиск по интернету и Text2SQL)
• Происходит выбор инструмента - в данном случае Text2SQLTool
• Инструмент получает на вход запрос пользователя и выдаёт в ответ сам SQL-запрос и результат его работы на естественном языке
• Формируется итоговый ответ на запрос пользователя

Зачем вообще всё это нужно?

Во-первых, для людей в компании, которым нужна какая-то статистическая информация, и при этом они не владеют SQL - маркетологи, проджект-менеджеры, топ-менеджмент. Да, для них можно делать какие-то дашборды, но не всё можно предусмотреть заранее.

Во-вторых, даже разработчикам и аналитикам иногда лень писать какие-то запросы, особенно сложные. Нужно вспоминать структуру БД, как называются колонки, какие у них типы, какие джойны нужно сделать...

Так или иначе, посмотрим, приживётся ли инструмент на практике. С Новым годом!