Фильтрация метаданных для Databricks RAG с использованием функций ai_extract и ai_prep_search

Пайплайн Databricks, построенный на функции SQL AI ai_prep_search, уже выполняет множество задач. Входят обработанные документы, а выходят семантически обогащенные фрагменты с учетом контекста, готовые к использованию в Vector Search. Процесс прост и четок: обработка документов, подготовка фрагментов, их объединение в таблицу формата Delta и создание индекса.

Однако что-то не так, когда этот процесс применяется к корпусу документов — в данном случае к набору руководств Orion Robotics. Сами фрагменты документов хороши: поле chunk_to_embed содержит необходимый контекст документа, а chunk_to_retrieve достаточно чист для передачи в LLM. Проблемы возникают при формулировке запросов, поскольку отсутствует надежный способ фильтрации. Запрос вроде "найдите для меня руководство по прошивке Orion Motion Controller" можно удовлетворить только с помощью семантического сходства, поскольку индекс не содержит структурированных столбцов вроде document_type, product_name или target_audience.

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

Почему важны фильтры метаданных

Гибридный поиск в RAG обычно подразумевает сочетание векторного поиска с поиском по ключевым словам, таким как BM25. Это полезно, но для структурированных коллекций документов более важен другой тип гибрида: векторный поиск в сочетании с фильтрацией метаданных.

Представьте агента, который уже знает, что пользователю нужно руководство по Orion Arm. Сортировка 10 000 фрагментов из всех видов документов была бы неэффективной. Лучший подход — сначала отсортировать данные, оставив только соответствующие документы, а затем использовать векторное сходство для определения лучших результатов в этом уменьшенном наборе. Databricks Vector Search позволяет это сделать благодаря колонкам-фильтрам, задаваемым при создании индекса.

Проблема заключается в том, что метаданные должны быть чистыми, последовательными и иметь правильный тип. Если в одной строке значение document_type равно «User Manual», в другой — «manual», а в третьей — ничего, фильтры не сработают корректно. В результате качество поиска становится плохим, и требуется много времени на отладку, хотя настоящей причиной являются несогласованные метаданные, созданные ранее.

ai_extract v2.1 решает эту проблему. Он обрабатывает структурированный объект VARIANT, сгенерированный функцией ai_parse_document, и возвращает метаданные с указанием типов, проверенные по заданной схеме. Поля-энумерации фиксируют определенные значения, поля-строки обеспечивают извлечение свободного текста, а массивы используются для полей с несколькими значениями. Поскольку процесс извлечения следует схеме, результат остается одинаковым во всех документах, даже если их структура отличается.

Три функции, три задачи

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

1. ai_parse_document v2.0 преобразует неструктурированные бинарные файлы (PDF, DOCX, изображения) в структурированный объект VARIANT, выявляя блоки текста, заголовки, таблицы, рисунки, области иерархии разделов. Она знает структурное устройство документа, но не имеет представления о его содержании.

ai_extract v2.1 берет обработанный VARIANT и, согласно схеме, предоставленной пользователем, извлекает метаданные на уровне документа. Он понимает суть содержимого только потому, что схема указывает, что искать. Результатом является новый VARIANT с заполненными полями такими как document_title, document_type, product_name и target_audience, а также любыми другими, определенными в схеме. Версия 2.1 ввела показатели уверенности и ссылки, и является рекомендуемой к использованию версией.

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

У каждой функции есть своя задача, и вся система состоит из объединения этих функций.

Шаг за шагом в ноутбуке

Ноутбук работает с корпусом данных Orion Robotics, который хранится в томе Unity Catalog и содержит PDF-файлы, файлы DOCX и изображения. Сначала производится парсинг:

SELECT
  path,
  ai_parse_document(
    content,
    MAP('version', '2.0')
  ) AS parsed
FROM READ_FILES(
  '/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/orion_docs/',
  format => 'binaryFile'
)

Каждый файл в томе считывается в двоичном формате и преобразуется в структурированный объект VARIANT. Разумно явно указать версию 2.0, поскольку в этой схеме содержатся координаты прямоугольника и описания фигур, от которых зависит функция ai_prep_search при принятии решений о разбиении на части. Если использовать значение по умолчанию, Databricks может изменить поведение в рамках данной обработки.

