Мультиагентный SQL-ассистент, часть 2: Создание RAG-менеджера

Автор: Дмитрий Иванов [Команда P9X]

~8 минут чтения

В своём блоге я исследовал, как создать мультиагентного помощника по SQL с помощью CrewAI и Streamlit. Пользователь мог запрашивать базу данных SQLite на естественном языке. Агенты ИИ генерировали SQL-запрос на основе ввода пользователя, проверяли его и контролировали соответствие перед выполнением в базе данных для получения результатов. Я также внедрил механизм контроля человека на каждом этапе для поддержания контроля и отображал затраты на LLM, связанные с каждым запросом, для прозрачности и контроля затрат.

Хотя прототип был отличным и давал хорошие результаты для моей небольшой демонстрационной базы данных, я знал, что этого будет недостаточно для реальных баз данных. В предыдущей настройке я отправлял всю схему базы данных в качестве контекста вместе с вводом пользователя. По мере увеличения размера схем баз данных передача полной схемы в LLM увеличивает использование токенов, замедляет время отклика и повышает вероятность галлюцинаций. Мне нужен был способ передавать в LLM только соответствующие фрагменты схемы. Здесь на помощь приходит RAG (Retrieval Augmented Generation).

Создание менеджера RAG

В этом блоге я создаю менеджера RAG и добавляю несколько стратегий RAG в своего помощника по SQL, чтобы сравнить их производительность по таким показателям, как время отклика и использование токенов. Помощник теперь поддерживает четыре стратегии RAG:

  • No RAG: передаёт полную схему (базовый уровень для сравнения).
  • Keyword RAG: использует сопоставление ключевых слов для выбора релевантных таблиц.
  • FAISS RAG: использует семантическое векторное сходство через FAISS с эмбеддингами all-MiniLM-L6-v2.
  • Chroma RAG: решение для постоянного векторного хранения с ChromaDB для масштабируемого поиска в производственных условиях.

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

No RAG Strategy

Этот подход аналогичен тому, который я использовал ранее, когда отправлял всю схему базы данных в качестве контекста в LLM без какой-либо фильтрации или оптимизации. Этот подход лучше всего подходит для очень небольших схем (предпочтительно менее 10 таблиц).

Keyword RAG Strategy

В подходе Keyword RAG я использую набор предопределённых ключевых слов, сопоставленных с каждой таблицей в схеме. Когда пользователь задаёт вопрос, система проверяет наличие совпадений ключевых слов в запросе и выбирает только наиболее релевантные таблицы. Таким образом, я не отправляю всю схему в LLM — экономя токены и ускоряя работу. Он хорошо работает, когда ваша схема знакома, а ваши запросы связаны с бизнесом или следуют общим шаблонам.

FAISS RAG Approach

Стратегия FAISS RAG делает вещи умнее. Вместо того чтобы сбрасывать всю схему, я встраиваю метаданные каждой таблицы (столбцы, отношения, бизнес-контекст) в векторы с помощью трансформера предложений. Когда пользователь задаёт вопрос, он также встраивается в этот запрос, и используется FAISS для выполнения семантического поиска по «смыслу», а не только по ключевым словам. Он идеально подходит для запросов, в которых пользователи не очень конкретны или когда таблицы имеют связанные термины. Мне нравится FAISS, потому что она бесплатна, работает локально и даёт довольно точные результаты, экономя при этом токены.

Chroma RAG Strategy

Chroma RAG — это более удобная для производства версия FAISS, потому что она предлагает постоянное хранилище. Вместо того чтобы хранить встраивания в памяти, Chroma хранит их локально, поэтому даже если я перезапущу приложение, векторный индекс всё равно будет там.

Сравнение различных стратегий RAG

Этот класс RAGManager является центром управления для переключения между различными стратегиями RAG. На основе пользовательского запроса он выбирает правильный подход, извлекает наиболее релевантную часть схемы и отслеживает такие показатели производительности, как время отклика, экономия токенов и количество таблиц. У него также есть функция сравнения, позволяющая сравнивать все RAG бок о бок, и хранит исторические показатели, чтобы вы могли анализировать, как каждая из них работает с течением времени. Очень удобно для тестирования того, что работает лучше всего, и оптимизации.

Все различные классы стратегий RAG инициализируются и хранятся в self.approaches. Каждый подход RAG — это класс, который наследуется от BaseRAG, поэтому все они имеют согласованный интерфейс (get_relevant_schema() и get_approach_info()). Это означает, что вы можете легко подключить новую стратегию (например, Pinecone или Weaviate), пока она расширяет BaseRAG.

Метод get_relevant_schema() возвращает схему, соответствующую этому запросу, на основе выбранной стратегии. Если передана недопустимая стратегия или по какой-то причине произошла ошибка, он разумно возвращается к стратегии 'Keyword RAG'.

Метод compare_approaches() запускает один и тот же запрос через все стратегии RAG. Он измеряет:

  • Длину результирующей схемы.
  • % Сокращение токенов по сравнению с полной схемой.
  • Время отклика.
  • Количество возвращаемых таблиц.

Это действительно полезно для сравнения стратегий бок о бок и выбора наиболее подходящей для вашего случая использования.