Создание агентного RAG для приложений Text-to-SQL

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

Видение агентного RAG для Text-to-SQL

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

1. Разбиваем задачи на более мелкие и управляемые инструменты в рамках агента
2. Улучшаем точность, назначая задачи специализированным инструментам
3. Увеличиваем прозрачность, отслеживая рассуждения и рабочий процесс каждого инструмента
4. Упрощаем масштабирование и отладку благодаря модульному дизайну

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

Обзор архитектуры

Структура состоит из агента, использующего инструменты в процессе Text-to-SQL. Процесс можно обобщить следующим образом:

Инструмент преобразования пользовательского запроса → Инструмент преобразования запроса → Инструмент подсказок Few Shot → Гибридный инструмент поиска → Инструмент повторного ранжирования → Инструмент поиска таблиц → Инструмент создания подсказок → Инструмент выполнения LLM → Инструмент выполнения SQL → Конечный результат

1. Инструмент преобразования пользовательского запроса

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

Улучшения

• Обработка временных ссылок. Сопоставление терминов, таких как “на сегодняшний день” или “до настоящего момента”, с явными датами.
• Замена неоднозначных слов. Например, “недавние” можно заменить на “за последние 7 дней”.
• Сопоставление сокращений или аббревиатур с их расшифровкой.

Пример

Ввод: “Показать недавние продажи MTD.”

Преобразованный запрос: “Извлечь данные о продажах за последние 7 дней (С начала месяца).”

from datetime import date, timedelta

def transform_query(user_query):
    # Обработка неявных временных ссылок
    today = date.today()
    transformations = {
        "as of today": f"up to {today}",
        "till now": f"up to {today}",
        "recent": "last 7 days",
        "last week": f"from {today - timedelta(days=7)} to {today}",
    }
    
    for key, value in transformations.items():
        user_query = user_query.replace(key, value)

    # Сопоставление общих сокращений
    abbreviations = {
        "MTD": "Month to Date",
        "YTD": "Year to Date",
    }
    for abbr, full_form in abbreviations.items():
        user_query = user_query.replace(abbr, full_form)

    return user_query

query_transform_tool = Tool(
    name="Query Transformer",
    func=transform_query,
    description="Refines user queries for clarity and specificity, handles abbreviations and open-ended terms."
)

2. Инструмент подсказок Few Shot

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

Пример рабочего процесса

1. Входящий вопрос: “Покажите мне общие продажи по продукту за 7 дней.”

2. Предопределенные шаблоны:

• “Показать продажи, сгруппированные по регионам.” → Пример SQL; SELECT region, SUM(sales) …
• “Показать общие продажи по продукту.” → Пример SQL; SELECT название_продукта, SUM(продажи) …

3. Самый похожий вопрос: “Показать общие продажи по продукту.”

4. Пример SQL вывода: SELECT название_продукта, SUM(продажи) FROM …

from langchain.chat_models import ChatOpenAI

llm = ChatOpenAI(model="gpt-4")

predefined_examples = {
    "Show sales grouped by region": "SELECT region, SUM(sales) FROM sales_data GROUP BY region;",
    "Show total sales by product": "SELECT product_name, SUM(sales) FROM sales_data GROUP BY product_name;",
}

def find_similar_question(user_query):
    prompt = "Find the most similar question type for the following user query:\n"
    prompt += f"User Query: {user_query}\n\nOptions:\n"
    for example in predefined_examples.keys():
        prompt += f"- {example}\n"
    prompt += "\nRespond with the closest match."

    response = llm.call_as_function(prompt)
    most_similar = response['content']
    return predefined_examples.get(most_similar, "")

few_shot_tool = Tool(
    name="Few-Shot Prompting",
    func=find_similar_question,
    description="Finds the most similar question type using an additional LLM call and retrieves the corresponding example SQL."
)

3. Гибридный инструмент поиска

Для надежного поиска этот инструмент объединяет семантический поиск, поиск по ключевым словам на основе BM25 и сопоставление на основе ключевых слов. Результаты поиска из этих методов объединяются с помощью объединения рангов взаимной связи.

Как это все сочетается?