Пример запуска дал следующий результат:

Следующий запрос извлекает метаданные из обработанного результата:

WITH parsed_documents AS (
  SELECT path, ai_parse_document(content, MAP('version', '2.0')) AS parsed
  FROM READ_FILES('/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/orion_docs/', format => 'binaryFile')
)
SELECT
  path,
  ai_extract(
    parsed,
    '{
      "document_title":  {"type": "string", "description": "Main title or heading of the document"},
      "document_type":   {"type": "enum", "labels": ["manual", "datasheet", "guide", "report", "specification", "policy", "other"]},
      "product_name":    {"type": "string", "description": "Primary product or system referenced"},
      "version":         {"type": "string", "description": "Document or product version if mentioned"},
      "key_topics":      {"type": "array", "items": {"type": "string"}},
      "target_audience": {"type": "enum", "labels": ["engineer", "operator", "manager", "customer", "general"]}
    }',
    MAP(
      'version', '2.1',
      'instructions', 'These are technical documents from Orion Robotics. Extract metadata that would help users filter and find relevant documents.'
    )
  ) AS metadata
FROM parsed_documents

Некоторые выборы требуют пояснения. Поля document_type и target_audience являются перечислениями, поскольку фильтры нуждаются в замкнутом наборе значений. В виде свободного текста они могут возвращаться в формате «manual», «User Manual» или «technical manual», что нарушает процесс сопоставления. Параметр instructions предоставляет дополнительную информацию — в данном случае о том, что файлы являются техническими документами от Orion Robotics, — что улучшает качество извлечения. Кроме того, функция ai_extract принимает данные VARIANT из ai_parse_document в их первоначальном виде, без преобразования в обычный текст.

Результат примерной работы:

Далее следует разбиение на части:

WITH parsed_documents AS ( ... )
SELECT
  path,
  ai_prep_search(parsed) AS chunked
FROM parsed_documents

ai_prep_search использует тот же самый разобранный ВАРИАНТ. Документы делятся на фрагменты, хранящиеся в массиве document.contents, причем каждый фрагмент содержит поля chunk_id, chunk_position, chunk_to_retrieve и chunk_to_embed. Последние два поля преднамеренно не содержат один и тот же текст.

chunk_to_embed включает дополнительный контекст, такой как название документа, заголовок раздела и детали страницы, чтобы модель встраивания могла правильно разместить фрагмент. chunk_to_retrieve — это более чистый текст, обычно тот, который отображается пользователям или передается в LLM.

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

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

Собирание всего вместе

Объединенный запрос — это место, где встречаются все компоненты:

WITH parsed_documents AS (
  SELECT path, ai_parse_document(content, MAP('version', '2.0')) AS parsed
  FROM READ_FILES('/Volumes/<YOUR_CATALOG>/<YOUR_SCHEMA>/orion_docs/', format => 'binaryFile')
),
doc_metadata AS (
  SELECT
    path,
    parsed,
    ai_extract(
      parsed,
      '{ ... }',
      MAP('version', '2.1', 'instructions', 'These are technical documents from Orion Robotics. Extract metadata for search filtering.')
    ) AS metadata
  FROM parsed_documents
),
doc_chunks AS (
  SELECT path, ai_prep_search(parsed) AS chunked
  FROM doc_metadata
)
SELECT
  chunk.value:chunk_id::STRING          AS chunk_id,
  chunk.value:chunk_position::INT       AS chunk_position,
  chunk.value:chunk_to_retrieve::STRING AS chunk_to_retrieve,
  chunk.value:chunk_to_embed::STRING    AS chunk_to_embed,
  doc_chunks.chunked:document.source_uri::STRING AS source_uri,
  dm.metadata:response:document_title::STRING    AS document_title,
  dm.metadata:response:document_type::STRING     AS document_type,
  dm.metadata:response:product_name::STRING      AS product_name,
  dm.metadata:response:version::STRING           AS doc_version,
  dm.metadata:response:target_audience::STRING   AS target_audience
FROM
  doc_chunks
  JOIN doc_metadata dm ON doc_chunks.path = dm.path,
  LATERAL variant_explode(doc_chunks.chunked:document.contents) AS chunk

