Фільтрація метаданих у Databricks RAG за допомогою ai_extract та ai_prep_search

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

Проте щось бракує, коли цей процес працює з набором документів — у цьому випадку з набором посібників Orion Robotics. Самі фрагменти документів є якісними. Поле chunk_to_embed містить належний контекст документа, а chunk_to_retrieve достатньо чисте для передачі до ШІ. Проблеми виникають під час формування запиту, оскільки немає надійного способу фільтрації. Запит на кшталт "Знайдіть для мене посібник з прошивкою для 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 використовує той самий розібраний VARIANT. Документи діляться на фрагменти, які зберігаються у масиві 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 — цей крок розширює вкладену масивну структуру chunk у рядки. Операція JOIN за шляхом додає метадані на рівні документа до кожного chunk цього документа. Результатом є по одному рядку на кожен chunk, який містить як поля для пошуку, так і стовпці метаданих, що піддаються фільтрації. Саме така структура зберігається в таблиці 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 залишається стабільним між запусками та складається з ідентифікатора документа та позиції його частини. Завдяки цьому інкрементальні оновлення за допомогою MERGE у таблиці 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 збільшує витрати та час обробки, оскільки для кожного документа потрібен додатковий виклик LLM. 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 це чистіший та простіший спосіб для початку. Крім того, частини все одно потрібно тестувати та оцінювати, і ніколи не варто довіряти їм сліпо.