Сопоставление таблиц и ключевых слов

Этот подход сопоставляет таблицы с ключевыми словами, содержащимися в запросе. Например:

• Наличие “продаж” приводит к отбору таблицы продаж.
• Наличие “продукта” приводит к отбору таблицы продуктов. 

Сопоставление перекрытия ключевых слов (BM25)

Это метод поиска на основе перекрытия ключевых слов, который отбирает таблицы в соответствии с релевантностью. Для этого мы применим технику BM25. Он сортирует статьи в порядке релевантности для поиска пользователя. Этот метод поиска учитывает насыщенность терминов, а также TF-IDF (Частота Термина – Обратная Частота Документа). 

Частота термина (TF) помогает измерить частоту термина в данном документе. Подход обратной частоты документа (IDF) подчеркивает слова, которые появляются в каждом документе, уменьшая их важность.

Нормализация учитывает длину документа, чтобы избежать предвзятости к более длинным статьям.

Учитывая:

• sales_data: Содержит термины “продажи,” “дата,” “продукт.”
• products: Содержит термины “продукт,” “категория.”
• orders: Содержит термины “заказ,” “дата,” “клиент.”
• financials: Содержит термины “выручка,” “прибыль,” “расход.”

Запрос пользователя: “Показать общие продажи по продукту.”

• Определите термины в запросе пользователя: [“продажи,” “продукт”].
• Отсортируйте каждый документ (на основе частоты и релевантности этих терминов) в DataBaseTable.

Релевантность документов:

• sales: Высокая релевантность из-за “продаж” и “продукт”
• products: Высокая релевантность из-за “продукт.”
• orders: Низкая релевантность из-за наличия только “продаж.”
• financials: Не релевантно.

Вывод:

Ранжированный список: [products, sales_data, orders, financials]

Семантический поиск

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

Слияние обратных рангов

Совмещает результаты BM25 и семантического поиска с использованием стратегии объединения обратного ранга, которая более подробно объяснена ниже:

Объединение обратного ранга (RRF) совмещение результатов BM25 и семантического поиска:

RRF – это метод объединения результатов из нескольких алгоритмов ранжирования (например, BM25 и семантический поиск). Он назначает оценку каждому документу на основе его ранга в индивидуальных методах, присваивая более высокие баллы документам, занимающим более высокие позиции по нескольким методам.

Формула RRF:

RRF(d) = Σ(r ∈ R) 1 / (k + r(d))

Где:

• d – документ
• R – набор ранжеров (методы поиска)
• k – константа (обычно 60)
• r(d) – ранг документа d в методе поиска r

Пошаговый пример

Входные данные.

1. Результаты ранжирования BM25:

• продукты (Ранг 1)
• данные о продажах (Ранг 2)
• заказы (Ранг 3)

2. Результаты ранжирования семантического поиска:

• данные о продажах (Ранг 1)
• финансы (Ранг 2)
• продукты (Ранг 3)

Объединение пошагово

Для каждой таблицы вычислите оценку:

1. данные о продажах

• Ранг BM25 = 2, Семантический Ранг = 1
• Оценка RRF = (1/60+2) + (1/60+1) = 0.03252

2. продукты

• Ранг BM25 = 1, Семантический Ранг = 3
• Оценка RRF = (1/60+1) + (1/60+3) = 0.03226

3. заказы

• BM25 Ранг = 3, Семантический Ранг = Не оценен
• RRF Оценка = (1/60+3)= 0.01587

4. финансовые данные

• BM25 Ранг = Не оценен, Семантический Ранг = 2
• RRF Оценка = (1/60+2)=0.01613

5. Сортировать по RRF оценке

• данные_по_продажам (высокая оценка из-за высокого ранга в семантическом поиске).
• продукты (высокая оценка от BM25).
• заказы (низкая релевантность в целом).
• финансовые данные (ограниченное совпадение).

Итоговый вывод: [‘данные_по_продажам’, ‘продукты,’ ‘финансовые данные,’ ‘заказы’]

Таблицы, полученные с помощью Ключевой таблицы отображения, всегда включаются.

from rank_bm25 import BM25Okapi