Шаг LATERAL variant_explode преобразует вложенный массив чанков в строки. Операция JOIN по пути привязывает метаданные на уровне документа к каждому чанку этого документа. В результате получается по одной строке на каждый чанк, в которой содержатся как поля для поиска, так и столбцы метаданных, подлежащие фильтрации. Именно такая структура хранится в таблице Delta и используется для пополнения индекса векторного поиска.

Результат примерной работы:

Таблица формируется вот так:

CREATE OR REPLACE TABLE <YOUR_CATALOG>.<YOUR_SCHEMA>.orion_docs_rag_chunks AS
-- [full CTE chain from Step 4]

После создания таблицы формируется индекс Mosaic AI Vector Search Delta Sync с следующими настройками:

1. Столбец исходных данных для встраивания: chunk_to_embed
2. Основной ключ: chunk_id
3. Поля для фильтрации: doc_version вместе с document_type, product_name и target_audience

chunk_id остается неизменным при каждой обработке и состоит из идентификатора документа и позиции его части. Благодаря этому при постепенных обновлениях таблицы Delta части документа, которые не изменились, остаются без изменений, вместо того чтобы их переиндексировать, что значительно экономит ресурсы при регулярном обновлении корпуса данных.

Преимущества при выполнении запросов

Полезность этого подхода проявляется при выполнении запроса пользователем, а не во время загрузки данных. Предположим, запрос звучит так: «Найти руководство по прошивке для Orion Motion Controller». Теперь поиск может выполняться в два этапа:

1. Определите соответствующие фильтры метаданных, такие как document_type = „guide“ и product_name = „Orion Motion Controller“.
2. Включите запрос и выполняйте поиск только среди документов, соответствующих этим фильтрам.

При отсутствии функции ai_extract агент может полагаться только на семантическое сходство. Для небольшого количества похожих документов это может подойти. Однако при обработке тысяч файлов, охватывающих множество продуктов, типов и аудиторий, такой подход становится ненадежным. Использование метаданных на основе перечислений решает эту проблему: агенту не нужно гадать, является ли документ руководством или инструкцией; достаточно просто использовать фильтр document_type = „manual“ для получения точных результатов.

Ограничения и затраты

ai_prep_search не позволяет контролировать процесс разбиения на части. Размер блоков, степень их перекрытия и алгоритм разбиения не могут быть настроены. Для многих случаев использования технологии RAG с документами это приемлемо, поскольку семантическое разбиение с добавлением контекста часто эффективнее простого разделения на фиксированные по размеру блоки. Однако для реализации структуры «родитель-потомок» или точного контроля перекрытия всё равно требуется пользовательская логика.

ai_extract увеличивает затраты и время обработки, поскольку для каждого документа требуется дополнительный вызов большой языковой модели. Функция Delta Sync может пропускать неизмененные документы, но первая полная обработка всё равно может быть затратной, поэтому разумно провести тестирование на выборке перед обработкой всего корпуса.

Используйте версию 2.0 функции ai_parse_document, когда её результат в формате VARIANT передается в функции ai_prep_search или ai_extract. Эта версия предоставляет информацию о рамках, типах элементов и описаниях изображений. Чистый текст можно передавать в функцию ai_extract, однако тогда теряются преимущества структурированной обработки.

Завершение работы над дизайном

Добавление ai_extract v2.1 требует немного дополнительного кода, но существенно меняет архитектуру, поскольку теперь фрагменты содержат чистые, типизированные и фильтруемые метаданные, а гибридный метод поиска становится более надежным. Процесс остается состоящим из трех четких этапов:

1. ai_parse_document определяет структуру документа.
2. ai_extract собирает полезные метаданные.
3. ai_prep_search формирует семантические фрагменты.

Такое разделение оказывается эффективным при работе с документами разных продуктов, версий, аудиторий и типов. При отсутствии фильтров метаданных поиск основан в основном на поиске по сходству. С их использованием агент сначала сужает круг возможных вариантов, а затем формирует рейтинги.

Этот подход не подойдет для всех случаев. Для работы с метаданными на уровне фрагментов, поиском родительских/дочерних элементов или пользовательским разделением на фрагменты всё равно требуется специальный код. Однако для многих пайплайнов RAG в Databricks это более чистый и простой способ начала работы. К тому же фрагменты всё равно необходимо тестировать и оценивать, нельзя слепо им доверять.