def hybrid_search(query):
    # Отображение на основе ключевых слов
    keyword_to_table = {
        "sales": "sales_data",
        "product": "products",
    }
    keyword_results = [table for keyword, table in keyword_to_table.items() if keyword in query.lower()]

    # Поиск BM25
    bm25 = BM25Okapi(["sales_data", "products", "orders", "financials"])
    bm25_results = bm25.get_top_n(query.split(), bm25.corpus, n=5)

    # Семантический поиск
    semantic_results = vector_store.similarity_search(query, k=5)

    # Объединение рангов
    def reciprocal_rank_fusion(results):
        rank_map = {}
        for rank, table in enumerate(results):
            rank_map[table] = rank_map.get(table, 0) + 1 / (1 + rank)
        return sorted(rank_map, key=rank_map.get, reverse=True)

    combined_results = reciprocal_rank_fusion(bm25_results + semantic_results)

    return list(set(keyword_results + combined_results))

hybrid_search_tool = Tool(
    name="Hybrid Search",
    func=hybrid_search,
    description="Combines keyword mapping, BM25, and semantic search with RRF for table retrieval."
)

4. Инструмент повторной оценки

Этот инструмент гарантирует приоритет наиболее релевантных таблиц.

Пример

• Входные таблицы: [“данные_по_продажам,” “продукты,” “финансовые данные”]
• Логика повторной оценки
  • Для каждой таблицы вычислите оценку релевантности, конкатенируя запрос и описание таблицы.
  • Сортируйте по оценке релевантности.

• Вывод: [“данные_по_продажам,” “продукты”]

Немного подробнее о логике повторного ранжирования:

Кросс–кодировщик вычисляет оценку релевантности, анализируя соединенный запрос и описание таблицы как единую пару входных данных. Этот процесс включает в себя:

• Парное ввод. Запрос и каждое описание таблицы объединяются и передаются как входные данные в кросс-кодировщик.
• Совместное кодирование. В отличие от отдельных кодировщиков (например, би-кодировщиков), кросс-кодировщик совместно кодирует пару, что позволяет лучше захватывать контекст и зависимости между запросом и описанием таблицы.
• Оценка. Модель выводит оценку релевантности для каждой пары, указывая, насколько хорошо таблица соответствует запросу.

from transformers import pipeline

reranker = pipeline("text-classification", model="cross-encoder/ms-marco-TinyBERT-L-2")

def re_rank_context(query, results):
    scores = [(doc, reranker(query + " " + doc)[0]['score']) for doc in results]
    return [doc for doc, score in sorted(scores, key=lambda x: x[1], reverse=True)]

re_rank_tool = Tool(
    name="Re-Ranker",
    func=re_rank_context,
    description="Re-ranks the retrieved tables based on relevance to the query."
)

5. Инструмент создания подсказок

Этот инструмент создает подробную подсказку для языковой модели, включая уточненный запрос пользователя, извлеченную схему и примеры из инструмента Few-Shot Prompting.

Предположим, вы человек, который хорошо разбирается в генерации SQL-запросов. Сгенерируйте SQL-запрос для: Получения общих продаж по продуктам за последние 7 дней.

Актуальные таблицы:

1. sales_data: Содержит столбцы [sales, date, product_id].
2. products: Содержит столбцы [product_id, product_name].

Пример SQL:

SELECT product_name, SUM(sales) FROM sales_data JOIN products ON sales_data.product_id = products.product_id GROUP BY product_name;

Будущие перспективы

Хотя эта система использует один агент с несколькими инструментами, чтобы упростить модульность и снизить сложность, в будущем можно рассмотреть многопользовательскую структуру. Мы могли бы рассмотреть следующее:

1. Посвященные агенты для извлечения контекста. Отдельные агенты для семантических и ключевых поисков.
2. Задачеориентированные агенты. Агенты, специализированные на проверке или оптимизации SQL-запросов.
3. Сотрудничество между агентами. Использование координационного агента для управления делегированием задач.

Такой подход может улучшить масштабируемость и позволить создавать более сложные рабочие процессы, особенно в корпоративных развертываниях.

Заключение

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